Sanitize key flags

pgpainless-core/src/main/kotlin/org/pgpainless/key/generation/OpenPgpKeyGenerator.kt

@@ -20,6 +20,7 @@ import org.pgpainless.algorithm.AlgorithmSuite
 import org.pgpainless.algorithm.CertificationType
 import org.pgpainless.algorithm.HashAlgorithm
 import org.pgpainless.algorithm.KeyFlag
+import org.pgpainless.algorithm.PublicKeyAlgorithm
 import org.pgpainless.implementation.ImplementationFactory
 import org.pgpainless.key.generation.DefinePrimaryKey.PrimaryKeyBuilder
 import org.pgpainless.key.generation.DefineSubkeys.SubkeyBuilder
@@ -169,7 +170,13 @@ internal constructor(val policy: Policy, val creationTime: Date, val preferences
         keyFlags: List<KeyFlag>? = listOf(KeyFlag.CERTIFY_OTHER),
         creationTime: Date = this.creationTime,
         block: PrimaryKeyBlock? = null
-    ): O = doSetPrimaryKey(type, keyFlags, creationTime, block)
+    ): O {
+        require(type.canCertify) {
+            "Primary key cannot use algorithm ${type.algorithm} because it needs to be " +
+                "signing capable."
+        }
+        return doSetPrimaryKey(type, keyFlags, creationTime, block)
+    }
 
     fun setPrimaryKey(type: KeyType, block: PrimaryKeyBlock?): O =
         setPrimaryKey(type, listOf(KeyFlag.CERTIFY_OTHER), this.creationTime, block)
@@ -200,6 +207,10 @@ internal constructor(val policy: Policy, val creationTime: Date, val preferences
         // Do nothing
     }
 
