Fix functional addressing with configurable ECU response addresses

doipclient/client.py

@@ -149,6 +149,13 @@ class DoIPClient:
     :type auto_reconnect_tcp: bool
     :param vm_specific: Optional 4 byte long int
     :type vm_specific: int, optional
+    :param allowed_response_addresses: Additional allowed source logical addresses for diagnostic responses.
+        Set this to ECU physical addresses when sending to a functional address. Alternatively, pass a callable
+        that takes a source address and returns whether it is allowed. Diagnostic responses may originate from
+        the current ecu_logical_address or any source accepted by this configuration. ACK/NACK messages may
+        originate from the request target or any source accepted by this configuration.
+        The response target must always match client_logical_address.
+    :type allowed_response_addresses: Union[Collection[int], Callable[[int], bool]], optional
 
     :raises ConnectionRefusedError: If the activation request fails
     :raises ValueError: If the IPAddress is neither an IPv4 nor an IPv6 address
@@ -167,8 +174,10 @@ def __init__(
         use_secure=False,
         auto_reconnect_tcp=False,
         vm_specific=None,
+        allowed_response_addresses=None,
     ):
         self._ecu_logical_address = ecu_logical_address
+        self._allowed_response_addresses = allowed_response_addresses
         self._client_logical_address = client_logical_address
         self._client_ip_address = client_ip_address
         self._use_secure = use_secure
@@ -728,6 +737,13 @@ def send_diagnostic(self, diagnostic_payload, timeout=A_PROCESSING_TIME):
             self._ecu_logical_address, diagnostic_payload, timeout
         )
 
+    def _is_allowed_response_address(self, source_address, address):
+        if source_address == address:
+            return True
+        if callable(self._allowed_response_addresses):
+            return self._allowed_response_addresses(source_address)
+        return source_address in (self._allowed_response_addresses or ())
+
     def send_diagnostic_to_address(
         self, address, diagnostic_payload, timeout=A_PROCESSING_TIME
     ):
