From 5cad6563f1b35cca9d8dafcccb35e95e7f718dcc Mon Sep 17 00:00:00 2001 From: Arnab Nandy Date: Thu, 25 Jun 2026 23:51:50 +0530 Subject: [PATCH] fix(iot-dev): Preserve response callback context --- .../iot/device/transport/IotHubTransport.java | 52 ++++++++++--------- .../device/transport/IotHubTransportTest.java | 48 +++++++++++++++++ 2 files changed, 75 insertions(+), 25 deletions(-) diff --git a/iothub/device/iot-device-client/src/main/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransport.java b/iothub/device/iot-device-client/src/main/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransport.java index 5c6ca241fb..2b59c65d2e 100644 --- a/iothub/device/iot-device-client/src/main/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransport.java +++ b/iothub/device/iot-device-client/src/main/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransport.java @@ -322,6 +322,16 @@ public void onMessageSent(Message message, String deviceId, TransportException e @Override public void onMessageReceived(IotHubTransportMessage message, TransportException e) { + CorrelationCallbackContext callbackContext = null; + if (message != null) + { + String correlationId = message.getCorrelationId(); + if (!correlationId.isEmpty()) + { + callbackContext = correlationCallbacks.get(correlationId); + } + } + if (message != null && e != null) { log.error("Exception encountered while receiving a message from service {}", message, e); @@ -338,36 +348,28 @@ else if (message != null) try { - if (message != null) + if (callbackContext != null && callbackContext.getCallback() != null) { - String correlationId = message.getCorrelationId(); - if (!correlationId.isEmpty()) + IotHubClientException clientException = null; + if (e != null) { - CorrelationCallbackContext callbackContext = correlationCallbacks.get(correlationId); - if (callbackContext != null && callbackContext.getCallback() != null) + // This case indicates that the transport layer failed to construct a valid message out of + // a message delivered by the service + clientException = e.toIotHubClientException(); + } + else + { + // This case indicates that the transport layer constructed a valid message out of a message + // delivered by the service, but that message may contain an unsuccessful status code in cases + // such as if an operation was rejected because it was badly formatted. + IotHubStatusCode statusCode = IotHubStatusCode.getIotHubStatusCode(Integer.parseInt(message.getStatus())); + if (!IotHubStatusCode.isSuccessful(statusCode)) { - IotHubClientException clientException = null; - if (e != null) - { - // This case indicates that the transport layer failed to construct a valid message out of - // a message delivered by the service - clientException = e.toIotHubClientException(); - } - else - { - // This case indicates that the transport layer constructed a valid message out of a message - // delivered by the service, but that message may contain an unsuccessful status code in cases - // such as if an operation was rejected because it was badly formatted. - IotHubStatusCode statusCode = IotHubStatusCode.getIotHubStatusCode(Integer.parseInt(message.getStatus())); - if (!IotHubStatusCode.isSuccessful(statusCode)) - { - clientException = new IotHubClientException(statusCode, "Received an unsuccessful operation error code from the service: " + statusCode); - } - } - - callbackContext.getCallback().onResponseReceived(message, callbackContext.getUserContext(), clientException); + clientException = new IotHubClientException(statusCode, "Received an unsuccessful operation error code from the service: " + statusCode); } } + + callbackContext.getCallback().onResponseReceived(message, callbackContext.getUserContext(), clientException); } } catch (Exception ex) diff --git a/iothub/device/iot-device-client/src/test/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransportTest.java b/iothub/device/iot-device-client/src/test/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransportTest.java index 0d6498f6f9..e3c2d26bae 100644 --- a/iothub/device/iot-device-client/src/test/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransportTest.java +++ b/iothub/device/iot-device-client/src/test/java/com/microsoft/azure/sdk/iot/device/transport/IotHubTransportTest.java @@ -225,6 +225,54 @@ public void onMessageReceivedWithMessageAndNoExceptionAddsToQueue() assertEquals(mockedTransportMessage, receivedMessagesQueue.poll()); } + //Tests_SRS_IOTHUBTRANSPORT_34_010: [If a correlation callback exists for a received message, this function shall call onResponseReceived even if the saved correlation callback is removed while queueing the received message.] + @Test + public void onMessageReceivedCallsResponseCallbackEvenIfCallbackIsRemovedWhileAddingMessageToQueue(@Mocked final CorrelatingMessageCallback mockedCorrelationCallback) + { + //arrange + final String correlationId = "1234"; + final Object callbackContext = new Object(); + final Map correlationCallbacks = new ConcurrentHashMap<>(); + correlationCallbacks.put(correlationId, new CorrelationCallbackContext(mockedCorrelationCallback, callbackContext, System.currentTimeMillis())); + + new Expectations() + { + { + mockedConfig.getDeviceId(); + result = "someDeviceId"; + + mockedTransportMessage.getCorrelationId(); + result = correlationId; + + mockedTransportMessage.getStatus(); + result = "200"; + } + }; + + new MockUp() + { + @Mock void addToReceivedMessagesQueue(IotHubTransportMessage message) + { + correlationCallbacks.remove(correlationId); + } + }; + + IotHubTransport transport = new IotHubTransport(mockedConfig, mockedIotHubConnectionStatusChangeCallback, false); + Deencapsulation.setField(transport, "correlationCallbacks", correlationCallbacks); + + //act + transport.onMessageReceived(mockedTransportMessage, null); + + //assert + new Verifications() + { + { + mockedCorrelationCallback.onResponseReceived(mockedTransportMessage, callbackContext, null); + times = 1; + } + }; + } + //Tests_SRS_IOTHUBTRANSPORT_34_014: [If the provided connectionId is associated with the current connection, This function shall invoke updateStatus with status CONNECTED, change reason CONNECTION_OK and a null throwable.] @Test public void onConnectionEstablishedCallsUpdateStatus()