Shaun Campbell
865f7c26c9
All checks were successful
Build and Push Image / docker-build-and-push (push) Successful in 2m28s
83 lines
2.3 KiB
Go
83 lines
2.3 KiB
Go
// Package proxy transports an opaque bidirectional byte stream.
|
|
package proxy
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
)
|
|
|
|
type copyResult struct {
|
|
err error
|
|
dst net.Conn
|
|
}
|
|
|
|
type closeWriter interface {
|
|
CloseWrite() error
|
|
}
|
|
|
|
// TrafficObserver receives each positive byte count as it is copied.
|
|
// clientToBackend identifies the PostgreSQL client-to-server direction.
|
|
type TrafficObserver func(clientToBackend bool, bytes int64)
|
|
|
|
// Bidirectional copies bytes in both directions. TCP peers retain half-close
|
|
// semantics; a TLS client is closed when its upstream direction has ended,
|
|
// because crypto/tls exposes no CloseWrite operation.
|
|
func Bidirectional(client, backend net.Conn) error {
|
|
return BidirectionalWithTraffic(client, backend, nil)
|
|
}
|
|
|
|
// BidirectionalWithTraffic copies bytes in both directions and reports every
|
|
// positive read to observer when one is provided.
|
|
func BidirectionalWithTraffic(client, backend net.Conn, observer TrafficObserver) error {
|
|
results := make(chan copyResult, 2)
|
|
copyStream := func(dst, src net.Conn, clientToBackend bool) {
|
|
_, err := io.Copy(dst, trafficReader{Reader: src, clientToBackend: clientToBackend, observer: observer})
|
|
results <- copyResult{err: err, dst: dst}
|
|
}
|
|
go copyStream(backend, client, true)
|
|
go copyStream(client, backend, false)
|
|
|
|
first := <-results
|
|
if writer, ok := first.dst.(closeWriter); ok {
|
|
_ = writer.CloseWrite()
|
|
} else {
|
|
// TLS cannot be half-closed. Closing unblocks the opposite copy and
|
|
// prevents an idle peer from retaining a handler goroutine.
|
|
_ = client.Close()
|
|
}
|
|
second := <-results
|
|
if writer, ok := second.dst.(closeWriter); ok {
|
|
_ = writer.CloseWrite()
|
|
}
|
|
|
|
return combine(first.err, second.err)
|
|
}
|
|
|
|
type trafficReader struct {
|
|
io.Reader
|
|
clientToBackend bool
|
|
observer TrafficObserver
|
|
}
|
|
|
|
func (r trafficReader) Read(buffer []byte) (int, error) {
|
|
bytes, err := r.Reader.Read(buffer)
|
|
if bytes > 0 && r.observer != nil {
|
|
r.observer(r.clientToBackend, int64(bytes))
|
|
}
|
|
return bytes, err
|
|
}
|
|
|
|
func combine(errs ...error) error {
|
|
var relevant []error
|
|
for _, err := range errs {
|
|
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
|
|
relevant = append(relevant, err)
|
|
}
|
|
}
|
|
return errors.Join(relevant...)
|
|
}
|
|
|
|
// Compile-time assertion: a TCP connection has practical half-close support.
|
|
var _ closeWriter = (*net.TCPConn)(nil)
|