tailcat: add Server.PeerKey to identify an accepted connection's peer

tailcat.go

@@ -1445,6 +1445,32 @@ func (b *locoBackend) peerConfig(k key.NodePublic) (_ wgcfg.PeerConfig, ok bool)
 	return withPSK(n.AllowedIPs), true
 }
 
+// PeerKey returns the node public key of the peer at remote, the
+// remote address of a connection or flow accepted by this server: the
+// net.Conn passed to [Server.OnTCP] or [Server.OnTCPForward], the
+// [ConnPacketConn] passed to [Server.OnUDP] or [Server.OnUDPForward],
+// or a connection from a [Server.Listen] listener.
+//
+// The tunnel has already authenticated the peer by this key, so a
+// caller can tell which peer it is serving, and can match it against
+// [Server.AllowedClients]. It reports ok=false if remote is not a
+// known peer's address.
+//
+// [Server.PeerEnv] reports the same key to served subprocesses as
+// TAILCAT_PEER_KEY.
+func (s *Server) PeerKey(remote net.Addr) (_ key.NodePublic, ok bool) {
+	var ap netip.AddrPort
+	switch a := remote.(type) {
+	case *net.TCPAddr:
+		ap = a.AddrPort()
+	case *net.UDPAddr:
+		ap = a.AddrPort()
+	default:
+		return key.NodePublic{}, false
+	}
+	return s.lb.peerByIP(ap.Addr().Unmap())
+}
+
 // peerByIP returns the public key of the peer that outbound packets
 // addressed to dst should be sent to (see
 // [wgengine.Engine.SetPeerByIPPacketFunc]).

tailcat_exec.go

@@ -49,10 +49,8 @@ func (s *Server) PeerEnv(local, remote net.Addr) []string {
 		"TAILCAT_REMOTE_ADDR=" + remote.String(),
 		"TAILCAT_LOCAL_ADDR=" + local.String(),
 	}
-	if ta, ok := remote.(*net.TCPAddr); ok {
-		if k, ok := s.lb.peerByIP(ta.AddrPort().Addr().Unmap()); ok {
-			env = append(env, "TAILCAT_PEER_KEY="+k.String())
-		}
+	if k, ok := s.PeerKey(remote); ok {
+		env = append(env, "TAILCAT_PEER_KEY="+k.String())
 	}
 	return env
 }

tailcat_ssh.go

@@ -24,7 +24,6 @@ import (
 
 	ssh "github.com/tailscale/gliderssh"
 	gossh "golang.org/x/crypto/ssh"
-	"tailscale.com/types/key"
 )
 
 const sshInteractiveMOTD = "🐈 Connected via tailcat SSH.\r\n"
@@ -127,7 +126,7 @@ func (s *Server) sessionHandler(sess ssh.Session) {
 	// authenticated the peer by this key (and --allow, if set, gated on
 	// it), so a shell wrapper can tell which allowed peer it's talking
 	// to. The value matches --allow's format ("nodekey:...").
-	if k, ok := s.peerKeyForSession(sess); ok {
+	if k, ok := s.PeerKey(sess.RemoteAddr()); ok {
 		cmd.Env = append(cmd.Env, "TAILCAT_PEER_KEY="+k.String())
 	}
 
@@ -168,19 +167,6 @@ func (s *Server) execSessionHandler(argv []string) ssh.Handler {
 	}
 }
 
-// peerKeyForSession returns the node public key of the peer on the
-// other end of sess. The session's remote address is the peer's
-// tailcat IP (derived from its key by tcAddrForKey); peerByIP reverses
-// that back to the key the tunnel authenticated. It returns ok=false
-// if the address can't be mapped to a known peer.
-func (s *Server) peerKeyForSession(sess ssh.Session) (key.NodePublic, bool) {
-	ta, ok := sess.RemoteAddr().(*net.TCPAddr)
-	if !ok {
-		return key.NodePublic{}, false
-	}
-	return s.lb.peerByIP(ta.AddrPort().Addr().Unmap())
-}
-
 // runWithPipes runs cmd with stdin/stdout/stderr pipes (no PTY).
 func runWithPipes(sess ssh.Session, cmd *exec.Cmd) {
 	stdinPipe, err := cmd.StdinPipe()

tailcat_test.go

@@ -1413,3 +1413,80 @@ func TestParseAddrRawKeepsNulls(t *testing.T) {
 		t.Errorf("Region = %v; want a single nil element", w.Region)
 	}
 }
+
+func TestPeerKey(t *testing.T) {
+	dm := integration.RunDERPAndSTUN(t, mkLogger(t, "derpstun"), "127.0.0.1")
+	reg := dm.Regions[1]
+	if reg == nil {
+		t.Fatal("no region 1 in derpmap")
+	}
+
+	type peerKey struct {
+		key key.NodePublic
+		ok  bool
+	}
+	tcpKey := make(chan peerKey, 1)
+	udpKey := make(chan peerKey, 1)
+
+	s := &Server{Logf: mkLogger(t, "server"), Region: reg}
+	t.Cleanup(func() { s.Close() })
+	s.OnTCP = func(port uint16) func(net.Conn) {
+		return func(c net.Conn) {
+			defer c.Close()
+			k, ok := s.PeerKey(c.RemoteAddr())
+			tcpKey <- peerKey{k, ok}
+		}
+	}
+	s.OnUDP = func(port uint16) func(ConnPacketConn) {
+		return func(c ConnPacketConn) {
+			defer c.Close()
+			k, ok := s.PeerKey(c.RemoteAddr())
+			udpKey <- peerKey{k, ok}
+		}
+	}
+	if err := s.Start(); err != nil {
+		t.Fatalf("server Start: %v", err)
+	}
+
+	c := &Client{Server: s.TailcatAddr(), Logf: mkLogger(t, "client")}
+	t.Cleanup(func() { c.Close() })
+	PingForTest(t, s, c)
+	want := c.PublicKey()
+
+	conn, err := c.DialTCPPort(t.Context(), 80)
+	if err != nil {
+		t.Fatalf("DialTCPPort: %v", err)
+	}
+	io.Copy(io.Discard, conn)
+	conn.Close()
+
+	pc, err := c.DialUDPPort(t.Context(), 53)
+	if err != nil {
+		t.Fatalf("DialUDPPort: %v", err)
+	}
+	defer pc.Close()
+	if _, err := pc.Write([]byte("hello")); err != nil {
+		t.Fatalf("UDP Write: %v", err)
+	}
+
+	for _, tt := range []struct {
+		proto string
+		ch    chan peerKey
+	}{
+		{"TCP", tcpKey},
+		{"UDP", udpKey},
+	} {
+		select {
+		case got := <-tt.ch:
+			if !got.ok {
+				t.Errorf("PeerKey on %s: ok=false; want the client's key", tt.proto)
+				continue
+			}
+			if got.key != want {
+				t.Errorf("PeerKey on %s = %v; want %v", tt.proto, got.key, want)
+			}
+		case <-time.After(30 * time.Second):
+			t.Errorf("timeout waiting for the %s handler", tt.proto)
+		}
+	}
+}