From e614269bc3b0ff11b9b7969691bc251b9662f0a4 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 14 Jun 2022 14:08:09 +0100
Subject: [PATCH] More strict keyable interface

---
 internal/caching/cache_serverkeys.go | 10 +++++++---
 internal/caching/caches.go           |  5 +++--
 internal/caching/impl_ristretto.go   |  2 +-
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/internal/caching/cache_serverkeys.go b/internal/caching/cache_serverkeys.go
index acab31bd2..836272502 100644
--- a/internal/caching/cache_serverkeys.go
+++ b/internal/caching/cache_serverkeys.go
@@ -1,6 +1,8 @@
 package caching
 
 import (
+	"fmt"
+
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
@@ -30,11 +32,12 @@ func (c Caches) GetServerKey(
 	request gomatrixserverlib.PublicKeyLookupRequest,
 	timestamp gomatrixserverlib.Timestamp,
 ) (gomatrixserverlib.PublicKeyLookupResult, bool) {
-	val, found := c.ServerKeys.Get(request)
+	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
+	val, found := c.ServerKeys.Get(key)
 	if found && !val.WasValidAt(timestamp, true) {
 		// The key wasn't valid at the requested timestamp so don't
 		// return it. The caller will have to work out what to do.
-		c.ServerKeys.Unset(request)
+		c.ServerKeys.Unset(key)
 		return gomatrixserverlib.PublicKeyLookupResult{}, false
 	}
 	return val, found
@@ -44,5 +47,6 @@ func (c Caches) StoreServerKey(
 	request gomatrixserverlib.PublicKeyLookupRequest,
 	response gomatrixserverlib.PublicKeyLookupResult,
 ) {
-	c.ServerKeys.Set(request, response)
+	key := fmt.Sprintf("%s/%s", request.ServerName, request.KeyID)
+	c.ServerKeys.Set(key, response)
 }
diff --git a/internal/caching/caches.go b/internal/caching/caches.go
index 624c8ffca..4d94e7321 100644
--- a/internal/caching/caches.go
+++ b/internal/caching/caches.go
@@ -12,7 +12,7 @@ import (
 // interface.
 type Caches struct {
 	RoomVersions       Cache[string, gomatrixserverlib.RoomVersion]
-	ServerKeys         Cache[gomatrixserverlib.PublicKeyLookupRequest, gomatrixserverlib.PublicKeyLookupResult]
+	ServerKeys         Cache[string, gomatrixserverlib.PublicKeyLookupResult]
 	RoomServerRoomNIDs Cache[string, types.RoomNID]
 	RoomServerRoomIDs  Cache[types.RoomNID, string]
 	RoomInfos          Cache[string, types.RoomInfo]
@@ -30,7 +30,8 @@ type Cache[K keyable, T any] interface {
 }
 
 type keyable interface {
-	comparable
+	// from https://github.com/dgraph-io/ristretto/blob/8e850b710d6df0383c375ec6a7beae4ce48fc8d5/z/z.go#L34
+	~uint64 | ~string | ~[]byte | ~byte | ~int | ~int32 | ~uint32 | ~int64
 }
 
 type costable interface {
diff --git a/internal/caching/impl_ristretto.go b/internal/caching/impl_ristretto.go
index 3d2d05bea..291b07335 100644
--- a/internal/caching/impl_ristretto.go
+++ b/internal/caching/impl_ristretto.go
@@ -33,7 +33,7 @@ func NewRistrettoCache(maxCost CacheSize, enablePrometheus bool) (*Caches, error
 			cache: cache,
 			Name:  "room_versions",
 		},
-		ServerKeys: &RistrettoCachePartition[gomatrixserverlib.PublicKeyLookupRequest, gomatrixserverlib.PublicKeyLookupResult]{
+		ServerKeys: &RistrettoCachePartition[string, gomatrixserverlib.PublicKeyLookupResult]{
 			cache:   cache,
 			Name:    "server_keys",
 			Mutable: true,