Add SP OBO Graph overage claim resolution support

auth/providers/azure/azure.go

@@ -267,6 +267,8 @@ func (s Authenticator) Check(ctx context.Context, token string) (*authv1.UserInf
 		return nil, err
 	}
 
+	isServicePrincipal := isAppToken(claims)
+
 	if s.Options.ResolveGroupMembershipOnlyOnOverageClaim {
 		groups, skipGraphAPI, err := getGroupsAndCheckOverage(claims)
 		if err != nil {
@@ -276,16 +278,16 @@ func (s Authenticator) Check(ctx context.Context, token string) (*authv1.UserInf
 			resp.Groups = groups
 			return resp, nil
 		}
-		// Service principal token with >200 groups. OBO flow does not
-		// support SPNs, so short-circuit instead of letting AAD fail with
-		// a cryptic AADSTS7000113.
+
+		// Service principal token with an overage claim (>200 groups) while SP
+		// group resolution is switched off: short-circuit instead of letting the
+		// user flow fail with a cryptic AADSTS7000113.
 		//
-		// StatusOK: K8s webhook authenticator only reads and logs
-		// TokenReview Status.Error on HTTP 200. Non-200 is treated as a
-		// transport error and the message is discarded. We need this error
-		// to surface in API server logs so operators can diagnose why the
-		// SPN authentication was rejected.
-		if isAppToken(claims) {
+		// StatusOK: K8s webhook authenticator only reads and logs TokenReview
+		// Status.Error on HTTP 200. Non-200 is treated as a transport error and
+		// the message is discarded. We need this error to surface in API server
+		// logs so operators can diagnose why the SPN was rejected.
+		if !s.Options.EnableSPGroupResolution && isServicePrincipal {
 			return nil, errutils.WithCode(
 				fmt.Errorf(
 					"service principal with group membership exceeding 200 is not supported. "+
@@ -299,11 +301,20 @@ func (s Authenticator) Check(ctx context.Context, token string) (*authv1.UserInf
 		if err := s.graphClient.RefreshToken(ctx, token); err != nil {
 			return nil, err
 		}
-		resp.Groups, err = s.graphClient.GetGroups(ctx, resp.Username, token)
+		principal := resp.Username
+		if isServicePrincipal {
+			// Service principals are resolved by object ID against the
+			// /servicePrincipals resource.
+			spOID, oidErr := claims.string(azureObjectIDClaim)
+			if oidErr != nil {
+				return nil, errors.Wrap(oidErr, "unable to get oid claim for service principal")
+			}
+			principal = spOID
+		}
+		resp.Groups, err = s.graphClient.GetGroups(ctx, principal, token, isServicePrincipal)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to get groups")
 		}
-
 	}
 	return resp, nil
 }

auth/providers/azure/azure_test.go

@@ -32,6 +32,7 @@ import (
 	"time"
 
 	"go.kubeguard.dev/guard/auth/providers/azure/graph"
+	errutils "go.kubeguard.dev/guard/util/error"
 
 	"github.com/coreos/go-oidc"
 	"github.com/go-chi/chi/v5"
@@ -216,6 +217,15 @@ func serverSetup(loginResp string, loginStatus int, jwkResp, groupIds, groupList
 		_, _ = w.Write(groupIds)
 	}))
 
+	m.Post("/api/servicePrincipals/{oid}/getMemberGroups", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if len(groupStatus) > 0 {
+			w.WriteHeader(groupStatus[0])
+		} else {
+			w.WriteHeader(http.StatusOK)
+		}
+		_, _ = w.Write(groupIds)
+	}))
+
 	m.Post("/api/directoryObjects/getByIds", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 		_, _ = w.Write(groupList)
@@ -456,7 +466,26 @@ func TestCheckAzureAuthenticationSPNWithOverage(t *testing.T) {
 		t.Fatalf("Error when creating signing key. reason : %v", err)
 	}
 
-	t.Run("SPN token with overage claim should return clear error", func(t *testing.T) {
+	t.Run("SPN token with overage claim resolves groups via servicePrincipals getMemberGroups", func(t *testing.T) {
+		srv, client := getServerAndClient(t, signKey, loginResp, 3, true, false, ClientCredentialAuthMode)
+		client.Options.ResolveGroupMembershipOnlyOnOverageClaim = true
+		client.Options.UseGroupUID = true
+		client.Options.EnableSPGroupResolution = true
+		defer srv.Close()
+
+		token, err := signKey.sign([]byte(fmt.Sprintf(accessTokenSPNWithOverage, srv.URL)))
+		if err != nil {
+			t.Fatalf("Error when signing token. reason: %v", err)
+		}
+
+		resp, err := client.Check(ctx, token)
+		assert.Nil(t, err)
+		assert.NotNil(t, resp)
+		assert.Equal(t, 3, len(resp.Groups))
+		assert.ElementsMatch(t, []string{"1", "2", "3"}, resp.Groups)
+	})
+
+	t.Run("SPN token with overage claim is rejected when SP group resolution is disabled", func(t *testing.T) {
 		srv, client := getServerAndClient(t, signKey, loginResp, 3, true, false, ClientCredentialAuthMode)
 		client.Options.ResolveGroupMembershipOnlyOnOverageClaim = true
 		client.Options.UseGroupUID = true
@@ -471,11 +500,29 @@ func TestCheckAzureAuthenticationSPNWithOverage(t *testing.T) {
 		assert.NotNil(t, err)
 		assert.Nil(t, resp)
 		assert.Contains(t, err.Error(), "service principal with group membership exceeding 200 is not supported")
-		assert.Contains(t, err.Error(), "kubelogin-authentication-in-aks-limitations")
 
-		codeErr, ok := err.(interface{ Code() int })
-		assert.True(t, ok, "error should implement Code()")
-		assert.Equal(t, http.StatusOK, codeErr.Code(), "error should HTTP 200 so the API server logs the message.")
+		codeErr, ok := err.(errutils.HttpStatusCode)
+		assert.True(t, ok, "error should carry an http status code")
+		assert.Equal(t, http.StatusOK, codeErr.Code())
+	})
+
+	t.Run("SPN token with overage claim surfaces graph failure", func(t *testing.T) {
+		srv, client := getServerAndClient(t, signKey, loginResp, 3, true, false, ClientCredentialAuthMode, http.StatusInternalServerError)
+		client.Options.ResolveGroupMembershipOnlyOnOverageClaim = true
+		client.Options.UseGroupUID = true
+		client.Options.EnableSPGroupResolution = true
+		defer srv.Close()
+
+		token, err := signKey.sign([]byte(fmt.Sprintf(accessTokenSPNWithOverage, srv.URL)))
+		if err != nil {
+			t.Fatalf("Error when signing token. reason: %v", err)
+		}
+
+		resp, err := client.Check(ctx, token)
+		assert.NotNil(t, err)
+		assert.Nil(t, resp)
+		assert.Contains(t, err.Error(), "failed to get groups")
+		assert.NotContains(t, err.Error(), "service principal with group membership exceeding 200 is not supported")
 	})
 
 	t.Run("SPN token with groups in token should succeed (no Graph call needed)", func(t *testing.T) {

auth/providers/azure/graph/graph.go

@@ -48,11 +48,18 @@ var (
 	expandedGroupsPerCall = 500
 	idtypClaim            = "idtyp"
 
+	// For getMemberGroups (user flow)
 	getMemberGroupsFailed = promauto.NewCounter(prometheus.CounterOpts{
 		Name: "guard_azure_graph_failure_total",
 		Help: "Azure graph getMemberGroups call failed.",
 	})
 
+	// For getMemberGroups (service principal flow)
+	getMemberGroupsForSPFailed = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "guard_azure_graph_sp_failure_total",
+		Help: "Azure graph getMemberGroups call for service principal failed.",
+	})
+
 	getMemberGroupsUsingARCOboServiceCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
 		Name: "guard_azure_arc_obo_request_total",
 		Help: "Total no of arc obo getMemberGroups calls by code.",
@@ -100,15 +107,23 @@ type UserInfo struct {
 	lock          sync.RWMutex
 }
 
-func (u *UserInfo) getGroupIDs(ctx context.Context, userPrincipal string) ([]string, error) {
-	// Create a new request for finding the user.
+// getGroupIDs returns the group IDs the given principal is a member of. Users and
+// service principals are distinct Graph resources, so the principal type selects
+// the resource to query: users are looked up by user principal name, service
+// principals by object ID.
+func (u *UserInfo) getGroupIDs(ctx context.Context, principal string, isServicePrincipal bool) ([]string, error) {
+	// Create a new request for finding the principal.
 	// Shallow copy of the base API URL
-	userSearchURL := *u.apiURL
+	searchURL := *u.apiURL
+	resource := "users"
+	if isServicePrincipal {
+		resource = "servicePrincipals"
+	}
 	// Append the path for the member list
-	userSearchURL.Path = path.Join(userSearchURL.Path, fmt.Sprintf("/users/%s/getMemberGroups", userPrincipal))
+	searchURL.Path = path.Join(searchURL.Path, fmt.Sprintf("/%s/%s/getMemberGroups", resource, principal))
 
 	// The body being sent makes sure that all groups are returned, not just security groups
-	req, err := http.NewRequest(http.MethodPost, userSearchURL.String(), strings.NewReader(`{"securityEnabledOnly": false}`))
+	req, err := http.NewRequest(http.MethodPost, searchURL.String(), strings.NewReader(`{"securityEnabledOnly": false}`))
 	if err != nil {
 		return nil, errors.Wrap(err, "error creating group IDs request")
 	}
@@ -117,8 +132,12 @@ func (u *UserInfo) getGroupIDs(ctx context.Context, userPrincipal string) ([]str
 
 	resp, err := u.client.Do(req.WithContext(ctx))
 	if err != nil {
-		getMemberGroupsFailed.Inc()
-		return nil, errors.Wrap(err, "error listing users")
+		if isServicePrincipal {
+			getMemberGroupsForSPFailed.Inc()
+		} else {
+			getMemberGroupsFailed.Inc()
+		}
+		return nil, errors.Wrapf(err, "error listing group memberships for %s", resource)
 	}
 	defer func() {
 		_ = resp.Body.Close()
@@ -363,10 +382,12 @@ func (u *UserInfo) isTokenExpired() bool {
 	return u.expires.Before(time.Now())
 }
 
-// GetGroups gets a list of all groups that the given user principal is part of
-// Generally in federated directories the email address is the userPrincipalName
-func (u *UserInfo) GetGroups(ctx context.Context, userPrincipal string, token string) ([]string, error) {
+// GetGroups gets a list of all groups that the given principal is part of.
+// For users the principal is generally the email address in federated directories
+// (the userPrincipalName), for service principals it is the object ID.
+func (u *UserInfo) GetGroups(ctx context.Context, principal string, token string, isServicePrincipal bool) ([]string, error) {
 	// use arc obo service to get groups if authn mode is arc
+	// the arc obo service rejects service principal tokens with a clear error
 	if u.authMode == arcAuthMode {
 		groupIds, err := u.getMemberGroupsUsingARCOboService(ctx, token)
 		if err != nil {
@@ -375,8 +396,8 @@ func (u *UserInfo) GetGroups(ctx context.Context, userPrincipal string, token st
 		return groupIds, nil
 	}
 
-	// Get the group IDs for the user
-	groupIDs, err := u.getGroupIDs(ctx, userPrincipal)
+	// Get the group IDs for the principal
+	groupIDs, err := u.getGroupIDs(ctx, principal, isServicePrincipal)
 	if err != nil {
 		return nil, err
 	}