From a33ff15f6d63a29d9a0cf4957835b65f48bb8a81 Mon Sep 17 00:00:00 2001
From: Ignasi Barrera <ignasi@tetrate.io>
Date: Wed, 17 Apr 2024 11:30:12 +0200
Subject: [PATCH] Cache well-known responses to avoid making too much calls to
 the IdP

Signed-off-by: Ignasi Barrera <ignasi@tetrate.io>
---
 internal/authz/oidc_test.go     |  2 ++
 internal/oidc/discovery.go      | 14 +++++++++++++-
 internal/oidc/discovery_test.go | 21 ++++++++++++++++++++-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/internal/authz/oidc_test.go b/internal/authz/oidc_test.go
index 1956e685..c68e1c97 100644
--- a/internal/authz/oidc_test.go
+++ b/internal/authz/oidc_test.go
@@ -1396,6 +1396,7 @@ func TestLoadWellKnownConfigMissingLogoutRedirectURI(t *testing.T) {
 	t.Cleanup(idpServer.Stop)
 
 	cfg := proto.Clone(dynamicOIDCConfig).(*oidcv1.OIDCConfig)
+	cfg.ConfigurationUri = "http://missing-logout/.well-known/openid-configuration"
 	require.ErrorIs(t, loadWellKnownConfig(idpServer.newHTTPClient(), cfg), ErrMissingLogoutRedirectURI)
 }
 
@@ -1403,6 +1404,7 @@ func TestLoadWellKnownConfigError(t *testing.T) {
 	clock := oidc.Clock{}
 	tlsPool := internal.NewTLSConfigPool(context.Background())
 	cfg := proto.Clone(dynamicOIDCConfig).(*oidcv1.OIDCConfig)
+	cfg.ConfigurationUri = "http://stopped-server/.well-known/openid-configuration"
 	sessions := &mockSessionStoreFactory{store: oidc.NewMemoryStore(&clock, time.Hour, time.Hour)}
 	_, err := NewOIDCHandler(cfg, tlsPool, oidc.NewJWKSProvider(newConfigFor(basicOIDCConfig), tlsPool),
 		sessions, clock, oidc.NewStaticGenerator(newSessionID, newNonce, newState))
diff --git a/internal/oidc/discovery.go b/internal/oidc/discovery.go
index 71f06823..8ea03e0a 100644
--- a/internal/oidc/discovery.go
+++ b/internal/oidc/discovery.go
@@ -40,8 +40,20 @@ type WellKnownConfig struct {
 	TokenRevocationEndpoint  string   `json:"token_revocation_endpoint"`
 }
 
+var (
+	// wellKnownConfigs is a map of issuer URL to the OIDC well-known configuration
+	// It is used to cache well-known configurations as they usually don't change. URLs are usually stable, and the only
+	// things that are subject to change are the signing keys, but those are already watched periodically by the JWKS fetcher.
+	wellKnownConfigs = make(map[string]WellKnownConfig)
+)
+
 // GetWellKnownConfig retrieves the OIDC well-known configuration from the given issuer URL.
 func GetWellKnownConfig(client *http.Client, url string) (WellKnownConfig, error) {
+	cfg, ok := wellKnownConfigs[url]
+	if ok {
+		return cfg, nil
+	}
+
 	// Make a GET request to the well-known configuration endpoint
 	response, err := client.Get(url)
 	if err != nil {
@@ -55,10 +67,10 @@ func GetWellKnownConfig(client *http.Client, url string) (WellKnownConfig, error
 	}
 
 	// Decode the JSON response into the OIDCConfig struct
-	var cfg WellKnownConfig
 	if err = json.NewDecoder(response.Body).Decode(&cfg); err != nil {
 		return WellKnownConfig{}, err
 	}
 
+	wellKnownConfigs[url] = cfg
 	return cfg, nil
 }
diff --git a/internal/oidc/discovery_test.go b/internal/oidc/discovery_test.go
index c9665783..10bc981a 100644
--- a/internal/oidc/discovery_test.go
+++ b/internal/oidc/discovery_test.go
@@ -53,7 +53,7 @@ func TestWellKnownConfig(t *testing.T) {
 		{"ok", "http://example.com/.well-known/openid-configuration", validWellKnownJSON, false, validWellKnown},
 		{"not-found", "http://example.com/not-found", validWellKnownJSON, true, WellKnownConfig{}},
 		{"invalid-url", "invalid", validWellKnownJSON, true, WellKnownConfig{}},
-		{"invalid-json", "http://example.com/.well-known/openid-configuration", invalidWellKnownJSON, true, WellKnownConfig{}},
+		{"invalid-json", "http://example2.com/.well-known/openid-configuration", invalidWellKnownJSON, true, WellKnownConfig{}},
 	}
 
 	for _, tt := range tests {
@@ -74,6 +74,25 @@ func TestWellKnownConfig(t *testing.T) {
 	}
 }
 
+func TestWellKnownConfigCache(t *testing.T) {
+	s := newServer()
+	s.wellKnownConfig = validWellKnownJSON
+	s.Start()
+	c := s.newHTTPClient()
+
+	got, err := GetWellKnownConfig(c, "http://example.com/.well-known/openid-configuration")
+	require.NoError(t, err)
+	require.Equal(t, validWellKnown, got)
+
+	// Stop the server and run the well-known request again.
+	// It should succeed and return the cached value.
+	s.Stop()
+
+	got, err = GetWellKnownConfig(c, "http://example.com/.well-known/openid-configuration")
+	require.NoError(t, err)
+	require.Equal(t, validWellKnown, got)
+}
+
 type idpServer struct {
 	server          *http.Server
 	listener        *bufconn.Listener