Merge pull request #292 from aojea/oidc

node: fix device flow polling against dex's non-RFC 401 responses

Author: aojea

cmd/sam-control-plane/main.go

@@ -36,6 +36,7 @@ var (
 	dbDSN                 string
 	dbDSNPath             string
 	oidcIssuer            string
+	oidcClientID          string
 	allowedAudiencesFlag  string
 	keyRotationInterval   time.Duration
 	keyGracePeriod        time.Duration
@@ -112,6 +113,7 @@ func main() {
 				DriverName:            dbDriver,
 				DataSourceName:        dbDSN,
 				OIDCIssuer:            oidcIssuer,
+				OIDCClientID:          oidcClientID,
 				AllowedAudiences:      auds,
 				LeaseDuration:         leaseDuration,
 				KeyRotationInterval:   keyRotationInterval,
@@ -146,6 +148,7 @@ func main() {
 	rootCmd.PersistentFlags().StringVar(&dbDSN, "db-dsn", "control-plane.db", "Database DSN/Connection URL (avoid for postgres: embeds a password; prefer --db-dsn-path or SAM_DB_DSN)")
 	rootCmd.PersistentFlags().StringVar(&dbDSNPath, "db-dsn-path", "", "Path to file containing the database DSN/Connection URL (overrides --db-dsn; or env SAM_DB_DSN)")
 	rootCmd.Flags().StringVar(&oidcIssuer, "issuer", "", "OIDC Issuer URL (comma-separated)")
+	rootCmd.Flags().StringVar(&oidcClientID, "oidc-client-id", "", "OAuth client ID advertised to joining nodes via /info (defaults to the first allowed audience)")
 	rootCmd.Flags().StringVar(&allowedAudiencesFlag, "allowed-audiences", api.DefaultAudience, "Comma-separated list of allowed OIDC audiences")
 	rootCmd.Flags().DurationVar(&keyRotationInterval, "key-rotation-interval", 24*time.Hour, "Key rotation interval (e.g. 24h). 0 disables rotation.")
 	rootCmd.Flags().DurationVar(&keyGracePeriod, "key-grace-period", 1*time.Hour, "Key grace period for rotated keys.")

internal/controlplane/config.go

@@ -25,6 +25,7 @@ type Options struct {
 	DriverName            string
 	DataSourceName        string
 	OIDCIssuer            string
+	OIDCClientID          string // OAuth client id advertised via /info; defaults to the first allowed audience
 	AllowedAudiences      []string
 	LeaseDuration         time.Duration
 	KeyRotationInterval   time.Duration

internal/controlplane/server.go

@@ -368,6 +368,13 @@ func (s *Server) HandleInfo(w http.ResponseWriter, r *http.Request) {
 		aud = s.config.AllowedAudiences[0]
 	}
 
+	// aud == client_id only holds for id_token-model IdPs (dex, Google); an
+	// explicit client id supports providers where the two differ.
+	clientID := s.config.OIDCClientID
+	if clientID == "" {
+		clientID = aud
+	}
+
 	// Fetch active routers
 	activeRouters, err := s.store.GetActiveRouters(r.Context())
 	if err != nil {
@@ -383,7 +390,7 @@ func (s *Server) HandleInfo(w http.ResponseWriter, r *http.Request) {
 
 	resp := &api.ControlPlaneInfoResponse{
 		OidcIssuer:      issuer,
-		ClientId:        aud,
+		ClientId:        clientID,
 		Audience:        aud,
 		RouterAddresses: routerAddrs, // Reused this field for back-compatibility with bootstrap routers list
 	}

internal/controlplane/server_test.go

@@ -202,6 +202,41 @@ func TestControlPlaneBasic(t *testing.T) {
 		t.Errorf("expected 0 active routers, got %d", len(info.RouterAddresses))
 	}
 
+	// With an explicit OIDC client id, /info must advertise it while the
+	// audience stays the first allowed audience.
+	t.Run("explicit oidc client id", func(t *testing.T) {
+		dbPath := filepath.Join(t.TempDir(), "cp-clientid.db")
+		st, err := storage.NewSQLStore("sqlite", dbPath)
+		if err != nil {
+			t.Fatalf("failed to create store: %v", err)
+		}
+		defer func() { _ = st.Close() }()
+
+		srv2, err := NewServer(Options{
+			DriverName:       "sqlite",
+			DataSourceName:   dbPath,
+			OIDCIssuer:       issuer,
+			OIDCClientID:     "sam-cli-app",
+			AllowedAudiences: []string{"sam-mesh-audience"},
+		}, st)
+		if err != nil {
+			t.Fatalf("failed to create server: %v", err)
+		}
+
+		rec := httptest.NewRecorder()
+		srv2.HandleInfo(rec, httptest.NewRequest(http.MethodGet, "/info", nil))
+		if rec.Code != http.StatusOK {
+			t.Fatalf("unexpected /info status: %d", rec.Code)
+		}
+		var info2 api.ControlPlaneInfoResponse
+		if err := proto.Unmarshal(rec.Body.Bytes(), &info2); err != nil {
+			t.Fatalf("failed to unmarshal ControlPlaneInfoResponse: %v", err)
+		}
+		if info2.ClientId != "sam-cli-app" || info2.Audience != "sam-mesh-audience" {
+			t.Errorf("unexpected client id/audience: %+v", &info2)
+		}
+	})
+
 	// 2. Test /keys
 	resp, err = client.Get(baseURL + "/keys")
 	if err != nil {

internal/identity/oidc.go

@@ -44,31 +44,39 @@ func VerifyJWT(ctx context.Context, jwtStr string, allowedAudiences []string, pr
 	}
 	iss, _ := claims["iss"].(string)
 
-	// 2. Extract the audience
-	var aud string
+	// 2. Extract all audiences; the aud claim may be a string or an array.
+	var auds []string
 	switch a := claims["aud"].(type) {
 	case string:
-		aud = a
+		auds = []string{a}
 	case []any:
-		if len(a) > 0 {
-			aud, _ = a[0].(string)
+		for _, v := range a {
+			if s, ok := v.(string); ok {
+				auds = append(auds, s)
+			}
 		}
 	}
 
-	if aud == "" {
+	if len(auds) == 0 {
 		return nil, nil, fmt.Errorf("missing aud claim")
 	}
 
-	// 3. Verify the audience matches one of your expected tenants/platforms
+	// 3. Accept if any audience matches an allowed one: a multi-audience token
+	// only needs to be intended for us, whatever else it names.
 	validAudience := false
-	for _, allowed := range allowedAudiences {
-		if aud == allowed {
-			validAudience = true
+	for _, aud := range auds {
+		for _, allowed := range allowedAudiences {
+			if aud == allowed {
+				validAudience = true
+				break
+			}
+		}
+		if validAudience {
 			break
 		}
 	}
 	if !validAudience {
-		return nil, nil, fmt.Errorf("untrusted audience: %s", aud)
+		return nil, nil, fmt.Errorf("untrusted audience(s): %s", strings.Join(auds, ", "))
 	}
 
 	// 4. Route to the correct provider

internal/identity/oidc_test.go

@@ -157,6 +157,39 @@ func TestVerifyJWT(t *testing.T) {
 		}
 	})
 
+	t.Run("allowed audience in any array position succeeds", func(t *testing.T) {
+		claims := validClaims()
+		claims["aud"] = []string{"some-other-audience", "sam-mesh-audience"}
+		tokenStr := signToken(t, key, testKID, claims)
+
+		_, _, err := VerifyJWT(ctx, tokenStr, allowedAudiences, providers)
+		if err != nil {
+			t.Fatalf("expected success for multi-audience token, got: %v", err)
+		}
+	})
+
+	t.Run("audience array with no allowed entry is rejected", func(t *testing.T) {
+		claims := validClaims()
+		claims["aud"] = []string{"some-other-audience", "yet-another-audience"}
+		tokenStr := signToken(t, key, testKID, claims)
+
+		_, _, err := VerifyJWT(ctx, tokenStr, allowedAudiences, providers)
+		if err == nil {
+			t.Fatal("expected error for audience array with no allowed entry")
+		}
+	})
+
+	t.Run("empty audience array is rejected", func(t *testing.T) {
+		claims := validClaims()
+		claims["aud"] = []string{}
+		tokenStr := signToken(t, key, testKID, claims)
+
+		_, _, err := VerifyJWT(ctx, tokenStr, allowedAudiences, providers)
+		if err == nil {
+			t.Fatal("expected error for empty audience array")
+		}
+	})
+
 	t.Run("unknown issuer is rejected", func(t *testing.T) {
 		claims := validClaims()
 		claims["iss"] = "https://unknown-issuer.example"

internal/node/oidc.go

@@ -383,6 +383,19 @@ func stdinIsInteractive() bool {
 	return isatty.IsTerminal(fd) || isatty.IsCygwinTerminal(fd)
 }
 
+// bodySnippet renders a response body for error messages, bounded so a
+// misbehaving server can't flood logs.
+func bodySnippet(body []byte) string {
+	if len(body) > 256 {
+		return strings.TrimSpace(string(body[:256])) + "..."
+	}
+	s := strings.TrimSpace(string(body))
+	if s == "" {
+		return "(empty response body)"
+	}
+	return s
+}
+
 // DeviceLogin performs OAuth 2.0 Device Authorization Grant (RFC 8628).
 func (n *SamNode) DeviceLogin(ctx context.Context, deviceAuthURL, tokenURL, clientID, audience string, requestRefresh bool) (string, error) {
 	if deviceAuthURL == "" {
@@ -514,7 +527,9 @@ func (n *SamNode) DeviceLogin(ctx context.Context, deviceAuthURL, tokenURL, clie
 			continue
 		}
 
-		body, readErr := io.ReadAll(tokenResp.Body)
+		// Cap the read: the poll target comes from discovery and a misbehaving
+		// server must not be able to exhaust memory.
+		body, readErr := io.ReadAll(io.LimitReader(tokenResp.Body, 1<<20))
 		if closeErr := tokenResp.Body.Close(); closeErr != nil {
 			logger.Errorf("Failed to close response body: %v", closeErr)
 		}
@@ -546,16 +561,29 @@ func (n *SamNode) DeviceLogin(ctx context.Context, deviceAuthURL, tokenURL, clie
 			return jwt, nil
 		}
 
-		if tokenResp.StatusCode != http.StatusBadRequest {
-			return "", fmt.Errorf("token request failed with status: %s", tokenResp.Status)
-		}
-
 		var errResp struct {
 			Error            string `json:"error"`
 			ErrorDescription string `json:"error_description"`
 		}
-		if err := json.Unmarshal(body, &errResp); err != nil {
-			return "", fmt.Errorf("failed to decode token polling error response: %w", err)
+		// Best-effort: non-JSON bodies (proxy HTML, empty) leave Error empty
+		// and are reported raw via bodySnippet below.
+		_ = json.Unmarshal(body, &errResp)
+		pending := errResp.Error == "authorization_pending" || errResp.Error == "slow_down"
+
+		// RFC 8628 §3.5 delivers polling errors as HTTP 400 with an OAuth error
+		// body. Known exception: dex returns the pending/slow_down signals with
+		// HTTP 401. Anything else is a real failure and is surfaced verbatim.
+		rfcError := tokenResp.StatusCode == http.StatusBadRequest && errResp.Error != ""
+		dexPending := tokenResp.StatusCode == http.StatusUnauthorized && pending
+		if !rfcError && !dexPending {
+			if errResp.Error != "" {
+				msg := errResp.Error
+				if errResp.ErrorDescription != "" {
+					msg += " - " + errResp.ErrorDescription
+				}
+				return "", fmt.Errorf("token request failed with status %s: %s", tokenResp.Status, msg)
+			}
+			return "", fmt.Errorf("token request failed with status %s: %s", tokenResp.Status, bodySnippet(body))
 		}
 
 		switch errResp.Error {

internal/node/oidc_test.go

@@ -510,6 +510,106 @@ func TestDeviceLoginRetriesOnTransientTokenError(t *testing.T) {
 	}
 }
 
+// TestDeviceLoginPending401 verifies polling survives providers (e.g. dex)
+// that return authorization_pending with HTTP 401 instead of the
+// RFC 8628-mandated 400: the OAuth error code in the body wins over the
+// HTTP status code.
+func TestDeviceLoginPending401(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/device", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]interface{}{
+			"device_code":      "dev_code_1",
+			"user_code":        "CCCC-DDDD",
+			"verification_uri": "https://example.com/device",
+			"expires_in":       60,
+			"interval":         1,
+		})
+	})
+
+	var pollCount int32
+	mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		if atomic.AddInt32(&pollCount, 1) == 1 {
+			w.WriteHeader(http.StatusUnauthorized)
+			_ = json.NewEncoder(w).Encode(map[string]string{"error": "authorization_pending"})
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]string{"id_token": "token_after_pending_401"})
+	})
+
+	server := httptest.NewServer(mux)
+	defer server.Close()
+
+	node := &SamNode{}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	token, err := node.DeviceLogin(ctx, server.URL+"/device", server.URL+"/token", "client_id_test", "sam-e2e", false)
+	if err != nil {
+		t.Fatalf("DeviceLogin failed: %v", err)
+	}
+	if token != "token_after_pending_401" {
+		t.Fatalf("Expected token_after_pending_401, got %q", token)
+	}
+	if got := atomic.LoadInt32(&pollCount); got < 2 {
+		t.Fatalf("expected at least 2 poll attempts, got %d", got)
+	}
+}
+
+// TestDeviceLoginFatalErrors verifies polling aborts on terminal responses:
+// a non-OAuth error body (regardless of status) and an OAuth error code
+// that is not a pending/slow_down signal.
+func TestDeviceLoginFatalErrors(t *testing.T) {
+	cases := []struct {
+		name    string
+		status  int
+		body    string
+		wantErr string
+	}{
+		{"non-oauth 401", http.StatusUnauthorized, `{"message":"nope"}`, "token request failed with status"},
+		{"invalid_client 401", http.StatusUnauthorized, `{"error":"invalid_client","error_description":"unknown client"}`, "invalid_client"},
+		{"access_denied 400", http.StatusBadRequest, `{"error":"access_denied"}`, "denied"},
+		{"oauth-shaped 500 is not a protocol error", http.StatusInternalServerError, `{"error":"server_error"}`, "server_error"},
+		{"html 502 surfaces the body", http.StatusBadGateway, `<html>bad gateway</html>`, "bad gateway"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			mux := http.NewServeMux()
+			mux.HandleFunc("/device", func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				_ = json.NewEncoder(w).Encode(map[string]interface{}{
+					"device_code":      "dev_code_1",
+					"user_code":        "EEEE-FFFF",
+					"verification_uri": "https://example.com/device",
+					"expires_in":       60,
+					"interval":         1,
+				})
+			})
+			mux.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(tc.status)
+				_, _ = w.Write([]byte(tc.body))
+			})
+
+			server := httptest.NewServer(mux)
+			defer server.Close()
+
+			node := &SamNode{}
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			_, err := node.DeviceLogin(ctx, server.URL+"/device", server.URL+"/token", "client_id_test", "sam-e2e", false)
+			if err == nil {
+				t.Fatal("expected DeviceLogin to fail")
+			}
+			if !strings.Contains(err.Error(), tc.wantErr) {
+				t.Fatalf("expected error containing %q, got %v", tc.wantErr, err)
+			}
+		})
+	}
+}
+
 func TestParseAuthMode(t *testing.T) {
 	cases := []struct {
 		in      string