Key-based locking in signature verification

1 files changed, 19 insertions, 8 deletions

diff --git a/src/main/java/com/juick/ActivityPubManager.java b/src/main/java/com/juick/ActivityPubManager.java
index 461f4731..33cc458b 100644
--- a/src/main/java/com/juick/ActivityPubManager.java
+++ b/src/main/java/com/juick/ActivityPubManager.java
@@ -31,10 +31,7 @@ import com.juick.service.activities.*;
 import com.juick.service.component.NotificationListener;
 import com.juick.service.component.PingEvent;
 import com.juick.service.component.SystemEvent;
-import com.juick.util.HttpBadRequestException;
-import com.juick.util.HttpNotFoundException;
-import com.juick.util.HttpUtils;
-import com.juick.util.MessageUtils;
+import com.juick.util.*;
 import com.juick.util.formatters.PlainTextFormatter;
 import com.juick.www.api.SystemActivity.ActivityType;
 import com.juick.www.api.activity.helpers.ProfileUriBuilder;
@@ -67,6 +64,8 @@ import java.security.NoSuchAlgorithmException;
 import java.security.SignatureException;
 import java.time.Instant;
 import java.util.*;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 
 public class ActivityPubManager implements ActivityListener, NotificationListener {
@@ -432,6 +431,8 @@ public class ActivityPubManager implements ActivityListener, NotificationListene
         }
     }
 
+    private static final ReentrantLockMap verificationLock = new ReentrantLockMap();
+
     public User verifyActor(String method, String path, Map<String, String> headers) {
         String signatureString = headers.get("signature");
         if (StringUtils.isNotEmpty(signatureString)) {
@@ -444,10 +445,20 @@ public class ActivityPubManager implements ActivityListener, NotificationListene
                 // local user
                 key = keystoreManager.getPublicKey();
             } else {
-                var context = activityPubService.get(keyId);
-                if (context.isPresent()) {
-                    actor = (Actor) context.get();
-                    key = KeystoreManager.publicKeyOf(actor);
+                ReentrantLock lock = null;
+                try {
+                    lock = verificationLock.getLock(keyId.toASCIIString());
+                    lock.lock();
+                    var context = activityPubService.get(keyId);
+                    if (context.isPresent()) {
+                        actor = (Actor) context.get();
+                        key = KeystoreManager.publicKeyOf(actor);
+                    }
+                } finally {
+                    if (lock != null) {
+                        lock.unlock();
+                        verificationLock.retainLock(keyId.toASCIIString());
+                    }
                 }
             }
             if (key != null) {