Track UDP connections

This will enable NAT to be performed on UDP packets that are sent
in response to packets sent by the stack.

This will also enable ICMP errors to be properly NAT-ed in response
to UDP packets (#5916).

Updates #5915.

PiperOrigin-RevId: 398373251

2 files changed, 292 insertions, 0 deletions

diff --git a/pkg/tcpip/tests/integration/BUILD b/pkg/tcpip/tests/integration/BUILD
index 181ef799e..7c998eaae 100644
--- a/pkg/tcpip/tests/integration/BUILD
+++ b/pkg/tcpip/tests/integration/BUILD
@@ -34,12 +34,16 @@ go_test(
         "//pkg/tcpip/checker",
         "//pkg/tcpip/header",
         "//pkg/tcpip/link/channel",
+        "//pkg/tcpip/network/arp",
         "//pkg/tcpip/network/ipv4",
         "//pkg/tcpip/network/ipv6",
         "//pkg/tcpip/stack",
         "//pkg/tcpip/tests/utils",
         "//pkg/tcpip/testutil",
+        "//pkg/tcpip/transport/tcp",
         "//pkg/tcpip/transport/udp",
+        "//pkg/waiter",
+        "@com_github_google_go_cmp//cmp:go_default_library",
     ],
 )
 
diff --git a/pkg/tcpip/tests/integration/iptables_test.go b/pkg/tcpip/tests/integration/iptables_test.go
index 28b49c6be..bdf4a64b9 100644
--- a/pkg/tcpip/tests/integration/iptables_test.go
+++ b/pkg/tcpip/tests/integration/iptables_test.go
@@ -15,19 +15,24 @@
 package iptables_test
 
 import (
+	"bytes"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/checker"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
+	"gvisor.dev/gvisor/pkg/tcpip/network/arp"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv6"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/tests/utils"
 	"gvisor.dev/gvisor/pkg/tcpip/testutil"
+	"gvisor.dev/gvisor/pkg/tcpip/transport/tcp"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
+	"gvisor.dev/gvisor/pkg/waiter"
 )
 
 type inputIfNameMatcher struct {
@@ -1156,3 +1161,286 @@ func TestInputHookWithLocalForwarding(t *testing.T) {
 		})
 	}
 }
