Key-based locking in signature verification

2 files changed, 66 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) {
diff --git a/src/main/java/com/juick/util/ReentrantLockMap.java b/src/main/java/com/juick/util/ReentrantLockMap.java
new file mode 100644
index 00000000..36aee218
--- /dev/null
+++ b/src/main/java/com/juick/util/ReentrantLockMap.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2008-2023, Juick
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.juick.util;
+
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
+
+public class ReentrantLockMap {
+
+    ConcurrentHashMap<String, ReentrantLock> mapStringLock;
+
+    public ReentrantLockMap() {
+        mapStringLock = new ConcurrentHashMap<>();
+    }
+
+    public ReentrantLock getLock(String key) {
+        ReentrantLock initValue = createIntanceLock();
+        ReentrantLock lock = mapStringLock.putIfAbsent(key, initValue);
+        if (lock == null) {
+            lock = initValue;
+        }
+        return lock;
+    }
+
+    public void retainLock(String key) {
+        mapStringLock.remove(key);
+    }
+
+    protected ReentrantLock createIntanceLock() {
+        return new ReentrantLock();
+    }
+}