Change HardwareSecurity DecryptionCallback to emit key-id

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/ConsumerOptions.java

@@ -49,7 +49,7 @@ public class ConsumerOptions {
 
     // Session key for decryption without passphrase/key
     private SessionKey sessionKey = null;
-    private Map<Set<Long>, PublicKeyDataDecryptorFactory> customPublicKeyDataDecryptorFactories = new HashMap<>();
+    private Map<Long, PublicKeyDataDecryptorFactory> customPublicKeyDataDecryptorFactories = new HashMap<>();
 
     private final Map<PGPSecretKeyRing, SecretKeyRingProtector> decryptionKeys = new HashMap<>();
     private final Set<Passphrase> decryptionPassphrases = new HashSet<>();
@@ -245,20 +245,15 @@ public class ConsumerOptions {
      * hardware-backed secret keys.
      * (See e.g. {@link org.pgpainless.decryption_verification.HardwareSecurity.HardwareDataDecryptorFactory}).
      *
-     * The set of key-ids determines, whether the decryptor factory shall be consulted to decrypt a given session key.
+     * @param factory decryptor factory
-     * See for example {@link HardwareSecurity#getIdsOfHardwareBackedKeys(PGPSecretKeyRing)}.
-     *
-     * @param keyIds set of key-ids for which the factory shall be consulted
-     * @param decryptorFactory decryptor factory
      * @return options
      */
-    public ConsumerOptions addCustomDecryptorFactory(
+    public ConsumerOptions addCustomDecryptorFactory(@Nonnull CustomPublicKeyDataDecryptorFactory factory) {
-            @Nonnull Set<Long> keyIds, @Nonnull PublicKeyDataDecryptorFactory decryptorFactory) {
+        this.customPublicKeyDataDecryptorFactories.put(factory.getKeyId(), factory);
-        this.customPublicKeyDataDecryptorFactories.put(keyIds, decryptorFactory);
         return this;
     }
 
-    Map<Set<Long>, PublicKeyDataDecryptorFactory> getCustomDecryptorFactories() {
+    Map<Long, PublicKeyDataDecryptorFactory> getCustomDecryptorFactories() {
         return new HashMap<>(customPublicKeyDataDecryptorFactories);
     }
 

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/CustomPublicKeyDataDecryptorFactory.java

@@ -0,0 +1,13 @@
+// SPDX-FileCopyrightText: 2022 Paul Schaub <vanitasvitae@fsfe.org>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package org.pgpainless.decryption_verification;
+
+import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
+
+public interface CustomPublicKeyDataDecryptorFactory extends PublicKeyDataDecryptorFactory {
+
+    long getKeyId();
+
+}

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/DecryptionStreamFactory.java

@@ -409,30 +409,28 @@ public final class DecryptionStreamFactory {
         }
 
         // Try custom PublicKeyDataDecryptorFactories (e.g. hardware-backed).
-        Map<Set<Long>, PublicKeyDataDecryptorFactory> customFactories = options.getCustomDecryptorFactories();
+        Map<Long, PublicKeyDataDecryptorFactory> customFactories = options.getCustomDecryptorFactories();
         for (PGPPublicKeyEncryptedData publicKeyEncryptedData : publicKeyProtected) {
             Long keyId = publicKeyEncryptedData.getKeyID();
-            for (Set<Long> keyIds : customFactories.keySet()) {
+            if (!customFactories.containsKey(keyId)) {
-                if (!keyIds.contains(keyId)) {
+                continue;
-                    continue;
+            }
-                }
 
-                PublicKeyDataDecryptorFactory decryptorFactory = customFactories.get(keyIds);
+            PublicKeyDataDecryptorFactory decryptorFactory = customFactories.get(keyId);
-                try {
+            try {
-                    InputStream decryptedDataStream = publicKeyEncryptedData.getDataStream(decryptorFactory);
+                InputStream decryptedDataStream = publicKeyEncryptedData.getDataStream(decryptorFactory);
-                    PGPSessionKey pgpSessionKey = publicKeyEncryptedData.getSessionKey(decryptorFactory);
+                PGPSessionKey pgpSessionKey = publicKeyEncryptedData.getSessionKey(decryptorFactory);
-                    SessionKey sessionKey = new SessionKey(pgpSessionKey);
+                SessionKey sessionKey = new SessionKey(pgpSessionKey);
-                    resultBuilder.setSessionKey(sessionKey);
+                resultBuilder.setSessionKey(sessionKey);
 
-                    throwIfAlgorithmIsRejected(sessionKey.getAlgorithm());
+                throwIfAlgorithmIsRejected(sessionKey.getAlgorithm());
 
-                    integrityProtectedEncryptedInputStream =
+                integrityProtectedEncryptedInputStream =
-                            new IntegrityProtectedInputStream(decryptedDataStream, publicKeyEncryptedData, options);
+                        new IntegrityProtectedInputStream(decryptedDataStream, publicKeyEncryptedData, options);
 
-                    return integrityProtectedEncryptedInputStream;
+                return integrityProtectedEncryptedInputStream;
-                } catch (PGPException e) {
+            } catch (PGPException e) {
-                    LOGGER.debug("Decryption with custom PublicKeyDataDecryptorFactory failed", e);
+                LOGGER.debug("Decryption with custom PublicKeyDataDecryptorFactory failed", e);
-                }
             }
         }
 

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/HardwareSecurity.java

@@ -28,13 +28,14 @@ public class HardwareSecurity {
          *
          * If decryption fails for some reason, a subclass of the {@link HardwareSecurityException} is thrown.
          *
+         * @param keyId id of the key
          * @param keyAlgorithm algorithm
          * @param sessionKeyData encrypted session key
          *
          * @return decrypted session key
          * @throws HardwareSecurityException exception
          */
-        byte[] decryptSessionKey(int keyAlgorithm, byte[] sessionKeyData)
+        byte[] decryptSessionKey(long keyId, int keyAlgorithm, byte[] sessionKeyData)
                 throws HardwareSecurityException;
 
     }
@@ -67,20 +68,22 @@ public class HardwareSecurity {
      * to a {@link DecryptionCallback}.
      * Users can provide such a callback to delegate decryption of messages to hardware security SDKs.
      */
-    public static class HardwareDataDecryptorFactory implements PublicKeyDataDecryptorFactory {
+    public static class HardwareDataDecryptorFactory implements CustomPublicKeyDataDecryptorFactory {
 
         private final DecryptionCallback callback;
         // luckily we can instantiate the BcPublicKeyDataDecryptorFactory with null as argument.
         private final PublicKeyDataDecryptorFactory factory =
                 new BcPublicKeyDataDecryptorFactory(null);
+        private long keyId;
 
         /**
          * Create a new {@link HardwareDataDecryptorFactory}.
          *
          * @param callback decryption callback
          */
-        public HardwareDataDecryptorFactory(DecryptionCallback callback) {
+        public HardwareDataDecryptorFactory(long keyId, DecryptionCallback callback) {
             this.callback = callback;
+            this.keyId = keyId;
         }
 
         @Override
@@ -88,7 +91,7 @@ public class HardwareSecurity {
                 throws PGPException {
             try {
                 // delegate decryption to the callback
-                return callback.decryptSessionKey(keyAlgorithm, secKeyData[0]);
+                return callback.decryptSessionKey(keyId, keyAlgorithm, secKeyData[0]);
             } catch (HardwareSecurityException e) {
                 throw new PGPException("Hardware-backed decryption failed.", e);
             }
@@ -99,10 +102,16 @@ public class HardwareSecurity {
                 throws PGPException {
             return factory.createDataDecryptor(withIntegrityPacket, encAlgorithm, key);
         }
+
+        @Override
+        public long getKeyId() {
+            return keyId;
+        }
     }
 
     public static class HardwareSecurityException
             extends Exception {
 
     }
+
 }

pgpainless-core/src/test/java/org/pgpainless/decryption_verification/CustomPublicKeyDataDecryptorFactoryTest.java

@@ -29,7 +29,6 @@ import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.NoSuchAlgorithmException;
-import java.util.Collections;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
@@ -55,7 +54,7 @@ public class CustomPublicKeyDataDecryptorFactoryTest {
 
         HardwareSecurity.DecryptionCallback hardwareDecryptionCallback = new HardwareSecurity.DecryptionCallback() {
             @Override
-            public byte[] decryptSessionKey(int keyAlgorithm, byte[] sessionKeyData)
+            public byte[] decryptSessionKey(long keyId, int keyAlgorithm, byte[] sessionKeyData)
                     throws HardwareSecurity.HardwareSecurityException {
                 // Emulate hardware decryption.
                 try {
@@ -74,8 +73,9 @@ public class CustomPublicKeyDataDecryptorFactoryTest {
                 .onInputStream(new ByteArrayInputStream(ciphertextOut.toByteArray()))
                 .withOptions(ConsumerOptions.get()
                         .addCustomDecryptorFactory(
-                                Collections.singleton(encryptionKey.getKeyID()),
+                                new HardwareSecurity.HardwareDataDecryptorFactory(
-                                new HardwareSecurity.HardwareDataDecryptorFactory(hardwareDecryptionCallback)));
+                                        encryptionKey.getKeyID(),
+                                        hardwareDecryptionCallback)));
 
         ByteArrayOutputStream decryptedOut = new ByteArrayOutputStream();
         Streams.pipeAll(decryptionStream, decryptedOut);