From a7b6d567081cc061d3ee5d8d92cd707aaaf6917f Mon Sep 17 00:00:00 2001 From: timedout Date: Sun, 17 May 2026 18:28:04 +0100 Subject: [PATCH 1/3] federation/cache: Add Context parameter to Cache --- federation/cache.go | 39 +++++++++++++++++++++------------------ 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/federation/cache.go b/federation/cache.go index 241549745..6fec7197c 100644 --- a/federation/cache.go +++ b/federation/cache.go @@ -7,6 +7,7 @@ package federation import ( + "context" "errors" "fmt" "math" @@ -16,17 +17,17 @@ import ( // ResolutionCache is an interface for caching resolved server names. type ResolutionCache interface { - StoreResolution(*ResolvedServerName) + StoreResolution(context.Context, *ResolvedServerName) // LoadResolution loads a resolved server name from the cache. // Expired entries MUST NOT be returned. - LoadResolution(serverName string) (*ResolvedServerName, error) + LoadResolution(ctx context.Context, serverName string) (*ResolvedServerName, error) } type KeyCache interface { - StoreKeys(*ServerKeyResponse) - StoreFetchError(serverName string, err error) - ShouldReQuery(serverName string) bool - LoadKeys(serverName string) (*ServerKeyResponse, error) + StoreKeys(context.Context, *ServerKeyResponse) + StoreFetchError(ctx context.Context, serverName string, err error) + ShouldReQuery(ctx context.Context, serverName string) bool + LoadKeys(ctx context.Context, serverName string) (*ServerKeyResponse, error) } type InMemoryCache struct { @@ -55,13 +56,13 @@ func NewInMemoryCache() *InMemoryCache { } } -func (c *InMemoryCache) StoreResolution(resolution *ResolvedServerName) { +func (c *InMemoryCache) StoreResolution(_ context.Context, resolution *ResolvedServerName) { c.resolutionsLock.Lock() defer c.resolutionsLock.Unlock() c.resolutions[resolution.ServerName] = resolution } -func (c *InMemoryCache) LoadResolution(serverName string) (*ResolvedServerName, error) { +func (c *InMemoryCache) LoadResolution(_ context.Context, serverName string) (*ResolvedServerName, error) { c.resolutionsLock.RLock() defer c.resolutionsLock.RUnlock() resolution, ok := c.resolutions[serverName] @@ -71,7 +72,7 @@ func (c *InMemoryCache) LoadResolution(serverName string) (*ResolvedServerName, return resolution, nil } -func (c *InMemoryCache) StoreKeys(keys *ServerKeyResponse) { +func (c *InMemoryCache) StoreKeys(_ context.Context, keys *ServerKeyResponse) { c.keysLock.Lock() defer c.keysLock.Unlock() c.keys[keys.ServerName] = keys @@ -93,7 +94,7 @@ func (rec *resolutionErrorCache) ShouldRetry() bool { var ErrRecentKeyQueryFailed = errors.New("last retry was too recent") -func (c *InMemoryCache) LoadKeys(serverName string) (*ServerKeyResponse, error) { +func (c *InMemoryCache) LoadKeys(_ context.Context, serverName string) (*ServerKeyResponse, error) { c.keysLock.RLock() defer c.keysLock.RUnlock() keys, ok := c.keys[serverName] @@ -112,7 +113,7 @@ func (c *InMemoryCache) LoadKeys(serverName string) (*ServerKeyResponse, error) return keys, nil } -func (c *InMemoryCache) StoreFetchError(serverName string, err error) { +func (c *InMemoryCache) StoreFetchError(_ context.Context, serverName string, err error) { c.keysLock.Lock() defer c.keysLock.Unlock() errorCache, ok := c.lastError[serverName] @@ -125,7 +126,7 @@ func (c *InMemoryCache) StoreFetchError(serverName string, err error) { } } -func (c *InMemoryCache) ShouldReQuery(serverName string) bool { +func (c *InMemoryCache) ShouldReQuery(_ context.Context, serverName string) bool { c.keysLock.Lock() defer c.keysLock.Unlock() lastQuery, ok := c.lastReQueryAt[serverName] @@ -138,12 +139,14 @@ func (c *InMemoryCache) ShouldReQuery(serverName string) bool { type noopCache struct{} -func (*noopCache) StoreKeys(_ *ServerKeyResponse) {} -func (*noopCache) LoadKeys(_ string) (*ServerKeyResponse, error) { return nil, nil } -func (*noopCache) StoreFetchError(_ string, _ error) {} -func (*noopCache) ShouldReQuery(_ string) bool { return true } -func (*noopCache) StoreResolution(_ *ResolvedServerName) {} -func (*noopCache) LoadResolution(_ string) (*ResolvedServerName, error) { return nil, nil } +func (*noopCache) StoreKeys(_ context.Context, _ *ServerKeyResponse) {} +func (*noopCache) LoadKeys(_ context.Context, _ string) (*ServerKeyResponse, error) { return nil, nil } +func (*noopCache) StoreFetchError(_ context.Context, _ string, _ error) {} +func (*noopCache) ShouldReQuery(_ context.Context, _ string) bool { return true } +func (*noopCache) StoreResolution(_ context.Context, _ *ResolvedServerName) {} +func (*noopCache) LoadResolution(_ context.Context, _ string) (*ResolvedServerName, error) { + return nil, nil +} var ( _ ResolutionCache = (*noopCache)(nil) From 064f629b9d945a87971800703fd0fdc5e1a63050 Mon Sep 17 00:00:00 2001 From: timedout Date: Sun, 17 May 2026 18:30:32 +0100 Subject: [PATCH 2/3] Fix build errors --- federation/httpclient.go | 4 ++-- federation/keyserver.go | 2 +- federation/serverauth.go | 10 +++++----- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/federation/httpclient.go b/federation/httpclient.go index 8375aab9f..2a70ae7dd 100644 --- a/federation/httpclient.go +++ b/federation/httpclient.go @@ -80,7 +80,7 @@ func (srt *ServerResolvingTransport) resolve(ctx context.Context, serverName str lock.Lock() defer lock.Unlock() - res, err := srt.cache.LoadResolution(serverName) + res, err := srt.cache.LoadResolution(ctx, serverName) if err != nil { return nil, fmt.Errorf("failed to read cache: %w", err) } else if res != nil { @@ -88,7 +88,7 @@ func (srt *ServerResolvingTransport) resolve(ctx context.Context, serverName str } else if res, err = ResolveServerName(ctx, serverName, srt.ResolveOpts); err != nil { return nil, err } else { - srt.cache.StoreResolution(res) + srt.cache.StoreResolution(ctx, res) return res, nil } } diff --git a/federation/keyserver.go b/federation/keyserver.go index d32ba5cf9..dd4b4fb5e 100644 --- a/federation/keyserver.go +++ b/federation/keyserver.go @@ -186,7 +186,7 @@ func (ks *KeyServer) GetQueryKeys(w http.ResponseWriter, r *http.Request) { resp.ServerKeys = append(resp.ServerKeys, key.GenerateKeyResponse(serverName, nil)) } } else if ks.OtherKeys != nil { - otherKey, err := ks.OtherKeys.LoadKeys(serverName) + otherKey, err := ks.OtherKeys.LoadKeys(r.Context(), serverName) if err != nil { mautrix.MUnknown.WithMessage("Failed to load keys from cache").Write(w) return diff --git a/federation/serverauth.go b/federation/serverauth.go index cd3003416..0bc06d388 100644 --- a/federation/serverauth.go +++ b/federation/serverauth.go @@ -103,7 +103,7 @@ func ParseXMatrixAuth(auth string) (xma XMatrixAuth) { } func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, keyID id.KeyID) (*ServerKeyResponse, error) { - res, err := sa.Keys.LoadKeys(serverName) + res, err := sa.Keys.LoadKeys(ctx, serverName) if err != nil { return nil, fmt.Errorf("failed to read cache: %w", err) } else if res.HasKey(keyID) { @@ -120,13 +120,13 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k lock.Lock() defer lock.Unlock() - res, err = sa.Keys.LoadKeys(serverName) + res, err = sa.Keys.LoadKeys(ctx, serverName) if err != nil { return nil, fmt.Errorf("failed to read cache: %w", err) } else if res != nil { if res.HasKey(keyID) { return res, nil - } else if !sa.Keys.ShouldReQuery(serverName) { + } else if !sa.Keys.ShouldReQuery(ctx, serverName) { zerolog.Ctx(ctx).Trace(). Str("server_name", serverName). Stringer("key_id", keyID). @@ -136,10 +136,10 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k } res, err = sa.Client.ServerKeys(ctx, serverName) if err != nil { - sa.Keys.StoreFetchError(serverName, err) + sa.Keys.StoreFetchError(ctx, serverName, err) return nil, err } - sa.Keys.StoreKeys(res) + sa.Keys.StoreKeys(ctx, res) return res, nil } From 17b1e8dbe0b1a348740fc74a9f291ac4faf93edc Mon Sep 17 00:00:00 2001 From: timedout Date: Sun, 17 May 2026 18:34:58 +0100 Subject: [PATCH 3/3] Allow returning an error in StoreResolution, StoreKeys, and ShouldReQuery --- federation/cache.go | 24 +++++++++++++----------- federation/serverauth.go | 7 +++++-- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/federation/cache.go b/federation/cache.go index 6fec7197c..c26786726 100644 --- a/federation/cache.go +++ b/federation/cache.go @@ -17,16 +17,16 @@ import ( // ResolutionCache is an interface for caching resolved server names. type ResolutionCache interface { - StoreResolution(context.Context, *ResolvedServerName) + StoreResolution(context.Context, *ResolvedServerName) error // LoadResolution loads a resolved server name from the cache. // Expired entries MUST NOT be returned. LoadResolution(ctx context.Context, serverName string) (*ResolvedServerName, error) } type KeyCache interface { - StoreKeys(context.Context, *ServerKeyResponse) + StoreKeys(context.Context, *ServerKeyResponse) error StoreFetchError(ctx context.Context, serverName string, err error) - ShouldReQuery(ctx context.Context, serverName string) bool + ShouldReQuery(ctx context.Context, serverName string) (bool, error) LoadKeys(ctx context.Context, serverName string) (*ServerKeyResponse, error) } @@ -56,10 +56,11 @@ func NewInMemoryCache() *InMemoryCache { } } -func (c *InMemoryCache) StoreResolution(_ context.Context, resolution *ResolvedServerName) { +func (c *InMemoryCache) StoreResolution(_ context.Context, resolution *ResolvedServerName) error { c.resolutionsLock.Lock() defer c.resolutionsLock.Unlock() c.resolutions[resolution.ServerName] = resolution + return nil } func (c *InMemoryCache) LoadResolution(_ context.Context, serverName string) (*ResolvedServerName, error) { @@ -72,11 +73,12 @@ func (c *InMemoryCache) LoadResolution(_ context.Context, serverName string) (*R return resolution, nil } -func (c *InMemoryCache) StoreKeys(_ context.Context, keys *ServerKeyResponse) { +func (c *InMemoryCache) StoreKeys(_ context.Context, keys *ServerKeyResponse) error { c.keysLock.Lock() defer c.keysLock.Unlock() c.keys[keys.ServerName] = keys delete(c.lastError, keys.ServerName) + return nil } type resolutionErrorCache struct { @@ -126,24 +128,24 @@ func (c *InMemoryCache) StoreFetchError(_ context.Context, serverName string, er } } -func (c *InMemoryCache) ShouldReQuery(_ context.Context, serverName string) bool { +func (c *InMemoryCache) ShouldReQuery(_ context.Context, serverName string) (bool, error) { c.keysLock.Lock() defer c.keysLock.Unlock() lastQuery, ok := c.lastReQueryAt[serverName] if ok && time.Since(lastQuery) < c.MinKeyRefetchDelay { - return false + return false, nil } c.lastReQueryAt[serverName] = time.Now() - return true + return true, nil } type noopCache struct{} -func (*noopCache) StoreKeys(_ context.Context, _ *ServerKeyResponse) {} +func (*noopCache) StoreKeys(_ context.Context, _ *ServerKeyResponse) error { return nil } func (*noopCache) LoadKeys(_ context.Context, _ string) (*ServerKeyResponse, error) { return nil, nil } func (*noopCache) StoreFetchError(_ context.Context, _ string, _ error) {} -func (*noopCache) ShouldReQuery(_ context.Context, _ string) bool { return true } -func (*noopCache) StoreResolution(_ context.Context, _ *ResolvedServerName) {} +func (*noopCache) ShouldReQuery(_ context.Context, _ string) (bool, error) { return true, nil } +func (*noopCache) StoreResolution(_ context.Context, _ *ResolvedServerName) error { return nil } func (*noopCache) LoadResolution(_ context.Context, _ string) (*ResolvedServerName, error) { return nil, nil } diff --git a/federation/serverauth.go b/federation/serverauth.go index 0bc06d388..4d6b3f003 100644 --- a/federation/serverauth.go +++ b/federation/serverauth.go @@ -126,8 +126,9 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k } else if res != nil { if res.HasKey(keyID) { return res, nil - } else if !sa.Keys.ShouldReQuery(ctx, serverName) { + } else if shouldQuery, err := sa.Keys.ShouldReQuery(ctx, serverName); !shouldQuery { zerolog.Ctx(ctx).Trace(). + Err(err). Str("server_name", serverName). Stringer("key_id", keyID). Msg("Not sending key request for missing key ID, last query was too recent") @@ -139,7 +140,9 @@ func (sa *ServerAuth) GetKeysWithCache(ctx context.Context, serverName string, k sa.Keys.StoreFetchError(ctx, serverName, err) return nil, err } - sa.Keys.StoreKeys(ctx, res) + if err := sa.Keys.StoreKeys(ctx, res); err != nil { + return nil, err + } return res, nil }