From d65fc88a149aa543c272f088febf21dad9c55cf8 Mon Sep 17 00:00:00 2001
From: Harsh Shandilya <me@msfjarvis.dev>
Date: Tue, 14 Jun 2022 13:00:54 +0530
Subject: Reimplement PGPainless encryption logic (#1955)

* crypto-pgpainless: reimplement encryption logic

* crypto-pgpainless: add an explicit error type for empty keyset
---
 .../msfjarvis/aps/crypto/errors/CryptoException.kt |  3 ++
 .../aps/crypto/PGPainlessCryptoHandler.kt          | 33 ++++++++++++++--------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/errors/CryptoException.kt b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/errors/CryptoException.kt
index aee23710..b795bbf1 100644
--- a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/errors/CryptoException.kt
+++ b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/errors/CryptoException.kt
@@ -37,5 +37,8 @@ public sealed class CryptoHandlerException(message: String? = null, cause: Throw
 /** The passphrase provided for decryption was incorrect. */
 public class IncorrectPassphraseException(cause: Throwable) : CryptoHandlerException(null, cause)
 
+/** No keys were provided for encryption. */
+public class NoKeysProvided(message: String?) : CryptoHandlerException(message, null)
+
 /** An unexpected error that cannot be mapped to a known type. */
 public class UnknownError(cause: Throwable) : CryptoHandlerException(null, cause)
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"
-- 
cgit v1.2.3