Enable AES-GCM and AES-CTR

Fixes #196.

Author: MatthiasValvekens

pkcs11/_pkcs11.pxd

@@ -255,6 +255,18 @@ cdef extern from '../extern/cryptoki.h':
         CK_ULONG ulPublicDataLen
         CK_BYTE *pPublicData
 
+    ctypedef struct CK_AES_CTR_PARAMS:
+        CK_ULONG ulCounterBits
+        CK_BYTE[16] cb
+
+    ctypedef struct CK_GCM_PARAMS:
+        CK_BYTE *pIv
+        CK_ULONG ulIvLen
+        CK_ULONG ulIvBits
+        CK_BYTE *pAAD
+        CK_ULONG ulAADLen
+        CK_ULONG ulTagBits
+
     ctypedef struct CK_KEY_DERIVATION_STRING_DATA:
         CK_BYTE *pData
         CK_ULONG ulLen

pkcs11/_pkcs11.pyx

@@ -142,12 +142,19 @@ cdef class MechanismWithParam:
     """The mechanism."""
     cdef void *param
     """Reference to a pointer we might need to free."""
+    cdef object _python_param
+    """
+    Hold a reference to the original parameter object so it doesn't get
+    GC'd before this one.
+    """
 
     def __cinit__(self, *args):
         self.data = <CK_MECHANISM *> PyMem_Malloc(sizeof(CK_MECHANISM))
         self.param = NULL
+        self._python_param = None
 
     def __init__(self, key_type, mapping, mechanism=None, param=None):
+        self._python_param = param
         if mechanism is None:
             try:
                 mechanism = mapping[key_type]
@@ -166,6 +173,8 @@ cdef class MechanismWithParam:
         cdef CK_ECDH1_DERIVE_PARAMS *ecdh1_params
         cdef CK_KEY_DERIVATION_STRING_DATA *aes_ecb_params
         cdef CK_AES_CBC_ENCRYPT_DATA_PARAMS *aes_cbc_params
+        cdef CK_GCM_PARAMS *gcm_params
+        cdef CK_AES_CTR_PARAMS *aes_ctr_params
 
         # Unpack mechanism parameters
 
@@ -264,6 +273,31 @@ cdef class MechanismWithParam:
             aes_cbc_params.pData = <CK_BYTE *> data
             aes_cbc_params.length = <CK_ULONG> len(data)
 
+        elif mechanism == Mechanism.AES_GCM:
+            paramlen = sizeof(CK_GCM_PARAMS)
+            if not isinstance(param, GCMParams):
+                raise TypeError
+            self.param = gcm_params = <CK_GCM_PARAMS *> PyMem_Malloc(paramlen)
+            gcm_params.pIv = <CK_BYTE *> param.nonce
+            gcm_params.ulIvLen = <CK_ULONG> len(param.nonce)
+            gcm_params.ulIvBits = <CK_ULONG> len(param.nonce) * 8
+            if param.aad is not None:
+                gcm_params.pAAD = <CK_BYTE *> param.aad
+                gcm_params.ulAADLen = <CK_ULONG> len(param.aad)
+            else:
+                gcm_params.pAAD = NULL
+                gcm_params.ulAADLen = 0
+            gcm_params.ulTagBits = <CK_ULONG> param.tag_bits
+
+        elif mechanism == Mechanism.AES_CTR:
+            paramlen = sizeof(CK_AES_CTR_PARAMS)
+            self.param = aes_ctr_params = <CK_AES_CTR_PARAMS *> PyMem_Malloc(paramlen)
+            # use a wrapper type to not break the forwards compat rule for params specified as `bytes`
+            if not isinstance(param, CTRParams):
+                raise TypeError
+            aes_ctr_params.ulCounterBits = (16 - len(param.nonce)) * 8
+            aes_ctr_params.cb = param.nonce + b"\x00" * (15 - len(param.nonce)) + b"\x01"
+
         elif param is None:
             self.data.pParameter = NULL
             paramlen = 0

pkcs11/mechanisms.py

@@ -1,5 +1,7 @@
 from enum import IntEnum
 
