From 44467ef90b6a5baa082fce47fc3ee18ad40733f0 Mon Sep 17 00:00:00 2001
From: Matthias Valvekens <matthias@mvalvekens.be>
Date: Thu, 14 Mar 2024 22:44:33 +0100
Subject: [PATCH] Fix regression in looking up PKCS#11 objects

Some config fallbacks were a little too aggressive.

Fixes #394.
---
 pyhanko/sign/pkcs11.py       | 15 ++++++---------
 pyhanko_tests/test_config.py | 16 ++++++++++++++++
 pyhanko_tests/test_pkcs11.py | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/pyhanko/sign/pkcs11.py b/pyhanko/sign/pkcs11.py
index 3f7caf4a..8ac4550a 100644
--- a/pyhanko/sign/pkcs11.py
+++ b/pyhanko/sign/pkcs11.py
@@ -47,7 +47,6 @@
         e,
     )
 
-
 __all__ = [
     'PKCS11Signer',
     'open_pkcs11_session',
@@ -165,7 +164,6 @@ class PKCS11SignatureOperationSpec:
     'sha512': Mechanism.SHA512_RSA_PKCS,
 }
 
-
 RSASSA_PSS_MECH_MAP = {
     'sha1': Mechanism.SHA1_RSA_PKCS_PSS,
     'sha224': Mechanism.SHA224_RSA_PKCS_PSS,
@@ -182,7 +180,6 @@ class PKCS11SignatureOperationSpec:
     'sha512': MGF.SHA512,
 }
 
-
 ECDSA_MECH_MAP = {
     'sha1': Mechanism.ECDSA_SHA1,
     'sha224': Mechanism.ECDSA_SHA224,
@@ -191,7 +188,6 @@ class PKCS11SignatureOperationSpec:
     'sha512': Mechanism.ECDSA_SHA512,
 }
 
-
 DSA_MECH_MAP = {
     'sha1': Mechanism.DSA_SHA1,
     'sha224': Mechanism.DSA_SHA224,
@@ -204,7 +200,6 @@ class PKCS11SignatureOperationSpec:
     'sha512': Mechanism.DSA_SHA512,
 }
 
-
 DIGEST_MECH_MAP = {
     'sha1': Mechanism.SHA_1,
     'sha224': Mechanism.SHA224,
@@ -517,10 +512,12 @@ def __init__(
         """
         Initialise a PKCS11 signer.
         """
-        self.cert_label = coalesce(cert_label, key_label)
-        self.key_id = coalesce(key_id, cert_id)
-        self.cert_id = coalesce(cert_id, key_id)
-        self.key_label = coalesce(key_label, cert_label)
+        self.cert_label = coalesce(
+            cert_label, key_label if not cert_id else None
+        )
+        self.key_id = coalesce(key_id, cert_id if not key_label else None)
+        self.cert_id = coalesce(cert_id, key_id if not cert_label else None)
+        self.key_label = coalesce(key_label, cert_label if not key_id else None)
         self.pkcs11_session = pkcs11_session
         self.other_certs = other_certs_to_pull
         self._other_certs_loaded = False
diff --git a/pyhanko_tests/test_config.py b/pyhanko_tests/test_config.py
index 28d5c70e..17741ec8 100644
--- a/pyhanko_tests/test_config.py
+++ b/pyhanko_tests/test_config.py
@@ -820,6 +820,22 @@ def test_read_pkcs11_config_key_label_from_cert_label():
     assert (cfg.key_label, cfg.key_id) == ("signer", None)
 
 
+def test_read_pkcs11_config_key_label_not_from_cert_label_if_key_id_defined():
+    cli_config = _parse_cli_config(
+        f"""
+        pkcs11-setups:
+            foo:
+                module-path: /path/to/libfoo.so
+                slot-no: 0
+                cert-label: "signer"
+                key-id: "deadbeef"
+        """
+    )
+    cfg = ModuleConfigWrapper(cli_config).get_pkcs11_config('foo')
+    assert (cfg.key_label, cfg.key_id) == (None, b"\xde\xad\xbe\xef")
+    assert (cfg.cert_label, cfg.cert_id) == ("signer", None)
+
+
 @pytest.mark.parametrize(
     'literal,exp_val',
     [
diff --git a/pyhanko_tests/test_pkcs11.py b/pyhanko_tests/test_pkcs11.py
index 4e7576d4..b61c2acd 100644
--- a/pyhanko_tests/test_pkcs11.py
+++ b/pyhanko_tests/test_pkcs11.py
@@ -24,6 +24,7 @@
 from pyhanko.sign import general, pkcs11, signers
 from pyhanko.sign.general import SigningError
 from pyhanko.sign.pkcs11 import (
+    PKCS11Signer,
     PKCS11SigningContext,
     TokenCriteria,
     criteria_satisfied_by,
@@ -831,3 +832,38 @@ def test_token_not_found(slot_list, token_lbl_query):
         token_criteria=TokenCriteria(label=token_lbl_query),
     )
     assert tok is None
+
+
+def _mirror(*cfgs):
+    for cfg in cfgs:
+        a = cfg
+        b = cfg[1], cfg[0], cfg[3], cfg[2]
+        yield a
+        if a != b:
+            yield b
+
+
+@pytest.mark.parametrize(
+    'key_in,cert_in,key_prop,cert_prop',
+    [
+        *_mirror(
+            ((None, None), (None, None), (None, None), (None, None)),
+            (('a', b"a"), (None, None), ('a', b"a"), ('a', b"a")),
+            (('a', None), (None, None), ('a', None), ('a', None)),
+            ((None, b"a"), (None, None), (None, b"a"), (None, b"a")),
+            (('a', None), (None, b"a"), ('a', None), (None, b"a")),
+        )
+    ],
+)
+def test_config_fallbacks(key_in, cert_in, key_prop, cert_prop):
+    # noinspection PyTypeChecker
+    signer = PKCS11Signer(
+        pkcs11_session=None,
+        cert_label=cert_in[0],
+        cert_id=cert_in[1],
+        key_label=key_in[0],
+        key_id=key_in[1],
+    )
+    actual_key_prop = signer.key_label, signer.key_id
+    actual_cert_prop = signer.cert_label, signer.cert_id
+    assert (actual_key_prop, actual_cert_prop) == (key_prop, cert_prop)