MQTT5Client disconnects and does not reconnect when connecting with authenticated Cognito Identity

[KamleshAtTray]

We connect to AWS IOT broker using authenticated access via AWS COgnito . we use the custom developer identity provider to authenticate our identities.

The overall process is done in 3 steps -

1. Using an IAM user with its accesskey and secretKey we call the GetOpenIDTokenForDeveloperIdentity API from the Cognito ecosystem to get IdentityId and token.
2. Then we call the AttachPolicy API to attach an IOT Core Policy to Cognito identity, this allows the identity to publish, connect, receive, Subscribe actions with IoT core.
3. We configure these identityId and token in the CognitoCredentialsProvider, then pass it as a config to the IoT Client builder via the WebSocketConfig, create an MQTT5Client, and start the client.

But we see that after connecting once the IoT client disconnects from the broker with error code - 1051 after roughly 15-20 mins. The client fails with errorcode 6146 everytime it attempts reconnection.

The MQTT5Client reconnected successfully when we connected using Cognito guess access.

We suspect that the OpenID token we get from Cognito System expires due to which the MQTT5Client is unable to reconnect to the IOT broker endpoint.

How can we achieve successful reconnection?

private void createMQTTClientUsingAuthenticatedCognitoRole(Mqtt5ClientOptions.PublishEvents publishEventsCallback) {

        IoTConfigurations configurations = util.getConfigurations();
        final String serverUniqueIdentifier = "IoTTester" + System.getProperty("tray.nodeId");

        //Fetch identityId and token from Cognito Server
        AuthenticationUtils.TrayCognitoAuthDetails trayCognitoAuthDetails = AuthenticationUtils.getCognitoIdentityAndToken(configurations.getAccessKey(),configurations.getSecretKey(),configurations.getCognitoDeveloperAuthenticatedIdentityProviderName(), configurations.getCognitoIdentityPoolId(), serverUniqueIdentifier);

        final String cognitoIdentityId = trayCognitoAuthDetails.getIdentityId();
        final String cognitoEndpoint = configurations.getCognitoEndpoint();
        final String ioTEndpoint = configurations.getIotBrokerEndpoint();
        final String token = trayCognitoAuthDetails.getToken();
        //While using CognitoCredentialsProvider use the below identityProviderName as key and token received from GetOpenIdTokenForDeveloperIdentity API.
        final String identityProviderName = "cognito-identity.amazonaws.com";
        final String serverIoTClientId = serverUniqueIdentifier + "_" + cognitoIdentityId;

        //Attach IoT Policy to the Cognito identity
        boolean attachPolicySuccessFlag = AuthenticationUtils.attachIoTPolicyToIdentity(serverUniqueIdentifier, cognitoIdentityId, configurations.getIotPolicyName());

        //Cognito TLS options
        TlsContextOptions cognitoTlsContextOptions = TlsContextOptions.createDefaultClient();
        ClientTlsContext cognitoTlsContext = new ClientTlsContext(cognitoTlsContextOptions);
        cognitoTlsContextOptions.close();

        //Cognito Credentials logins map
        CognitoCredentialsProvider.CognitoLoginTokenPair cognitoLoginTokenPair = new CognitoCredentialsProvider.CognitoLoginTokenPair(identityProviderName, token);

        CognitoCredentialsProvider.CognitoCredentialsProviderBuilder cognitoBuilder = new CognitoCredentialsProvider.CognitoCredentialsProviderBuilder();

        cognitoBuilder.withLogin(cognitoLoginTokenPair)
                .withIdentity(cognitoIdentityId)
                .withEndpoint(cognitoEndpoint)
                .withClientBootstrap(ClientBootstrap.getOrCreateStaticDefault())
                .withTlsContext(cognitoTlsContext);

        //IoT WebSocket Config
        AwsIotMqtt5ClientBuilder.WebsocketSigv4Config websocketSigv4Config = new AwsIotMqtt5ClientBuilder.WebsocketSigv4Config();
        websocketSigv4Config.credentialsProvider = cognitoBuilder.build();

        //Setting up the IoT Client configuration
        ConnectPacket.ConnectPacketBuilder connectProperties = new ConnectPacket.ConnectPacketBuilder()
                .withClientId(serverIoTClientId)
                .withSessionExpiryIntervalSeconds(configurations.getSessionExpiryIntervalInSeconds());  //Set for persistent sessions
        IoTClientLifeCycle lifecycleEvents = new IoTClientLifeCycle();

        try(
                AwsIotMqtt5ClientBuilder mqtt5ClientBuilder = AwsIotMqtt5ClientBuilder
                        .newWebsocketMqttBuilderWithSigv4Auth(ioTEndpoint, websocketSigv4Config)
                        .withConnectProperties(connectProperties)
                        .withLifeCycleEvents(lifecycleEvents)
                        .withPublishEvents(publishEventsCallback);
        ) {
            mqtt5Client = mqtt5ClientBuilder.build();
            if (mqtt5Client == null) {
                throw new RuntimeException("Failed to construct MQTT5Client.");
            } else {
                mqtt5Client.start();
            }
        }  catch (Exception e) {
            log.error("Exception while starting MQTT5Client", e);
            throw e;
        }
    }

