From 8b0795dd34acb1a22811b78acb63316acc290a52 Mon Sep 17 00:00:00 2001
From: aeneasr <3372410+aeneasr@users.noreply.github.com>
Date: Fri, 5 Jan 2024 09:49:10 +0100
Subject: [PATCH] fix: do not read outdated trust relations

---
 persistence/sql/persister_grant_jwk.go |  5 ++-
 persistence/sql/persister_nid_test.go  | 45 ++++++++++++++++++--------
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/persistence/sql/persister_grant_jwk.go b/persistence/sql/persister_grant_jwk.go
index 7341cefe001..ac3c8b6f880 100644
--- a/persistence/sql/persister_grant_jwk.go
+++ b/persistence/sql/persister_grant_jwk.go
@@ -135,9 +135,12 @@ func (p *Persister) GetPublicKeys(ctx context.Context, issuer string, subject st
 
 	grantsData := make([]trust.SQLData, 0)
 	query := p.QueryWithNetwork(ctx).
+		Select("key_set", "key_id").
+		Where("expires_at > NOW()").
 		Where("issuer = ?", issuer).
 		Where("(subject = ? OR allow_any_subject IS TRUE)", subject).
-		Where("nid = ?", p.NetworkID(ctx))
+		Order("created_at DESC").
+		Limit(100) // Load maximum of 100 keys
 
 	if err := query.All(&grantsData); err != nil {
 		return nil, sqlcon.HandleError(err)
diff --git a/persistence/sql/persister_nid_test.go b/persistence/sql/persister_nid_test.go
index 2ca0f672be2..0d778fef44f 100644
--- a/persistence/sql/persister_nid_test.go
+++ b/persistence/sql/persister_nid_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
+	"github.com/ory/x/sqlcon"
 	"testing"
 	"time"
 
@@ -1341,22 +1342,38 @@ func (s *PersisterTestSuite) TestGetPublicKeys() {
 	t := s.T()
 	for k, r := range s.registries {
 		t.Run(k, func(t *testing.T) {
-			ks := newKeySet("ks-id", "use")
-			grant := trust.Grant{
-				ID:        uuid.Must(uuid.NewV4()).String(),
-				ExpiresAt: time.Now().Add(time.Hour),
-				PublicKey: trust.PublicKey{Set: "ks-id", KeyID: ks.Keys[0].KeyID},
-			}
-			require.NoError(t, r.Persister().AddKeySet(s.t1, "ks-id", ks))
-			require.NoError(t, r.Persister().CreateGrant(s.t1, grant, ks.Keys[0]))
+			t.Run("get key", func(t *testing.T) {
+				ks := newKeySet("ks-id", "use")
+				grant := trust.Grant{
+					ID:        uuid.Must(uuid.NewV4()).String(),
+					ExpiresAt: time.Now().Add(time.Hour),
+					PublicKey: trust.PublicKey{Set: "ks-id", KeyID: ks.Keys[0].KeyID},
+				}
+				require.NoError(t, r.Persister().AddKeySet(s.t1, "ks-id", ks))
+				require.NoError(t, r.Persister().CreateGrant(s.t1, grant, ks.Keys[0]))
 
-			actual, err := r.Persister().GetPublicKeys(s.t2, grant.Issuer, grant.Subject)
-			require.NoError(t, err)
-			require.Nil(t, actual.Keys)
+				actual, err := r.Persister().GetPublicKeys(s.t2, grant.Issuer, grant.Subject)
+				require.NoError(t, err)
+				require.Nil(t, actual.Keys)
 
-			actual, err = r.Persister().GetPublicKeys(s.t1, grant.Issuer, grant.Subject)
-			require.NoError(t, err)
-			require.NotNil(t, actual.Keys)
+				actual, err = r.Persister().GetPublicKeys(s.t1, grant.Issuer, grant.Subject)
+				require.NoError(t, err)
+				require.NotNil(t, actual.Keys)
+			})
+
+			t.Run("get expired key fails", func(t *testing.T) {
+				ks := newKeySet("ks-id", "use")
+				grant := trust.Grant{
+					ID:        uuid.Must(uuid.NewV4()).String(),
+					ExpiresAt: time.Now().Add(-time.Hour),
+					PublicKey: trust.PublicKey{Set: "ks-id", KeyID: ks.Keys[0].KeyID},
+				}
+				require.NoError(t, r.Persister().AddKeySet(s.t1, "ks-id", ks))
+				require.NoError(t, r.Persister().CreateGrant(s.t1, grant, ks.Keys[0]))
+
+				_, err := r.Persister().GetPublicKeys(s.t2, grant.Issuer, grant.Subject)
+				require.ErrorIs(t, err, sqlcon.ErrNoRows)
+			})
 		})
 	}
 }