From f74ee80abe8819f3d3bfd9138056a46820f4fc54 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 20 Apr 2022 13:53:01 -0700
Subject: [PATCH] ssh/tailssh: support expansions in public key fetch URL too

Updates #3802

Change-Id: I5aa98bdab14fd1c1c00ba63b93f8d7e670f72437
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 ssh/tailssh/tailssh.go      | 18 +++++++++++++++++-
 ssh/tailssh/tailssh_test.go | 20 ++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/ssh/tailssh/tailssh.go b/ssh/tailssh/tailssh.go
index 247217a21..94a21c805 100644
--- a/ssh/tailssh/tailssh.go
+++ b/ssh/tailssh/tailssh.go
@@ -541,6 +541,22 @@ func (srv *server) expandDelegateURL(ci *sshConnInfo, lu *user.User, actionURL s
 	).Replace(actionURL)
 }
 
+func (ci *sshConnInfo) expandPublicKeyURL(pubKeyURL string) string {
+	if !strings.Contains(pubKeyURL, "$") {
+		return pubKeyURL
+	}
+	var localPart string
+	var loginName string
+	if ci.uprof != nil {
+		loginName = ci.uprof.LoginName
+		localPart, _, _ = strings.Cut(loginName, "@")
+	}
+	return strings.NewReplacer(
+		"$LOGINNAME_EMAIL", loginName,
+		"$LOGINNAME_LOCALPART", localPart,
+	).Replace(pubKeyURL)
+}
+
 // sshSession is an accepted Tailscale SSH session.
 type sshSession struct {
 	ssh.Session
@@ -1011,7 +1027,7 @@ func principalMatchesPubKey(p *tailcfg.SSHPrincipal, ci *sshConnInfo, clientPubK
 			return false, fmt.Errorf("no public key fetcher")
 		}
 		var err error
-		knownKeys, err = ci.fetchPublicKeysURL(knownKeys[0])
+		knownKeys, err = ci.fetchPublicKeysURL(ci.expandPublicKeyURL(knownKeys[0]))
 		if err != nil {
 			return false, err
 		}
diff --git a/ssh/tailssh/tailssh_test.go b/ssh/tailssh/tailssh_test.go
index 2d8ad255f..93a82700f 100644
--- a/ssh/tailssh/tailssh_test.go
+++ b/ssh/tailssh/tailssh_test.go
@@ -406,3 +406,23 @@ func TestPublicKeyFetching(t *testing.T) {
 	}
 
 }
+
+func TestExpandPublicKeyURL(t *testing.T) {
+	ci := &sshConnInfo{
+		uprof: &tailcfg.UserProfile{
+			LoginName: "bar@baz.tld",
+		},
+	}
+	if got, want := ci.expandPublicKeyURL("foo"), "foo"; got != want {
+		t.Errorf("basic: got %q; want %q", got, want)
+	}
+	if got, want := ci.expandPublicKeyURL("https://example.com/$LOGINNAME_LOCALPART.keys"), "https://example.com/bar.keys"; got != want {
+		t.Errorf("localpart: got %q; want %q", got, want)
+	}
+	if got, want := ci.expandPublicKeyURL("https://example.com/keys?email=$LOGINNAME_EMAIL"), "https://example.com/keys?email=bar@baz.tld"; got != want {
+		t.Errorf("email: got %q; want %q", got, want)
+	}
+	if got, want := new(sshConnInfo).expandPublicKeyURL("https://example.com/keys?email=$LOGINNAME_EMAIL"), "https://example.com/keys?email="; got != want {
+		t.Errorf("on empty: got %q; want %q", got, want)
+	}
+}