improve readability of ports package

Lots of small changes:
- simplify package API via Reservation type
- rename some single-letter variable names that were hard to follow
- rename some types

PiperOrigin-RevId: 362442366

pkg/tcpip/ports/BUILD

@@ -4,7 +4,10 @@ package(licenses = ["notice"])
 
 go_library(
     name = "ports",
-    srcs = ["ports.go"],
+    srcs = [
+        "flags.go",
+        "ports.go",
+    ],
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/sync",

pkg/tcpip/ports/flags.go

@@ -0,0 +1,150 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ports
+
+// Flags represents the type of port reservation.
+//
+// +stateify savable
+type Flags struct {
+	// MostRecent represents UDP SO_REUSEADDR.
+	MostRecent bool
+
+	// LoadBalanced indicates SO_REUSEPORT.
+	//
+	// LoadBalanced takes precidence over MostRecent.
+	LoadBalanced bool
+
+	// TupleOnly represents TCP SO_REUSEADDR.
+	TupleOnly bool
+}
+
+// Bits converts the Flags to their bitset form.
+func (f Flags) Bits() BitFlags {
+	var rf BitFlags
+	if f.MostRecent {
+		rf |= MostRecentFlag
+	}
+	if f.LoadBalanced {
+		rf |= LoadBalancedFlag
+	}
+	if f.TupleOnly {
+		rf |= TupleOnlyFlag
+	}
+	return rf
+}
+
+// Effective returns the effective behavior of a flag config.
+func (f Flags) Effective() Flags {
+	e := f
+	if e.LoadBalanced && e.MostRecent {
+		e.MostRecent = false
+	}
+	return e
+}
+
+// BitFlags is a bitset representation of Flags.
+type BitFlags uint32
+
+const (
+	// MostRecentFlag represents Flags.MostRecent.
+	MostRecentFlag BitFlags = 1 << iota
+
+	// LoadBalancedFlag represents Flags.LoadBalanced.
+	LoadBalancedFlag
+
+	// TupleOnlyFlag represents Flags.TupleOnly.
+	TupleOnlyFlag
+
+	// nextFlag is the value that the next added flag will have.
+	//
+	// It is used to calculate FlagMask below. It is also the number of
+	// valid flag states.
+	nextFlag
+
+	// FlagMask is a bit mask for BitFlags.
+	FlagMask = nextFlag - 1
+
+	// MultiBindFlagMask contains the flags that allow binding the same
+	// tuple multiple times.
+	MultiBindFlagMask = MostRecentFlag | LoadBalancedFlag
+)
+
+// ToFlags converts the bitset into a Flags struct.
+func (f BitFlags) ToFlags() Flags {
+	return Flags{
+		MostRecent:   f&MostRecentFlag != 0,
+		LoadBalanced: f&LoadBalancedFlag != 0,
+		TupleOnly:    f&TupleOnlyFlag != 0,
+	}
+}
+
+// FlagCounter counts how many references each flag combination has.
+type FlagCounter struct {
+	// refs stores the count for each possible flag combination, (0 though
+	// FlagMask).
+	refs [nextFlag]int
+}
+
+// AddRef increases the reference count for a specific flag combination.
+func (c *FlagCounter) AddRef(flags BitFlags) {
+	c.refs[flags]++
+}
+
+// DropRef decreases the reference count for a specific flag combination.
+func (c *FlagCounter) DropRef(flags BitFlags) {
+	c.refs[flags]--
+}
+
+// TotalRefs calculates the total number of references for all flag
+// combinations.
+func (c FlagCounter) TotalRefs() int {
+	var total int
+	for _, r := range c.refs {
+		total += r
+	}
+	return total
+}
+
+// FlagRefs returns the number of references with all specified flags.
+func (c FlagCounter) FlagRefs(flags BitFlags) int {
+	var total int
+	for i, r := range c.refs {
+		if BitFlags(i)&flags == flags {
+			total += r
+		}
+	}
+	return total
+}
+
+// AllRefsHave returns if all references have all specified flags.
+func (c FlagCounter) AllRefsHave(flags BitFlags) bool {
+	for i, r := range c.refs {
+		if BitFlags(i)&flags != flags && r > 0 {
+			return false
+		}
+	}
+	return true
+}
+
+// SharedFlags returns the set of flags shared by all references.
+func (c FlagCounter) SharedFlags() BitFlags {
+	intersection := FlagMask
+	for i, r := range c.refs {
+		if r > 0 {
+			intersection &= BitFlags(i)
+		}
+	}
+	return intersection
+}

pkg/tcpip/ports/ports_test.go

@@ -331,15 +331,33 @@ func TestPortReservation(t *testing.T) {
 			for _, test := range test.actions {
 				first, _ := pm.PortRange()
 				if test.release {
-					pm.ReleasePort(net, fakeTransNumber, test.ip, test.port, test.flags, test.device, test.dest)
+					portRes := Reservation{
+						Networks:     net,
+						Transport:    fakeTransNumber,
+						Addr:         test.ip,
+						Port:         test.port,
+						Flags:        test.flags,
+						BindToDevice: test.device,
+						Dest:         test.dest,
+					}
+					pm.ReleasePort(portRes)
 					continue
 				}
-				gotPort, err := pm.ReservePort(net, fakeTransNumber, test.ip, test.port, test.flags, test.device, test.dest, nil /* testPort */)
+				portRes := Reservation{
+					Networks:     net,
+					Transport:    fakeTransNumber,
+					Addr:         test.ip,
+					Port:         test.port,
+					Flags:        test.flags,
+					BindToDevice: test.device,
+					Dest:         test.dest,
+				}
+				gotPort, err := pm.ReservePort(portRes, nil /* testPort */)
 				if diff := cmp.Diff(test.want, err); diff != "" {
-					t.Fatalf("unexpected error from ReservePort(.., .., %s, %d, %+v, %d, %v), (-want, +got):\n%s", test.ip, test.port, test.flags, test.device, test.dest, diff)
+					t.Fatalf("unexpected error from ReservePort(%+v, _), (-want, +got):\n%s", portRes, diff)
 				}
 				if test.port == 0 && (gotPort == 0 || gotPort < first) {
-					t.Fatalf("ReservePort(.., .., .., 0, ..) = %d, want port number >= %d to be picked", gotPort, first)
+					t.Fatalf("ReservePort(%+v, _) = %d, want port number >= %d to be picked", portRes, gotPort, first)
 				}
 			}
 		})

