github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/tcpip/ports/ports_test.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ports
    16  
    17  import (
    18  	"math"
    19  	"math/rand"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/google/go-cmp/cmp"
    24  	"github.com/SagerNet/gvisor/pkg/tcpip"
    25  	"github.com/SagerNet/gvisor/pkg/tcpip/testutil"
    26  )
    27  
    28  const (
    29  	fakeTransNumber   tcpip.TransportProtocolNumber = 1
    30  	fakeNetworkNumber tcpip.NetworkProtocolNumber   = 2
    31  )
    32  
    33  var (
    34  	fakeIPAddress  = testutil.MustParse4("8.8.8.8")
    35  	fakeIPAddress1 = testutil.MustParse4("8.8.8.9")
    36  )
    37  
    38  type portReserveTestAction struct {
    39  	port    uint16
    40  	ip      tcpip.Address
    41  	want    tcpip.Error
    42  	flags   Flags
    43  	release bool
    44  	device  tcpip.NICID
    45  	dest    tcpip.FullAddress
    46  }
    47  
    48  func TestPortReservation(t *testing.T) {
    49  	for _, test := range []struct {
    50  		tname   string
    51  		actions []portReserveTestAction
    52  	}{
    53  		{
    54  			tname: "bind to ip",
    55  			actions: []portReserveTestAction{
    56  				{port: 80, ip: fakeIPAddress, want: nil},
    57  				{port: 80, ip: fakeIPAddress1, want: nil},
    58  				/* N.B. Order of tests matters! */
    59  				{port: 80, ip: anyIPAddress, want: &tcpip.ErrPortInUse{}},
    60  				{port: 80, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}, flags: Flags{LoadBalanced: true}},
    61  			},
    62  		},
    63  		{
    64  			tname: "bind to inaddr any",
    65  			actions: []portReserveTestAction{
    66  				{port: 22, ip: anyIPAddress, want: nil},
    67  				{port: 22, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}},
    68  				/* release fakeIPAddress, but anyIPAddress is still inuse */
    69  				{port: 22, ip: fakeIPAddress, release: true},
    70  				{port: 22, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}},
    71  				{port: 22, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}, flags: Flags{LoadBalanced: true}},
    72  				/* Release port 22 from any IP address, then try to reserve fake IP address on 22 */
    73  				{port: 22, ip: anyIPAddress, want: nil, release: true},
    74  				{port: 22, ip: fakeIPAddress, want: nil},
    75  			},
    76  		}, {
    77  			tname: "bind to zero port",
    78  			actions: []portReserveTestAction{
    79  				{port: 00, ip: fakeIPAddress, want: nil},
    80  				{port: 00, ip: fakeIPAddress, want: nil},
    81  				{port: 00, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    82  			},
    83  		}, {
    84  			tname: "bind to ip with reuseport",
    85  			actions: []portReserveTestAction{
    86  				{port: 25, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    87  				{port: 25, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    88  
    89  				{port: 25, ip: fakeIPAddress, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
    90  				{port: 25, ip: anyIPAddress, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
    91  
    92  				{port: 25, ip: anyIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    93  			},
    94  		}, {
    95  			tname: "bind to inaddr any with reuseport",
    96  			actions: []portReserveTestAction{
    97  				{port: 24, ip: anyIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    98  				{port: 24, ip: anyIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
    99  
   100  				{port: 24, ip: anyIPAddress, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
   101  				{port: 24, ip: fakeIPAddress, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
   102  
   103  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
   104  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, release: true, want: nil},
   105  
   106  				{port: 24, ip: anyIPAddress, flags: Flags{LoadBalanced: true}, release: true},
   107  				{port: 24, ip: anyIPAddress, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
   108  
   109  				{port: 24, ip: anyIPAddress, flags: Flags{LoadBalanced: true}, release: true},
   110  				{port: 24, ip: anyIPAddress, flags: Flags{}, want: nil},
   111  			},
   112  		}, {
   113  			tname: "bind twice with device fails",
   114  			actions: []portReserveTestAction{
   115  				{port: 24, ip: fakeIPAddress, device: 3, want: nil},
   116  				{port: 24, ip: fakeIPAddress, device: 3, want: &tcpip.ErrPortInUse{}},
   117  			},
   118  		}, {
   119  			tname: "bind to device",
   120  			actions: []portReserveTestAction{
   121  				{port: 24, ip: fakeIPAddress, device: 1, want: nil},
   122  				{port: 24, ip: fakeIPAddress, device: 2, want: nil},
   123  			},
   124  		}, {
   125  			tname: "bind to device and then without device",
   126  			actions: []portReserveTestAction{
   127  				{port: 24, ip: fakeIPAddress, device: 123, want: nil},
   128  				{port: 24, ip: fakeIPAddress, device: 0, want: &tcpip.ErrPortInUse{}},
   129  			},
   130  		}, {
   131  			tname: "bind without device",
   132  			actions: []portReserveTestAction{
   133  				{port: 24, ip: fakeIPAddress, want: nil},
   134  				{port: 24, ip: fakeIPAddress, device: 123, want: &tcpip.ErrPortInUse{}},
   135  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   136  				{port: 24, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}},
   137  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   138  			},
   139  		}, {
   140  			tname: "bind with device",
   141  			actions: []portReserveTestAction{
   142  				{port: 24, ip: fakeIPAddress, device: 123, want: nil},
   143  				{port: 24, ip: fakeIPAddress, device: 123, want: &tcpip.ErrPortInUse{}},
   144  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   145  				{port: 24, ip: fakeIPAddress, device: 0, want: &tcpip.ErrPortInUse{}},
   146  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   147  				{port: 24, ip: fakeIPAddress, device: 456, flags: Flags{LoadBalanced: true}, want: nil},
   148  				{port: 24, ip: fakeIPAddress, device: 789, want: nil},
   149  				{port: 24, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}},
   150  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   151  			},
   152  		}, {
   153  			tname: "bind with reuseport",
   154  			actions: []portReserveTestAction{
   155  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
   156  				{port: 24, ip: fakeIPAddress, device: 123, want: &tcpip.ErrPortInUse{}},
   157  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   158  				{port: 24, ip: fakeIPAddress, device: 0, want: &tcpip.ErrPortInUse{}},
   159  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: nil},
   160  			},
   161  		}, {
   162  			tname: "binding with reuseport and device",
   163  			actions: []portReserveTestAction{
   164  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   165  				{port: 24, ip: fakeIPAddress, device: 123, want: &tcpip.ErrPortInUse{}},
   166  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   167  				{port: 24, ip: fakeIPAddress, device: 0, want: &tcpip.ErrPortInUse{}},
   168  				{port: 24, ip: fakeIPAddress, device: 456, flags: Flags{LoadBalanced: true}, want: nil},
   169  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: nil},
   170  				{port: 24, ip: fakeIPAddress, device: 789, flags: Flags{LoadBalanced: true}, want: nil},
   171  				{port: 24, ip: fakeIPAddress, device: 999, want: &tcpip.ErrPortInUse{}},
   172  			},
   173  		}, {
   174  			tname: "mixing reuseport and not reuseport by binding to device",
   175  			actions: []portReserveTestAction{
   176  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   177  				{port: 24, ip: fakeIPAddress, device: 456, want: nil},
   178  				{port: 24, ip: fakeIPAddress, device: 789, flags: Flags{LoadBalanced: true}, want: nil},
   179  				{port: 24, ip: fakeIPAddress, device: 999, want: nil},
   180  			},
   181  		}, {
   182  			tname: "can't bind to 0 after mixing reuseport and not reuseport",
   183  			actions: []portReserveTestAction{
   184  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   185  				{port: 24, ip: fakeIPAddress, device: 456, want: nil},
   186  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   187  			},
   188  		}, {
   189  			tname: "bind and release",
   190  			actions: []portReserveTestAction{
   191  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{LoadBalanced: true}, want: nil},
   192  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: nil},
   193  				{port: 24, ip: fakeIPAddress, device: 345, flags: Flags{}, want: &tcpip.ErrPortInUse{}},
   194  				{port: 24, ip: fakeIPAddress, device: 789, flags: Flags{LoadBalanced: true}, want: nil},
   195  
   196  				// Release the bind to device 0 and try again.
   197  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: nil, release: true},
   198  				{port: 24, ip: fakeIPAddress, device: 345, flags: Flags{}, want: nil},
   199  			},
   200  		}, {
   201  			tname: "bind twice with reuseport once",
   202  			actions: []portReserveTestAction{
   203  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{}, want: nil},
   204  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   205  			},
   206  		}, {
   207  			tname: "release an unreserved device",
   208  			actions: []portReserveTestAction{
   209  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{}, want: nil},
   210  				{port: 24, ip: fakeIPAddress, device: 456, flags: Flags{}, want: nil},
   211  				// The below don't exist.
   212  				{port: 24, ip: fakeIPAddress, device: 345, flags: Flags{}, want: nil, release: true},
   213  				{port: 9999, ip: fakeIPAddress, device: 123, flags: Flags{}, want: nil, release: true},
   214  				// Release all.
   215  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{}, want: nil, release: true},
   216  				{port: 24, ip: fakeIPAddress, device: 456, flags: Flags{}, want: nil, release: true},
   217  			},
   218  		}, {
   219  			tname: "bind with reuseaddr",
   220  			actions: []portReserveTestAction{
   221  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: nil},
   222  				{port: 24, ip: fakeIPAddress, device: 123, want: &tcpip.ErrPortInUse{}},
   223  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{MostRecent: true}, want: nil},
   224  				{port: 24, ip: fakeIPAddress, device: 0, want: &tcpip.ErrPortInUse{}},
   225  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{MostRecent: true}, want: nil},
   226  			},
   227  		}, {
   228  			tname: "bind twice with reuseaddr once",
   229  			actions: []portReserveTestAction{
   230  				{port: 24, ip: fakeIPAddress, device: 123, flags: Flags{}, want: nil},
   231  				{port: 24, ip: fakeIPAddress, device: 0, flags: Flags{MostRecent: true}, want: &tcpip.ErrPortInUse{}},
   232  			},
   233  		}, {
   234  			tname: "bind with reuseaddr and reuseport",
   235  			actions: []portReserveTestAction{
   236  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   237  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   238  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   239  			},
   240  		}, {
   241  			tname: "bind with reuseaddr and reuseport, and then reuseaddr",
   242  			actions: []portReserveTestAction{
   243  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   244  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: nil},
   245  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   246  			},
   247  		}, {
   248  			tname: "bind with reuseaddr and reuseport, and then reuseport",
   249  			actions: []portReserveTestAction{
   250  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   251  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
   252  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: &tcpip.ErrPortInUse{}},
   253  			},
   254  		}, {
   255  			tname: "bind with reuseaddr and reuseport twice, and then reuseaddr",
   256  			actions: []portReserveTestAction{
   257  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   258  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   259  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: nil},
   260  			},
   261  		}, {
   262  			tname: "bind with reuseaddr and reuseport twice, and then reuseport",
   263  			actions: []portReserveTestAction{
   264  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   265  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   266  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
   267  			},
   268  		}, {
   269  			tname: "bind with reuseaddr, and then reuseaddr and reuseport",
   270  			actions: []portReserveTestAction{
   271  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: nil},
   272  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   273  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: &tcpip.ErrPortInUse{}},
   274  			},
   275  		}, {
   276  			tname: "bind with reuseport, and then reuseaddr and reuseport",
   277  			actions: []portReserveTestAction{
   278  				{port: 24, ip: fakeIPAddress, flags: Flags{LoadBalanced: true}, want: nil},
   279  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true, LoadBalanced: true}, want: nil},
   280  				{port: 24, ip: fakeIPAddress, flags: Flags{MostRecent: true}, want: &tcpip.ErrPortInUse{}},
   281  			},
   282  		}, {
   283  			tname: "bind tuple with reuseaddr, and then wildcard with reuseaddr",
   284  			actions: []portReserveTestAction{
   285  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: nil},
   286  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{}, want: nil},
   287  			},
   288  		}, {
   289  			tname: "bind tuple with reuseaddr, and then wildcard",
   290  			actions: []portReserveTestAction{
   291  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: nil},
   292  				{port: 24, ip: fakeIPAddress, want: &tcpip.ErrPortInUse{}},
   293  			},
   294  		}, {
   295  			tname: "bind wildcard with reuseaddr, and then tuple with reuseaddr",
   296  			actions: []portReserveTestAction{
   297  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{}, want: nil},
   298  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: nil},
   299  			},
   300  		}, {
   301  			tname: "bind tuple with reuseaddr, and then wildcard",
   302  			actions: []portReserveTestAction{
   303  				{port: 24, ip: fakeIPAddress, want: nil},
   304  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: &tcpip.ErrPortInUse{}},
   305  			},
   306  		}, {
   307  			tname: "bind two tuples with reuseaddr",
   308  			actions: []portReserveTestAction{
   309  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: nil},
   310  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 25}, want: nil},
   311  			},
   312  		}, {
   313  			tname: "bind two tuples",
   314  			actions: []portReserveTestAction{
   315  				{port: 24, ip: fakeIPAddress, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: nil},
   316  				{port: 24, ip: fakeIPAddress, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 25}, want: nil},
   317  			},
   318  		}, {
   319  			tname: "bind wildcard, and then tuple with reuseaddr",
   320  			actions: []portReserveTestAction{
   321  				{port: 24, ip: fakeIPAddress, dest: tcpip.FullAddress{}, want: nil},
   322  				{port: 24, ip: fakeIPAddress, flags: Flags{TupleOnly: true}, dest: tcpip.FullAddress{Addr: fakeIPAddress, Port: 24}, want: &tcpip.ErrPortInUse{}},
   323  			},
   324  		}, {
   325  			tname: "bind wildcard twice with reuseaddr",
   326  			actions: []portReserveTestAction{
   327  				{port: 24, ip: anyIPAddress, flags: Flags{TupleOnly: true}, want: nil},
   328  				{port: 24, ip: anyIPAddress, flags: Flags{TupleOnly: true}, want: nil},
   329  			},
   330  		},
   331  	} {
   332  		t.Run(test.tname, func(t *testing.T) {
   333  			pm := NewPortManager()
   334  			net := []tcpip.NetworkProtocolNumber{fakeNetworkNumber}
   335  			rng := rand.New(rand.NewSource(time.Now().UnixNano()))
   336  
   337  			for _, test := range test.actions {
   338  				first, _ := pm.PortRange()
   339  				if test.release {
   340  					portRes := Reservation{
   341  						Networks:     net,
   342  						Transport:    fakeTransNumber,
   343  						Addr:         test.ip,
   344  						Port:         test.port,
   345  						Flags:        test.flags,
   346  						BindToDevice: test.device,
   347  						Dest:         test.dest,
   348  					}
   349  					pm.ReleasePort(portRes)
   350  					continue
   351  				}
   352  				portRes := Reservation{
   353  					Networks:     net,
   354  					Transport:    fakeTransNumber,
   355  					Addr:         test.ip,
   356  					Port:         test.port,
   357  					Flags:        test.flags,
   358  					BindToDevice: test.device,
   359  					Dest:         test.dest,
   360  				}
   361  				gotPort, err := pm.ReservePort(rng, portRes, nil /* testPort */)
   362  				if diff := cmp.Diff(test.want, err); diff != "" {
   363  					t.Fatalf("unexpected error from ReservePort(%+v, _), (-want, +got):\n%s", portRes, diff)
   364  				}
   365  				if test.port == 0 && (gotPort == 0 || gotPort < first) {
   366  					t.Fatalf("ReservePort(%+v, _) = %d, want port number >= %d to be picked", portRes, gotPort, first)
   367  				}
   368  			}
   369  		})
   370  	}
   371  }
   372  
   373  func TestPickEphemeralPort(t *testing.T) {
   374  	const (
   375  		firstEphemeral    = 32000
   376  		numEphemeralPorts = 1000
   377  	)
   378  
   379  	for _, test := range []struct {
   380  		name     string
   381  		f        func(port uint16) (bool, tcpip.Error)
   382  		wantErr  tcpip.Error
   383  		wantPort uint16
   384  	}{
   385  		{
   386  			name: "no-port-available",
   387  			f: func(port uint16) (bool, tcpip.Error) {
   388  				return false, nil
   389  			},
   390  			wantErr: &tcpip.ErrNoPortAvailable{},
   391  		},
   392  		{
   393  			name: "port-tester-error",
   394  			f: func(port uint16) (bool, tcpip.Error) {
   395  				return false, &tcpip.ErrBadBuffer{}
   396  			},
   397  			wantErr: &tcpip.ErrBadBuffer{},
   398  		},
   399  		{
   400  			name: "only-port-16042-available",
   401  			f: func(port uint16) (bool, tcpip.Error) {
   402  				if port == firstEphemeral+42 {
   403  					return true, nil
   404  				}
   405  				return false, nil
   406  			},
   407  			wantPort: firstEphemeral + 42,
   408  		},
   409  		{
   410  			name: "only-port-under-16000-available",
   411  			f: func(port uint16) (bool, tcpip.Error) {
   412  				if port < firstEphemeral {
   413  					return true, nil
   414  				}
   415  				return false, nil
   416  			},
   417  			wantErr: &tcpip.ErrNoPortAvailable{},
   418  		},
   419  	} {
   420  		t.Run(test.name, func(t *testing.T) {
   421  			pm := NewPortManager()
   422  			rng := rand.New(rand.NewSource(time.Now().UnixNano()))
   423  			if err := pm.SetPortRange(firstEphemeral, firstEphemeral+numEphemeralPorts); err != nil {
   424  				t.Fatalf("failed to set ephemeral port range: %s", err)
   425  			}
   426  			port, err := pm.PickEphemeralPort(rng, test.f)
   427  			if diff := cmp.Diff(test.wantErr, err); diff != "" {
   428  				t.Fatalf("unexpected error from PickEphemeralPort(..), (-want, +got):\n%s", diff)
   429  			}
   430  			if port != test.wantPort {
   431  				t.Errorf("got PickEphemeralPort(..) = (%d, nil); want (%d, nil)", port, test.wantPort)
   432  			}
   433  		})
   434  	}
   435  }
   436  
   437  func TestPickEphemeralPortStable(t *testing.T) {
   438  	const (
   439  		firstEphemeral    = 32000
   440  		numEphemeralPorts = 1000
   441  	)
   442  
   443  	for _, test := range []struct {
   444  		name     string
   445  		f        func(port uint16) (bool, tcpip.Error)
   446  		wantErr  tcpip.Error
   447  		wantPort uint16
   448  	}{
   449  		{
   450  			name: "no-port-available",
   451  			f: func(port uint16) (bool, tcpip.Error) {
   452  				return false, nil
   453  			},
   454  			wantErr: &tcpip.ErrNoPortAvailable{},
   455  		},
   456  		{
   457  			name: "port-tester-error",
   458  			f: func(port uint16) (bool, tcpip.Error) {
   459  				return false, &tcpip.ErrBadBuffer{}
   460  			},
   461  			wantErr: &tcpip.ErrBadBuffer{},
   462  		},
   463  		{
   464  			name: "only-port-16042-available",
   465  			f: func(port uint16) (bool, tcpip.Error) {
   466  				if port == firstEphemeral+42 {
   467  					return true, nil
   468  				}
   469  				return false, nil
   470  			},
   471  			wantPort: firstEphemeral + 42,
   472  		},
   473  		{
   474  			name: "only-port-under-16000-available",
   475  			f: func(port uint16) (bool, tcpip.Error) {
   476  				if port < firstEphemeral {
   477  					return true, nil
   478  				}
   479  				return false, nil
   480  			},
   481  			wantErr: &tcpip.ErrNoPortAvailable{},
   482  		},
   483  	} {
   484  		t.Run(test.name, func(t *testing.T) {
   485  			pm := NewPortManager()
   486  			if err := pm.SetPortRange(firstEphemeral, firstEphemeral+numEphemeralPorts); err != nil {
   487  				t.Fatalf("failed to set ephemeral port range: %s", err)
   488  			}
   489  			portOffset := uint32(rand.Int31n(int32(numEphemeralPorts)))
   490  			port, err := pm.PickEphemeralPortStable(portOffset, test.f)
   491  			if diff := cmp.Diff(test.wantErr, err); diff != "" {
   492  				t.Fatalf("unexpected error from PickEphemeralPort(..), (-want, +got):\n%s", diff)
   493  			}
   494  			if port != test.wantPort {
   495  				t.Errorf("got PickEphemeralPort(..) = (%d, nil); want (%d, nil)", port, test.wantPort)
   496  			}
   497  		})
   498  	}
   499  }
   500  
   501  // TestOverflow addresses b/183593432, wherein an overflowing uint16 causes a
   502  // port allocation failure.
   503  func TestOverflow(t *testing.T) {
   504  	// Use a small range and start at offsets that will cause an overflow.
   505  	count := uint16(50)
   506  	for offset := uint32(math.MaxUint16 - count); offset < math.MaxUint16; offset++ {
   507  		reservedPorts := make(map[uint16]struct{})
   508  		// Ensure we can reserve everything in the allowed range.
   509  		for i := uint16(0); i < count; i++ {
   510  			port, err := pickEphemeralPort(offset, firstEphemeral, count, func(port uint16) (bool, tcpip.Error) {
   511  				if _, ok := reservedPorts[port]; !ok {
   512  					reservedPorts[port] = struct{}{}
   513  					return true, nil
   514  				}
   515  				return false, nil
   516  			})
   517  			if err != nil {
   518  				t.Fatalf("port picking failed at iteration %d, for offset %d, len(reserved): %+v", i, offset, len(reservedPorts))
   519  			}
   520  			if port < firstEphemeral || port > firstEphemeral+count {
   521  				t.Fatalf("reserved port %d, which is not in range [%d, %d]", port, firstEphemeral, firstEphemeral+count-1)
   522  			}
   523  		}
   524  	}
   525  }