From c02832a3e92569f64f180dec1555056dc8f8c3e3 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 26 Apr 2022 09:58:06 +0100
Subject: [PATCH] Refactor `QuerySignatures`

---
 keyserver/internal/cross_signing.go | 46 ++++++++++++++---------------
 1 file changed, 23 insertions(+), 23 deletions(-)

diff --git a/keyserver/internal/cross_signing.go b/keyserver/internal/cross_signing.go
index 08bbfedb8..2cdcff8bd 100644
--- a/keyserver/internal/cross_signing.go
+++ b/keyserver/internal/cross_signing.go
@@ -521,7 +521,7 @@ func (a *KeyInternalAPI) crossSigningKeysFromDatabase(
 }
 
 func (a *KeyInternalAPI) QuerySignatures(ctx context.Context, req *api.QuerySignaturesRequest, res *api.QuerySignaturesResponse) {
-	for targetUserID, forTargetUser := range req.TargetIDs {
+	for targetUserID := range req.TargetIDs {
 		keyMap, err := a.DB.CrossSigningKeysForUser(ctx, targetUserID)
 		if err != nil && err != sql.ErrNoRows {
 			res.Error = &api.KeyError{
@@ -552,31 +552,31 @@ func (a *KeyInternalAPI) QuerySignatures(ctx context.Context, req *api.QuerySign
 			}
 		}
 
-		for _, targetKeyID := range forTargetUser {
-			// Get own signatures only.
-			sigMap, err := a.DB.CrossSigningSigsForTarget(ctx, targetUserID, targetUserID, targetKeyID)
-			if err != nil && err != sql.ErrNoRows {
-				res.Error = &api.KeyError{
-					Err: fmt.Sprintf("a.DB.CrossSigningSigsForTarget: %s", err),
+		for _, targetSection := range []map[string]gomatrixserverlib.CrossSigningKey{
+			res.MasterKeys,
+			res.SelfSigningKeys,
+			res.UserSigningKeys,
+		} {
+			for targetKeyID, section := range targetSection {
+				// Get own signatures only.
+				sigMap, err := a.DB.CrossSigningSigsForTarget(ctx, targetUserID, targetUserID, gomatrixserverlib.KeyID(targetKeyID))
+				if err != nil && err != sql.ErrNoRows {
+					res.Error = &api.KeyError{
+						Err: fmt.Sprintf("a.DB.CrossSigningSigsForTarget: %s", err),
+					}
+					return
 				}
-				return
-			}
 
-			for sourceUserID, forSourceUser := range sigMap {
-				for sourceKeyID, sourceSig := range forSourceUser {
-					if res.Signatures == nil {
-						res.Signatures = map[string]map[gomatrixserverlib.KeyID]types.CrossSigningSigMap{}
+				for sourceUserID, forSourceUser := range sigMap {
+					for sourceKeyID, sourceSig := range forSourceUser {
+						if section.Signatures == nil {
+							section.Signatures = map[string]map[gomatrixserverlib.KeyID]gomatrixserverlib.Base64Bytes{}
+						}
+						if _, ok := res.Signatures[sourceUserID]; !ok {
+							section.Signatures[sourceUserID] = map[gomatrixserverlib.KeyID]gomatrixserverlib.Base64Bytes{}
+						}
+						section.Signatures[sourceUserID][sourceKeyID] = sourceSig
 					}
-					if _, ok := res.Signatures[targetUserID]; !ok {
-						res.Signatures[targetUserID] = map[gomatrixserverlib.KeyID]types.CrossSigningSigMap{}
-					}
-					if _, ok := res.Signatures[targetUserID][targetKeyID]; !ok {
-						res.Signatures[targetUserID][targetKeyID] = types.CrossSigningSigMap{}
-					}
-					if _, ok := res.Signatures[targetUserID][targetKeyID][sourceUserID]; !ok {
-						res.Signatures[targetUserID][targetKeyID][sourceUserID] = map[gomatrixserverlib.KeyID]gomatrixserverlib.Base64Bytes{}
-					}
-					res.Signatures[targetUserID][targetKeyID][sourceUserID][sourceKeyID] = sourceSig
 				}
 			}
 		}