Update logic as per PR comments

Author: nkshah2

recipe/session/recipeImplementation.go

@@ -45,10 +45,19 @@ var JWKRefreshRateLimit = 500
 // Maintains a map of the core path to the result
 var jwksCache *sessmodels.GetJWKSResult = nil
 
-func getJWKS() []sessmodels.GetJWKSFunctionObject {
-	result := []sessmodels.GetJWKSFunctionObject{}
+func getJWKS() sessmodels.GetJWKSResult {
 	corePaths := supertokens.GetAllCoreUrlsForPath("/.well-known/jwks.json")
 
+	if len(corePaths) == 0 {
+		return sessmodels.GetJWKSResult{
+			JWKS:        nil,
+			Error:       defaultErrors.New("No SuperTokens core available to query. Please pass supertokens > connectionURI to the init function, or override all the functions of the recipe you are using."),
+			LastFetched: 0,
+		}
+	}
+
+	var lastError error
+
 	for _, path := range corePaths {
 		// Here we dont need to check if cached result had an error because we only add to cache
 		// if the JWKS result was successful
@@ -57,86 +66,71 @@ func getJWKS() []sessmodels.GetJWKSFunctionObject {
 			// We check if we need to refresh before returning
 			currentTime := time.Now().UnixNano() / int64(time.Millisecond)
 
-			// This means that the value in cache is not expired, in this case we return a function that simply
-			// returns the cached keys
+			// This means that the value in cache is not expired, in this case we return the cached value
 			//
 			// Note that this also means that the SDK will not try to query any other Core (if there are multiple)
-			// if it has even a single valid cache entry from one of the core URLs. It will only attempt to fetch
+			// if it has a valid cache entry from one of the core URLs. It will only attempt to fetch
 			// from the cores again after the entry in the cache is expired
 			if (currentTime - jwksCache.LastFetched) < JWKCacheMaxAgeInMs {
-				finalResult := []sessmodels.GetJWKSFunctionObject{}
-
-				finalResult = append(finalResult, sessmodels.GetJWKSFunctionObject{
-					Fn: func(_ string) sessmodels.GetJWKSResult {
-						if supertokens.IsRunningInTestMode() {
-							returnedFromCache = true
-						}
-
-						return *jwksCache
-					},
-					Path: path,
-				})
+				if supertokens.IsRunningInTestMode() {
+					returnedFromCache = true
+				}
 
-				return finalResult
+				return *jwksCache
 			}
 
 			// This means that the value in cache is expired, we clear from cache and proceed
 			// as if it was never cached because that would be the equivalent of refreshing
 			//
 			// This has the added benefit where if there are multiple cores [Core1, Core2] and initially
-			// Core1 was down (so the cache only has a result for Core2). When Core2's cache expires the SDK
+			// Core1 was down (so the cache only has the result from Core2). When the cache expires the SDK
 			// will try to re-fetch for Core1 and will return that result (and save to cache) if Core1 is now up
 			jwksCache = nil
 			if supertokens.IsRunningInTestMode() {
 				deleteFromCacheCount++
 			}
 		}
 
-		// We need to also save the path of the JWKS request and pass it to the function this way because
-		// golang only saves a reference to the last value set to path when this function actually is called
-		// so all the functions in the slice would end up calling the last core host
-		//
-		// Doing it this way makes sure that all individual hosts are called
-		result = append(result, sessmodels.GetJWKSFunctionObject{
-			Fn: func(inputPath string) sessmodels.GetJWKSResult {
-				if supertokens.IsRunningInTestMode() {
-					urlsAttemptedForJWKSFetch = append(urlsAttemptedForJWKSFetch, inputPath)
-				}
+		if supertokens.IsRunningInTestMode() {
+			urlsAttemptedForJWKSFetch = append(urlsAttemptedForJWKSFetch, path)
+		}
 
-				// RefreshUnknownKID - Fetch JWKS again if the kid in the header of the JWT does not match any in
-				// the keyfunc library's cache
-				jwks, err := keyfunc.Get(inputPath, keyfunc.Options{
-					RefreshUnknownKID: true,
-				})
+		// RefreshUnknownKID - Fetch JWKS again if the kid in the header of the JWT does not match any in
+		// the keyfunc library's cache
+		jwks, jwksError := keyfunc.Get(path, keyfunc.Options{
+			RefreshUnknownKID: true,
+		})
 
-				jwksResult := sessmodels.GetJWKSResult{
-					JWKS:        jwks,
-					Error:       err,
-					LastFetched: time.Now().UnixNano() / int64(time.Millisecond),
-				}
+		if jwksError == nil {
+			jwksResult := sessmodels.GetJWKSResult{
+				JWKS:        jwks,
+				Error:       jwksError,
+				LastFetched: time.Now().UnixNano() / int64(time.Millisecond),
+			}
 
-				// Dont add to cache if there is an error to keep the logic of checking cache simple
-				// This means that for multiple cores, the only item we add to cache would be the first
-				// core that returned keys without an error
-				//
-				// This also has the added benefit where if initially the request failed because the core
-				// was down and then it comes back up, the next time it will try to request that core again
-				// after the cache has expired
-				if err == nil {
-					jwksCache = &jwksResult
-				}
+			// Dont add to cache if there is an error to keep the logic of checking cache simple
+			//
+			// This also has the added benefit where if initially the request failed because the core
+			// was down and then it comes back up, the next time it will try to request that core again
+			// after the cache has expired
+			jwksCache = &jwksResult
 
-				if supertokens.IsRunningInTestMode() {
-					returnedFromCache = false
-				}
+			if supertokens.IsRunningInTestMode() {
+				returnedFromCache = false
+			}
 
-				return jwksResult
-			},
-			Path: path,
-		})
+			return jwksResult
+		}
+
+		lastError = jwksError
 	}
 
-	return result
+	// This means that fetching from all cores failed
+	return sessmodels.GetJWKSResult{
+		JWKS:        nil,
+		Error:       lastError,
+		LastFetched: 0,
+	}
 }
 
 /**
@@ -147,29 +141,17 @@ Every core instance a backend is connected to is expected to connect to the same
 token verification. Otherwise, the result of session verification would depend on which core is currently available.
 */
 func GetCombinedJWKS() (*keyfunc.JWKS, error) {
-	var lastError error
-	jwksObjects := getJWKS()
-
 	if supertokens.IsRunningInTestMode() {
 		urlsAttemptedForJWKSFetch = []string{}
 	}
 
-	if len(jwksObjects) == 0 {
-		return nil, defaultErrors.New("No SuperTokens core available to query. Please pass supertokens > connectionURI to the init function, or override all the functions of the recipe you are using.")
-	}
-
-	for _, jwkObject := range jwksObjects {
-		jwksResult := jwkObject.Fn(jwkObject.Path)
-		err := jwksResult.Error
+	jwksResult := getJWKS()
 
-		if err != nil {
-			lastError = err
-		} else {
-			return jwksResult.JWKS, nil
-		}
+	if jwksResult.Error != nil {
+		return nil, jwksResult.Error
 	}
 
-	return nil, lastError
+	return jwksResult.JWKS, nil
 }
 
 func MakeRecipeImplementation(querier supertokens.Querier, config sessmodels.TypeNormalisedInput, appInfo supertokens.NormalisedAppinfo) sessmodels.RecipeInterface {

recipe/session/session_test.go

@@ -1335,13 +1335,13 @@ func TestThatJWKSResultIsRefreshedProperly(t *testing.T) {
 		t.Error(err.Error())
 	}
 
-	jwksBefore := getJWKS()[0]
-	beforeKids := jwksBefore.Fn(jwksBefore.Path).JWKS.KIDs()
+	jwksBefore := getJWKS()
+	beforeKids := jwksBefore.JWKS.KIDs()
 
 	time.Sleep(3 * time.Second)
 
-	jwksAfter := getJWKS()[0]
-	afterKids := jwksAfter.Fn(jwksAfter.Path).JWKS.KIDs()
+	jwksAfter := getJWKS()
+	afterKids := jwksAfter.JWKS.KIDs()
 	var newKeys []string
 
 	for _, key := range afterKids {