1 files changed, 97 insertions, 0 deletions
diff --git a/pkg/tcpip/stack/stack_test.go b/pkg/tcpip/stack/stack_test.go
index 57de5b93a..c46e91241 100644
--- a/pkg/tcpip/stack/stack_test.go
+++ b/pkg/tcpip/stack/stack_test.go
@@ -19,6 +19,7 @@ package stack_test
 
 import (
 	"math"
+	"strings"
 	"testing"
 
 	"gvisor.googlesource.com/gvisor/pkg/tcpip"
@@ -763,6 +764,102 @@ func TestNetworkOptions(t *testing.T) {
 	}
 }
 
+func TestSubnetAddRemove(t *testing.T) {
+	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
+	id, _ := channel.New(10, defaultMTU, "")
+	if err := s.CreateNIC(1, id); err != nil {
+		t.Fatalf("CreateNIC failed: %v", err)
+	}
+
+	addr := tcpip.Address("\x01\x01\x01\x01")
+	mask := tcpip.AddressMask(strings.Repeat("\xff", len(addr)))
+	subnet, err := tcpip.NewSubnet(addr, mask)
+	if err != nil {
+		t.Fatalf("NewSubnet failed: %v", err)
+	}
+
+	if contained, err := s.ContainsSubnet(1, subnet); err != nil {
+		t.Fatalf("ContainsSubnet failed: %v", err)
+	} else if contained {
+		t.Fatal("got s.ContainsSubnet(...) = true, want = false")
+	}
+
+	if err := s.AddSubnet(1, fakeNetNumber, subnet); err != nil {
+		t.Fatalf("AddSubnet failed: %v", err)
+	}
+
+	if contained, err := s.ContainsSubnet(1, subnet); err != nil {
+		t.Fatalf("ContainsSubnet failed: %v", err)
+	} else if !contained {
+		t.Fatal("got s.ContainsSubnet(...) = false, want = true")
+	}
+
+	if err := s.RemoveSubnet(1, subnet); err != nil {
+		t.Fatalf("RemoveSubnet failed: %v", err)
+	}
+
+	if contained, err := s.ContainsSubnet(1, subnet); err != nil {
+		t.Fatalf("ContainsSubnet failed: %v", err)
+	} else if contained {
+		t.Fatal("got s.ContainsSubnet(...) = true, want = false")
+	}
+}
+
+func TestGetMainNICAddress(t *testing.T) {
+	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
+	id, _ := channel.New(10, defaultMTU, "")
+	if err := s.CreateNIC(1, id); err != nil {
+		t.Fatalf("CreateNIC failed: %v", err)
+	}
+
+	for _, tc := range []struct {
+		name    string
+		address tcpip.Address
+	}{
+		{"IPv4", "\x01\x01\x01\x01"},
+		{"IPv6", "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			address := tc.address
+			mask := tcpip.AddressMask(strings.Repeat("\xff", len(address)))
+			subnet, err := tcpip.NewSubnet(address, mask)
+			if err != nil {
+				t.Fatalf("NewSubnet failed: %v", err)
+			}
+
+			if err := s.AddAddress(1, fakeNetNumber, address); err != nil {
+				t.Fatalf("AddAddress failed: %v", err)
+			}
+
+			if err := s.AddSubnet(1, fakeNetNumber, subnet); err != nil {
+				t.Fatalf("AddSubnet failed: %v", err)
+			}
+
+			// Check that we get the right initial address and subnet.
+			if gotAddress, gotSubnet, err := s.GetMainNICAddress(1, fakeNetNumber); err != nil {
+				t.Fatalf("GetMainNICAddress failed: %v", err)
+			} else if gotAddress != address {
+				t.Fatalf("got GetMainNICAddress = (%v, ...), want = (%v, ...)", gotAddress, address)
+			} else if gotSubnet != subnet {
+				t.Fatalf("got GetMainNICAddress = (..., %v), want = (..., %v)", gotSubnet, subnet)
+			}
+
+			if err := s.RemoveSubnet(1, subnet); err != nil {
+				t.Fatalf("RemoveSubnet failed: %v", err)
+			}
+
+			if err := s.RemoveAddress(1, address); err != nil {
+				t.Fatalf("RemoveAddress failed: %v", err)
+			}
+
+			// Check that we get an error after removal.
+			if _, _, err := s.GetMainNICAddress(1, fakeNetNumber); err != tcpip.ErrNoLinkAddress {
+				t.Fatalf("got s.GetMainNICAddress(...) = %v, want = %v", err, tcpip.ErrNoLinkAddress)
+			}
+		})
+	}
+}
+
 func init() {
 	stack.RegisterNetworkProtocolFactory("fakeNet", func() stack.NetworkProtocol {
 		return &fakeNetworkProtocol{}