diff --git a/connector/masque/association.go b/connector/masque/association.go new file mode 100644 index 00000000..0beb5e61 --- /dev/null +++ b/connector/masque/association.go @@ -0,0 +1,328 @@ +package masque + +import ( + "context" + "errors" + "net" + "os" + "sync" + "time" + + "github.com/go-gost/core/logger" +) + +const udpAssociationBufferSize = 64 * 1024 + +var errMissingDestination = errors.New("masque: destination address required") + +type udpAssociationDialFunc func(ctx context.Context, addr net.Addr) (net.PacketConn, error) + +type udpAssociationResult struct { + data []byte + addr net.Addr +} + +type udpAssociationTunnel struct { + conn net.PacketConn +} + +type udpAssociationConn struct { + ctx context.Context + cancel context.CancelFunc + localAddr net.Addr + dial udpAssociationDialFunc + closeIdle func() error + log logger.Logger + timeout time.Duration + + closed chan struct{} + closeOnce sync.Once + closeErr error + results chan udpAssociationResult + + mu sync.Mutex + tunnels map[string]*udpAssociationTunnel + readDeadline time.Time + readDeadlineChanged chan struct{} +} + +func newUDPAssociationConn( + ctx context.Context, + localAddr net.Addr, + timeout time.Duration, + dial udpAssociationDialFunc, + closeIdle func() error, + log logger.Logger, +) *udpAssociationConn { + ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + c := &udpAssociationConn{ + ctx: ctx, + cancel: cancel, + localAddr: localAddr, + dial: dial, + closeIdle: closeIdle, + log: log, + timeout: timeout, + closed: make(chan struct{}), + results: make(chan udpAssociationResult, 32), + tunnels: make(map[string]*udpAssociationTunnel), + readDeadlineChanged: make(chan struct{}), + } + return c +} + +func (c *udpAssociationConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { + for { + select { + case <-c.closed: + return 0, nil, net.ErrClosed + default: + } + + c.mu.Lock() + deadline := c.readDeadline + deadlineChanged := c.readDeadlineChanged + c.mu.Unlock() + + var ( + timer *time.Timer + timeout <-chan time.Time + ) + if !deadline.IsZero() { + remaining := time.Until(deadline) + if remaining <= 0 { + return 0, nil, os.ErrDeadlineExceeded + } + timer = time.NewTimer(remaining) + timeout = timer.C + } + + select { + case result := <-c.results: + if timer != nil { + timer.Stop() + } + return copy(b, result.data), result.addr, nil + case <-deadlineChanged: + if timer != nil { + timer.Stop() + } + continue + case <-timeout: + return 0, nil, os.ErrDeadlineExceeded + case <-c.closed: + if timer != nil { + timer.Stop() + } + return 0, nil, net.ErrClosed + } + } +} + +func (c *udpAssociationConn) WriteTo(b []byte, addr net.Addr) (int, error) { + if addr == nil { + return 0, errMissingDestination + } + + tunnel, err := c.tunnel(addr) + if err != nil { + if c.isClosed() { + return 0, net.ErrClosed + } + c.logError(addr, err) + return len(b), nil + } + n, err := tunnel.conn.WriteTo(b, addr) + if err == nil { + return n, nil + } + if c.isClosed() { + return 0, net.ErrClosed + } + if c.removeTunnel(addr.String(), tunnel) { + c.logError(addr, err) + } + return len(b), nil +} + +func (c *udpAssociationConn) tunnel(addr net.Addr) (*udpAssociationTunnel, error) { + key := addr.String() + + c.mu.Lock() + select { + case <-c.closed: + c.mu.Unlock() + return nil, net.ErrClosed + default: + } + + if tunnel := c.tunnels[key]; tunnel != nil { + c.mu.Unlock() + return tunnel, nil + } + c.mu.Unlock() + + ctx := c.ctx + if c.timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, c.timeout) + defer cancel() + } + conn, err := c.dial(ctx, addr) + if err != nil { + return nil, err + } + if conn == nil { + return nil, errors.New("masque: nil UDP tunnel") + } + + tunnel := &udpAssociationTunnel{conn: conn} + c.mu.Lock() + select { + case <-c.closed: + c.mu.Unlock() + conn.Close() + return nil, net.ErrClosed + default: + } + if current := c.tunnels[key]; current != nil { + c.mu.Unlock() + conn.Close() + return current, nil + } + c.tunnels[key] = tunnel + c.mu.Unlock() + go c.readTunnel(key, tunnel) + return tunnel, nil +} + +func (c *udpAssociationConn) readTunnel(key string, tunnel *udpAssociationTunnel) { + buf := make([]byte, udpAssociationBufferSize) + for { + n, addr, err := tunnel.conn.ReadFrom(buf) + if err != nil { + if c.isClosed() { + return + } + if c.removeTunnel(key, tunnel) { + c.logError(key, err) + } + return + } + + data := append([]byte(nil), buf[:n]...) + select { + case c.results <- udpAssociationResult{data: data, addr: addr}: + case <-c.closed: + return + } + } +} + +func (c *udpAssociationConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +func (c *udpAssociationConn) logError(addr any, err error) { + if c.log != nil { + c.log.Warnf("masque: UDP tunnel %v: %v", addr, err) + } +} + +func (c *udpAssociationConn) removeTunnel(key string, tunnel *udpAssociationTunnel) bool { + c.mu.Lock() + removed := c.tunnels[key] == tunnel + if c.tunnels[key] == tunnel { + delete(c.tunnels, key) + } + c.mu.Unlock() + tunnel.conn.Close() + return removed +} + +func (c *udpAssociationConn) Read(b []byte) (int, error) { + n, _, err := c.ReadFrom(b) + return n, err +} + +func (c *udpAssociationConn) Write(b []byte) (int, error) { + return 0, errMissingDestination +} + +func (c *udpAssociationConn) Close() error { + c.closeOnce.Do(func() { + close(c.closed) + c.cancel() + + c.mu.Lock() + close(c.readDeadlineChanged) + tunnels := make([]*udpAssociationTunnel, 0, len(c.tunnels)) + for _, tunnel := range c.tunnels { + tunnels = append(tunnels, tunnel) + } + clear(c.tunnels) + c.mu.Unlock() + + var errs []error + for _, tunnel := range tunnels { + if err := tunnel.conn.Close(); err != nil { + errs = append(errs, err) + } + } + if c.closeIdle != nil { + if err := c.closeIdle(); err != nil { + errs = append(errs, err) + } + } + c.closeErr = errors.Join(errs...) + }) + return c.closeErr +} + +func (c *udpAssociationConn) LocalAddr() net.Addr { + return c.localAddr +} + +func (c *udpAssociationConn) RemoteAddr() net.Addr { + return &net.UDPAddr{} +} + +func (c *udpAssociationConn) SetDeadline(t time.Time) error { + if err := c.SetReadDeadline(t); err != nil { + return err + } + return c.SetWriteDeadline(t) +} + +func (c *udpAssociationConn) SetReadDeadline(t time.Time) error { + c.mu.Lock() + defer c.mu.Unlock() + select { + case <-c.closed: + return net.ErrClosed + default: + } + c.readDeadline = t + close(c.readDeadlineChanged) + c.readDeadlineChanged = make(chan struct{}) + return nil +} + +func (c *udpAssociationConn) SetWriteDeadline(t time.Time) error { + select { + case <-c.closed: + return net.ErrClosed + default: + } + return nil +} + +var ( + _ net.Conn = (*udpAssociationConn)(nil) + _ net.PacketConn = (*udpAssociationConn)(nil) +) diff --git a/connector/masque/association_test.go b/connector/masque/association_test.go new file mode 100644 index 00000000..ebd983ba --- /dev/null +++ b/connector/masque/association_test.go @@ -0,0 +1,321 @@ +package masque + +import ( + "context" + "errors" + "net" + "sync" + "testing" + "time" +) + +func TestUDPAssociationConnRoutesByDestination(t *testing.T) { + type dialCall struct { + addr net.Addr + conn *fakePacketConn + } + + var ( + mu sync.Mutex + calls []dialCall + ) + dial := func(ctx context.Context, addr net.Addr) (net.PacketConn, error) { + conn := newFakePacketConn() + mu.Lock() + calls = append(calls, dialCall{addr: addr, conn: conn}) + mu.Unlock() + return conn, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + conn := newUDPAssociationConn(ctx, &net.UDPAddr{}, 0, dial, nil, nil) + cancel() // The route dial context ends before the UDP association is used. + t.Cleanup(func() { conn.Close() }) + + addr1 := &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 53} + addr2 := &net.UDPAddr{IP: net.ParseIP("8.8.8.8"), Port: 53} + if _, err := conn.WriteTo([]byte("first"), addr1); err != nil { + t.Fatalf("first WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("second"), addr1); err != nil { + t.Fatalf("second WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("third"), addr2); err != nil { + t.Fatalf("third WriteTo failed: %v", err) + } + + mu.Lock() + gotCalls := append([]dialCall(nil), calls...) + mu.Unlock() + if len(gotCalls) != 2 { + t.Fatalf("expected one tunnel per destination, got %d", len(gotCalls)) + } + if gotCalls[0].addr.String() != addr1.String() || gotCalls[1].addr.String() != addr2.String() { + t.Fatalf("unexpected tunnel destinations: %s, %s", gotCalls[0].addr, gotCalls[1].addr) + } + if got := gotCalls[0].conn.writtenData(); len(got) != 2 || string(got[0]) != "first" || string(got[1]) != "second" { + t.Fatalf("unexpected writes for first destination: %q", got) + } + if got := gotCalls[1].conn.writtenData(); len(got) != 1 || string(got[0]) != "third" { + t.Fatalf("unexpected writes for second destination: %q", got) + } + + replyAddr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 1} + gotCalls[1].conn.queueRead([]byte("reply"), replyAddr) + buf := make([]byte, 32) + n, addr, err := conn.ReadFrom(buf) + if err != nil { + t.Fatalf("ReadFrom failed: %v", err) + } + if string(buf[:n]) != "reply" { + t.Fatalf("unexpected reply: %q", buf[:n]) + } + if addr.String() != replyAddr.String() { + t.Fatalf("expected reply from %s, got %s", replyAddr, addr) + } +} + +func TestUDPAssociationConnClose(t *testing.T) { + upstream := newFakePacketConn() + idleClosed := 0 + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 0, + func(context.Context, net.Addr) (net.PacketConn, error) { return upstream, nil }, + func() error { + idleClosed++ + return nil + }, + nil, + ) + + if _, err := conn.WriteTo([]byte("request"), &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 53}); err != nil { + t.Fatalf("WriteTo failed: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + if err := conn.Close(); err != nil { + t.Fatalf("second Close failed: %v", err) + } + if idleClosed != 1 { + t.Fatalf("expected idle stream to be closed once, got %d", idleClosed) + } + if !upstream.isClosed() { + t.Fatal("expected destination tunnel to be closed") + } + if _, err := conn.WriteTo([]byte("request"), &net.UDPAddr{}); !errors.Is(err, net.ErrClosed) { + t.Fatalf("expected net.ErrClosed after Close, got %v", err) + } +} + +func TestUDPAssociationConnTunnelFailureIsIsolated(t *testing.T) { + var ( + mu sync.Mutex + calls []*fakePacketConn + ) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 0, + func(context.Context, net.Addr) (net.PacketConn, error) { + pc := newFakePacketConn() + mu.Lock() + calls = append(calls, pc) + mu.Unlock() + return pc, nil + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + addr1 := &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443} + addr2 := &net.UDPAddr{IP: net.ParseIP("8.8.8.8"), Port: 443} + if _, err := conn.WriteTo([]byte("first"), addr1); err != nil { + t.Fatalf("first WriteTo failed: %v", err) + } + if _, err := conn.WriteTo([]byte("second"), addr2); err != nil { + t.Fatalf("second WriteTo failed: %v", err) + } + + mu.Lock() + first, second := calls[0], calls[1] + mu.Unlock() + first.queueReadError(errors.New("stream reset")) + select { + case <-first.closed: + case <-time.After(time.Second): + t.Fatal("failed tunnel was not removed") + } + + second.queueRead([]byte("reply"), addr2) + buf := make([]byte, 32) + n, addr, err := conn.ReadFrom(buf) + if err != nil { + t.Fatalf("healthy tunnel failed after peer reset: %v", err) + } + if string(buf[:n]) != "reply" || addr.String() != addr2.String() { + t.Fatalf("unexpected healthy tunnel reply %q from %v", buf[:n], addr) + } + + if _, err := conn.WriteTo([]byte("retry"), addr1); err != nil { + t.Fatalf("retry WriteTo failed: %v", err) + } + mu.Lock() + callCount := len(calls) + mu.Unlock() + if callCount != 3 { + t.Fatalf("expected failed destination to redial, got %d dials", callCount) + } +} + +func TestUDPAssociationConnDialTimeout(t *testing.T) { + dialErr := make(chan error, 1) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 20*time.Millisecond, + func(ctx context.Context, _ net.Addr) (net.PacketConn, error) { + <-ctx.Done() + dialErr <- ctx.Err() + return nil, ctx.Err() + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + payload := []byte("request") + n, err := conn.WriteTo(payload, &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443}) + if err != nil || n != len(payload) { + t.Fatalf("timed out datagram was not dropped cleanly: n=%d err=%v", n, err) + } + if err := <-dialErr; !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected dial deadline, got %v", err) + } +} + +func TestUDPAssociationConnDialDoesNotHoldAssociationLock(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + conn := newUDPAssociationConn( + context.Background(), + &net.UDPAddr{}, + 0, + func(context.Context, net.Addr) (net.PacketConn, error) { + close(started) + <-release + return newFakePacketConn(), nil + }, + nil, + nil, + ) + t.Cleanup(func() { conn.Close() }) + + writeDone := make(chan struct{}) + go func() { + conn.WriteTo([]byte("request"), &net.UDPAddr{IP: net.ParseIP("1.1.1.1"), Port: 443}) + close(writeDone) + }() + <-started + + deadlineDone := make(chan error, 1) + go func() { + deadlineDone <- conn.SetReadDeadline(time.Now().Add(time.Second)) + }() + select { + case err := <-deadlineDone: + if err != nil { + t.Fatalf("SetReadDeadline failed: %v", err) + } + case <-time.After(time.Second): + t.Fatal("dial held the association lock") + } + + close(release) + <-writeDone +} + +type fakePacket struct { + data []byte + addr net.Addr + err error +} + +type fakePacketConn struct { + reads chan fakePacket + closed chan struct{} + closeOnce sync.Once + mu sync.Mutex + writes [][]byte +} + +func newFakePacketConn() *fakePacketConn { + return &fakePacketConn{ + reads: make(chan fakePacket, 1), + closed: make(chan struct{}), + } +} + +func (c *fakePacketConn) queueRead(data []byte, addr net.Addr) { + c.reads <- fakePacket{data: append([]byte(nil), data...), addr: addr} +} + +func (c *fakePacketConn) queueReadError(err error) { + c.reads <- fakePacket{err: err} +} + +func (c *fakePacketConn) writtenData() [][]byte { + c.mu.Lock() + defer c.mu.Unlock() + result := make([][]byte, len(c.writes)) + for i, data := range c.writes { + result[i] = append([]byte(nil), data...) + } + return result +} + +func (c *fakePacketConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +func (c *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) { + select { + case packet := <-c.reads: + if packet.err != nil { + return 0, nil, packet.err + } + return copy(b, packet.data), packet.addr, nil + case <-c.closed: + return 0, nil, net.ErrClosed + } +} + +func (c *fakePacketConn) WriteTo(b []byte, _ net.Addr) (int, error) { + select { + case <-c.closed: + return 0, net.ErrClosed + default: + } + c.mu.Lock() + c.writes = append(c.writes, append([]byte(nil), b...)) + c.mu.Unlock() + return len(b), nil +} + +func (c *fakePacketConn) Close() error { + c.closeOnce.Do(func() { close(c.closed) }) + return nil +} + +func (c *fakePacketConn) LocalAddr() net.Addr { return &net.UDPAddr{} } +func (c *fakePacketConn) SetDeadline(time.Time) error { return nil } +func (c *fakePacketConn) SetReadDeadline(time.Time) error { return nil } +func (c *fakePacketConn) SetWriteDeadline(time.Time) error { return nil } diff --git a/connector/masque/connector.go b/connector/masque/connector.go index 8040bced..c586d60f 100644 --- a/connector/masque/connector.go +++ b/connector/masque/connector.go @@ -10,6 +10,7 @@ import ( "net/url" "strconv" "strings" + "sync" "time" "github.com/go-gost/core/connector" @@ -19,6 +20,7 @@ import ( masque_dialer "github.com/go-gost/x/dialer/http3/masque" masque_util "github.com/go-gost/x/internal/util/masque" "github.com/go-gost/x/registry" + "github.com/quic-go/quic-go/http3" ) func init() { @@ -73,6 +75,16 @@ func (c *masqueConnector) Connect(ctx context.Context, conn net.Conn, network, a return nil, fmt.Errorf("%w: %s", ErrUnsupportedNetwork, network) } +func newTCPConnectRequest(address string) *http.Request { + return &http.Request{ + Method: http.MethodConnect, + URL: &url.URL{}, + Host: address, + Header: http.Header{}, + ProtoMajor: 3, + } +} + func (c *masqueConnector) connectTCP(ctx context.Context, conn net.Conn, address string, log logger.Logger) (net.Conn, error) { log.Debugf("connect-tcp %s", address) @@ -103,13 +115,7 @@ func (c *masqueConnector) connectTCP(ctx context.Context, conn net.Conn, address // Create standard HTTP/3 CONNECT request (RFC 9114) // No :protocol pseudo-header for standard CONNECT - req := &http.Request{ - Method: http.MethodConnect, - URL: &url.URL{}, - Host: address, // Target address goes in :authority - Header: http.Header{}, - Proto: "HTTP/3.0", - } + req := newTCPConnectRequest(address) // Add proxy authentication if configured if c.options.Auth != nil { @@ -172,9 +178,69 @@ func (c *masqueConnector) connectUDP(ctx context.Context, conn net.Conn, address return nil, ErrInvalidConnection } - // Get pre-opened stream from dialer (stream opening happens there for dead connection detection) - reqStream := masqueConn.GetRequestStream() - proxyHost := masqueConn.GetHost() + if address == "" { + return c.connectUDPAssociation(ctx, masqueConn, log), nil + } + + return c.connectUDPTarget( + ctx, + conn.LocalAddr(), + masqueConn.GetHost(), + masqueConn.GetRequestStream(), + address, + nil, + log, + ) +} + +func (c *masqueConnector) connectUDPAssociation(ctx context.Context, conn *masque_dialer.MasqueConn, log logger.Logger) net.Conn { + var ( + mu sync.Mutex + firstStream = conn.GetRequestStream() + ) + + nextStream := func(ctx context.Context) (*http3.RequestStream, error) { + mu.Lock() + defer mu.Unlock() + if firstStream != nil { + stream := firstStream + firstStream = nil + return stream, nil + } + return conn.OpenRequestStream(ctx) + } + + closeIdle := func() error { + mu.Lock() + defer mu.Unlock() + if firstStream == nil { + return nil + } + err := firstStream.Close() + firstStream = nil + return err + } + + dial := func(ctx context.Context, addr net.Addr) (net.PacketConn, error) { + stream, err := nextStream(ctx) + if err != nil { + return nil, err + } + return c.connectUDPTarget(ctx, conn.LocalAddr(), conn.GetHost(), stream, addr.String(), addr, log) + } + + return newUDPAssociationConn(ctx, conn.LocalAddr(), c.md.connectTimeout, dial, closeIdle, log) +} + +func (c *masqueConnector) connectUDPTarget( + ctx context.Context, + localAddr net.Addr, + proxyHost string, + reqStream *http3.RequestStream, + address string, + remoteAddr net.Addr, + log logger.Logger, +) (*masque_util.DatagramConn, error) { // Apply connect timeout to the actual stream if c.md.connectTimeout > 0 { @@ -259,14 +325,15 @@ func (c *masqueConnector) connectUDP(ctx context.Context, conn net.Conn, address // Get the underlying HTTP/3 stream for datagrams stream := reqStream - // Resolve target address - raddr, err := net.ResolveUDPAddr("udp", address) - if err != nil { - return nil, err + if remoteAddr == nil { + remoteAddr, err = net.ResolveUDPAddr("udp", address) + if err != nil { + return nil, err + } } // Create datagram connection wrapping the request stream - datagramConn := masque_util.NewDatagramConnFromRequestStream(stream, conn.LocalAddr(), raddr) + datagramConn := masque_util.NewDatagramConnFromRequestStream(stream, localAddr, remoteAddr) success = true // Prevent defer from closing stream - datagramConn now owns it return datagramConn, nil diff --git a/connector/masque/connector_test.go b/connector/masque/connector_test.go new file mode 100644 index 00000000..96571fb5 --- /dev/null +++ b/connector/masque/connector_test.go @@ -0,0 +1,24 @@ +package masque + +import ( + "net/http" + "testing" +) + +func TestNewTCPConnectRequestUsesStandardConnect(t *testing.T) { + req := newTCPConnectRequest("example.com:443") + + if req.Method != http.MethodConnect { + t.Fatalf("expected CONNECT method, got %s", req.Method) + } + if req.Host != "example.com:443" { + t.Fatalf("expected target authority, got %s", req.Host) + } + if req.Proto != "" { + // quic-go treats any non-empty Proto as an extended CONNECT :protocol. + t.Fatalf("standard CONNECT must not set :protocol, got %q", req.Proto) + } + if req.ProtoMajor != 3 { + t.Fatalf("expected HTTP/3 request, got HTTP/%d", req.ProtoMajor) + } +} diff --git a/dialer/http3/masque/client.go b/dialer/http3/masque/client.go index c482afd6..98f8d0a3 100644 --- a/dialer/http3/masque/client.go +++ b/dialer/http3/masque/client.go @@ -102,6 +102,15 @@ func (c *MasqueConn) GetRequestStream() *http3.RequestStream { return c.reqStream } +// OpenRequestStream opens another request stream on the shared HTTP/3 connection. +// CONNECT-UDP associations use one stream per destination. +func (c *MasqueConn) OpenRequestStream(ctx context.Context) (*http3.RequestStream, error) { + if c.clientConn == nil { + return nil, errors.New("masque: connection is closed") + } + return c.clientConn.OpenRequestStream(ctx) +} + // GetHost returns the proxy host. func (c *MasqueConn) GetHost() string { return c.host