Pass CRT error, if it exists, into completion function

awslabs · Apr 29, 2024 · d8c01ae · d8c01ae
1 parent c11a554
commit d8c01ae
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/awscrt/mqtt5.py b/awscrt/mqtt5.py
@@ -1391,11 +1391,14 @@ def __init__(self, client_options: ClientOptions):
 
     def _ws_handshake_transform(self, http_request_binding, http_headers_binding, native_userdata):
         if self._ws_handshake_transform_cb is None:
-            _awscrt.mqtt5_ws_handshake_transform_complete(None, native_userdata)
+            _awscrt.mqtt5_ws_handshake_transform_complete(None, native_userdata, 0)
             return
 
         def _on_complete(f):
-            _awscrt.mqtt5_ws_handshake_transform_complete(f.exception(), native_userdata)
+            error_code = 0
+            if isinstance(f, exceptions.AwsCrtError):
+                error_code = f.code
+            _awscrt.mqtt5_ws_handshake_transform_complete(f.exception(), native_userdata, error_code)
 
         future = Future()
         future.add_done_callback(_on_complete)

diff --git a/source/mqtt5_client.c b/source/mqtt5_client.c
@@ -711,13 +711,13 @@ PyObject *aws_py_mqtt5_ws_handshake_transform_complete(PyObject *self, PyObject
 
     PyObject *exception_py;
     PyObject *ws_transform_capsule;
-    if (!PyArg_ParseTuple(args, "OO", &exception_py, &ws_transform_capsule)) {
+    int error_code = AWS_ERROR_SUCCESS;
+    if (!PyArg_ParseTuple(args, "OOi", &exception_py, &ws_transform_capsule, &error_code)) {
         return NULL;
     }
 
-    int error_code = AWS_ERROR_SUCCESS;
-    if (exception_py != Py_None) {
-        /* TODO: Translate Python exception to aws error. In the meantime here's a catch-all. */
+    if (exception_py != Py_None && error_code == AWS_ERROR_SUCCESS) {
+        /* Fallback code for if the error source was outside the CRT native implementation */
         error_code = AWS_ERROR_HTTP_CALLBACK_FAILURE;
     }