diff --git a/ccbt/security/mse_handshake.py b/ccbt/security/mse_handshake.py index cec4e2f..f85d134 100644 --- a/ccbt/security/mse_handshake.py +++ b/ccbt/security/mse_handshake.py @@ -182,12 +182,7 @@ async def initiate_as_initiator( None, f"Failed to read RKEYE message ({rke_failure.value})", ) - legacy_type = None - if len(rke_message) > 0: - try: - legacy_type = MSEHandshakeType(rke_message[0]) - except ValueError: - legacy_type = None + legacy_type = self._legacy_message_type(rke_message) if legacy_type in { MSEHandshakeType.SKEYE, MSEHandshakeType.CRYPTO, @@ -254,12 +249,7 @@ async def initiate_as_initiator( None, f"Failed to read CRYPTO message ({crypto_failure.value})", ) - legacy_crypto_type = None - if len(crypto_response) > 0: - try: - legacy_crypto_type = MSEHandshakeType(crypto_response[0]) - except ValueError: - legacy_crypto_type = None + legacy_crypto_type = self._legacy_message_type(crypto_response) if ( legacy_crypto_type is not None and legacy_crypto_type != MSEHandshakeType.CRYPTO @@ -352,12 +342,7 @@ async def respond_as_receiver( None, f"Failed to read SKEYE message ({ske_failure.value})", ) - legacy_type = None - if len(ske_message) > 0: - try: - legacy_type = MSEHandshakeType(ske_message[0]) - except ValueError: - legacy_type = None + legacy_type = self._legacy_message_type(ske_message) if legacy_type in { MSEHandshakeType.RKEYE, MSEHandshakeType.CRYPTO, @@ -413,12 +398,7 @@ async def respond_as_receiver( None, f"Failed to read CRYPTO message ({crypto_failure.value})", ) - legacy_crypto_type = None - if len(crypto_message) > 0: - try: - legacy_crypto_type = MSEHandshakeType(crypto_message[0]) - except ValueError: - legacy_crypto_type = None + legacy_crypto_type = self._legacy_message_type(crypto_message) if ( legacy_crypto_type is not None and legacy_crypto_type != MSEHandshakeType.CRYPTO @@ -620,6 +600,29 @@ def _dh_public_key_length_for_size(size: int) -> int: def _dh_public_key_length(self) -> int: return self._dh_public_key_length_for_size(self.dh_exchange.key_size) + def _legacy_message_type(self, payload: bytes) -> Optional[MSEHandshakeType]: + """Return a legacy MSE type only for short type-prefixed payloads. + + Modern BEP-style frames carry raw DH key material or RC4 ciphertext + without a leading type byte. Their first byte can coincidentally match + ``SKEYE``/``RKEYE``/``CRYPTO`` enum values and must not be rejected. + """ + if not payload: + return None + try: + msg_type = MSEHandshakeType(payload[0]) + except ValueError: + return None + if msg_type == MSEHandshakeType.CRYPTO: + # Legacy CRYPTO is exactly type (1) + cipher (1). + return msg_type if len(payload) == 2 else None + if msg_type in {MSEHandshakeType.SKEYE, MSEHandshakeType.RKEYE}: + # Legacy SKEYE/RKEYE is type + DH key (+ optional leading 0x00). + dh_len = self._dh_public_key_length() + if len(payload) in {dh_len + 1, dh_len + 2}: + return msg_type + return None + @staticmethod def _select_handshake_padding() -> bytes: """Select handshake padding for PadC/PadD. diff --git a/tests/integration/test_mse_tcp_server_pe_first.py b/tests/integration/test_mse_tcp_server_pe_first.py index ed963c7..79e4f9a 100644 --- a/tests/integration/test_mse_tcp_server_pe_first.py +++ b/tests/integration/test_mse_tcp_server_pe_first.py @@ -50,7 +50,7 @@ async def _run_loopback_mse_handshake( outbound_payload: bytes, port: int, *, - attempts: int = 2, + attempts: int = 3, accept_mock: AsyncMock | None = None, ) -> None: """Run an outbound MSE handshake against an already-started loopback server.""" @@ -72,6 +72,8 @@ async def _run_loopback_mse_handshake( finally: writer.close() await writer.wait_closed() + # Only stop retrying once the server has accepted — a failed initiator + # attempt that never reached accept should be retried. if accept_mock is not None and accept_mock.await_count > 0: break if attempt + 1 < attempts: diff --git a/tests/unit/security/test_mse_handshake.py b/tests/unit/security/test_mse_handshake.py index 42f17f9..0e449ba 100644 --- a/tests/unit/security/test_mse_handshake.py +++ b/tests/unit/security/test_mse_handshake.py @@ -63,6 +63,22 @@ def test_init_custom_allowed_ciphers(self): assert CipherType.AES not in handshake.allowed_ciphers +class TestMSEHandshakeLegacyTypeDetection: + """Legacy type detection must not misfire on modern BEP payloads.""" + + def test_modern_ciphertext_starting_with_skeye_byte_is_not_legacy(self): + """RC4 crypto-select blobs may start with 0x02 without being SKEYE.""" + handshake = MSEHandshake() + modern_payload = bytes([int(MSEHandshakeType.SKEYE)]) + b"\x00" * 32 + assert handshake._legacy_message_type(modern_payload) is None + + def test_legacy_crypto_two_byte_payload_is_detected(self): + """Legacy CRYPTO remains type + cipher byte.""" + handshake = MSEHandshake() + payload = bytes([int(MSEHandshakeType.CRYPTO), int(CipherType.RC4)]) + assert handshake._legacy_message_type(payload) == MSEHandshakeType.CRYPTO + + class TestMSEHandshakeMessageEncoding: """Tests for message encoding/decoding."""