diff --git a/transport/transport.go b/transport/transport.go index 82f5f09..d2319d0 100644 --- a/transport/transport.go +++ b/transport/transport.go @@ -4,6 +4,10 @@ import ( "net" ) +// Transport wraps and unwraps a proxy protocol frame on a net.Conn. +// Wrap writes the protocol header (including target address) to the connection. +// Unwrap reads and validates the protocol header, returning the target address. +// Neither method closes the connection; the caller is responsible for that. type Transport interface { Unwrap(conn net.Conn) (net.Addr, error) Wrap(conn net.Conn, tgtHost string, tgtPort uint16) error diff --git a/transport/transport_test.go b/transport/transport_test.go index e0c3453..5756caa 100644 --- a/transport/transport_test.go +++ b/transport/transport_test.go @@ -2,10 +2,10 @@ package transport import ( "net" - "strings" "testing" "github.com/sower-proxy/conns/reread" + "github.com/sower-proxy/sower/transport/socks5" "github.com/sower-proxy/sower/transport/sower" "github.com/sower-proxy/sower/transport/trojan" ) @@ -22,14 +22,32 @@ func testPipe(tran Transport) (net.Addr, error) { return tran.Unwrap(reread.New(r)) } +// testPipeDirect is like testPipe but passes the raw conn to Unwrap +// (no reread wrapper), needed for transports where Unwrap writes back. +func testPipeDirect(tran Transport) (net.Addr, error) { + r, w := net.Pipe() + defer r.Close() + + go func(w net.Conn) { + defer w.Close() + tran.Wrap(w, "sower", 443) + }(w) + + return tran.Unwrap(r) +} + func Test_Transports(t *testing.T) { - if addr, err := testPipe(newSower()); err != nil || strings.TrimSpace(addr.String()) != "sower:443" { + if addr, err := testPipe(newSower()); err != nil || addr.String() != "sower:443" { t.Errorf("test sower, unexpected address: %s, err: %s", addr, err) } - if addr, err := testPipe(newTrojan()); err != nil || strings.TrimSpace(addr.String()) != "sower:443" { + if addr, err := testPipe(newTrojan()); err != nil || addr.String() != "sower:443" { t.Errorf("test trojan, unexpected address: %s, err: %s", addr, err) } + + if addr, err := testPipeDirect(socks5.New()); err != nil || addr.String() != "sower:443" { + t.Errorf("test socks5, unexpected address: %s, err: %s", addr, err) + } } func newSower() *sower.Sower {