add pagination for graph api

tuunit · jjlakis · commit dfe9b5336244 · 2024-12-24T01:25:44.000+01:00
diff --git a/providers/ms_entra_id.go b/providers/ms_entra_id.go
@@ -117,29 +117,39 @@ func (p *MicrosoftEntraIDProvider) addGraphGroupsToSession(ctx context.Context,
 	groupsHeaders := makeAuthorizationHeader(tokenTypeBearer, s.AccessToken, nil)
 	groupsHeaders.Add("ConsistencyLevel", "eventual")
 
-	groupsURL := fmt.Sprintf("%s/transitiveMemberOf?$select=id&$top=999", p.microsoftGraphURL)
+	var allGroups []string
+	var nextLink string
 
-	jsonRequest, err := requests.New(groupsURL).
-		WithContext(ctx).
-		WithHeaders(groupsHeaders).
-		Do().
-		UnmarshalSimpleJSON()
+	for {
+		if nextLink == "" {
+			nextLink = fmt.Sprintf("%s/transitiveMemberOf?$select=id&$top=100", p.microsoftGraphURL)
+		}
 
-	if err != nil {
-		logger.Errorf("invalid response from microsoft graph, no groups added to session: %v", err)
-		return nil
-	}
+		response, err := requests.New(nextLink).
+			WithContext(ctx).
+			WithHeaders(groupsHeaders).
+			Do().
+			UnmarshalSimpleJSON()
 
-	reqGroups := jsonRequest.Get("value").MustArray()
-	groups := make([]string, len(reqGroups))
+		if err != nil {
+			return fmt.Errorf("invalid response from microsoft graph, no groups added to session: %v", err)
+		}
+		reqGroups := response.Get("value").MustArray()
 
-	for i := range reqGroups {
-		value := jsonRequest.Get("value").GetIndex(i).Get("id").MustString()
-		groups = append(groups, value)
-	}
+		for i := range reqGroups {
+			value := response.Get("value").GetIndex(i).Get("id").MustString()
+			allGroups = append(allGroups, value)
+		}
+
+		// https://learn.microsoft.com/en-us/graph/paging?view=graph-rest-1.0&tabs=http#how-paging-works
+		nextLink = response.Get("@odata.nextLink").MustString()
 
-	s.Groups = util.RemoveDuplicateStr(append(s.Groups, groups...))
+		if nextLink == "" {
+			break
+		}
+	}
 
+	s.Groups = util.RemoveDuplicateStr(append(s.Groups, allGroups...))
 	return nil
 }
 
diff --git a/providers/ms_entra_id_test.go b/providers/ms_entra_id_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -77,6 +78,7 @@ func TestAzureEntraOIDCProviderEnrichSessionGroupOverage(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Contains(t, session.Groups, "85d7d600-7804-4d92-8d43-9c33c21c130c")
 	assert.Contains(t, session.Groups, "916f0604-8a3b-4a69-bda9-06db11a8f0cd")
+	assert.Contains(t, session.Groups, "b1aef995-6b55-4ac6-bbfe-e829810e9352", "Pagination using $skiptoken failed")
 }
 
 func TestAzureEntraOIDCProviderValidateSessionAllowedTenants(t *testing.T) {
@@ -133,9 +135,23 @@ func mockGraphAPI(noGroupMemberPermissions bool) *httptest.Server {
 		func(w http.ResponseWriter, r *http.Request) {
 			if noGroupMemberPermissions {
 				w.WriteHeader(401)
-			} else if r.URL.Path == groupsPath && r.Method == http.MethodGet {
+			} else if r.URL.Path == groupsPath && r.Method == http.MethodGet && len(r.URL.Query()["$skiptoken"]) > 0 {
+				// Second page (pagination)
 				w.Write([]byte(`{
 					"@odata.context": "https://graph.microsoft.com/v1.0/$metadata#directoryObjects(id)",
+					"value": [
+						{
+							"@odata.type": "#microsoft.graph.group",
+							"id": "b1aef995-6b55-4ac6-bbfe-e829810e9352"
+						}
+					]
+				}`))
+
+			} else if r.URL.Path == groupsPath && r.Method == http.MethodGet {
+				// First page (pagination)
+				w.Write([]byte(fmt.Sprintf(`{
+					"@odata.context": "https://graph.microsoft.com/v1.0/$metadata#directoryObjects(id)",
+					"@odata.nextLink": "http://%s/v1.0/me/transitiveMemberOf?$select=id&$top=2&$skiptoken=TEST_TOKEN",
 					"value": [
 						{
 							"@odata.type": "#microsoft.graph.group",
@@ -146,7 +162,7 @@ func mockGraphAPI(noGroupMemberPermissions bool) *httptest.Server {
 							"id": "916f0604-8a3b-4a69-bda9-06db11a8f0cd"
 						  }
 					]
-				}`))
+				}`, r.Host)))
 			}
 		},
 	))
