tsnet: allow for automatic ID token generation

Allow for optionally specifiying an audience for tsnet. This is passed
to the underlying identity federation logic to allow for tsnet auth to
use automatic ID token generation for authentication.

Updates https://github.com/tailscale/corp/issues/33316

Signed-off-by: Mario Minardi <mario@tailscale.com>

tsnet/tsnet.go

@@ -139,6 +139,14 @@ type Server struct {
 	// field is not used.
 	IDToken string
 
+	// Audience, if non-empty, is the audience to use when requesting
+	// an ID token from a well-known identity provider to exchange
+	// with the control server for workload identity federation. It
+	// will be preferred over the TS_AUDIENCE environment variable. If
+	// the node is already created (from state previously stored in Store),
+	// then this field is not used.
+	Audience string
+
 	// ControlURL optionally specifies the coordination server URL.
 	// If empty, the Tailscale default is used.
 	ControlURL string
@@ -567,6 +575,13 @@ func (s *Server) getIDToken() string {
 	return os.Getenv("TS_ID_TOKEN")
 }
 
+func (s *Server) getAudience() string {
+	if v := s.Audience; v != "" {
+		return v
+	}
+	return os.Getenv("TS_AUDIENCE")
+}
+
 func (s *Server) start() (reterr error) {
 	var closePool closeOnErrorPool
 	defer closePool.closeAllIfError(&reterr)
@@ -805,13 +820,22 @@ func (s *Server) resolveAuthKey() (string, error) {
 	if wifOk && authKey == "" {
 		clientID := s.getClientID()
 		idToken := s.getIDToken()
-		if clientID != "" && idToken == "" {
-			return "", fmt.Errorf("client ID for workload identity federation found, but ID token is empty")
+		audience := s.getAudience()
+		if clientID != "" && idToken == "" && audience == "" {
+			return "", fmt.Errorf("client ID for workload identity federation found, but ID token and audience are empty")
+		}
+		if idToken != "" && audience != "" {
+			return "", fmt.Errorf("only one of ID token and audience should be for workload identity federation")
 		}
-		if clientID == "" && idToken != "" {
-			return "", fmt.Errorf("ID token for workload identity federation found, but client ID is empty")
+		if clientID == "" {
+			if idToken != "" {
+				return "", fmt.Errorf("ID token for workload identity federation found, but client ID is empty")
+			}
+			if audience != "" {
+				return "", fmt.Errorf("audience for workload identity federation found, but client ID is empty")
+			}
 		}
-		authKey, err = resolveViaWIF(s.shutdownCtx, s.ControlURL, clientID, idToken, "", s.AdvertiseTags)
+		authKey, err = resolveViaWIF(s.shutdownCtx, s.ControlURL, clientID, idToken, audience, s.AdvertiseTags)
 		if err != nil {
 			return "", err
 		}

tsnet/tsnet_test.go

@@ -1503,6 +1503,7 @@ func TestResolveAuthKey(t *testing.T) {
 		clientSecret    string
 		clientID        string
 		idToken         string
+		audience        string
 		oauthAvailable  bool
 		wifAvailable    bool
 		resolveViaOAuth func(ctx context.Context, clientSecret string, tags []string) (string, error)
@@ -1550,6 +1551,23 @@ func TestResolveAuthKey(t *testing.T) {
 			wantAuthKey:     "tskey-auth-via-wif",
 			wantErrContains: "",
 		},
+		{
+			name:         "successful resolution via federated audience",
+			clientID:     "client-id-123",
+			audience:     "api.tailscale.com",
+			wifAvailable: true,
+			resolveViaWIF: func(ctx context.Context, baseURL, clientID, idToken, audience string, tags []string) (string, error) {
+				if clientID != "client-id-123" {
+					return "", fmt.Errorf("unexpected client ID: %s", clientID)
+				}
+				if audience != "api.tailscale.com" {
+					return "", fmt.Errorf("unexpected ID token: %s", idToken)
+				}
+				return "tskey-auth-via-wif", nil
+			},
+			wantAuthKey:     "tskey-auth-via-wif",
+			wantErrContains: "",
+		},
 		{
 			name:         "failing resolution via federated ID token",
 			clientID:     "client-id-123",
@@ -1561,7 +1579,7 @@ func TestResolveAuthKey(t *testing.T) {
 			wantErrContains: "resolution failed",
 		},
 		{
-			name:         "empty client ID",
+			name:         "empty client ID with ID token",
 			clientID:     "",
 			idToken:      "id-token-456",
 			wifAvailable: true,
@@ -1570,6 +1588,16 @@ func TestResolveAuthKey(t *testing.T) {
 			},
 			wantErrContains: "empty",
 		},
+		{
+			name:         "empty client ID with audience",
+			clientID:     "",
+			audience:     "api.tailscale.com",
+			wifAvailable: true,
+			resolveViaWIF: func(ctx context.Context, baseURL, clientID, idToken, audience string, tags []string) (string, error) {
+				return "", fmt.Errorf("should not be called")
+			},
+			wantErrContains: "empty",
+		},
 		{
 			name:         "empty ID token",
 			clientID:     "client-id-123",
@@ -1580,6 +1608,17 @@ func TestResolveAuthKey(t *testing.T) {
 			},
 			wantErrContains: "empty",
 		},
+		{
+			name:         "audience with ID token",
+			clientID:     "client-id-123",
+			idToken:      "id-token-456",
+			audience:     "api.tailscale.com",
+			wifAvailable: true,
+			resolveViaWIF: func(ctx context.Context, baseURL, clientID, idToken, audience string, tags []string) (string, error) {
+				return "", fmt.Errorf("should not be called")
+			},
+			wantErrContains: "only one of ID token and audience",
+		},
 		{
 			name:           "workload identity resolution skipped if resolution via OAuth token succeeds",
 			clientSecret:   "tskey-client-secret-123",
@@ -1665,6 +1704,7 @@ func TestResolveAuthKey(t *testing.T) {
 				ClientSecret: tt.clientSecret,
 				ClientID:     tt.clientID,
 				IDToken:      tt.idToken,
+				Audience:     tt.audience,
 				ControlURL:   "https://control.example.com",
 			}
 			s.shutdownCtx = context.Background()