tailcfg,all: add/plumb Node.IsJailed

This adds a new bool that can be sent down from control
to do jailing on the client side. Previously this would
only be done from control by modifying the packet filter
we sent down to clients. This would result in a lot of
additional work/CPU on control, we could instead just
do this on the client. This has always been a TODO which
we keep putting off, might as well do it now.

Updates tailscale/corp#19623

Signed-off-by: Maisem Ali <maisem@tailscale.com>

tstest/integration/integration_test.go

@@ -14,6 +14,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/netip"
@@ -29,6 +30,7 @@ import (
 	"time"
 
 	"go4.org/mem"
+	"tailscale.com/client/tailscale"
 	"tailscale.com/clientupdate"
 	"tailscale.com/cmd/testwrapper/flakytest"
 	"tailscale.com/ipn"
@@ -722,6 +724,121 @@ func TestOneNodeUpWindowsStyle(t *testing.T) {
 	d1.MustCleanShutdown(t)
 }
 
+// TestClientSideJailing tests that when one node is jailed for another, the
+// jailed node cannot initiate connections to the other node however the other
+// node can initiate connections to the jailed node.
+func TestClientSideJailing(t *testing.T) {
+	tstest.Shard(t)
+	tstest.Parallel(t)
+	env := newTestEnv(t)
+	registerNode := func() (*testNode, key.NodePublic) {
+		n := newTestNode(t, env)
+		n.StartDaemon()
+		n.AwaitListening()
+		n.MustUp()
+		n.AwaitRunning()
+		k := n.MustStatus().Self.PublicKey
+		return n, k
+	}
+	n1, k1 := registerNode()
+	n2, k2 := registerNode()
+
+	ln, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	port := uint16(ln.Addr().(*net.TCPAddr).Port)
+
+	lc1 := &tailscale.LocalClient{
+		Socket:        n1.sockFile,
+		UseSocketOnly: true,
+	}
+	lc2 := &tailscale.LocalClient{
+		Socket:        n2.sockFile,
+		UseSocketOnly: true,
+	}
+
+	ip1 := n1.AwaitIP4()
+	ip2 := n2.AwaitIP4()
+
+	tests := []struct {
+		name          string
+		n1JailedForN2 bool
+		n2JailedForN1 bool
+	}{
+		{
+			name:          "not_jailed",
+			n1JailedForN2: false,
+			n2JailedForN1: false,
+		},
+		{
+			name:          "uni_jailed",
+			n1JailedForN2: true,
+			n2JailedForN1: false,
+		},
+		{
+			name:          "bi_jailed", // useless config?
+			n1JailedForN2: true,
+			n2JailedForN1: true,
+		},
+	}
+
+	testDial := func(t *testing.T, lc *tailscale.LocalClient, ip netip.Addr, port uint16, shouldFail bool) {
+		t.Helper()
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		c, err := lc.DialTCP(ctx, ip.String(), port)
+		failed := err != nil
+		if failed != shouldFail {
+			t.Errorf("failed = %v; want %v", failed, shouldFail)
+		}
+		if c != nil {
+			c.Close()
+		}
+	}
+
+	b1, err := lc1.WatchIPNBus(context.Background(), 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	b2, err := lc2.WatchIPNBus(context.Background(), 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	waitPeerIsJailed := func(t *testing.T, b *tailscale.IPNBusWatcher, jailed bool) {
+		t.Helper()
+		for {
+			n, err := b.Next()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if n.NetMap == nil {
+				continue
+			}
+			if len(n.NetMap.Peers) == 0 {
+				continue
+			}
+			if j := n.NetMap.Peers[0].IsJailed(); j == jailed {
+				break
+			}
+		}
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			env.Control.SetJailed(k1, k2, tc.n2JailedForN1)
+			env.Control.SetJailed(k2, k1, tc.n1JailedForN2)
+
+			// Wait for the jailed status to propagate.
+			waitPeerIsJailed(t, b1, tc.n2JailedForN1)
+			waitPeerIsJailed(t, b2, tc.n1JailedForN2)
+
+			testDial(t, lc1, ip2, port, tc.n1JailedForN2)
+			testDial(t, lc2, ip1, port, tc.n2JailedForN1)
+		})
+	}
+}
+
 // TestNATPing creates two nodes, n1 and n2, sets up masquerades for both and
 // tries to do bi-directional pings between them.
 func TestNATPing(t *testing.T) {

tstest/integration/testcontrol/testcontrol.go

@@ -71,6 +71,9 @@ type Server struct {
 	// by the specified node.
 	nodeSubnetRoutes map[key.NodePublic][]netip.Prefix
 
+	// peerIsJailed is the set of peers that are jailed for a node.
+	peerIsJailed map[key.NodePublic]map[key.NodePublic]bool // node => peer => isJailed
+
 	// masquerades is the set of masquerades that should be applied to
 	// MapResponses sent to clients. It is keyed by the requesting nodes
 	// public key, and then the peer node's public key. The value is the
@@ -379,6 +382,20 @@ type MasqueradePair struct {
 	NodeMasqueradesAs netip.Addr
 }
 
+// SetJailed sets b to be jailed when it is a peer of a.
+func (s *Server) SetJailed(a, b key.NodePublic, jailed bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.peerIsJailed == nil {
+		s.peerIsJailed = map[key.NodePublic]map[key.NodePublic]bool{}
+	}
+	if s.peerIsJailed[a] == nil {
+		s.peerIsJailed[a] = map[key.NodePublic]bool{}
+	}
+	s.peerIsJailed[a][b] = jailed
+	s.updateLocked("SetJailed", s.nodeIDsLocked(0))
+}
+
 // SetMasqueradeAddresses sets the masquerade addresses for the server.
 // See MasqueradePair for more details.
 func (s *Server) SetMasqueradeAddresses(pairs []MasqueradePair) {
@@ -945,6 +962,7 @@ func (s *Server) MapResponse(req *tailcfg.MapRequest) (res *tailcfg.MapResponse,
 
 	s.mu.Lock()
 	nodeMasqs := s.masquerades[node.Key]
+	jailed := maps.Clone(s.peerIsJailed[node.Key])
 	s.mu.Unlock()
 	for _, p := range s.AllNodes() {
 		if p.StableID == node.StableID {
@@ -957,6 +975,7 @@ func (s *Server) MapResponse(req *tailcfg.MapRequest) (res *tailcfg.MapResponse,
 				p.SelfNodeV4MasqAddrForThisPeer = ptr.To(masqIP)
 			}
 		}
+		p.IsJailed = jailed[p.Key]
 
 		s.mu.Lock()
 		peerAddress := s.masquerades[p.Key][node.Key]