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>

control/controlclient/map.go

@@ -736,6 +736,10 @@ func peerChangeDiff(was tailcfg.NodeView, n *tailcfg.Node) (_ *tailcfg.PeerChang
 			if was.IsWireGuardOnly() != n.IsWireGuardOnly {
 				return nil, false
 			}
+		case "IsJailed":
+			if was.IsJailed() != n.IsJailed {
+				return nil, false
+			}
 		case "Expired":
 			if was.Expired() != n.Expired {
 				return nil, false

net/tstun/wrap.go

@@ -739,18 +739,16 @@ func peerConfigTableFromWGConfig(wcfg *wgcfg.Config) *peerConfigTable {
 			}
 		}
 
-		if !addrToUse4.IsValid() && !addrToUse6.IsValid() {
+		if !addrToUse4.IsValid() && !addrToUse6.IsValid() && !p.IsJailed {
 			// NAT not required for this peer.
 			continue
 		}
 
-		const peerIsJailed = false // TODO: implement jailed peers
-
 		// Use the same peer configuration for each address of the peer.
 		pc := &peerConfig{
 			dstMasqAddr4: addrToUse4,
 			dstMasqAddr6: addrToUse6,
-			jailed:       peerIsJailed,
+			jailed:       p.IsJailed,
 		}
 
 		// Insert an entry into our routing table for each allowed IP.

tailcfg/tailcfg.go

@@ -134,7 +134,8 @@ type CapabilityVersion int
 //   - 91: 2024-04-24: Client understands PeerCapabilityTaildriveSharer.
 //   - 92: 2024-05-06: Client understands NodeAttrUserDialUseRoutes.
 //   - 93: 2024-05-06: added support for stateful firewalling.
-const CurrentCapabilityVersion CapabilityVersion = 93
+//   - 94: 2024-05-06: Client understands Node.IsJailed.
+const CurrentCapabilityVersion CapabilityVersion = 94
 
 type StableID string
 
@@ -416,6 +417,11 @@ type Node struct {
 	// order to be reachable.
 	IsWireGuardOnly bool `json:",omitempty"`
 
+	// IsJailed indicates that this node is jailed and should not be allowed
+	// initiate connections, however outbound connections to it should still be
+	// allowed.
+	IsJailed bool `json:",omitempty"`
+
 	// ExitNodeDNSResolvers is the list of DNS servers that should be used when this
 	// node is marked IsWireGuardOnly and being used as an exit node.
 	ExitNodeDNSResolvers []*dnstype.Resolver `json:",omitempty"`
@@ -2046,7 +2052,8 @@ func (n *Node) Equal(n2 *Node) bool {
 		n.Expired == n2.Expired &&
 		eqPtr(n.SelfNodeV4MasqAddrForThisPeer, n2.SelfNodeV4MasqAddrForThisPeer) &&
 		eqPtr(n.SelfNodeV6MasqAddrForThisPeer, n2.SelfNodeV6MasqAddrForThisPeer) &&
-		n.IsWireGuardOnly == n2.IsWireGuardOnly
+		n.IsWireGuardOnly == n2.IsWireGuardOnly &&
+		n.IsJailed == n2.IsJailed
 }
 
 func eqPtr[T comparable](a, b *T) bool {

tailcfg/tailcfg_clone.go

@@ -118,6 +118,7 @@ var _NodeCloneNeedsRegeneration = Node(struct {
 	SelfNodeV4MasqAddrForThisPeer *netip.Addr
 	SelfNodeV6MasqAddrForThisPeer *netip.Addr
 	IsWireGuardOnly               bool
+	IsJailed                      bool
 	ExitNodeDNSResolvers          []*dnstype.Resolver
 }{})
 

tailcfg/tailcfg_test.go

@@ -363,7 +363,7 @@ func TestNodeEqual(t *testing.T) {
 		"UnsignedPeerAPIOnly",
 		"ComputedName", "computedHostIfDifferent", "ComputedNameWithHost",
 		"DataPlaneAuditLogID", "Expired", "SelfNodeV4MasqAddrForThisPeer",
-		"SelfNodeV6MasqAddrForThisPeer", "IsWireGuardOnly", "ExitNodeDNSResolvers",
+		"SelfNodeV6MasqAddrForThisPeer", "IsWireGuardOnly", "IsJailed", "ExitNodeDNSResolvers",
 	}
 	if have := fieldsOf(reflect.TypeFor[Node]()); !reflect.DeepEqual(have, nodeHandles) {
 		t.Errorf("Node.Equal check might be out of sync\nfields: %q\nhandled: %q\n",
@@ -607,6 +607,16 @@ func TestNodeEqual(t *testing.T) {
 			},
 			false,
 		},
+		{
+			&Node{IsJailed: true},
+			&Node{IsJailed: true},
+			true,
+		},
+		{
+			&Node{IsJailed: false},
+			&Node{IsJailed: true},
+			false,
+		},
 	}
 	for i, tt := range tests {
 		got := tt.a.Equal(tt.b)

tailcfg/tailcfg_view.go

@@ -195,6 +195,7 @@ func (v NodeView) SelfNodeV6MasqAddrForThisPeer() *netip.Addr {
 }
 
 func (v NodeView) IsWireGuardOnly() bool { return v.ж.IsWireGuardOnly }
+func (v NodeView) IsJailed() bool        { return v.ж.IsJailed }
 func (v NodeView) ExitNodeDNSResolvers() views.SliceView[*dnstype.Resolver, dnstype.ResolverView] {
 	return views.SliceOfViews[*dnstype.Resolver, dnstype.ResolverView](v.ж.ExitNodeDNSResolvers)
 }
@@ -235,6 +236,7 @@ var _NodeViewNeedsRegeneration = Node(struct {
 	SelfNodeV4MasqAddrForThisPeer *netip.Addr
 	SelfNodeV6MasqAddrForThisPeer *netip.Addr
 	IsWireGuardOnly               bool
+	IsJailed                      bool
 	ExitNodeDNSResolvers          []*dnstype.Resolver
 }{})
 

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]

wgengine/wgcfg/config.go

@@ -41,6 +41,7 @@ type Peer struct {
 	AllowedIPs          []netip.Prefix
 	V4MasqAddr          *netip.Addr // if non-nil, masquerade IPv4 traffic to this peer using this address
 	V6MasqAddr          *netip.Addr // if non-nil, masquerade IPv6 traffic to this peer using this address
+	IsJailed            bool        // if true, this peer is jailed and cannot initiate connections
 	PersistentKeepalive uint16      // in seconds between keep-alives; 0 to disable
 	// wireguard-go's endpoint for this peer. It should always equal Peer.PublicKey.
 	// We represent it explicitly so that we can detect if they diverge and recover.

wgengine/wgcfg/nmcfg/nmcfg.go

@@ -110,6 +110,7 @@ func WGCfg(nm *netmap.NetworkMap, logf logger.Logf, flags netmap.WGConfigFlags,
 		didExitNodeWarn := false
 		cpeer.V4MasqAddr = peer.SelfNodeV4MasqAddrForThisPeer()
 		cpeer.V6MasqAddr = peer.SelfNodeV6MasqAddrForThisPeer()
+		cpeer.IsJailed = peer.IsJailed()
 		for i := range peer.AllowedIPs().Len() {
 			allowedIP := peer.AllowedIPs().At(i)
 			if allowedIP.Bits() == 0 && peer.StableID() != exitNode {

wgengine/wgcfg/wgcfg_clone.go

@@ -74,6 +74,7 @@ var _PeerCloneNeedsRegeneration = Peer(struct {
 	AllowedIPs          []netip.Prefix
 	V4MasqAddr          *netip.Addr
 	V6MasqAddr          *netip.Addr
+	IsJailed            bool
 	PersistentKeepalive uint16
 	WGEndpoint          key.NodePublic
 }{})