Custom auth websocket support

.github/workflows/ci.yml

@@ -247,6 +247,9 @@ jobs:
       - name: run MQTT5 CustomAuthorizerConnect sample
         run: |
           python3 ${{ env.CI_UTILS_FOLDER }}/run_sample_ci.py --language Python --sample_file "${{ env.CI_SAMPLES_FOLDER }}/mqtt5_custom_authorizer_connect.py" --sample_region ${{ env.AWS_DEFAULT_REGION }} --sample_secret_endpoint 'ci/endpoint' --sample_secret_custom_authorizer_name 'ci/CustomAuthorizer/name' --sample_secret_custom_authorizer_password 'ci/CustomAuthorizer/password'
+      - name: run MQTT5 CustomAuthorizerConnect sample (websockets)
+        run: |
+          python3 ${{ env.CI_UTILS_FOLDER }}/run_sample_ci.py --language Python --sample_file "${{ env.CI_SAMPLES_FOLDER }}/mqtt5_custom_authorizer_connect.py" --sample_region ${{ env.AWS_DEFAULT_REGION }} --sample_secret_endpoint 'ci/endpoint' --sample_secret_custom_authorizer_name 'ci/CustomAuthorizer/name' --sample_secret_custom_authorizer_password 'ci/CustomAuthorizer/password' --sample_arguments '--use_websockets true'
       - name: configure AWS credentials (Shadow)
         uses: aws-actions/configure-aws-credentials@v1
         with:

awsiot/mqtt5_client_builder.py

@@ -622,6 +622,77 @@ def direct_with_custom_authorizer(
                     **kwargs)
 
 
+def websockets_with_custom_authorizer(
+        auth_username=None,
+        auth_authorizer_name=None,
+        auth_authorizer_signature=None,
+        auth_password=None,
+        websocket_proxy_options=None,
+        **kwargs) -> awscrt.mqtt5.Client:
+    """
+    This builder creates an :class:`awscrt.mqtt5.Client`, configured for an MQTT5 Client using a custom
+    authorizer through websockets. This function will set the username, port, and TLS options.
+
+    This function takes all :mod:`common arguments<awsiot.mqtt5_client_builder>`
+    described at the top of this doc, as well as...
+
+    Keyword Args:
+        auth_username (`str`): The username to use with the custom authorizer.
+            If provided, the username given will be passed when connecting to the custom authorizer.
+            If not provided, it will check to see if a username has already been set (via username="example")
+            and will use that instead.
+            If no username has been set then no username will be sent with the MQTT connection.
+
+        auth_authorizer_name (`str`):  The name of the custom authorizer.
+            If not provided, then "x-amz-customauthorizer-name" will not be added with the MQTT connection.
+
+        auth_authorizer_signature (`str`):  The signature of the custom authorizer.
+            If not provided, then "x-amz-customauthorizer-name" will not be added with the MQTT connection.
+
+        auth_password (`str`):  The password to use with the custom authorizer.
+            If not provided, then no passord will be set.
+
+        websocket_proxy_options (awscrt.http.HttpProxyOptions): Deprecated,
+            for proxy settings use `http_proxy_options` (described in
+            :mod:`common arguments<awsiot.mqtt5_client_builder>`)
+    """
+
+    _check_required_kwargs(**kwargs)
+    username_string = ""
+
+    if auth_username is None:
+        if not _get(kwargs, "username") is None:
+            username_string += _get(kwargs, "username")
+    else:
+        username_string += auth_username
+
+    if auth_authorizer_name is not None:
+        username_string = _add_to_username_parameter(
+            username_string, auth_authorizer_name, "x-amz-customauthorizer-name=")
+    if auth_authorizer_signature is not None:
+        username_string = _add_to_username_parameter(
+            username_string, auth_authorizer_signature, "x-amz-customauthorizer-signature=")
+
+    kwargs["username"] = username_string
+    kwargs["password"] = auth_password
+
+    tls_ctx_options = awscrt.io.TlsContextOptions()
+
+    def _sign_websocket_handshake_request(transform_args, **kwargs):
+        # transform_args need to know when transform is done
+        try:
+            transform_args.set_done()
+        except Exception as e:
+            transform_args.set_done(e)
+
+    return _builder(tls_ctx_options=tls_ctx_options,
+                    use_websockets=True,
+                    use_custom_authorizer=True,
+                    websocket_handshake_transform=_sign_websocket_handshake_request,
+                    websocket_proxy_options=websocket_proxy_options,
+                    **kwargs)
+
+
 def new_default_builder(**kwargs) -> awscrt.mqtt5.Client:
     """
     This builder creates an :class:`awscrt.mqtt5.Client`, without any configuration besides the default TLS context options.

awsiot/mqtt_connection_builder.py

@@ -456,7 +456,7 @@ def direct_with_custom_authorizer(
         **kwargs) -> awscrt.mqtt.Connection:
     """
     This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection using a custom
