Handle increased message buffer in YubiHSM 2.4.0

Author: elibon99

tests/device/test_basic.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from yubihsm.core import MAX_MSG_SIZE
 from yubihsm.defs import ALGORITHM, CAPABILITY, OBJECT, COMMAND, ORIGIN, FIPS_STATUS
 from yubihsm.objects import (
     YhsmObject,
@@ -141,8 +140,9 @@ def test_get_pseudo_random(self, session):
         assert len(data2) == 10
         assert data != data2
 
-    def test_send_too_big(self, hsm):
-        buf = os.urandom(MAX_MSG_SIZE - 3 + 1)  # Message 1 byte too large
+    def test_send_too_big(self, hsm, session):
+        max_msg_size = hsm._msg_buf_size - 1
+        buf = os.urandom(max_msg_size - 3 + 1)  # Message 1 byte too large
         with pytest.raises(YubiHsmInvalidRequestError):
             hsm.send_cmd(COMMAND.ECHO, buf)
 

tests/device/test_opaque.py

@@ -66,7 +66,7 @@ def test_put_too_big(session):
             1,
             CAPABILITY.NONE,
             ALGORITHM.OPAQUE_DATA,
-            os.urandom(1976),
+            os.urandom(3064),
         )
 
     # Make sure our session is still working

tests/test_core.py

@@ -54,6 +54,7 @@ def get_mocked_session(patch):
     """
     mocked_backend = get_backend()
     mocked_backend.transceive.side_effect = [
+        _TRANSCEIVE_DEVICE_INFO,  # get_device_info is called during initialization
         b"\x83\x00\x11\x00\x05MV1\xc9\x18o\x802%\xed\x8a2$\xf2\xcf",
         b"\x84\x00\x00",
     ]
@@ -266,9 +267,12 @@ def test_create_session_derived(self, item):
         expect_enc = b"\t\x0bG\xdb\xedYVT\x90\x1d\xee\x1c\xc6U\xe4 "
         expect_mac = b"Y/\xd4\x83\xf7Y\xe2\x99\t\xa0LE\x05\xd2\xce\n"
 
-        # Note: backend doesn't do anything here; it's just required by the
-        # function's signature
-        hsm = YubiHsm(backend=None)
+        # Note: get_device_info gets called during initialization
+        # which is why we mock the transceive function.
+        backend = get_backend()
+        backend.transceive.return_value = _TRANSCEIVE_DEVICE_INFO
+
+        hsm = YubiHsm(backend)
         hsm.create_session_derived(auth_key_id, password)
 
         hsm.create_session.assert_called_once_with(auth_key_id, expect_enc, expect_mac)
@@ -284,7 +288,12 @@ def test_get_device_info_mock_transceive(self):
         hsm = YubiHsm(backend)
 
         info = hsm.get_device_info()
-        hsm._backend.transceive.assert_called_once_with(b"\x06\x00\x00")
+        hsm._backend.transceive.assert_has_calls(
+            [
+                call(b"\x06\x00\x00"),  # first call during YubiHSM::__init__
+                call(b"\x06\x00\x00"),
+            ]
+        )
 
         self.assertEqual(info.version, (2, 0, 0))
         self.assertEqual(info.serial, 7550140)

yubihsm/core.py

@@ -54,8 +54,6 @@
 CARD_CRYPTOGRAM = 0x00
 HOST_CRYPTOGRAM = 0x01
 
-MAX_MSG_SIZE = 2048 - 1
-
 
 def _derive(key: bytes, t: int, context: bytes, L: int = 0x80) -> bytes:
     # this only supports aes128
@@ -272,6 +270,12 @@ def __init__(self, backend: YhsmBackend):
         """
         self._backend: YhsmBackend = backend
 
+        # Initialize the message buffer size to 2048 bytes. This may be updated
+        # depending on the YubiHSM FW version (in 2.4.0 or higher the
+        # buffer size is 3136) in get_device_info.
+        self._msg_buf_size = 2048
+        self.get_device_info()
+
     def __enter__(self):
         return self
 
@@ -285,7 +289,7 @@ def close(self) -> None:
             self._backend = _ClosedBackend()
 
     def _transceive(self, msg: bytes) -> bytes:
-        if len(msg) > MAX_MSG_SIZE:
+        if len(msg) > self._msg_buf_size - 1:
             raise YubiHsmInvalidRequestError("Message too long.")
         return self._backend.transceive(msg)
 
@@ -307,6 +311,8 @@ def get_device_info(self) -> DeviceInfo:
         first_page = self.send_cmd(COMMAND.DEVICE_INFO)
         device_info = DeviceInfo.parse(first_page)
         if device_info.version >= (2, 4, 0):
+            # Update maximum message buffer size
+            self._msg_buf_size = 3136
             # fetch next page
             second_page = self.send_cmd(COMMAND.DEVICE_INFO, struct.pack("!B", 1))
             device_info = DeviceInfo.parse(first_page, second_page)