pkg/tcpip/stack/transport_demuxer.go

@@ -359,7 +359,7 @@ func selectEndpoint(id TransportEndpointID, mpep *multiPortEndpoint, seed uint32
 		return mpep.endpoints[0]
 	}
 
-	if mpep.flags.IntersectionRefs().ToFlags().Effective().MostRecent {
+	if mpep.flags.SharedFlags().ToFlags().Effective().MostRecent {
 		return mpep.endpoints[len(mpep.endpoints)-1]
 	}
 
@@ -410,7 +410,7 @@ func (ep *multiPortEndpoint) singleRegisterEndpoint(t TransportEndpoint, flags p
 
 	if len(ep.endpoints) != 0 {
 		// If it was previously bound, we need to check if we can bind again.
-		if ep.flags.TotalRefs() > 0 && bits&ep.flags.IntersectionRefs() == 0 {
+		if ep.flags.TotalRefs() > 0 && bits&ep.flags.SharedFlags() == 0 {
 			return &tcpip.ErrPortInUse{}
 		}
 	}
@@ -429,7 +429,7 @@ func (ep *multiPortEndpoint) singleCheckEndpoint(flags ports.Flags) tcpip.Error
 
 	if len(ep.endpoints) != 0 {
 		// If it was previously bound, we need to check if we can bind again.
-		if ep.flags.TotalRefs() > 0 && bits&ep.flags.IntersectionRefs() == 0 {
+		if ep.flags.TotalRefs() > 0 && bits&ep.flags.SharedFlags() == 0 {
 			return &tcpip.ErrPortInUse{}
 		}
 	}

pkg/tcpip/transport/tcp/accept.go

@@ -27,6 +27,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/pkg/tcpip/ports"
 	"gvisor.dev/gvisor/pkg/tcpip/seqnum"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/waiter"
@@ -432,15 +433,16 @@ func (e *endpoint) propagateInheritableOptionsLocked(n *endpoint) {
 // * e.mu is held.
 func (e *endpoint) reserveTupleLocked() bool {
 	dest := tcpip.FullAddress{Addr: e.ID.RemoteAddress, Port: e.ID.RemotePort}
-	if !e.stack.ReserveTuple(
-		e.effectiveNetProtos,
-		ProtocolNumber,
-		e.ID.LocalAddress,
-		e.ID.LocalPort,
-		e.boundPortFlags,
-		e.boundBindToDevice,
-		dest,
-	) {
+	portRes := ports.Reservation{
+		Networks:     e.effectiveNetProtos,
+		Transport:    ProtocolNumber,
+		Addr:         e.ID.LocalAddress,
+		Port:         e.ID.LocalPort,
+		Flags:        e.boundPortFlags,
+		BindToDevice: e.boundBindToDevice,
+		Dest:         dest,
+	}
+	if !e.stack.ReserveTuple(portRes) {
 		return false
 	}
 

pkg/tcpip/transport/tcp/endpoint.go

@@ -1097,7 +1097,16 @@ func (e *endpoint) closeNoShutdownLocked() {
 			e.isRegistered = false
 		}
 
-		e.stack.ReleasePort(e.effectiveNetProtos, ProtocolNumber, e.ID.LocalAddress, e.ID.LocalPort, e.boundPortFlags, e.boundBindToDevice, e.boundDest)
+		portRes := ports.Reservation{
+			Networks:     e.effectiveNetProtos,
+			Transport:    ProtocolNumber,
+			Addr:         e.ID.LocalAddress,
+			Port:         e.ID.LocalPort,
+			Flags:        e.boundPortFlags,
+			BindToDevice: e.boundBindToDevice,
+			Dest:         e.boundDest,
+		}
+		e.stack.ReleasePort(portRes)
 		e.isPortReserved = false
 		e.boundBindToDevice = 0
 		e.boundPortFlags = ports.Flags{}
@@ -1172,7 +1181,16 @@ func (e *endpoint) cleanupLocked() {
 	}
 
 	if e.isPortReserved {
-		e.stack.ReleasePort(e.effectiveNetProtos, ProtocolNumber, e.ID.LocalAddress, e.ID.LocalPort, e.boundPortFlags, e.boundBindToDevice, e.boundDest)
+		portRes := ports.Reservation{
+			Networks:     e.effectiveNetProtos,
+			Transport:    ProtocolNumber,
+			Addr:         e.ID.LocalAddress,
+			Port:         e.ID.LocalPort,
+			Flags:        e.boundPortFlags,
+			BindToDevice: e.boundBindToDevice,
+			Dest:         e.boundDest,
+		}
+		e.stack.ReleasePort(portRes)
 		e.isPortReserved = false
 	}
 	e.boundBindToDevice = 0
@@ -2242,7 +2260,16 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) tcp
 			if sameAddr && p == e.ID.RemotePort {
 				return false, nil
 			}
-			if _, err := e.stack.ReservePort(netProtos, ProtocolNumber, e.ID.LocalAddress, p, e.portFlags, bindToDevice, addr, nil /* testPort */); err != nil {
+			portRes := ports.Reservation{
+				Networks:     netProtos,
+				Transport:    ProtocolNumber,
+				Addr:         e.ID.LocalAddress,
+				Port:         p,
+				Flags:        e.portFlags,
+				BindToDevice: bindToDevice,
+				Dest:         addr,
+			}
+			if _, err := e.stack.ReservePort(portRes, nil /* testPort */); err != nil {
 				if _, ok := err.(*tcpip.ErrPortInUse); !ok || !reuse {
 					return false, nil
 				}
@@ -2280,7 +2307,16 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) tcp
 				tcpEP.notifyProtocolGoroutine(notifyAbort)
 				tcpEP.UnlockUser()
 				// Now try and Reserve again if it fails then we skip.
-				if _, err := e.stack.ReservePort(netProtos, ProtocolNumber, e.ID.LocalAddress, p, e.portFlags, bindToDevice, addr, nil /* testPort */); err != nil {
+				portRes := ports.Reservation{
+					Networks:     netProtos,
+					Transport:    ProtocolNumber,
+					Addr:         e.ID.LocalAddress,
+					Port:         p,
+					Flags:        e.portFlags,
+					BindToDevice: bindToDevice,
+					Dest:         addr,
+				}
+				if _, err := e.stack.ReservePort(portRes, nil /* testPort */); err != nil {
 					return false, nil
 				}
 			}
@@ -2288,7 +2324,16 @@ func (e *endpoint) connect(addr tcpip.FullAddress, handshake bool, run bool) tcp
 			id := e.ID
 			id.LocalPort = p
 			if err := e.stack.RegisterTransportEndpoint(netProtos, ProtocolNumber, id, e, e.portFlags, bindToDevice); err != nil {
-				e.stack.ReleasePort(netProtos, ProtocolNumber, e.ID.LocalAddress, p, e.portFlags, bindToDevice, addr)
+				portRes := ports.Reservation{
+					Networks:     netProtos,
+					Transport:    ProtocolNumber,
+					Addr:         e.ID.LocalAddress,
+					Port:         p,
+					Flags:        e.portFlags,
+					BindToDevice: bindToDevice,
+					Dest:         addr,
+				}
+				e.stack.ReleasePort(portRes)
 				if _, ok := err.(*tcpip.ErrPortInUse); ok {
 					return false, nil
 				}
@@ -2604,7 +2649,16 @@ func (e *endpoint) bindLocked(addr tcpip.FullAddress) (err tcpip.Error) {
 	}
 
 	bindToDevice := tcpip.NICID(e.ops.GetBindToDevice())
-	port, err := e.stack.ReservePort(netProtos, ProtocolNumber, addr.Addr, addr.Port, e.portFlags, bindToDevice, tcpip.FullAddress{}, func(p uint16) bool {
+	portRes := ports.Reservation{
+		Networks:     netProtos,
+		Transport:    ProtocolNumber,
+		Addr:         addr.Addr,
+		Port:         addr.Port,
+		Flags:        e.portFlags,
+		BindToDevice: bindToDevice,
+		Dest:         tcpip.FullAddress{},
+	}
+	port, err := e.stack.ReservePort(portRes, func(p uint16) (bool, tcpip.Error) {
 		id := e.ID
 		id.LocalPort = p
 		// CheckRegisterTransportEndpoint should only return an error if there is a
@@ -2616,9 +2670,9 @@ func (e *endpoint) bindLocked(addr tcpip.FullAddress) (err tcpip.Error) {
 		// address/port. Hence this will only return an error if there is a matching
 		// listening endpoint.
 		if err := e.stack.CheckRegisterTransportEndpoint(netProtos, ProtocolNumber, id, e.portFlags, bindToDevice); err != nil {
-			return false
+			return false, nil
 		}
-		return true
+		return true, nil
 	})
 	if err != nil {
 		return err

pkg/tcpip/transport/tcp/endpoint_state.go

@@ -22,6 +22,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/pkg/tcpip/ports"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
@@ -208,7 +209,16 @@ func (e *endpoint) Resume(s *stack.Stack) {
 		if err != nil {
 			panic("unable to parse BindAddr: " + err.String())
 		}
-		if ok := e.stack.ReserveTuple(e.effectiveNetProtos, ProtocolNumber, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest); !ok {
+		portRes := ports.Reservation{
+			Networks:     e.effectiveNetProtos,
+			Transport:    ProtocolNumber,
+			Addr:         addr.Addr,
+			Port:         addr.Port,
+			Flags:        e.boundPortFlags,
+			BindToDevice: e.boundBindToDevice,
+			Dest:         e.boundDest,
+		}
+		if ok := e.stack.ReserveTuple(portRes); !ok {
 			panic(fmt.Sprintf("unable to re-reserve tuple (%v, %q, %d, %+v, %d, %v)", e.effectiveNetProtos, addr.Addr, addr.Port, e.boundPortFlags, e.boundBindToDevice, e.boundDest))
 		}
 		e.isPortReserved = true

pkg/tcpip/transport/udp/endpoint.go

@@ -245,7 +245,16 @@ func (e *endpoint) Close() {
 	switch e.EndpointState() {
 	case StateBound, StateConnected:
 		e.stack.UnregisterTransportEndpoint(e.effectiveNetProtos, ProtocolNumber, e.ID, e, e.boundPortFlags, e.boundBindToDevice)
-		e.stack.ReleasePort(e.effectiveNetProtos, ProtocolNumber, e.ID.LocalAddress, e.ID.LocalPort, e.boundPortFlags, e.boundBindToDevice, tcpip.FullAddress{})
+		portRes := ports.Reservation{
+			Networks:     e.effectiveNetProtos,
+			Transport:    ProtocolNumber,
+			Addr:         e.ID.LocalAddress,
+			Port:         e.ID.LocalPort,
+			Flags:        e.boundPortFlags,
+			BindToDevice: e.boundBindToDevice,
+			Dest:         tcpip.FullAddress{},
+		}
+		e.stack.ReleasePort(portRes)
 		e.boundBindToDevice = 0
 		e.boundPortFlags = ports.Flags{}
 	}
@@ -920,7 +929,16 @@ func (e *endpoint) Disconnect() tcpip.Error {
 	} else {
 		if e.ID.LocalPort != 0 {
 			// Release the ephemeral port.
-			e.stack.ReleasePort(e.effectiveNetProtos, ProtocolNumber, e.ID.LocalAddress, e.ID.LocalPort, boundPortFlags, e.boundBindToDevice, tcpip.FullAddress{})
+			portRes := ports.Reservation{
+				Networks:     e.effectiveNetProtos,
+				Transport:    ProtocolNumber,
+				Addr:         e.ID.LocalAddress,
+				Port:         e.ID.LocalPort,
+				Flags:        boundPortFlags,
+				BindToDevice: e.boundBindToDevice,
+				Dest:         tcpip.FullAddress{},
+			}
+			e.stack.ReleasePort(portRes)
 			e.boundPortFlags = ports.Flags{}
 		}
 		e.setEndpointState(StateInitial)
@@ -1072,7 +1090,16 @@ func (*endpoint) Accept(*tcpip.FullAddress) (tcpip.Endpoint, *waiter.Queue, tcpi
 func (e *endpoint) registerWithStack(netProtos []tcpip.NetworkProtocolNumber, id stack.TransportEndpointID) (stack.TransportEndpointID, tcpip.NICID, tcpip.Error) {
 	bindToDevice := tcpip.NICID(e.ops.GetBindToDevice())
 	if e.ID.LocalPort == 0 {
-		port, err := e.stack.ReservePort(netProtos, ProtocolNumber, id.LocalAddress, id.LocalPort, e.portFlags, bindToDevice, tcpip.FullAddress{}, nil /* testPort */)
+		portRes := ports.Reservation{
+			Networks:     netProtos,
+			Transport:    ProtocolNumber,
+			Addr:         id.LocalAddress,
+			Port:         id.LocalPort,
+			Flags:        e.portFlags,
+			BindToDevice: bindToDevice,
+			Dest:         tcpip.FullAddress{},
+		}
+		port, err := e.stack.ReservePort(portRes, nil /* testPort */)
 		if err != nil {
 			return id, bindToDevice, err
 		}
@@ -1082,7 +1109,16 @@ func (e *endpoint) registerWithStack(netProtos []tcpip.NetworkProtocolNumber, id
 
 	err := e.stack.RegisterTransportEndpoint(netProtos, ProtocolNumber, id, e, e.boundPortFlags, bindToDevice)
 	if err != nil {
-		e.stack.ReleasePort(netProtos, ProtocolNumber, id.LocalAddress, id.LocalPort, e.boundPortFlags, bindToDevice, tcpip.FullAddress{})
+		portRes := ports.Reservation{
+			Networks:     netProtos,
+			Transport:    ProtocolNumber,
+			Addr:         id.LocalAddress,
+			Port:         id.LocalPort,
+			Flags:        e.boundPortFlags,
+			BindToDevice: bindToDevice,
+			Dest:         tcpip.FullAddress{},
+		}
+		e.stack.ReleasePort(portRes)
 		e.boundPortFlags = ports.Flags{}
 	}
 	return id, bindToDevice, err