+
+func TestSNAT(t *testing.T) {
+	const listenPort = 8080
+
+	type endpointAndAddresses struct {
+		serverEP         tcpip.Endpoint
+		serverAddr       tcpip.Address
+		serverReadableCH chan struct{}
+
+		clientEP         tcpip.Endpoint
+		clientAddr       tcpip.Address
+		clientReadableCH chan struct{}
+
+		nattedClientAddr tcpip.Address
+	}
+
+	newEP := func(t *testing.T, s *stack.Stack, transProto tcpip.TransportProtocolNumber, netProto tcpip.NetworkProtocolNumber) (tcpip.Endpoint, chan struct{}) {
+		t.Helper()
+		var wq waiter.Queue
+		we, ch := waiter.NewChannelEntry(nil)
+		wq.EventRegister(&we, waiter.ReadableEvents)
+		t.Cleanup(func() {
+			wq.EventUnregister(&we)
+		})
+
+		ep, err := s.NewEndpoint(transProto, netProto, &wq)
+		if err != nil {
+			t.Fatalf("s.NewEndpoint(%d, %d, _): %s", transProto, netProto, err)
+		}
+		t.Cleanup(ep.Close)
+
+		return ep, ch
+	}
+
+	tests := []struct {
+		name       string
+		epAndAddrs func(t *testing.T, host1Stack, routerStack, host2Stack *stack.Stack, proto tcpip.TransportProtocolNumber) endpointAndAddresses
+	}{
+		{
+			name: "IPv4 host1 server with host2 client",
+			epAndAddrs: func(t *testing.T, host1Stack, routerStack, host2Stack *stack.Stack, proto tcpip.TransportProtocolNumber) endpointAndAddresses {
+				t.Helper()
+
+				ipt := routerStack.IPTables()
+				filter := ipt.GetTable(stack.NATID, false /* ipv6 */)
+				ruleIdx := filter.BuiltinChains[stack.Postrouting]
+				filter.Rules[ruleIdx].Filter = stack.IPHeaderFilter{OutputInterface: utils.RouterNIC1Name}
+				filter.Rules[ruleIdx].Target = &stack.SNATTarget{NetworkProtocol: ipv4.ProtocolNumber, Addr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address}
+				// Make sure the packet is not dropped by the next rule.
+				filter.Rules[ruleIdx+1].Target = &stack.AcceptTarget{}
+				if err := ipt.ReplaceTable(stack.NATID, filter, false /* ipv6 */); err != nil {
+					t.Fatalf("ipt.ReplaceTable(%d, _, %t): %s", stack.NATID, false, err)
+				}
+
+				ep1, ep1WECH := newEP(t, host1Stack, proto, ipv4.ProtocolNumber)
+				ep2, ep2WECH := newEP(t, host2Stack, proto, ipv4.ProtocolNumber)
+				return endpointAndAddresses{
+					serverEP:         ep1,
+					serverAddr:       utils.Host1IPv4Addr.AddressWithPrefix.Address,
+					serverReadableCH: ep1WECH,
+
+					clientEP:         ep2,
+					clientAddr:       utils.Host2IPv4Addr.AddressWithPrefix.Address,
+					clientReadableCH: ep2WECH,
+
+					nattedClientAddr: utils.RouterNIC1IPv4Addr.AddressWithPrefix.Address,
+				}
+			},
+		},
+		{
+			name: "IPv6 host1 server with host2 client",
+			epAndAddrs: func(t *testing.T, host1Stack, routerStack, host2Stack *stack.Stack, proto tcpip.TransportProtocolNumber) endpointAndAddresses {
+				t.Helper()
+
+				ipt := routerStack.IPTables()
+				filter := ipt.GetTable(stack.NATID, true /* ipv6 */)
+				ruleIdx := filter.BuiltinChains[stack.Postrouting]
+				filter.Rules[ruleIdx].Filter = stack.IPHeaderFilter{OutputInterface: utils.RouterNIC1Name}
+				filter.Rules[ruleIdx].Target = &stack.SNATTarget{NetworkProtocol: ipv6.ProtocolNumber, Addr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address}
+				// Make sure the packet is not dropped by the next rule.
+				filter.Rules[ruleIdx+1].Target = &stack.AcceptTarget{}
+				if err := ipt.ReplaceTable(stack.NATID, filter, true /* ipv6 */); err != nil {
+					t.Fatalf("ipt.ReplaceTable(%d, _, %t): %s", stack.NATID, true, err)
+				}
+
+				ep1, ep1WECH := newEP(t, host1Stack, proto, ipv6.ProtocolNumber)
+				ep2, ep2WECH := newEP(t, host2Stack, proto, ipv6.ProtocolNumber)
+				return endpointAndAddresses{
+					serverEP:         ep1,
+					serverAddr:       utils.Host1IPv6Addr.AddressWithPrefix.Address,
+					serverReadableCH: ep1WECH,
+
+					clientEP:         ep2,
+					clientAddr:       utils.Host2IPv6Addr.AddressWithPrefix.Address,
+					clientReadableCH: ep2WECH,
+
+					nattedClientAddr: utils.RouterNIC1IPv6Addr.AddressWithPrefix.Address,
+				}
+			},
+		},
+	}
+
+	subTests := []struct {
+		name               string
+		proto              tcpip.TransportProtocolNumber
+		expectedConnectErr tcpip.Error
+		setupServer        func(t *testing.T, ep tcpip.Endpoint)
+		setupServerConn    func(t *testing.T, ep tcpip.Endpoint, ch <-chan struct{}, clientAddr tcpip.FullAddress) (tcpip.Endpoint, chan struct{})
+		needRemoteAddr     bool
+	}{
+		{
+			name:               "UDP",
+			proto:              udp.ProtocolNumber,
+			expectedConnectErr: nil,
+			setupServerConn: func(t *testing.T, ep tcpip.Endpoint, _ <-chan struct{}, clientAddr tcpip.FullAddress) (tcpip.Endpoint, chan struct{}) {
+				t.Helper()
+
+				if err := ep.Connect(clientAddr); err != nil {
+					t.Fatalf("ep.Connect(%#v): %s", clientAddr, err)
+				}
+				return nil, nil
+			},
+			needRemoteAddr: true,
+		},
+		{
+			name:               "TCP",
+			proto:              tcp.ProtocolNumber,
+			expectedConnectErr: &tcpip.ErrConnectStarted{},
+			setupServer: func(t *testing.T, ep tcpip.Endpoint) {
+				t.Helper()
+
+				if err := ep.Listen(1); err != nil {
+					t.Fatalf("ep.Listen(1): %s", err)
+				}
+			},
+			setupServerConn: func(t *testing.T, ep tcpip.Endpoint, ch <-chan struct{}, clientAddr tcpip.FullAddress) (tcpip.Endpoint, chan struct{}) {
+				t.Helper()
+
+				var addr tcpip.FullAddress
+				for {
+					newEP, wq, err := ep.Accept(&addr)
+					if _, ok := err.(*tcpip.ErrWouldBlock); ok {
+						<-ch
+						continue
+					}
+					if err != nil {
+						t.Fatalf("ep.Accept(_): %s", err)
+					}
+					if diff := cmp.Diff(clientAddr, addr, checker.IgnoreCmpPath(
+						"NIC",
+					)); diff != "" {
+						t.Errorf("accepted address mismatch (-want +got):\n%s", diff)
+					}
+
+					we, newCH := waiter.NewChannelEntry(nil)
+					wq.EventRegister(&we, waiter.ReadableEvents)
+					return newEP, newCH
+				}
+			},
+			needRemoteAddr: false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			for _, subTest := range subTests {
+				t.Run(subTest.name, func(t *testing.T) {
+					stackOpts := stack.Options{
+						NetworkProtocols:   []stack.NetworkProtocolFactory{arp.NewProtocol, ipv4.NewProtocol, ipv6.NewProtocol},
+						TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, tcp.NewProtocol},
+					}
+
+					host1Stack := stack.New(stackOpts)
+					routerStack := stack.New(stackOpts)
+					host2Stack := stack.New(stackOpts)
+					utils.SetupRoutedStacks(t, host1Stack, routerStack, host2Stack)
+
+					epsAndAddrs := test.epAndAddrs(t, host1Stack, routerStack, host2Stack, subTest.proto)
+					serverAddr := tcpip.FullAddress{Addr: epsAndAddrs.serverAddr, Port: listenPort}
+					if err := epsAndAddrs.serverEP.Bind(serverAddr); err != nil {
+						t.Fatalf("epsAndAddrs.serverEP.Bind(%#v): %s", serverAddr, err)
+					}
+					clientAddr := tcpip.FullAddress{Addr: epsAndAddrs.clientAddr}
+					if err := epsAndAddrs.clientEP.Bind(clientAddr); err != nil {
+						t.Fatalf("epsAndAddrs.clientEP.Bind(%#v): %s", clientAddr, err)
+					}
+
+					if subTest.setupServer != nil {
+						subTest.setupServer(t, epsAndAddrs.serverEP)
+					}
+					{
+						err := epsAndAddrs.clientEP.Connect(serverAddr)
+						if diff := cmp.Diff(subTest.expectedConnectErr, err); diff != "" {
+							t.Fatalf("unexpected error from epsAndAddrs.clientEP.Connect(%#v), (-want, +got):\n%s", serverAddr, diff)
+						}
+					}
+					nattedClientAddr := tcpip.FullAddress{Addr: epsAndAddrs.nattedClientAddr}
+					if addr, err := epsAndAddrs.clientEP.GetLocalAddress(); err != nil {
+						t.Fatalf("epsAndAddrs.clientEP.GetLocalAddress(): %s", err)
+					} else {
+						nattedClientAddr.Port = addr.Port
+					}
+
+					serverEP := epsAndAddrs.serverEP
+					serverCH := epsAndAddrs.serverReadableCH
+					if ep, ch := subTest.setupServerConn(t, serverEP, serverCH, nattedClientAddr); ep != nil {
+						defer ep.Close()
+						serverEP = ep
+						serverCH = ch
+					}
+
+					write := func(ep tcpip.Endpoint, data []byte) {
+						t.Helper()
+
+						var r bytes.Reader
+						r.Reset(data)
+						var wOpts tcpip.WriteOptions
+						n, err := ep.Write(&r, wOpts)
+						if err != nil {
+							t.Fatalf("ep.Write(_, %#v): %s", wOpts, err)
+						}
+						if want := int64(len(data)); n != want {
+							t.Fatalf("got ep.Write(_, %#v) = (%d, _), want = (%d, _)", wOpts, n, want)
+						}
+					}
+
+					read := func(ch chan struct{}, ep tcpip.Endpoint, data []byte, expectedFrom tcpip.FullAddress) {
+						t.Helper()
+
+						var buf bytes.Buffer
+						var res tcpip.ReadResult
+						for {
+							var err tcpip.Error
+							opts := tcpip.ReadOptions{NeedRemoteAddr: subTest.needRemoteAddr}
+							res, err = ep.Read(&buf, opts)
+							if _, ok := err.(*tcpip.ErrWouldBlock); ok {
+								<-ch
+								continue
+							}
+							if err != nil {
+								t.Fatalf("ep.Read(_, %d, %#v): %s", len(data), opts, err)
+							}
+							break
+						}
+
+						readResult := tcpip.ReadResult{
+							Count: len(data),
+							Total: len(data),
+						}
+						if subTest.needRemoteAddr {
+							readResult.RemoteAddr = expectedFrom
+						}
+						if diff := cmp.Diff(readResult, res, checker.IgnoreCmpPath(
+							"ControlMessages",
+							"RemoteAddr.NIC",
+						)); diff != "" {
+							t.Errorf("ep.Read: unexpected result (-want +got):\n%s", diff)
+						}
+						if diff := cmp.Diff(buf.Bytes(), data); diff != "" {
+							t.Errorf("received data mismatch (-want +got):\n%s", diff)
+						}
+
+						if t.Failed() {
+							t.FailNow()
+						}
+					}
+
+					{
+						data := []byte{1, 2, 3, 4}
+						write(epsAndAddrs.clientEP, data)
+						read(serverCH, serverEP, data, nattedClientAddr)
+					}
+
+					{
+						data := []byte{5, 6, 7, 8, 9, 10, 11, 12}
+						write(serverEP, data)
+						read(epsAndAddrs.clientReadableCH, epsAndAddrs.clientEP, data, serverAddr)
+					}
+				})
+			}
+		})
+	}
+}