From 984c174ef697c8756841399aabc049d84d6c0aae Mon Sep 17 00:00:00 2001
From: Paul Schaub <vanitasvitae@fsfe.org>
Date: Fri, 23 Sep 2022 14:51:06 +0200
Subject: [PATCH] Implement caching PublicKeyDataDecryptorFactory

---
 .../CachingPublicKeyDataDecryptorFactory.java | 75 +++++++++++++++++++
 .../java/org/bouncycastle/package-info.java   |  8 ++
 2 files changed, 83 insertions(+)
 create mode 100644 pgpainless-core/src/main/java/org/bouncycastle/CachingPublicKeyDataDecryptorFactory.java
 create mode 100644 pgpainless-core/src/main/java/org/bouncycastle/package-info.java

diff --git a/pgpainless-core/src/main/java/org/bouncycastle/CachingPublicKeyDataDecryptorFactory.java b/pgpainless-core/src/main/java/org/bouncycastle/CachingPublicKeyDataDecryptorFactory.java
new file mode 100644
index 00000000..1498b6f2
--- /dev/null
+++ b/pgpainless-core/src/main/java/org/bouncycastle/CachingPublicKeyDataDecryptorFactory.java
@@ -0,0 +1,75 @@
+// SPDX-FileCopyrightText: 2022 Paul Schaub <vanitasvitae@fsfe.org>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package org.bouncycastle;
+
+import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.operator.PGPDataDecryptor;
+import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
+import org.bouncycastle.util.encoders.Base64;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Implementation of the {@link PublicKeyDataDecryptorFactory} which caches decrypted session keys.
+ * That way, if a message needs to be decrypted multiple times, expensive private key operations can be omitted.
+ *
+ * This implementation changes the behavior or {@link #recoverSessionData(int, byte[][])} to first return any
+ * cache hits.
+ * If no hit is found, the method call is delegated to the underlying {@link PublicKeyDataDecryptorFactory}.
+ * The result of that is then placed in the cache and returned.
+ *
+ * TODO: Do we also cache invalid session keys?
+ */
+public class CachingPublicKeyDataDecryptorFactory implements PublicKeyDataDecryptorFactory {
+
+    private final Map<String, byte[]> cachedSessionKeys = new HashMap<>();
+    private final PublicKeyDataDecryptorFactory factory;
+
+    public CachingPublicKeyDataDecryptorFactory(PublicKeyDataDecryptorFactory factory) {
+        this.factory = factory;
+    }
+
+    @Override
+    public byte[] recoverSessionData(int keyAlgorithm, byte[][] secKeyData) throws PGPException {
+        byte[] sessionKey = lookup(secKeyData);
+        if (sessionKey == null) {
+            sessionKey = factory.recoverSessionData(keyAlgorithm, secKeyData);
+            cache(secKeyData, sessionKey);
+        }
+        return sessionKey;
+    }
+
+    private byte[] lookup(byte[][] secKeyData) {
+        byte[] sk = secKeyData[0];
+        String key = Base64.toBase64String(sk);
+        byte[] sessionKey = cachedSessionKeys.get(key);
+        return copy(sessionKey);
+    }
+
+    private void cache(byte[][] secKeyData, byte[] sessionKey) {
+        byte[] sk = secKeyData[0];
+        String key = Base64.toBase64String(sk);
+        cachedSessionKeys.put(key, copy(sessionKey));
+    }
+
+    private static byte[] copy(byte[] bytes) {
+        if (bytes == null) {
+            return null;
+        }
+        byte[] copy = new byte[bytes.length];
+        System.arraycopy(bytes, 0, copy, 0, copy.length);
+        return copy;
+    }
+
+    public void clear() {
+        cachedSessionKeys.clear();
+    }
+
+    @Override
+    public PGPDataDecryptor createDataDecryptor(boolean withIntegrityPacket, int encAlgorithm, byte[] key) throws PGPException {
+        return null;
+    }
+}
diff --git a/pgpainless-core/src/main/java/org/bouncycastle/package-info.java b/pgpainless-core/src/main/java/org/bouncycastle/package-info.java
new file mode 100644
index 00000000..565bb5f4
--- /dev/null
+++ b/pgpainless-core/src/main/java/org/bouncycastle/package-info.java
@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: 2022 Paul Schaub <vanitasvitae@fsfe.org>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Classes which could be upstreamed to BC at some point.
+ */
+package org.bouncycastle;