@@ -754,8 +770,10 @@ def send_diagnostic_to_address(
                 result = self.read_doip()
             if type(result) == DiagnosticMessageNegativeAcknowledgement:
                 if (
-                    result.source_address == address
-                    and result.target_address == self._client_logical_address
+                    result.target_address == self._client_logical_address
+                    and self._is_allowed_response_address(
+                        result.source_address, address
+                    )
                 ):
                     raise IOError(
                         "Diagnostic request rejected with negative acknowledge code: {}".format(
@@ -768,8 +786,10 @@ def send_diagnostic_to_address(
                     )
             elif type(result) == DiagnosticMessagePositiveAcknowledgement:
                 if (
-                    result.source_address == address
-                    and result.target_address == self._client_logical_address
+                    result.target_address == self._client_logical_address
+                    and self._is_allowed_response_address(
+                        result.source_address, address
+                    )
                 ):
                     return
                 else:
@@ -800,17 +820,20 @@ def receive_diagnostic(self, timeout=None):
             else:
                 result = self.read_doip()
             if type(result) == DiagnosticMessage:
+                source_address_allowed = self._is_allowed_response_address(
+                    result.source_address, self._ecu_logical_address
+                )
                 if (
-                    result.source_address == self._ecu_logical_address
-                    and result.target_address == self._client_logical_address
+                    result.target_address == self._client_logical_address
+                    and source_address_allowed
                 ):
                     return result.user_data
                 elif (
-                    result.source_address != self._ecu_logical_address
-                    and result.target_address == self._client_logical_address
+                    result.target_address == self._client_logical_address
+                    and not source_address_allowed
                 ):
                     logger.warning(
-                        "Received DiagnosticMessage with expected target address, but source address doesn't match expected ECU logical address. Ignoring."
+                        "Received DiagnosticMessage with expected target address, but source address is not allowed. Ignoring."
                     )
                     start_time = time.time()
                 else:

tests/test_client.py

@@ -730,9 +730,16 @@ def test_send_diagnostic_to_address_times_out_when_only_invalid_target_ack_nack(
     assert mock_socket.tx_queue[-1] == diagnostic_request_to_address
 
 
-def test_receive_diagnostic(mock_socket):
-    sut = DoIPClient(test_ip, test_logical_address)
-    mock_socket.rx_queue.append(diagnostic_result)
+@pytest.mark.parametrize("allowed_addresses", [None, {}, {0x1234}, lambda address: False])
+@pytest.mark.parametrize("ecu_address", [1, 2])
+def test_receive_diagnostic(mock_socket, allowed_addresses, ecu_address):
+    sut = DoIPClient(
+        test_ip, test_logical_address, allowed_response_addresses=allowed_addresses
+    )
+    sut._ecu_logical_address = ecu_address
+    response = bytearray(diagnostic_result)
+    response[8:10] = ecu_address.to_bytes(2, "big")
+    mock_socket.rx_queue.append(response)
     result = sut.receive_diagnostic()
     assert result == bytearray([0, 1, 2, 3])
 
@@ -743,6 +750,61 @@ def test_receive_diagnostic_timeout(mock_socket):
         sut.receive_diagnostic(timeout=0.01)
 
 
+@pytest.mark.parametrize("source_address", [1, 2])
+@pytest.mark.parametrize("allowed_addresses", [{1, 2}, lambda address: address in (1, 2)])
+def test_functional_diagnostic_accepts_allowed_responses(
+    mock_socket, source_address, allowed_addresses
+):
+    sut = DoIPClient(test_ip, 0xE400, allowed_response_addresses=allowed_addresses)
+    ack, response = bytearray(diagnostic_ack), bytearray(diagnostic_result)
+    ack[8:10] = response[8:10] = source_address.to_bytes(2, "big")
+    mock_socket.rx_queue.extend([ack, response])
+    sut.send_diagnostic(bytearray([0, 1, 2]), timeout=0.01)
+    assert mock_socket.tx_queue[-1][10:12] == b"\xe4\x00"
+    assert sut.receive_diagnostic(timeout=0.01) == bytearray([0, 1, 2, 3])
+
+
+@pytest.mark.parametrize("allowed_addresses", [{1}, lambda address: address == 1])
+def test_functional_diagnostic_accepts_allowed_nack(mock_socket, allowed_addresses):
+    sut = DoIPClient(test_ip, 0xE400, allowed_response_addresses=allowed_addresses)
+    mock_socket.rx_queue.append(diagnostic_nack)
+    with pytest.raises(IOError, match="Diagnostic request rejected"):
+        sut.send_diagnostic(bytearray([0, 1, 2]), timeout=0.01)
+
+
+def test_functional_diagnostic_rejects_unlisted_ecu_ack(mock_socket):
+    sut = DoIPClient(test_ip, test_logical_address)
+    sut._ecu_logical_address = 2
+    mock_socket.rx_queue.append(diagnostic_ack_invalid_source)
+    with pytest.raises(TimeoutError):
+        sut.send_diagnostic_to_address(0xE400, bytearray([0, 1, 2]), timeout=0.01)
+
+
+@pytest.mark.parametrize("allowed_addresses", [{2}, lambda address: address == 2])
+def test_receive_diagnostic_uses_allowed_addresses_and_checks_target(
+    mock_socket, allowed_addresses
+):
+    sut = DoIPClient(
+        test_ip, test_logical_address, allowed_response_addresses=allowed_addresses
+    )
+    sut._ecu_logical_address = 3
+    valid_response = bytearray(diagnostic_result)
+    valid_response[8:10] = b"\x00\x02"
+    wrong_target = bytearray(valid_response)
+    wrong_target[10:12] = b"\x0e\x01"
+    mock_socket.rx_queue.extend([diagnostic_result, wrong_target, valid_response])
+    assert sut.receive_diagnostic(timeout=0.01) == bytearray([0, 1, 2, 3])
+    assert mock_socket.rx_queue == []
+
+
+def test_diagnostic_ack_accepts_request_target_with_callable(mock_socket):
+    sut = DoIPClient(
+        test_ip, test_logical_address, allowed_response_addresses=lambda address: False
+    )
+    mock_socket.rx_queue.append(diagnostic_ack_to_address)
+    assert sut.send_diagnostic_to_address(0x1234, bytearray([0, 1, 2]), timeout=0.01) is None
+
+
 def test_request_vehicle_identification(mock_socket):
     sut = DoIPClient(test_ip, test_logical_address)
     mock_socket.rx_queue.append(vehicle_identification_response)