[bretambrose]

This looks like a serious API defect; there is no way to dynamically modify the login key value pairs over time. I'm not sure what the right change would be -- possibly a callback to let you add pairs every time the cognito http request is made, in that case you'd move the first few lines into the callback and have it inject your provider-token pair there. I'll try to get some traction on this internally, but I can't give a timeline yet.

[bretambrose]

Current tentative proposal is that we will add a new configuration option that allows you to set a callback function that will be invoked right before credentials are fetched from cognito:

fn CallbackFunction() -> CompletableFuture<List<CognitoLoginTokenPair>>

If this callback is set, then the credential fetch will be deferred until the returned future completes. The pairs in the future's result will be added to the cognito API call to fetch credentials. In your case, you'd leave the login token pairs empty in the provider constructor (since the only pair you want is session-limited) and you'd implement your callback to asynchronously fetch the token pair you need and complete the returned future with it.

You would also need to wrap the cognito provider with a caching provider (https://github.com/awslabs/aws-crt-java/blob/main/src/main/java/software/amazon/awssdk/crt/auth/credentials/CachedCredentialsProvider.java), where the CachedCredentialsProvider's caching duration is a value comfortably less than your sourced token's expiration time.

[KamleshAtTray]

@bretambrose Thank you for taking the time to answer this.

Can you give us a rough idea of when this callback feature might be available? We rely heavily on authenticated Cognito identities, which is really important.

In the meantime, is there any manual workaround without an API to force a token refresh? We're looking for a way to ensure that offline devices can reconnect properly using persistent sessions.

[bretambrose]

I cannot give a time estimate at this time.

For a workaround:

The CRT provides a generic credentials provider that will source credentials from whatever method you would like: https://github.com/awslabs/aws-crt-java/blob/main/src/main/java/software/amazon/awssdk/crt/auth/credentials/DelegateCredentialsProvider.java

So you could use this provider and in your callback, you could source the session credentials directly from cognito using GetCredentialsForIdentity: https://docs.aws.amazon.com/cognitoidentity/latest/APIReference/API_GetCredentialsForIdentity.html

Before calling GetCredentialsForIdentity, you would source your id provider token and configure the Logins key-value set appropriately.

The downside of this solution is that it will block the event loop the connection is seated on, since the DelegateProvider's API is not asynchronous.

[KamleshAtTray]

Hi @bretambrose

We implemented the workaround, the MQTT client is unable to reconnect to the IoT broker especially when the machine on which this code runs comes back online after being disconnected for few minutes. What are we doing wrong here ?

For testing purposes, we've set the cognito token expiration is 60 seconds and cache_duration to 58 seconds.

Code for your reference -

`Callback -

public class CognitoCredentialsCallBack implements DelegateCredentialsHandler {

private static final Logger logger = LogManager.getLogger(CognitoCredentialsCallBack.class);
private AWSIoTHelper iotHelper = new AWSIoTHelper();

@Override
public Credentials getCredentials() {

    final String serverUniqueIdentifier = "Server" + System.getProperty("nodeId");

    logger.info("getCredentials called");

    GetOpenIdTokenForDeveloperIdentityResponse identityIdAndTokenResponse = getIdentityAndToken(serverUniqueIdentifier);
    GetCredentialsForIdentityResponse credentialsResponse = getCredentialsForIdentityIdAndToken(identityIdAndTokenResponse.identityId(), identityIdAndTokenResponse.token());
    byte[] accessKey = credentialsResponse.credentials().accessKeyId().getBytes(CharsetUtil.UTF_8);
    byte[] secretKey = credentialsResponse.credentials().secretKey().getBytes(CharsetUtil.UTF_8);
    byte[] session = credentialsResponse.credentials().sessionToken().getBytes(CharsetUtil.UTF_8);
    return new Credentials(accessKey, secretKey, session);
}

private GetCredentialsForIdentityResponse getCredentialsForIdentityIdAndToken(String identityId, String token) {

    logger.info("Called getCredentialsForIdentityIdAndToken from thread {}", Thread.currentThread().getName());

    try (CognitoIdentityClient client = CognitoIdentityClient.builder().region(Region.US_WEST_2).build();){

        GetCredentialsForIdentityRequest request = GetCredentialsForIdentityRequest.builder()
                .identityId(identityId)
                .logins(Map.of("cognito-identity.amazonaws.com", token)).build();

        GetCredentialsForIdentityResponse response  = client.getCredentialsForIdentity(request);

        String maskedAccessKeyId = "xxxx" + response.credentials().accessKeyId().substring(response.credentials().accessKeyId().length()-2);

        logger.info("Fetched credentials accessKey={}", maskedAccessKeyId);

        return response;

    }
}

private GetOpenIdTokenForDeveloperIdentityResponse getIdentityAndToken (String id) {

    IoTConfigurations configurations = iotHelper.getConfigurations();

    logger.info("Called getIdentityAndToken from thread {}", Thread.currentThread().getName());

    String accessKey = configurations.getAccessKey();
    String secretAccessKey = configurations.getSecretKey();

    try (CognitoIdentityClient client = CognitoIdentityClient.builder()
            .region(Region.US_WEST_2)
            .credentialsProvider(StaticCredentialsProvider
                    .create(AwsBasicCredentials.create(accessKey, secretAccessKey)))
            .build();
    ) {

        Map<String, String> map = Map.of(configurations.getCognitoDeveloperAuthenticatedIdentityProviderName(), id);

        //request
        GetOpenIdTokenForDeveloperIdentityRequest request = GetOpenIdTokenForDeveloperIdentityRequest.builder()
                .identityPoolId(configurations.getCognitoIdentityPoolId())
                .logins(map)
                //TODO: Make configurable
                .tokenDuration(60L) // 1 min token expiry
                .build();

        //response
        GetOpenIdTokenForDeveloperIdentityResponse response = client.getOpenIdTokenForDeveloperIdentity(request);

        logger.info("Received Response, identityId={} and token={}", response.identityId(), "xxxxx" + response.token().substring(response.token().length()-4));

        //attach Policy
        AuthenticationUtils.attachIoTPolicyToIdentity(configurations.getTrayUniqueId(),
                response.identityId(),
                configurations.getIotPolicyName());

        return response;
    }
}

}

IOT CLient code

CognitoCredentialsCallBack cognitoCredentialsCallBack = new CognitoCredentialsCallBack();
try (DelegateCredentialsProvider delegateCredentialsProvider = new DelegateCredentialsProvider.DelegateCredentialsProviderBuilder()
.withHandler(cognitoCredentialsCallBack)
.build();

         CachedCredentialsProvider iotCachedCredentialsProvider = new CachedCredentialsProvider.CachedCredentialsProviderBuilder()
                 .withCachedProvider(delegateCredentialsProvider)
                 .withCachingDurationInSeconds(59)
                 .build();
    ) {

        //IoT WebSocket Config
        AwsIotMqtt5ClientBuilder.WebsocketSigv4Config websocketSigv4Config = new AwsIotMqtt5ClientBuilder.WebsocketSigv4Config();
        websocketSigv4Config.credentialsProvider = iotCachedCredentialsProvider;
        websocketSigv4Config.region = Region.US_WEST_2.toString();

        //Setting up the IoT Client configuration
        ConnectPacket.ConnectPacketBuilder connectProperties = new ConnectPacket.ConnectPacketBuilder()
                .withClientId(serverInstanceUniqueIdentifier)
                .withSessionExpiryIntervalSeconds(configurations.getSessionExpiryIntervalInSeconds());  //Set for persistent sessions

        IoTClientLifeCycle lifecycleEvents = new IoTClientLifeCycle();

        try(
                AwsIotMqtt5ClientBuilder mqtt5ClientBuilder = AwsIotMqtt5ClientBuilder
                        .newWebsocketMqttBuilderWithSigv4Auth(configurations.getIotBrokerEndpoint(), websocketSigv4Config)
                        .withConnectProperties(connectProperties)
                        .withLifeCycleEvents(lifecycleEvents)
                        .withPublishEvents(publishEventsCallback);
        ) {
            mqtt5Client = mqtt5ClientBuilder.build();
            if (mqtt5Client == null) {
                throw new RuntimeException("Failed to construct MQTT5Client.");
            } else {
                mqtt5Client.start();
            }
        }  catch (Exception e) {
            log.error("Exception while starting MQTT5Client", e);
            throw e;
        }
    }

`

[bretambrose]

Honestly, this is something you're going to have to debug yourself.

Logs can be helpful in seeing what's going on. You can enable CRT logging via system properties using the instructions here: https://github.com/awslabs/aws-crt-java?tab=readme-ov-file#system-properties

[KamleshAtTray]

@bretambrose
When the internet disconnects, the MQTT client detects the loss of connectivity and triggers the onDisconnection callback from LifeCycleEvents, as well as the getCredentials() callback method implemented in our DelegateCredentialsHandler (in our case, CognitoCredentialsCallback).

Within the credentials callback, we invoke the following Cognito and IoT APIs:

1. GetOpenIdTokenForDeveloperIdentity: Fetches a new token for the given identityId.
2. GetCredentialsForIdentity: Retrieves credentials for the token and identityId.
3. AttachIoTPolicy: Attaches the appropriate policy to the identity.

We have observed that reconnection does not occur immediately upon losing the internet connection. Instead, there is a variable time interval after which the MQTT5Client attempts to reconnect to the IoT broker by calling the getCredentials() callback. Because the internet is disconnected, the call to GetOpenIdTokenForDeveloperIdentity fails and is retried three times before the MQTT5Client stops attempting to reconnect.

How can we ensure that the MQTT5Client keeps calling the getCredentials() callback until the device is back online?

To enforce this reconnection in a loop, we added the following to the onConnectionFailure() and onDisconnection() LifeCycleEvent's callback method, but this is not enforcing the MQTT5Client to keep trying the reconnection until the device comes back online or the client reconnects successfully.

//code added to onDisconnection() and onConnection() lifeCycleEvents callback class -

try {
            log.info("Attempting to reconnect...");
            attemptReconnect(client);
        } catch (Exception e) {
            log.error("Failed to reconnect", e);
        }

private void attemptReconnect(Mqtt5Client client) {
        int retryCount = 0;
        while (!client.getIsConnected() && !shutdownFlag) { // Keep retrying indefinitely
            try {
                // Check internet connectivity before reconnecting
                if (isInternetAvailable()) {
                    log.info("Internet is available. Attempting to reconnect (attempt {})...", retryCount+1);
                    client.start();
                    log.info("Reconnection successful.");
                } else {
                    log.warn("No internet connection. Waiting before retrying...");
                }
            } catch (Exception e) {
                log.error("Reconnection attempt {} failed", retryCount+1, e);
            }

            // Calculate delay using exponential backoff
            long delay = (long) Math.pow(2, retryCount) * INITIAL_RETRY_DELAY_MS;
            log.info("Waiting {} ms before retrying...", delay);

            try {
                Thread.sleep(delay); // Wait before retrying
            } catch (InterruptedException ie) {
                log.warn("Reconnection thread interrupted", ie);
                Thread.currentThread().interrupt(); // Restore the interrupted status
                break; // Exit the loop if the thread is interrupted
            }
            retryCount++;
        }
    }

    private boolean isInternetAvailable() {
        try {
            URL url = new URL("https://www.google.com");
            HttpURLConnection urlConnect = (HttpURLConnection) url.openConnection();
            urlConnect.setConnectTimeout(5000); // Timeout in milliseconds
            urlConnect.connect();
            return urlConnect.getResponseCode() == 200;
        } catch (IOException e) {
            System.out.format("Internet connectivity check failed: %s\n", e.getMessage());
            return false;
        }
    }

Cognito callback failing while calling the GetOpenIdTokenForDeveloperIdentity
IOT_Cognito_Reconnectivity_Issue.txt

AWS CRT logs
AWS_IOT_Logs.txt