-    authorizer. This function will set the username, port, and TLS options.
+    authorizer using a direct MQTT connection. This function will set the username, port, and TLS options.
 
     This function takes all :mod:`common arguments<awsiot.mqtt_connection_builder>`
     described at the top of this doc, as well as...
@@ -478,6 +478,73 @@ def direct_with_custom_authorizer(
             If not provided, then no passord will be set.
     """
 
+    return _with_custom_authorizer(
+        auth_username=auth_username,
+        auth_authorizer_name=auth_authorizer_name,
+        auth_authorizer_signature=auth_authorizer_signature,
+        auth_password=auth_password,
+        use_websockets=False,
+        **kwargs)
+
+def websockets_with_custom_authorizer(
+        region,
+        credentials_provider,
+        auth_username=None,
+        auth_authorizer_name=None,
+        auth_authorizer_signature=None,
+        auth_password=None,
+        **kwargs) -> awscrt.mqtt.Connection:
+    """
+    This builder creates an :class:`awscrt.mqtt.Connection`, configured for an MQTT connection using a custom
+    authorizer using websockets. This function will set the username, port, and TLS options.
+
+    This function takes all :mod:`common arguments<awsiot.mqtt_connection_builder>`
+    described at the top of this doc, as well as...
+
+    Keyword Args:
+        region (str): AWS region to use when signing.
+
+        credentials_provider (awscrt.auth.AwsCredentialsProvider): Source of AWS credentials to use when signing.
+
+        auth_username (`str`): The username to use with the custom authorizer.
+            If provided, the username given will be passed when connecting to the custom authorizer.
+            If not provided, it will check to see if a username has already been set (via username="example")
+            and will use that instead.
+            If no username has been set then no username will be sent with the MQTT connection.
+
+        auth_authorizer_name (`str`):  The name of the custom authorizer.
+            If not provided, then "x-amz-customauthorizer-name" will not be added with the MQTT connection.
+
+        auth_authorizer_signature (`str`):  The signature of the custom authorizer.
+            If not provided, then "x-amz-customauthorizer-name" will not be added with the MQTT connection.
+
+        auth_password (`str`):  The password to use with the custom authorizer.
+            If not provided, then no passord will be set.
+    """
+
+    return _with_custom_authorizer(
+        auth_username=auth_username,
+        auth_authorizer_name=auth_authorizer_name,
+        auth_authorizer_signature=auth_authorizer_signature,
+        auth_password=auth_password,
+        use_websockets=True,
+        websockets_region=region,
+        websockets_credentials_provider=credentials_provider,
+        **kwargs)
+
+
+def _with_custom_authorizer(auth_username=None,
+        auth_authorizer_name=None,
+        auth_authorizer_signature=None,
+        auth_password=None,
+        use_websockets=False,
+        websockets_credentials_provider=None,
+        websockets_region=None,
+        **kwargs) -> awscrt.mqtt.Connection:
+    """
+    Helper function that contains the setup needed for custom authorizers
+    """
+
     _check_required_kwargs(**kwargs)
     username_string = ""
 
@@ -496,14 +563,33 @@ def direct_with_custom_authorizer(
 
     kwargs["username"] = username_string
     kwargs["password"] = auth_password
-    kwargs["port"] = 443
 
     tls_ctx_options = awscrt.io.TlsContextOptions()
-    tls_ctx_options.alpn_list = ["mqtt"]
+    if use_websockets == False:
+        kwargs["port"] = 443
+        tls_ctx_options.alpn_list = ["mqtt"]
+
+    def _sign_websocket_handshake_request(transform_args, **kwargs):
+        # transform_args need to know when transform is done
+        try:
+            signing_config = awscrt.auth.AwsSigningConfig(
+                algorithm=awscrt.auth.AwsSigningAlgorithm.V4,
+                signature_type=awscrt.auth.AwsSignatureType.HTTP_REQUEST_QUERY_PARAMS,
+                credentials_provider=websockets_credentials_provider,
+                region=websockets_region,
+                service='iotdevicegateway',
+                omit_session_token=True,  # IoT is weird and does not sign X-Amz-Security-Token
+            )
+
+            signing_future = awscrt.auth.aws_sign_request(transform_args.http_request, signing_config)
+            signing_future.add_done_callback(lambda x: transform_args.set_done(x.exception()))
+        except Exception as e:
+            transform_args.set_done(e)
 
     return _builder(tls_ctx_options=tls_ctx_options,
-                    use_websockets=False,
+                    use_websockets=use_websockets,
                     use_custom_authorizer=True,
+                    websocket_handshake_transform=_sign_websocket_handshake_request if use_websockets else None,
                     **kwargs)
 
 

samples/mqtt5_custom_authorizer_connect.py

@@ -18,13 +18,15 @@
 cmdUtils.register_command("client_id", "<str>",
                           "Client ID to use for MQTT connection (optional, default='test-*').",
                           default="test-" + str(uuid4()))
+cmdUtils.register_command("use_websockets", "<str>", "If set, websockets will be used (optional, do not set to use direct MQTT)")
 cmdUtils.register_command("is_ci", "<str>", "If present the sample will run in CI mode (optional, default='None')")
 # Needs to be called so the command utils parse the commands
 cmdUtils.get_args()
 
 future_stopped = Future()
 future_connection_success = Future()
 is_ci = cmdUtils.get_command("is_ci", None) != None
+use_websockets = cmdUtils.get_command("use_websockets", None) != None
 
 # Callback for the lifecycle event Stopped
 def on_lifecycle_stopped(lifecycle_stopped_data: mqtt5.LifecycleStoppedData):
@@ -42,16 +44,29 @@ def on_lifecycle_connection_success(lifecycle_connect_success_data: mqtt5.Lifecy
 if __name__ == '__main__':
 
     # Create MQTT5 Client with a custom authorizer
-    client = mqtt5_client_builder.direct_with_custom_authorizer(
-        endpoint=cmdUtils.get_command_required(cmdUtils.m_cmd_endpoint),
-        ca_filepath=cmdUtils.get_command(cmdUtils.m_cmd_ca_file),
-        auth_username=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_username),
-        auth_authorizer_name=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_name),
-        auth_authorizer_signature=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_signature),
-        auth_password=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_password),
-        on_lifecycle_stopped=on_lifecycle_stopped,
-        on_lifecycle_connection_success=on_lifecycle_connection_success,
-        client_id=cmdUtils.get_command("client_id"))
+    client = None
+    if use_websockets == None:
+        client = mqtt5_client_builder.direct_with_custom_authorizer(
+            endpoint=cmdUtils.get_command_required(cmdUtils.m_cmd_endpoint),
+            ca_filepath=cmdUtils.get_command(cmdUtils.m_cmd_ca_file),
+            auth_username=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_username),
+            auth_authorizer_name=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_name),
+            auth_authorizer_signature=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_signature),
+            auth_password=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_password),
+            on_lifecycle_stopped=on_lifecycle_stopped,
+            on_lifecycle_connection_success=on_lifecycle_connection_success,
+            client_id=cmdUtils.get_command("client_id"))
+    else:
+        client = mqtt5_client_builder.websockets_with_custom_authorizer(
+            endpoint=cmdUtils.get_command_required(cmdUtils.m_cmd_endpoint),
+            ca_filepath=cmdUtils.get_command(cmdUtils.m_cmd_ca_file),
+            auth_username=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_username),
+            auth_authorizer_name=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_name),
+            auth_authorizer_signature=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_authorizer_signature),
+            auth_password=cmdUtils.get_command(cmdUtils.m_cmd_custom_auth_password),
+            on_lifecycle_stopped=on_lifecycle_stopped,
+            on_lifecycle_connection_success=on_lifecycle_connection_success,
+            client_id=cmdUtils.get_command("client_id"))
 
     if is_ci == False:
         print("Connecting to {} with client ID '{}'...".format(
@@ -61,7 +76,7 @@ def on_lifecycle_connection_success(lifecycle_connect_success_data: mqtt5.Lifecy
 
     client.start()
     future_connection_success.result(TIMEOUT)
-    print("Clint Connected")
+    print("Client Connected")
 
     print("Stopping Client")
     client.stop()