1 files changed, 21 insertions, 12 deletions
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
index 416f4bb4..3ebdb44c 100644
--- a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
@@ -10,12 +10,12 @@ import com.github.michaelbull.result.mapError
 import com.github.michaelbull.result.runCatching
 import dev.msfjarvis.aps.crypto.errors.CryptoHandlerException
 import dev.msfjarvis.aps.crypto.errors.IncorrectPassphraseException
+import dev.msfjarvis.aps.crypto.errors.NoKeysProvided
 import dev.msfjarvis.aps.crypto.errors.UnknownError
 import java.io.ByteArrayInputStream
 import java.io.InputStream
 import java.io.OutputStream
 import javax.inject.Inject
-import org.bouncycastle.bcpg.ArmoredInputStream
 import org.bouncycastle.openpgp.PGPSecretKeyRingCollection
 import org.pgpainless.PGPainless
 import org.pgpainless.decryption_verification.ConsumerOptions
@@ -64,24 +64,33 @@ public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler<PGPKe
     outputStream: OutputStream,
   ): Result<Unit, CryptoHandlerException> =
     runCatching {
+        if (keys.isEmpty()) throw NoKeysProvided("No keys provided for encryption")
         val armoredKeys = keys.map { key -> key.contents.decodeToString() }
         val pubKeysStream = ByteArrayInputStream(armoredKeys.joinToString("\n").toByteArray())
         val publicKeyRingCollection =
-          pubKeysStream.use {
-            ArmoredInputStream(it).use { armoredInputStream ->
-              PGPainless.readKeyRing().publicKeyRingCollection(armoredInputStream)
-            }
+          pubKeysStream.use { PGPainless.readKeyRing().publicKeyRingCollection(pubKeysStream) }
+        val encryptionOptions =
+          EncryptionOptions.encryptCommunications()
+            .addRecipients(publicKeyRingCollection.asIterable())
+        val producerOptions = ProducerOptions.encrypt(encryptionOptions).setAsciiArmor(true)
+        val encryptor =
+          PGPainless.encryptAndOrSign().onOutputStream(outputStream).withOptions(producerOptions)
+        plaintextStream.copyTo(encryptor)
+        encryptor.close()
+        val result = encryptor.result
+        publicKeyRingCollection.keyRings.forEach { keyRing ->
+          require(result.isEncryptedFor(keyRing)) {
+            "Stream should be encrypted for ${keyRing.publicKey.keyID} but wasn't"
           }
-        val encOpt =
-          EncryptionOptions().apply { publicKeyRingCollection.forEach { addRecipient(it) } }
-        val prodOpt = ProducerOptions.encrypt(encOpt).setAsciiArmor(true)
-        PGPainless.encryptAndOrSign().onOutputStream(outputStream).withOptions(prodOpt).use {
-          encryptionStream ->
-          plaintextStream.copyTo(encryptionStream)
         }
         return@runCatching
       }
-      .mapError { error -> UnknownError(error) }
+      .mapError { error ->
+        when (error) {
+          is CryptoHandlerException -> error
+          else -> UnknownError(error)
+        }
+      }
 
   public override fun canHandle(fileName: String): Boolean {
     return fileName.split('.').lastOrNull() == "gpg"