3 files changed, 302 insertions, 0 deletions
diff --git a/pkg/tcpip/ports/BUILD b/pkg/tcpip/ports/BUILD
new file mode 100644
index 000000000..e0140cea6
--- /dev/null
+++ b/pkg/tcpip/ports/BUILD
@@ -0,0 +1,20 @@
+package(licenses = ["notice"])  # BSD
+
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "ports",
+    srcs = ["ports.go"],
+    importpath = "gvisor.googlesource.com/gvisor/pkg/tcpip/ports",
+    visibility = ["//:sandbox"],
+    deps = ["//pkg/tcpip"],
+)
+
+go_test(
+    name = "ports_test",
+    srcs = ["ports_test.go"],
+    embed = [":ports"],
+    deps = [
+        "//pkg/tcpip",
+    ],
+)
diff --git a/pkg/tcpip/ports/ports.go b/pkg/tcpip/ports/ports.go
new file mode 100644
index 000000000..24f3095d6
--- /dev/null
+++ b/pkg/tcpip/ports/ports.go
@@ -0,0 +1,148 @@
+// Copyright 2016 The Netstack Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package ports provides PortManager that manages allocating, reserving and releasing ports.
+package ports
+
+import (
+	"math"
+	"math/rand"
+	"sync"
+
+	"gvisor.googlesource.com/gvisor/pkg/tcpip"
+)
+
+const (
+	// firstEphemeral is the first ephemeral port.
+	firstEphemeral uint16 = 16000
+
+	anyIPAddress = tcpip.Address("")
+)
+
+type portDescriptor struct {
+	network   tcpip.NetworkProtocolNumber
+	transport tcpip.TransportProtocolNumber
+	port      uint16
+}
+
+// PortManager manages allocating, reserving and releasing ports.
+type PortManager struct {
+	mu             sync.RWMutex
+	allocatedPorts map[portDescriptor]bindAddresses
+}
+
+// bindAddresses is a set of IP addresses.
+type bindAddresses map[tcpip.Address]struct{}
+
+// isAvailable checks whether an IP address is available to bind to.
+func (b bindAddresses) isAvailable(addr tcpip.Address) bool {
+	if addr == anyIPAddress {
+		return len(b) == 0
+	}
+
+	// If all addresses for this portDescriptor are already bound, no
+	// address is available.
+	if _, ok := b[anyIPAddress]; ok {
+		return false
+	}
+
+	if _, ok := b[addr]; ok {
+		return false
+	}
+	return true
+}
+
+// NewPortManager creates new PortManager.
+func NewPortManager() *PortManager {
+	return &PortManager{allocatedPorts: make(map[portDescriptor]bindAddresses)}
+}
+
+// PickEphemeralPort randomly chooses a starting point and iterates over all
+// possible ephemeral ports, allowing the caller to decide whether a given port
+// is suitable for its needs, and stopping when a port is found or an error
+// occurs.
+func (s *PortManager) PickEphemeralPort(testPort func(p uint16) (bool, *tcpip.Error)) (port uint16, err *tcpip.Error) {
+	count := uint16(math.MaxUint16 - firstEphemeral + 1)
+	offset := uint16(rand.Int31n(int32(count)))
+
+	for i := uint16(0); i < count; i++ {
+		port = firstEphemeral + (offset+i)%count
+		ok, err := testPort(port)
+		if err != nil {
+			return 0, err
+		}
+
+		if ok {
+			return port, nil
+		}
+	}
+
+	return 0, tcpip.ErrNoPortAvailable
+}
+
+// ReservePort marks a port/IP combination as reserved so that it cannot be
+// 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(network []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) (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(network, transport, addr, port) {
+			return 0, tcpip.ErrPortInUse
+		}
+		return port, nil
+	}
+
+	// A port wasn't specified, so try to find one.
+	return s.PickEphemeralPort(func(p uint16) (bool, *tcpip.Error) {
+		return s.reserveSpecificPort(network, transport, addr, p), nil
+	})
+}
+
+// reserveSpecificPort tries to reserve the given port on all given protocols.
+func (s *PortManager) reserveSpecificPort(network []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) bool {
+	// Check that the port is available on all network protocols.
+	desc := portDescriptor{0, transport, port}
+	for _, n := range network {
+		desc.network = n
+		if addrs, ok := s.allocatedPorts[desc]; ok {
+			if !addrs.isAvailable(addr) {
+				return false
+			}
+		}
+	}
+
+	// Reserve port on all network protocols.
+	for _, n := range network {
+		desc.network = n
+		m, ok := s.allocatedPorts[desc]
+		if !ok {
+			m = make(bindAddresses)
+			s.allocatedPorts[desc] = m
+		}
+		m[addr] = struct{}{}
+	}
+
+	return true
+}
+
+// ReleasePort releases the reservation on a port/IP combination so that it can
+// be reserved by other endpoints.
+func (s *PortManager) ReleasePort(network []tcpip.NetworkProtocolNumber, transport tcpip.TransportProtocolNumber, addr tcpip.Address, port uint16) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for _, n := range network {
+		desc := portDescriptor{n, transport, port}
+		m := s.allocatedPorts[desc]
+		delete(m, addr)
+		if len(m) == 0 {
+			delete(s.allocatedPorts, desc)
+		}
+	}
+}
diff --git a/pkg/tcpip/ports/ports_test.go b/pkg/tcpip/ports/ports_test.go
new file mode 100644
index 000000000..9a4c702b2
--- /dev/null
+++ b/pkg/tcpip/ports/ports_test.go
@@ -0,0 +1,134 @@
+// Copyright 2016 The Netstack Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+package ports
+
+import (
+	"testing"
+
+	"gvisor.googlesource.com/gvisor/pkg/tcpip"
+)
+
+const (
+	fakeTransNumber   tcpip.TransportProtocolNumber = 1
+	fakeNetworkNumber tcpip.NetworkProtocolNumber   = 2
+
+	fakeIPAddress  = tcpip.Address("\x08\x08\x08\x08")
+	fakeIPAddress1 = tcpip.Address("\x08\x08\x08\x09")
+)
+
+func TestPortReservation(t *testing.T) {
+	pm := NewPortManager()
+	net := []tcpip.NetworkProtocolNumber{fakeNetworkNumber}
+
+	for _, test := range []struct {
+		port uint16
+		ip   tcpip.Address
+		want *tcpip.Error
+	}{
+		{
+			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,
+		},
+		{
+			port: 0,
+			ip:   fakeIPAddress,
+			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)
+		}
+	}
+
+	// Release port 22 from any IP address, then try to reserve fake IP
+	// address on 22.
+	pm.ReleasePort(net, fakeTransNumber, anyIPAddress, 22)
+
+	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)
+	}
+}
+
+func TestPickEphemeralPort(t *testing.T) {
+	pm := NewPortManager()
+	customErr := &tcpip.Error{}
+	for _, test := range []struct {
+		name     string
+		f        func(port uint16) (bool, *tcpip.Error)
+		wantErr  *tcpip.Error
+		wantPort uint16
+	}{
+		{
+			name: "no-port-available",
+			f: func(port uint16) (bool, *tcpip.Error) {
+				return false, nil
+			},
+			wantErr: tcpip.ErrNoPortAvailable,
+		},
+		{
+			name: "port-tester-error",
+			f: func(port uint16) (bool, *tcpip.Error) {
+				return false, customErr
+			},
+			wantErr: customErr,
+		},
+		{
+			name: "only-port-16042-available",
+			f: func(port uint16) (bool, *tcpip.Error) {
+				if port == firstEphemeral+42 {
+					return true, nil
+				}
+				return false, nil
+			},
+			wantPort: firstEphemeral + 42,
+		},
+		{
+			name: "only-port-under-16000-available",
+			f: func(port uint16) (bool, *tcpip.Error) {
+				if port < firstEphemeral {
+					return true, nil
+				}
+				return false, nil
+			},
+			wantErr: tcpip.ErrNoPortAvailable,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			if port, err := pm.PickEphemeralPort(test.f); port != test.wantPort || err != test.wantErr {
+				t.Errorf("PickEphemeralPort(..) = (port %d, err %v); want (port %d, err %v)", port, err, test.wantPort, test.wantErr)
+			}
+		})
+	}
+}