From 4878d8246340b35dba44f9682eaeca401c06e0ea Mon Sep 17 00:00:00 2001
From: Laura Hausmann <laura@hausmann.dev>
Date: Sat, 16 Mar 2024 17:17:43 +0100
Subject: [PATCH] [backend/federation] Fix sporadic key fetch failures
 (ISH-194)

---
 .../Middleware/AuthorizedFetchMiddleware.cs    |  3 +++
 Iceshrimp.Backend/Core/Services/UserService.cs | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/Iceshrimp.Backend/Core/Middleware/AuthorizedFetchMiddleware.cs b/Iceshrimp.Backend/Core/Middleware/AuthorizedFetchMiddleware.cs
index d38506f2..9d774f1d 100644
--- a/Iceshrimp.Backend/Core/Middleware/AuthorizedFetchMiddleware.cs
+++ b/Iceshrimp.Backend/Core/Middleware/AuthorizedFetchMiddleware.cs
@@ -67,6 +67,9 @@ public class AuthorizedFetchMiddleware(
 						var user = await userResolver.ResolveAsync(sig.KeyId).WaitAsync(ct);
 						key = await db.UserPublickeys.Include(p => p.User)
 						              .FirstOrDefaultAsync(p => p.User == user, ct);
+
+						// If the key is still null here, we have a data consistency issue and need to update the key manually
+						key ??= await userSvc.UpdateUserPublicKeyAsync(user).WaitAsync(ct);
 					}
 					catch (Exception e)
 					{
diff --git a/Iceshrimp.Backend/Core/Services/UserService.cs b/Iceshrimp.Backend/Core/Services/UserService.cs
index 07991682..42cb4a59 100644
--- a/Iceshrimp.Backend/Core/Services/UserService.cs
+++ b/Iceshrimp.Backend/Core/Services/UserService.cs
@@ -432,6 +432,24 @@ public class UserService(
 		return key;
 	}
 
+	public async Task<UserPublickey> UpdateUserPublicKeyAsync(User user)
+	{
+		var uri   = user.Uri ?? throw new Exception("Can't update public key of user without Uri");
+		var actor = await fetchSvc.FetchActorAsync(uri);
+
+		if (actor.PublicKey?.PublicKey == null)
+			throw new Exception("Failed to update user public key: Invalid or missing public key");
+
+		var key = await db.UserPublickeys.FirstOrDefaultAsync(p => p.User == user) ?? new UserPublickey { User = user };
+
+		key.KeyId  = actor.PublicKey.Id;
+		key.KeyPem = actor.PublicKey.PublicKey;
+
+		db.Update(key);
+		await db.SaveChangesAsync();
+		return key;
+	}
+
 	public async Task DeleteUserAsync(ASActor actor)
 	{
 		var user = await db.Users