+    protected open fun sanitizeKeyFlags(algorithm: PublicKeyAlgorithm, keyFlags: List<KeyFlag>?) {
+        // Do nothing
+    }
+
     /**
      * Function that can be applied to the primary key.
      *
@@ -407,14 +418,23 @@ internal constructor(
     @JvmOverloads
     fun addSubkey(
         type: KeyType,
+        flags: List<KeyFlag>? = null,
         creationTime: Date = this.creationTime,
         block: SubkeyBlock? = null
     ): B =
         apply {
+            sanitizeKeyFlags(type.algorithm, flags)
             sanitizeSubkeyCreationTime(creationTime, primaryKey)
 
             var subkey = generateSubkey(type, creationTime)
-            subkey = invokeOnSubkey(subkey, block)
+            val subkeyBlock =
+                block
+                    ?: {
+                        addBindingSignature(
+                            SelfSignatureSubpackets.applyHashed { flags?.let { setKeyFlags(it) } },
+                            bindingTime = creationTime)
+                    }
+            subkey = invokeOnSubkey(subkey, subkeyBlock)
             subkeys.add(subkey)
         }
             as B
@@ -490,6 +510,10 @@ internal constructor(
         // Do nothing
     }
 
+    protected open fun sanitizeKeyFlags(algorithm: PublicKeyAlgorithm, keyFlags: List<KeyFlag>?) {
+        // Do nothing
+    }
+
     /**
      * Function that can be applied to subkeys.
      *
@@ -723,8 +747,6 @@ internal constructor(policy: Policy, creationTime: Date, preferences: AlgorithmS
         creationTime: Date,
         block: PrimaryKeyBlock?
     ): OpinionatedDefineSubkeysV4 {
-        // Check algorithm is signing capable
-        require(type.algorithm.isSigningCapable()) { "Primary Key MUST be capable of signing." }
 
         // Check key strength
         require(policy.publicKeyAlgorithmPolicy.isAcceptable(type.algorithm, type.bitStrength)) {
@@ -732,6 +754,9 @@ internal constructor(policy: Policy, creationTime: Date, preferences: AlgorithmS
                 " for the current public key algorithm policy."
         }
 
+        // Sanitize key flags
+        sanitizeKeyFlags(type.algorithm, keyFlags)
+
         // Remember flags for DK and UID signatures
         this.keyFlags = keyFlags
 
@@ -762,6 +787,27 @@ internal constructor(policy: Policy, creationTime: Date, preferences: AlgorithmS
                 "key was created at ${primaryKey.publicKey.creationTime.formatUTC()}."
         }
     }
+
+    override fun sanitizeKeyFlags(algorithm: PublicKeyAlgorithm, keyFlags: List<KeyFlag>?) {
+        keyFlags?.forEach { flag ->
+            when (flag) {
+                KeyFlag.CERTIFY_OTHER,
+                KeyFlag.SIGN_DATA,
+                KeyFlag.AUTHENTICATION ->
+                    require(algorithm.isSigningCapable()) {
+                        "Primary key cannot carry key flag $flag because the " +
+                            "algorithm $algorithm is not signing capable."
+                    }
+                KeyFlag.ENCRYPT_COMMS,
+                KeyFlag.ENCRYPT_STORAGE ->
+                    require(algorithm.isEncryptionCapable()) {
+                        "Primary key cannot carry key flag $flag because the " +
+                            "algorithm $algorithm is not encryption capable."
+                    }
+                else -> {} // no special requirements for SPLIT and SHARED
+            }
+        }
+    }
 }
 
 /**
@@ -842,11 +888,10 @@ internal constructor(primaryKey: PGPKeyPair, policy: Policy, creationTime: Date)
     fun addSigningSubkey(
         type: KeyType,
         creationTime: Date = this.creationTime,
-        block: SubkeyBlock? = {
+        block: SubkeyBlock? = null
-            addBindingSignature(
+    ): OpinionatedDefineSubkeysV4 {
-                SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.SIGN_DATA) })
+        return addSubkey(type, listOf(KeyFlag.SIGN_DATA), creationTime, block)
-        }
+    }
-    ) = addSubkey(type, creationTime, block)
 
     /**
      * Add a subkey for signing messages to the OpenPGP key.
@@ -868,13 +913,11 @@ internal constructor(primaryKey: PGPKeyPair, policy: Policy, creationTime: Date)
     fun addEncryptionSubkey(
         type: KeyType,
         creationTime: Date = this.creationTime,
-        block: SubkeyBlock? = {
+        block: SubkeyBlock? = null,
-            addBindingSignature(
+    ): OpinionatedDefineSubkeysV4 {
-                SelfSignatureSubpackets.applyHashed {
+        return addSubkey(
-                    setKeyFlags(KeyFlag.ENCRYPT_COMMS, KeyFlag.ENCRYPT_STORAGE)
+            type, listOf(KeyFlag.ENCRYPT_COMMS, KeyFlag.ENCRYPT_STORAGE), creationTime, block)
-                })
+    }
-        }
-    ) = addSubkey(type, creationTime, block)
 
     /**
      * Add a subkey for message encryption to the OpenPGP key.
@@ -907,6 +950,27 @@ internal constructor(primaryKey: PGPKeyPair, policy: Policy, creationTime: Date)
                 "Primary key was created at ${primaryKey.publicKey.creationTime.formatUTC()}."
         }
     }
+
+    override fun sanitizeKeyFlags(algorithm: PublicKeyAlgorithm, keyFlags: List<KeyFlag>?) {
+        keyFlags?.forEach { flag ->
+            when (flag) {
+                KeyFlag.CERTIFY_OTHER,
+                KeyFlag.SIGN_DATA,
+                KeyFlag.AUTHENTICATION ->
+                    require(algorithm.isSigningCapable()) {
+                        "Subkey cannot carry key flag $flag because the " +
+                            "algorithm $algorithm is not signing capable."
+                    }
+                KeyFlag.ENCRYPT_COMMS,
+                KeyFlag.ENCRYPT_STORAGE ->
+                    require(algorithm.isEncryptionCapable()) {
+                        "Subkey cannot carry key flag $flag because the " +
+                            "algorithm $algorithm is not encryption capable."
+                    }
+                else -> {} // no special requirements for SPLIT and SHARED
+            }
+        }
+    }
 }
 
 /**
@@ -1078,17 +1142,9 @@ class OpenPgpKeyTemplates private constructor() {
                         SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.CERTIFY_OTHER) })
                 }
                 // singing key
-                .addSubkey(KeyType.EDDSA(EdDSACurve._Ed25519)) {
+                .addSigningSubkey(KeyType.EDDSA(EdDSACurve._Ed25519))
-                    addBindingSignature(
-                        SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.SIGN_DATA) })
-                }
                 // encryption key
-                .addSubkey(KeyType.XDH(XDHSpec._X25519)) {
+                .addEncryptionSubkey(KeyType.XDH(XDHSpec._X25519))
-                    addBindingSignature(
-                        SelfSignatureSubpackets.applyHashed {
-                            setKeyFlags(KeyFlag.ENCRYPT_COMMS, KeyFlag.ENCRYPT_STORAGE)
-                        })
-                }
                 .build()
 
         /**
@@ -1122,17 +1178,9 @@ class OpenPgpKeyTemplates private constructor() {
                         SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.CERTIFY_OTHER) })
                 }
                 // signing key
-                .addSubkey(KeyType.RSA(length)) {
+                .addSigningSubkey(KeyType.RSA(length))
-                    addBindingSignature(
-                        SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.SIGN_DATA) })
-                }
                 // encryption key
-                .addSubkey(KeyType.RSA(length)) {
+                .addEncryptionSubkey(KeyType.RSA(length))
-                    addBindingSignature(
-                        SelfSignatureSubpackets.applyHashed {
-                            setKeyFlags(KeyFlag.ENCRYPT_COMMS, KeyFlag.ENCRYPT_STORAGE)
-                        })
-                }
                 .build()
 
         /**

pgpainless-core/src/test/kotlin/org/pgpainless/key/generation/OpenPgpKeyGeneratorTest.kt

@@ -109,10 +109,7 @@ class OpenPgpKeyGeneratorTest {
         val key =
             OpenPgpKeyGenerator.buildV4Key()
                 .setPrimaryKey(KeyType.EDDSA(EdDSACurve._Ed25519))
-                .addSubkey(KeyType.EDDSA(EdDSACurve._Ed25519)) {
+                .addSubkey(KeyType.EDDSA(EdDSACurve._Ed25519), listOf(KeyFlag.SIGN_DATA))
-                    addBindingSignature(
-                        SelfSignatureSubpackets.applyHashed { setKeyFlags(KeyFlag.SIGN_DATA) })
-                }
                 .build()
 
         assertFalse(
@@ -376,7 +373,7 @@ class OpenPgpKeyGeneratorTest {
         assertThrows<IllegalArgumentException> {
             OpenPgpKeyGenerator.buildV4Key(policy, t1)
                 .setPrimaryKey(KeyType.EDDSA(EdDSACurve._Ed25519))
-                .addSubkey(KeyType.XDH(XDHSpec._X25519), t0)
+                .addSubkey(KeyType.XDH(XDHSpec._X25519), null, t0)
         }
     }
 
@@ -389,7 +386,7 @@ class OpenPgpKeyGeneratorTest {
         OpenPgpKeyGenerator.buildV4Key(policy, t1)
             .unopinionated()
             .setPrimaryKey(KeyType.EDDSA(EdDSACurve._Ed25519))
-            .addSubkey(KeyType.XDH(XDHSpec._X25519), t0)
+            .addSubkey(KeyType.XDH(XDHSpec._X25519), null, t0)
     }
 
     @Test
@@ -442,6 +439,50 @@ class OpenPgpKeyGeneratorTest {
             }
     }
 
+    @Test
+    fun `opinionated set primary key to encryption key fails`() {
+        val policy = Policy()
+
+        assertThrows<IllegalArgumentException> {
+            OpenPgpKeyGenerator.buildV4Key(policy).setPrimaryKey(KeyType.XDH(XDHSpec._X25519))
+        }
+    }
+
+    @Test
+    fun `opinionated set primary key to sign-only algorithm but with encryption flag fails`() {
+        val policy = Policy()
+
+        assertThrows<IllegalArgumentException> {
+            OpenPgpKeyGenerator.buildV4Key(policy)
+                .setPrimaryKey(
+                    KeyType.EDDSA(EdDSACurve._Ed25519),
+                    listOf(KeyFlag.CERTIFY_OTHER, KeyFlag.ENCRYPT_STORAGE))
+        }
+    }
+
+    @Test
+    fun `unopinionated set primary key to sign-only algorithm but with encryption flag is okay`() {
+        val policy = Policy()
+
+        OpenPgpKeyGenerator.buildV4Key(policy)
+            .unopinionated()
+            .setPrimaryKey(
+                KeyType.EDDSA(EdDSACurve._Ed25519),
+                listOf(KeyFlag.CERTIFY_OTHER, KeyFlag.ENCRYPT_STORAGE))
+    }
+
+    @Test
+    fun `opinionated add encryption-only subkey with additional sign flag fails`() {
+        val policy = Policy()
+
+        assertThrows<IllegalArgumentException> {
+            OpenPgpKeyGenerator.buildV4Key(policy)
+                .setPrimaryKey(KeyType.EDDSA(EdDSACurve._Ed25519))
+                .addSubkey(
+                    KeyType.XDH(XDHSpec._X25519), listOf(KeyFlag.ENCRYPT_COMMS, KeyFlag.SIGN_DATA))
+        }
+    }
+
     @Test
     fun `add image attribute to key`() {
         // smallest JPEG according to https://stackoverflow.com/a/2349470/11150851