Implement SO_REUSEPORT for TCP and UDP sockets

This option allows multiple sockets to be bound to the same port.

Incoming packets are distributed to sockets using a hash based on source and
destination addresses. This means that all packets from one sender will be
received by the same server socket.

PiperOrigin-RevId: 227153413
Change-Id: I59b6edda9c2209d5b8968671e9129adb675920cf

3 files changed, 142 insertions, 70 deletions

diff --git a/pkg/tcpip/ports/BUILD b/pkg/tcpip/ports/BUILD
index c69fc0744..a2fa9b84a 100644
--- a/pkg/tcpip/ports/BUILD
+++ b/pkg/tcpip/ports/BUILD
@@ -7,7 +7,9 @@ go_library(
     srcs = ["ports.go"],
     importpath = "gvisor.googlesource.com/gvisor/pkg/tcpip/ports",
     visibility = ["//:sandbox"],
-    deps = ["//pkg/tcpip"],
+    deps = [
+        "//pkg/tcpip",
+    ],
 )
 
 go_test(
diff --git a/pkg/tcpip/ports/ports.go b/pkg/tcpip/ports/ports.go
index 41ef32921..d212a5792 100644
--- a/pkg/tcpip/ports/ports.go
+++ b/pkg/tcpip/ports/ports.go
@@ -42,23 +42,47 @@ type PortManager struct {
 	allocatedPorts map[portDescriptor]bindAddresses
 }
 
+type portNode struct {
+	reuse bool
+	refs  int
+}
+
 // bindAddresses is a set of IP addresses.
-type bindAddresses map[tcpip.Address]struct{}
+type bindAddresses map[tcpip.Address]portNode
 
 // isAvailable checks whether an IP address is available to bind to.
-func (b bindAddresses) isAvailable(addr tcpip.Address) bool {
+func (b bindAddresses) isAvailable(addr tcpip.Address, reuse bool) bool {
 	if addr == anyIPAddress {
-		return len(b) == 0
+		if len(b) == 0 {
+			return true
+		}
+		if !reuse {
+			return false
+		}
+		for _, n := range b {
+			if !n.reuse {
+				return false
+			}
+		}
+		return true
 	}
 
 	// If all addresses for this portDescriptor are already bound, no
 	// address is available.
-	if _, ok := b[anyIPAddress]; ok {
-		return false
+	if n, ok := b[anyIPAddress]; ok {
+		if !reuse {
+			return false
+		}
+		if !n.reuse {
+			return false
+		}
 	}
 
-	if _, ok := b[addr]; ok {
-		return false
+	if n, ok := b[addr]; ok {
+		if !reuse {
+			return false
+		}
+		return n.reuse
 	}
 	return true
 }
@@ -92,17 +116,17 @@ func (s *PortManager) PickEphemeralPort(testPort func(p uint16) (bool, *tcpip.Er
 }
 
 // IsPortAvailable tests if the given port is available on all given protocols.
-func (s *PortManager) IsPortAvailable(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) bool {
+func (s *PortManager) IsPortAvailable(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16, reuse bool) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	return s.isPortAvailableLocked(networks, transport, addr, port)
+	return s.isPortAvailableLocked(networks, transport, addr, port, reuse)
 }
 
-func (s *PortManager) isPortAvailableLocked(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) bool {
+func (s *PortManager) isPortAvailableLocked(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16, reuse bool) bool {
 	for _, network := range networks {
 		desc := portDescriptor{network, transport, port}
 		if addrs, ok := s.allocatedPorts[desc]; ok {
-			if !addrs.isAvailable(addr) {
+			if !addrs.isAvailable(addr, reuse) {
 				return false
 			}
 		}
@@ -114,14 +138,14 @@ func (s *PortManager) isPortAvailableLocked(networks []tcpip.NetworkProtocolNumb
 // reserved by another endpoint. If port is zero, ReservePort will search for
 // an unreserved ephemeral port and reserve it, returning its value in the
 // "port" return value.
-func (s *PortManager) ReservePort(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) (reservedPort uint16, err *tcpip.Error) {
+func (s *PortManager) ReservePort(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16, reuse bool) (reservedPort uint16, err *tcpip.Error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	// If a port is specified, just try to reserve it for all network
 	// protocols.
 	if port != 0 {
-		if !s.reserveSpecificPort(networks, transport, addr, port) {
+		if !s.reserveSpecificPort(networks, transport, addr, port, reuse) {
 			return 0, tcpip.ErrPortInUse
 		}
 		return port, nil
@@ -129,13 +153,13 @@ func (s *PortManager) ReservePort(networks []tcpip.NetworkProtocolNumber, transp
 
 	// A port wasn't specified, so try to find one.
 	return s.PickEphemeralPort(func(p uint16) (bool, *tcpip.Error) {
-		return s.reserveSpecificPort(networks, transport, addr, p), nil
+		return s.reserveSpecificPort(networks, transport, addr, p, reuse), nil
 	})
 }
 
 // reserveSpecificPort tries to reserve the given port on all given protocols.
-func (s *PortManager) reserveSpecificPort(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) bool {
-	if !s.isPortAvailableLocked(networks, transport, addr, port) {
+func (s *PortManager) reserveSpecificPort(networks []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16, reuse bool) bool {
+	if !s.isPortAvailableLocked(networks, transport, addr, port, reuse) {
 		return false
 	}
 
@@ -147,7 +171,12 @@ func (s *PortManager) reserveSpecificPort(networks []tcpip.NetworkProtocolNumber
 			m = make(bindAddresses)
 			s.allocatedPorts[desc] = m
 		}
-		m[addr] = struct{}{}
+		if n, ok := m[addr]; ok {
+			n.refs++
+			m[addr] = n
+		} else {
+			m[addr] = portNode{reuse: reuse, refs: 1}
+		}
 	}
 
 	return true
@@ -162,7 +191,16 @@ func (s *PortManager) ReleasePort(networks []tcpip.NetworkProtocolNumber, transp
 	for _, network := range networks {
 		desc := portDescriptor{network, transport, port}
 		if m, ok := s.allocatedPorts[desc]; ok {
-			delete(m, addr)
+			n, ok := m[addr]
+			if !ok {
+				continue
+			}
+			n.refs--
+			if n.refs == 0 {
+				delete(m, addr)
+			} else {
+				m[addr] = n
+			}
 			if len(m) == 0 {
 				delete(s.allocatedPorts, desc)
 			}
diff --git a/pkg/tcpip/ports/ports_test.go b/pkg/tcpip/ports/ports_test.go
index 72577dfcb..01e7320b4 100644
--- a/pkg/tcpip/ports/ports_test.go
+++ b/pkg/tcpip/ports/ports_test.go
@@ -28,67 +28,99 @@ const (
 	fakeIPAddress1 = tcpip.Address("\x08\x08\x08\x09")
 )
 
-func TestPortReservation(t *testing.T) {
-	pm := NewPortManager()
-	net := []tcpip.NetworkProtocolNumber{fakeNetworkNumber}
+type portReserveTestAction struct {
+	port    uint16
+	ip      tcpip.Address
+	want    *tcpip.Error
+	reuse   bool
+	release bool
+}
 
+func TestPortReservation(t *testing.T) {
 	for _, test := range []struct {
-		port uint16
-		ip   tcpip.Address
-		want *tcpip.Error
+		tname   string
+		actions []portReserveTestAction
 	}{
 		{
-			port: 80,
-			ip:   fakeIPAddress,
-			want: nil,
-		},
-		{
-			port: 80,
-			ip:   fakeIPAddress1,
-			want: nil,
-		},
-		{
-			/* N.B. Order of tests matters! */
-			port: 80,
-			ip:   anyIPAddress,
-			want: tcpip.ErrPortInUse,
-		},
-		{
-			port: 22,
-			ip:   anyIPAddress,
-			want: nil,
-		},
-		{
-			port: 22,
-			ip:   fakeIPAddress,
-			want: tcpip.ErrPortInUse,
-		},
-		{
-			port: 0,
-			ip:   fakeIPAddress,
-			want: nil,
+			tname: "bind to ip",
+			actions: []portReserveTestAction{
+				{port: 80, ip: fakeIPAddress, want: nil},
+				{port: 80, ip: fakeIPAddress1, want: nil},
+				/* N.B. Order of tests matters! */
+				{port: 80, ip: anyIPAddress, want: tcpip.ErrPortInUse},
+				{port: 80, ip: fakeIPAddress, want: tcpip.ErrPortInUse, reuse: true},
+			},
 		},
 		{
-			port: 0,
-			ip:   fakeIPAddress,
-			want: nil,
+			tname: "bind to inaddr any",
+			actions: []portReserveTestAction{
+				{port: 22, ip: anyIPAddress, want: nil},
+				{port: 22, ip: fakeIPAddress, want: tcpip.ErrPortInUse},
+				/* release fakeIPAddress, but anyIPAddress is still inuse */
+				{port: 22, ip: fakeIPAddress, release: true},
+				{port: 22, ip: fakeIPAddress, want: tcpip.ErrPortInUse},
+				{port: 22, ip: fakeIPAddress, want: tcpip.ErrPortInUse, reuse: true},
+				/* Release port 22 from any IP address, then try to reserve fake IP address on 22 */
+				{port: 22, ip: anyIPAddress, want: nil, release: true},
+				{port: 22, ip: fakeIPAddress, want: nil},
+			},
+		}, {
+			tname: "bind to zero port",
+			actions: []portReserveTestAction{
+				{port: 00, ip: fakeIPAddress, want: nil},
+				{port: 00, ip: fakeIPAddress, want: nil},
+				{port: 00, ip: fakeIPAddress, reuse: true, want: nil},
+			},
+		}, {
+			tname: "bind to ip with reuseport",
+			actions: []portReserveTestAction{
+				{port: 25, ip: fakeIPAddress, reuse: true, want: nil},
+				{port: 25, ip: fakeIPAddress, reuse: true, want: nil},
+
+				{port: 25, ip: fakeIPAddress, reuse: false, want: tcpip.ErrPortInUse},
+				{port: 25, ip: anyIPAddress, reuse: false, want: tcpip.ErrPortInUse},
+
+				{port: 25, ip: anyIPAddress, reuse: true, want: nil},
+			},
+		}, {
+			tname: "bind to inaddr any with reuseport",
+			actions: []portReserveTestAction{
+				{port: 24, ip: anyIPAddress, reuse: true, want: nil},
+				{port: 24, ip: anyIPAddress, reuse: true, want: nil},
+
+				{port: 24, ip: anyIPAddress, reuse: false, want: tcpip.ErrPortInUse},
+				{port: 24, ip: fakeIPAddress, reuse: false, want: tcpip.ErrPortInUse},
+
+				{port: 24, ip: fakeIPAddress, reuse: true, want: nil},
+				{port: 24, ip: fakeIPAddress, release: true, want: nil},
+
+				{port: 24, ip: anyIPAddress, release: true},
+				{port: 24, ip: anyIPAddress, reuse: false, want: tcpip.ErrPortInUse},
+
+				{port: 24, ip: anyIPAddress, release: true},
+				{port: 24, ip: anyIPAddress, reuse: false, want: nil},
+			},
 		},
 	} {
-		gotPort, err := pm.ReservePort(net, fakeTransNumber, test.ip, test.port)
-		if err != test.want {
-			t.Fatalf("ReservePort(.., .., %s, %d) = %v, want %v", test.ip, test.port, err, test.want)
-		}
-		if test.port == 0 && (gotPort == 0 || gotPort < FirstEphemeral) {
-			t.Fatalf("ReservePort(.., .., .., 0) = %d, want port number >= %d to be picked", gotPort, FirstEphemeral)
-		}
-	}
+		t.Run(test.tname, func(t *testing.T) {
+			pm := NewPortManager()
+			net := []tcpip.NetworkProtocolNumber{fakeNetworkNumber}
 
-	// Release port 22 from any IP address, then try to reserve fake IP
-	// address on 22.
-	pm.ReleasePort(net, fakeTransNumber, anyIPAddress, 22)
+			for _, test := range test.actions {
+				if test.release {
+					pm.ReleasePort(net, fakeTransNumber, test.ip, test.port)
+					continue
+				}
+				gotPort, err := pm.ReservePort(net, fakeTransNumber, test.ip, test.port, test.reuse)
+				if err != test.want {
+					t.Fatalf("ReservePort(.., .., %s, %d, %t) = %v, want %v", test.ip, test.port, test.release, err, test.want)
+				}
+				if test.port == 0 && (gotPort == 0 || gotPort < FirstEphemeral) {
+					t.Fatalf("ReservePort(.., .., .., 0) = %d, want port number >= %d to be picked", gotPort, FirstEphemeral)
+				}
+			}
+		})
 
-	if port, err := pm.ReservePort(net, fakeTransNumber, fakeIPAddress, 22); port != 22 || err != nil {
-		t.Fatalf("ReservePort(.., .., .., %d) = (port %d, err %v), want (22, nil); failed to reserve port after it should have been released", 22, port, err)
 	}
 }