diff --git a/internal/amneziawgnet/portfwd.go b/internal/amneziawgnet/portfwd.go index 5dff9e474..25019345d 100644 --- a/internal/amneziawgnet/portfwd.go +++ b/internal/amneziawgnet/portfwd.go @@ -24,7 +24,6 @@ package amneziawgnet import ( "context" "fmt" - "io" "net" "net/netip" "sync" @@ -327,10 +326,7 @@ func relayTCPForward(gstack *stack.Stack, conn net.Conn, inboundID int, key port } defer tunnelConn.Close() - done := make(chan struct{}, 2) - go func() { _, _ = io.Copy(tunnelConn, conn); done <- struct{}{} }() - go func() { _, _ = io.Copy(conn, tunnelConn); done <- struct{}{} }() - <-done + pipeBothWays(conn, tunnelConn) } // Close stops accepting new connections. Already-relaying connections are diff --git a/internal/amneziawgnet/relay.go b/internal/amneziawgnet/relay.go index 33f3b7662..80982c8a9 100644 --- a/internal/amneziawgnet/relay.go +++ b/internal/amneziawgnet/relay.go @@ -15,6 +15,7 @@ import ( "net" "net/netip" "sync" + "sync/atomic" "time" "golang.org/x/net/proxy" @@ -60,7 +61,7 @@ func SocksInboundSettings(emails []string, password string) ([]byte, error) { } // RelayTCP dials r.Addr, authenticates as email, issues a SOCKS5 CONNECT to -// dest, and pipes bytes both ways until either side closes or errors. +// dest, and pipes bytes both ways until both directions end. // Blocks until the relay ends; meant to be called from (or as) an // AttachTCPForwarder handler, which already runs each connection on its own // goroutine. @@ -80,10 +81,57 @@ func (r SocksRelay) RelayTCP(conn *gonet.TCPConn, email string, dest netip.AddrP } defer upstream.Close() - done := make(chan struct{}, 2) - go func() { _, _ = io.Copy(upstream, conn); done <- struct{}{} }() - go func() { _, _ = io.Copy(conn, upstream); done <- struct{}{} }() - <-done + pipeBothWays(conn, upstream) +} + +// halfCloseIdle bounds how long the surviving direction of a half-closed pair +// may sit idle, so a peer that vanished mid-transfer cannot pin it forever. +const halfCloseIdle = 2 * time.Minute + +// closeWriter is the half-close half of *net.TCPConn and *gonet.TCPConn. +type closeWriter interface{ CloseWrite() error } + +// guardedReader reads one side of a relayed pair, re-arming its read deadline +// on every read once armed, so the bound is an idle window, not a total one. +type guardedReader struct { + conn net.Conn + armed atomic.Bool +} + +func (r *guardedReader) Read(p []byte) (int, error) { + if r.armed.Load() { + _ = r.conn.SetReadDeadline(time.Now().Add(halfCloseIdle)) + } + return r.conn.Read(p) +} + +// arm bounds this side's remaining reads, including one already in flight. +func (r *guardedReader) arm() { + r.armed.Store(true) + _ = r.conn.SetReadDeadline(time.Now().Add(halfCloseIdle)) +} + +// pipeBothWays copies a and b into each other until BOTH directions end, +// half-closing each far side in turn so a half-closed peer still gets its reply. +func pipeBothWays(a, b net.Conn) { + ga, gb := &guardedReader{conn: a}, &guardedReader{conn: b} + var wg sync.WaitGroup + wg.Add(2) + // Arming dst bounds the direction still reading from it -- the one this + // copy just signalled EOF to. + pipe := func(dst, src *guardedReader) { + defer wg.Done() + _, _ = io.Copy(dst.conn, src) + if cw, ok := dst.conn.(closeWriter); ok { + _ = cw.CloseWrite() + } else { + _ = dst.conn.Close() + } + dst.arm() + } + go pipe(gb, ga) + go pipe(ga, gb) + wg.Wait() } // socks5UDPSession is one established SOCKS5 UDP ASSOCIATE session: udpConn @@ -278,36 +326,45 @@ func (s *socks5UDPSession) receive(buf []byte) (netip.AddrPort, []byte, error) { if len(data) < 4 { return netip.AddrPort{}, nil, fmt.Errorf("amneziawgnet: short SOCKS5 UDP reply (%d bytes)", n) } - atyp := data[3] - data = data[4:] - addr, err := readSocks5Addr(bytesReader{data}, atyp) + addr, rest, err := splitSocks5Addr(data[4:], data[3]) if err != nil { return netip.AddrPort{}, nil, err } - switch atyp { - case 0x01: - data = data[4:] - case 0x04: - data = data[16:] - } - if len(data) < 2 { + if len(rest) < 2 { return netip.AddrPort{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 UDP reply port") } - port := binary.BigEndian.Uint16(data[:2]) - return netip.AddrPortFrom(addr, port), data[2:], nil + return netip.AddrPortFrom(addr, binary.BigEndian.Uint16(rest[:2])), rest[2:], nil } -// bytesReader is the minimal io.Reader readSocks5Addr needs, over an -// in-memory slice that's already fully available (a received UDP -// datagram) -- avoids pulling in bytes.Reader just for this. -type bytesReader struct{ b []byte } - -func (r bytesReader) Read(p []byte) (int, error) { - n := copy(p, r.b) - if n < len(p) { - return n, io.ErrUnexpectedEOF +// splitSocks5Addr decodes the address at the head of b for address type atyp +// and returns it with whatever follows, length-checked at every step. +func splitSocks5Addr(b []byte, atyp byte) (netip.Addr, []byte, error) { + switch atyp { + case 0x01: + if len(b) < 4 { + return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 IPv4 reply address") + } + return netip.AddrFrom4([4]byte(b[:4])), b[4:], nil + case 0x04: + if len(b) < 16 { + return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 IPv6 reply address") + } + return netip.AddrFrom16([16]byte(b[:16])), b[16:], nil + case 0x03: + // Resolving here would block the receive loop on DNS, and a datagram's + // own source is an address already -- so only a literal is accepted. + if len(b) < 1 || len(b) < 1+int(b[0]) { + return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: truncated SOCKS5 domain reply address") + } + name := string(b[1 : 1+int(b[0])]) + addr, err := netip.ParseAddr(name) + if err != nil { + return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: SOCKS5 UDP reply from non-literal address %q", name) + } + return addr, b[1+int(b[0]):], nil + default: + return netip.Addr{}, nil, fmt.Errorf("amneziawgnet: unsupported SOCKS5 address type %d", atyp) } - return n, nil } // UDPRelay tracks one SOCKS5 UDP ASSOCIATE session per source (tunnel- @@ -318,13 +375,15 @@ type UDPRelay struct { relay SocksRelay gstack *stack.Stack + // Keyed by the comparable netip.AddrPort, like udpForwardListener's own + // session map: src.String() would allocate on every relayed datagram. mu sync.Mutex - sessions map[string]*socks5UDPSession + sessions map[netip.AddrPort]*socks5UDPSession } // NewUDPRelay creates a UDPRelay for one embedded AmneziaWG Device's stack. func NewUDPRelay(relay SocksRelay, gstack *stack.Stack) *UDPRelay { - return &UDPRelay{relay: relay, gstack: gstack, sessions: map[string]*socks5UDPSession{}} + return &UDPRelay{relay: relay, gstack: gstack, sessions: map[netip.AddrPort]*socks5UDPSession{}} } // Handle relays one packet from src (the peer's tunnel-internal source) to @@ -334,20 +393,28 @@ func NewUDPRelay(relay SocksRelay, gstack *stack.Stack) *UDPRelay { // reusing it for subsequent packets from the same src. func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte) { u.mu.Lock() - sess, ok := u.sessions[src.String()] + sess, ok := u.sessions[src] u.mu.Unlock() if !ok { - var err error - sess, err = newSocks5UDPSession(u.relay.Addr, email, u.relay.Password) + fresh, err := newSocks5UDPSession(u.relay.Addr, email, u.relay.Password) if err != nil { logger.Warningf("amneziawgnet: UDPRelay: SOCKS5 associate for %q: %v", email, err) return } u.mu.Lock() - u.sessions[src.String()] = sess - u.mu.Unlock() - go u.pump(src, sess) + // Associating happens off-lock, so a concurrent Handle for the same src + // may already have published one; keep it, so the key has a single pump. + if existing, dup := u.sessions[src]; dup { + u.mu.Unlock() + fresh.Close() + sess = existing + } else { + u.sessions[src] = fresh + u.mu.Unlock() + sess = fresh + go u.pump(src, fresh) + } } if err := sess.sendTo(dst, payload); err != nil { logger.Warningf("amneziawgnet: UDPRelay: send to %s: %v", dst, err) @@ -360,7 +427,11 @@ func (u *UDPRelay) Handle(src, dst netip.AddrPort, email string, payload []byte) func (u *UDPRelay) pump(src netip.AddrPort, sess *socks5UDPSession) { defer func() { u.mu.Lock() - delete(u.sessions, src.String()) + // Only retire our own entry: a delete by key alone would evict whichever + // session currently holds src, orphaning a live one. + if u.sessions[src] == sess { + delete(u.sessions, src) + } u.mu.Unlock() sess.Close() }() diff --git a/internal/amneziawgnet/relay_test.go b/internal/amneziawgnet/relay_test.go new file mode 100644 index 000000000..4437e223d --- /dev/null +++ b/internal/amneziawgnet/relay_test.go @@ -0,0 +1,257 @@ +package amneziawgnet + +import ( + "io" + "net" + "net/netip" + "testing" + "time" +) + +// newDeadUDPSession builds a socks5UDPSession over real but already-closed +// sockets, so pump's receive fails immediately and its teardown runs at once. +func newDeadUDPSession(t *testing.T) *socks5UDPSession { + t.Helper() + peer, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + t.Fatalf("listen udp: %v", err) + } + t.Cleanup(func() { _ = peer.Close() }) + udpConn, err := net.DialUDP("udp", nil, peer.LocalAddr().(*net.UDPAddr)) + if err != nil { + t.Fatalf("dial udp: %v", err) + } + ctrlLn, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen tcp: %v", err) + } + t.Cleanup(func() { _ = ctrlLn.Close() }) + ctrl, err := net.Dial("tcp", ctrlLn.Addr().String()) + if err != nil { + t.Fatalf("dial tcp: %v", err) + } + _ = udpConn.Close() + return &socks5UDPSession{ctrl: ctrl, udpConn: udpConn} +} + +// TestUDPRelayPumpOnlyRetiresItsOwnSession pins the flow that survives a +// duplicate associate: a losing pump must not evict the published session. +func TestUDPRelayPumpOnlyRetiresItsOwnSession(t *testing.T) { + relay := NewUDPRelay(SocksRelay{Addr: "127.0.0.1:1", Password: "x"}, nil) + src := netip.MustParseAddrPort("10.8.1.5:51820") + + live := newDeadUDPSession(t) + superseded := newDeadUDPSession(t) + relay.sessions[src] = live + + // Returns as soon as receive fails on the closed socket, so no wait is needed. + relay.pump(src, superseded) + + got, ok := relay.sessions[src] + if !ok { + t.Fatal("live session was evicted: a retiring pump deleted src's entry regardless of which session held it") + } + if got != live { + t.Fatalf("sessions[%v] = %p, want the live session %p", src, got, live) + } +} + +// TestUDPRelayCloseDropsEverySession keeps Close's contract explicit now that +// pump's teardown is conditional on still owning the key. +func TestUDPRelayCloseDropsEverySession(t *testing.T) { + relay := NewUDPRelay(SocksRelay{Addr: "127.0.0.1:1", Password: "x"}, nil) + for _, s := range []string{"10.8.1.5:51820", "10.8.1.6:2000"} { + relay.sessions[netip.MustParseAddrPort(s)] = newDeadUDPSession(t) + } + + relay.Close() + + if n := len(relay.sessions); n != 0 { + t.Fatalf("Close left %d sessions behind, want 0", n) + } +} + +// tcpPair returns a connected pair of real loopback TCP conns; net.Pipe would +// not do, since these tests turn on CloseWrite, which it does not implement. +func tcpPair(t *testing.T) (client, server net.Conn) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer ln.Close() + + type accepted struct { + conn net.Conn + err error + } + ch := make(chan accepted, 1) + go func() { + c, err := ln.Accept() + ch <- accepted{c, err} + }() + client, err = net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + got := <-ch + if got.err != nil { + t.Fatalf("accept: %v", got.err) + } + t.Cleanup(func() { _ = client.Close(); _ = got.conn.Close() }) + return client, got.conn +} + +// TestPipeBothWaysDeliversReplyAfterHalfClose is the half-close regression: a +// client that shuts down its write side must still receive the full response. +func TestPipeBothWaysDeliversReplyAfterHalfClose(t *testing.T) { + const request = "GET / HTTP/1.0\r\n\r\n" + const response = "the reply that arrives only after the request is complete" + + client, a := tcpPair(t) + b, server := tcpPair(t) + go pipeBothWays(a, b) + + if _, err := client.Write([]byte(request)); err != nil { + t.Fatalf("client write: %v", err) + } + // The half-close the old relay treated as "tear the whole pair down". + if err := client.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatalf("client CloseWrite: %v", err) + } + + _ = server.SetReadDeadline(time.Now().Add(10 * time.Second)) + gotReq, err := io.ReadAll(server) + if err != nil { + t.Fatalf("server read: %v", err) + } + if string(gotReq) != request { + t.Fatalf("server got request %q, want %q", gotReq, request) + } + + if _, err := server.Write([]byte(response)); err != nil { + t.Fatalf("server write: %v", err) + } + if err := server.(*net.TCPConn).CloseWrite(); err != nil { + t.Fatalf("server CloseWrite: %v", err) + } + + _ = client.SetReadDeadline(time.Now().Add(10 * time.Second)) + gotResp, err := io.ReadAll(client) + if err != nil { + t.Fatalf("client read: %v", err) + } + if string(gotResp) != response { + t.Fatalf("client got response %q, want %q: the reply was cut off by the half-close", gotResp, response) + } +} + +// TestPipeBothWaysClosesWhenBothSidesFinish keeps the teardown contract: both +// directions ending must return, not hang on the idle bound. +func TestPipeBothWaysClosesWhenBothSidesFinish(t *testing.T) { + client, a := tcpPair(t) + b, server := tcpPair(t) + + done := make(chan struct{}) + go func() { defer close(done); pipeBothWays(a, b) }() + + _ = client.(*net.TCPConn).CloseWrite() + _, _ = io.ReadAll(server) + _ = server.(*net.TCPConn).CloseWrite() + _, _ = io.ReadAll(client) + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("pipeBothWays did not return after both directions ended") + } +} + +// liveUDPSession returns a session whose udpConn is connected to the returned +// peer, so a test can hand receive() one exact reply datagram. +func liveUDPSession(t *testing.T) (*socks5UDPSession, *net.UDPConn) { + t.Helper() + peer, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + t.Fatalf("listen udp: %v", err) + } + t.Cleanup(func() { _ = peer.Close() }) + udpConn, err := net.DialUDP("udp", nil, peer.LocalAddr().(*net.UDPAddr)) + if err != nil { + t.Fatalf("dial udp: %v", err) + } + t.Cleanup(func() { _ = udpConn.Close() }) + return &socks5UDPSession{udpConn: udpConn}, peer +} + +// sendReply delivers one raw datagram to sess's socket. +func sendReply(t *testing.T, sess *socks5UDPSession, peer *net.UDPConn, datagram []byte) { + t.Helper() + if _, err := peer.WriteToUDP(datagram, sess.udpConn.LocalAddr().(*net.UDPAddr)); err != nil { + t.Fatalf("write reply: %v", err) + } + _ = sess.udpConn.SetReadDeadline(time.Now().Add(5 * time.Second)) +} + +// TestSocks5ReceiveDecodesReplyAddressTypes covers all three ATYP forms; the +// domain form used to misread its own length byte and never skip the name. +func TestSocks5ReceiveDecodesReplyAddressTypes(t *testing.T) { + tests := []struct { + name string + addrPart []byte + wantAddr string + }{ + {"IPv4", []byte{0x01, 10, 0, 0, 7}, "10.0.0.7"}, + {"IPv6", append([]byte{0x04}, netip.MustParseAddr("2001:db8::5").AsSlice()...), "2001:db8::5"}, + {"domain holding a literal", append([]byte{0x03, 8}, []byte("10.0.0.9")...), "10.0.0.9"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sess, peer := liveUDPSession(t) + payload := []byte("the-actual-datagram-payload") + datagram := append([]byte{0x00, 0x00, 0x00}, tt.addrPart...) + datagram = append(datagram, 0x1f, 0x90) // port 8080 + datagram = append(datagram, payload...) + sendReply(t, sess, peer, datagram) + + buf := make([]byte, 4096) + from, got, err := sess.receive(buf) + if err != nil { + t.Fatalf("receive: %v", err) + } + want := netip.AddrPortFrom(netip.MustParseAddr(tt.wantAddr), 8080) + if from != want { + t.Errorf("source = %v, want %v", from, want) + } + if string(got) != string(payload) { + t.Errorf("payload = %q, want %q", got, payload) + } + }) + } +} + +// TestSocks5ReceiveRejectsTruncatedReplies pins that a short datagram is an +// error, not a slice-bounds panic in the relay's own pump goroutine. +func TestSocks5ReceiveRejectsTruncatedReplies(t *testing.T) { + tests := []struct { + name string + datagram []byte + }{ + {"header only, IPv4 announced", []byte{0x00, 0x00, 0x00, 0x01}}, + {"IPv4 address cut short", []byte{0x00, 0x00, 0x00, 0x01, 10, 0}}, + {"IPv6 address cut short", []byte{0x00, 0x00, 0x00, 0x04, 0x20, 0x01}}, + {"domain length past the end", []byte{0x00, 0x00, 0x00, 0x03, 40, 'a', 'b'}}, + {"address complete but port missing", []byte{0x00, 0x00, 0x00, 0x01, 10, 0, 0, 7}}, + {"unsupported address type", []byte{0x00, 0x00, 0x00, 0x09, 1, 2, 3, 4, 0, 80}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sess, peer := liveUDPSession(t) + sendReply(t, sess, peer, tt.datagram) + buf := make([]byte, 4096) + if _, _, err := sess.receive(buf); err == nil { + t.Fatal("receive accepted a malformed datagram instead of returning an error") + } + }) + } +}