From 7a794161e05e78cf6e91a46cb1169958a287c3a5 Mon Sep 17 00:00:00 2001
From: FrancescoL96 <Francesco29.fl@gmail.com>
Date: Fri, 28 Nov 2025 12:00:25 +0100
Subject: [PATCH] fix: match partial objects from RESTAction in OIDC solver

---
 internal/routes/auth/oidc/login.go   |  8 ++++--
 internal/routes/auth/oidc/support.go | 43 ++++++++++++++++++++++++++++
 testdata/oidc-azure-pagination.yaml  |  8 +++---
 3 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/internal/routes/auth/oidc/login.go b/internal/routes/auth/oidc/login.go
index 495778b..61f28e2 100644
--- a/internal/routes/auth/oidc/login.go
+++ b/internal/routes/auth/oidc/login.go
@@ -94,9 +94,11 @@ func (r *loginRoute) Handler() http.HandlerFunc {
 		log.Debug().Str("name", name).Msg("resolving restaction")
 		if cfg.RESTActionRef != nil {
 			additionalFieldstoReplace, err := restaction.Resolve(r.ctx, r.rc, cfg.RESTActionRef, idToken.email, idToken.bearerToken)
-			value, ok := additionalFieldstoReplace["name"]
-			_, okk := additionalFieldstoReplace["name"].(string)
-			if err != nil || !ok || !okk || value == nil {
+			ok, errr := checkKeys(additionalFieldstoReplace)
+			if err != nil || !ok || errr != nil {
+				if errr != nil {
+					log.Err(err).Str("name", name).Msgf("restaction parsing has found the following error: %v", errr)
+				}
 				log.Err(err).Str("name", name).Msg("unable to resolve restaction, retrying with legacy resolve (copy)")
 				additionalFieldstoReplace, err = restaction.LegacyResolve(r.ctx, r.rc, cfg.RESTActionRef, idToken.email, idToken.bearerToken)
 				if err != nil {
diff --git a/internal/routes/auth/oidc/support.go b/internal/routes/auth/oidc/support.go
index 3bb02de..a8bc330 100644
--- a/internal/routes/auth/oidc/support.go
+++ b/internal/routes/auth/oidc/support.go
@@ -270,3 +270,46 @@ func updateConfig(config idToken, additionalFieldstoReplace map[string]interface
 	}
 	return config, nil
 }
+
+func checkKeys(additionalFieldstoReplace map[string]interface{}) (bool, error) {
+	for key := range additionalFieldstoReplace {
+		if additionalFieldstoReplace[key] != nil {
+			switch key {
+			case "name":
+				_, ok := additionalFieldstoReplace[key].(string)
+				if !ok {
+					return false, fmt.Errorf("key name not string")
+				}
+			case "email":
+				_, ok := additionalFieldstoReplace[key].(string)
+				if !ok {
+					return false, fmt.Errorf("key email not string")
+				}
+			case "preferredUsername":
+				_, ok := additionalFieldstoReplace[key].(string)
+				if !ok {
+					return false, fmt.Errorf("key preferredUsername not string")
+				}
+			case "groups":
+				v, ok := additionalFieldstoReplace[key].([]interface{})
+				if !ok {
+					return false, fmt.Errorf("key groups not []any")
+				}
+				for _, i := range v {
+					if _, okk := i.(string); !okk {
+						return false, fmt.Errorf("key in groups not string")
+					}
+				}
+			case "avatarURL":
+				_, ok := additionalFieldstoReplace[key].(string)
+				if !ok {
+					return false, fmt.Errorf("key avatarURL not string")
+				}
+			// If a key outside of those allowed is found, then error
+			default:
+				return false, fmt.Errorf("found unexpected key: %s", key)
+			}
+		}
+	}
+	return true, nil
+}
diff --git a/testdata/oidc-azure-pagination.yaml b/testdata/oidc-azure-pagination.yaml
index e8f6566..7b546ba 100644
--- a/testdata/oidc-azure-pagination.yaml
+++ b/testdata/oidc-azure-pagination.yaml
@@ -56,19 +56,19 @@ spec:
     headers:
     - 'Accept: application/json'
     - "${ \"Authorization: Bearer \" + .token }"
-    path: ${ .next // "v1.0/me/memberOf?$select=displayName&" }
+    path: ${ .next // "v1.0/me/memberOf?$select=id" }
     endpointRef:
       name: azure-entra
       namespace: krateo-system
     filter: |
       .groups | if has("@odata.nextLink") then
       {
-        res: (.value | map(.displayName)),
+        res: (.value | map(.id)),
         next: (."@odata.nextLink" | sub("https://graph.microsoft.com"; ""))
       }
       else
       {
-        res: (.value | map(.displayName))
+        res: (.value | map(.id))
       }
       end
   - name: recursive
@@ -82,7 +82,7 @@ spec:
         (.counter // 0) as $counter | .token as $token | if $counter < 10 then ( .groups | if has("next") then [{next: .next, token: $token, counter: $counter}] else [] end ) else [] end
     headers:
     - 'Accept: application/json'
-    path: |
+    path: | # set your restaction name and namespace in the URI to match 
       ${ "/call?apiVersion=templates.krateo.io/v1&resource=restactions&name=test-rest-action&namespace=krateo-system&extras=" + ({token: .token, next: .next, counter: (.counter+1)} | tostring | @uri ) }
     endpointRef:
       name: snowplow-endpoint