+from pkcs11.exceptions import ArgumentsBad
+
 
 class KeyType(IntEnum):
     """
@@ -794,3 +796,22 @@ class MGF(IntEnum):
 
     def __repr__(self):
         return "<MGF.%s>" % self.name
+
+
+class GCMParams:
+    def __init__(self, nonce, aad=None, tag_bits=128):
+        if len(nonce) > 12:
+            raise ArgumentsBad("IV must be less than 12 bytes")
+        self.nonce = nonce
+        self.aad = aad
+        self.tag_bits = tag_bits
+
+
+class CTRParams:
+    def __init__(self, nonce):
+        if len(nonce) >= 16:
+            raise ArgumentsBad(
+                f"{nonce.hex()} is too long to serve as a CTR nonce, must be 15 bytes or less "
+                f"to leave room for the block counter."
+            )
+        self.nonce = nonce

tests/test_aes.py

@@ -5,7 +5,7 @@
 from parameterized import parameterized
 
 import pkcs11
-from pkcs11 import Mechanism
+from pkcs11 import ArgumentsBad, CTRParams, GCMParams, Mechanism, PKCS11Error
 
 from . import FIXME, TestCase, requires
 
@@ -462,3 +462,152 @@ def test_encrypt_with_key_derived_using_cbc_encrypt(
         text = self.key.decrypt(crypttext, mechanism_param=iv)
 
         self.assertEqual(text, data)
+
+    @requires(Mechanism.AES_GCM)
+    def test_encrypt_gcm(self):
+        data = b"INPUT DATA"
+        nonce = b"0" * 12
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce)
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce)
+        )
+        self.assertEqual(data, text)
+
+    def test_gcm_nonce_size_limit(self):
+        def _inst():
+            return GCMParams(nonce=b"0" * 13)
+
+        self.assertRaises(ArgumentsBad, _inst)
+
+    @requires(Mechanism.AES_GCM)
+    def test_encrypt_gcm_with_aad(self):
+        data = b"INPUT DATA"
+        nonce = b"0" * 12
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"foo")
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"foo")
+        )
+        self.assertEqual(data, text)
+
+    @requires(Mechanism.AES_GCM)
+    def test_encrypt_gcm_with_mismatching_nonces(self):
+        data = b"INPUT DATA"
+        nonce1 = b"0" * 12
+        nonce2 = b"1" * 12
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce1, b"foo")
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        # This should be EncryptedDataInvalid, but in practice not all tokens support this
+        with self.assertRaises(PKCS11Error):
+            self.key.decrypt(
+                crypttext, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce2, b"foo")
+            )
+
+    @requires(Mechanism.AES_GCM)
+    def test_encrypt_gcm_with_mismatching_aad(self):
+        data = b"INPUT DATA"
+        nonce = b"0" * 12
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"foo")
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        with self.assertRaises(PKCS11Error):
+            self.key.decrypt(
+                crypttext, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"bar")
+            )
+
+    @requires(Mechanism.AES_GCM)
+    def test_encrypt_gcm_with_custom_tag_length(self):
+        data = b"INPUT DATA"
+        nonce = b"0" * 12
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"foo", 120)
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_GCM, mechanism_param=GCMParams(nonce, b"foo", 120)
+        )
+        self.assertEqual(data, text)
+
+        # This should be EncryptedDataInvalid, but in practice not all tokens support this
+        with self.assertRaises(PKCS11Error):
+            text = self.key.decrypt(
+                crypttext,
+                mechanism=Mechanism.AES_GCM,
+                mechanism_param=GCMParams(nonce, b"foo", 128),
+            )
+
+    @parameterized.expand(
+        [
+            (b""),
+            (b"0" * 12),
+            (b"0" * 15),
+        ]
+    )
+    @requires(Mechanism.AES_CTR)
+    @FIXME.opencryptoki  # opencryptoki incorrectly forces AES-CTR input to be padded
+    def test_encrypt_ctr(self, nonce):
+        data = b"INPUT DATA SEVERAL BLOCKS LONG SO THE COUNTER GOES UP A FEW TIMES" * 20
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(nonce)
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(nonce)
+        )
+        self.assertEqual(data, text)
+
+    @requires(Mechanism.AES_CTR)
+    def test_encrypt_ctr_exactly_padded(self):
+        # let's still verify the "restricted" AES-CTR supported by opencryptoki
+        data = b"PADDED INPUT DATA TO MAKE OPENCRYPTOKI HAPPY" * 16
+        nonce = b"0" * 15
+
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(nonce)
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(nonce)
+        )
+        self.assertEqual(data, text)
+
+    def test_ctr_nonce_size_limit(self):
+        def _inst():
+            return CTRParams(nonce=b"0" * 16)
+
+        self.assertRaises(ArgumentsBad, _inst)
+
+    @requires(Mechanism.AES_CTR)
+    @FIXME.opencryptoki  # opencryptoki incorrectly forces AES-CTR input to be padded
+    def test_encrypt_ctr_nonce_mismatch(self):
+        data = b"INPUT DATA SEVERAL BLOCKS LONG SO THE COUNTER GOES UP A FEW TIMES" * 20
+        crypttext = self.key.encrypt(
+            data, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(b"0" * 12)
+        )
+        self.assertIsInstance(crypttext, bytes)
+        self.assertNotEqual(data, crypttext)
+        text = self.key.decrypt(
+            crypttext, mechanism=Mechanism.AES_CTR, mechanism_param=CTRParams(b"1" * 12)
+        )
+        self.assertNotEqual(data, text)