63 lines
1.7 KiB
Go
63 lines
1.7 KiB
Go
package meta
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
)
|
|
|
|
// proxy brokers a connection from src to dst
|
|
func proxy(dst, src *net.TCPConn) error {
|
|
// channels to wait on the close event for each connection
|
|
serverClosed := make(chan struct{}, 1)
|
|
clientClosed := make(chan struct{}, 1)
|
|
errors := make(chan error, 1)
|
|
|
|
go broker(dst, src, clientClosed, errors)
|
|
go broker(src, dst, serverClosed, errors)
|
|
|
|
// wait for one half of the proxy to exit, then trigger a shutdown of the
|
|
// other half by calling CloseRead(). This will break the read loop in the
|
|
// broker and allow us to fully close the connection cleanly without a
|
|
// "use of closed network connection" error.
|
|
var waitFor chan struct{}
|
|
select {
|
|
case <-clientClosed:
|
|
// the client closed first and any more packets from the server aren't
|
|
// useful, so we can optionally SetLinger(0) here to recycle the port
|
|
// faster.
|
|
dst.SetLinger(0)
|
|
dst.CloseRead()
|
|
waitFor = serverClosed
|
|
case <-serverClosed:
|
|
src.CloseRead()
|
|
waitFor = clientClosed
|
|
case err := <-errors:
|
|
src.CloseRead()
|
|
dst.SetLinger(0)
|
|
dst.CloseRead()
|
|
return err
|
|
}
|
|
|
|
// Wait for the other connection to close.
|
|
<-waitFor
|
|
return nil
|
|
}
|
|
|
|
// This does the actual data transfer.
|
|
// The broker only closes the Read side.
|
|
func broker(dst, src net.Conn, srcClosed chan struct{}, errors chan error) {
|
|
// We can handle errors in a finer-grained manner by inlining io.Copy (it's
|
|
// simple, and we drop the ReaderFrom or WriterTo checks for
|
|
// net.Conn->net.Conn transfers, which aren't needed). This would also let
|
|
// us adjust buffersize.
|
|
_, err := io.Copy(dst, src)
|
|
|
|
if err != nil {
|
|
errors <- err
|
|
}
|
|
if err := src.Close(); err != nil {
|
|
errors <- err
|
|
}
|
|
srcClosed <- struct{}{}
|
|
}
|