Remove Tailscale policies after client disconnects

Author: cedws

pkg/sshgate/policy.go

@@ -5,8 +5,8 @@ import "sync"
 type principalType string
 
 const (
-	principalTypeFingerprint = "fingerprint"
-	principalTypeTailscale   = "tailscale"
+	principalTypeFingerprint   = "fingerprint"
+	principalTypeTailscaleNode = "tailscale_node"
 )
 
 func newPolicyEngine() *policyEngine {
@@ -29,7 +29,7 @@ func (p *policyEngine) AddPolicy(principalType principalType, principal string,
 	switch principalType {
 	case principalTypeFingerprint:
 		p.fingerprints[principal] = append(p.fingerprints[principal], rules...)
-	case principalTypeTailscale:
+	case principalTypeTailscaleNode:
 		p.tailscaleNodes[principal] = append(p.tailscaleNodes[principal], rules...)
 	}
 }
@@ -41,7 +41,7 @@ func (p *policyEngine) RemovePolicy(principalType principalType, principal strin
 	switch principalType {
 	case principalTypeFingerprint:
 		delete(p.fingerprints, principal)
-	case principalTypeTailscale:
+	case principalTypeTailscaleNode:
 		delete(p.tailscaleNodes, principal)
 	}
 }
@@ -56,7 +56,7 @@ func (p *policyEngine) Principal(principalType principalType, principal string)
 	switch principalType {
 	case principalTypeFingerprint:
 		rules, ok = p.fingerprints[principal]
-	case principalTypeTailscale:
+	case principalTypeTailscaleNode:
 		rules, ok = p.tailscaleNodes[principal]
 	}
 

pkg/sshgate/sshgate.go

@@ -30,7 +30,7 @@ const tsnetCapability = "github.com/cedws/sshgate"
 
 const (
 	fingerprintKey     = "fingerprint"
-	tailscaleNodeIDKey = "tailscale_user"
+	tailscaleNodeIDKey = "tailscale_node_id"
 )
 
 type rejectionError struct {
@@ -410,7 +410,7 @@ func (s *Server) addCapMapPolicy(logger *slog.Logger, id string, whois *apitype.
 		rules = append(rules, capData)
 	}
 
-	s.policyEngine.AddPolicy(principalTypeTailscale, id, rules)
+	s.policyEngine.AddPolicy(principalTypeTailscaleNode, id, rules)
 
 	return nil
 }
@@ -435,6 +435,14 @@ func (s *Server) handleConnection(ctx context.Context, c net.Conn, config *ssh.S
 	s.conns.Add(1)
 	defer s.conns.Add(-1)
 
+	defer func() {
+		// Clean up policies from Tailscale capabilities associated with this client
+		nodeID, ok := sshConn.Permissions.Extensions[tailscaleNodeIDKey]
+		if ok {
+			s.policyEngine.RemovePolicy(principalTypeTailscaleNode, nodeID)
+		}
+	}()
+
 	go func() {
 		for req := range requests {
 			req.Reply(false, nil)