python credential provider (#215)

AwsCredentialsProvider.new_delegate(...) allows custom credentials provider to be defined with python code

Author: TingDaoK

awscrt/auth.py

@@ -101,44 +101,29 @@ def __deepcopy__(self, memo):
 
 
 class AwsCredentialsProviderBase(NativeResource):
-    """
-    Base class for providers that source the AwsCredentials needed to sign an authenticated AWS request.
-
-    NOTE: Custom subclasses of AwsCredentialsProviderBase are not yet supported.
-    """
-    __slots__ = ()
-
-    def __init__(self, binding=None):
-        super().__init__()
-
-        if binding is None:
-            # TODO: create binding type that lets native code call into python subclass
-            raise NotImplementedError("Custom subclasses of AwsCredentialsProviderBase are not yet supported")
-
-        self._binding = binding
-
-    def get_credentials(self):
-        """
-        Asynchronously fetch AwsCredentials.
-
-        Returns:
-            concurrent.futures.Future: A Future which will contain
-            :class:`AwsCredentials` (or an exception) when the operation completes.
-            The operation may complete on a different thread.
-        """
-        raise NotImplementedError()
+    # Pointless base class, kept for backwards compatibility.
+    # AwsCredentialsProvider is (and always will be) the only subclass.
+    #
+    # Originally created with the thought that, when we supported
+    # custom python providers, they would inherit from this class.
+    # We ended up supporting custom python providers via
+    # AwsCredentialsProvider.new_delegate() instead.
+    pass
 
 
 class AwsCredentialsProvider(AwsCredentialsProviderBase):
     """
     Credentials providers source the AwsCredentials needed to sign an authenticated AWS request.
 
-    Base class: AwsCredentialsProviderBase
-
     This class provides `new()` functions for several built-in provider types.
+    To define a custom provider, use the `new_delegate()` function.
     """
     __slots__ = ()
 
+    def __init__(self, binding):
+        super().__init__()
+        self._binding = binding
+
     @classmethod
     def new_default_chain(cls, client_bootstrap):
         """
@@ -289,7 +274,35 @@ def new_chain(cls, providers):
         binding = _awscrt.credentials_provider_new_chain(providers)
         return cls(binding)
 
+    @classmethod
+    def new_delegate(cls, get_credentials):
+        """
+        Creates a provider that sources credentials from a custom
+        synchronous callback.
+
+        Args:
+            get_credentials: Callable which takes no arguments and returns
+                :class:`AwsCredentials`.
+
+        Returns:
+            AwsCredentialsProvider:
+        """
+        # TODO: support async delegates
+
+        assert callable(get_credentials)
+
+        binding = _awscrt.credentials_provider_new_delegate(get_credentials)
+        return cls(binding)
+
     def get_credentials(self):
+        """
+        Asynchronously fetch AwsCredentials.
+
+        Returns:
+            concurrent.futures.Future: A Future which will contain
+            :class:`AwsCredentials` (or an exception) when the operation completes.
+            The operation may complete on a different thread.
+        """
         future = Future()
 
         def _on_complete(error_code, binding):
@@ -306,7 +319,7 @@ def _on_complete(error_code, binding):
         try:
             _awscrt.credentials_provider_get_credentials(self._binding, _on_complete)
         except Exception as e:
-            future.set_result(e)
+            future.set_exception(e)
 
         return future
 
@@ -383,7 +396,7 @@ class AwsSigningConfig(NativeResource):
         signature_type (AwsSignatureType): Which sort of signature should be
             computed from the signable.
 
-        credentials_provider (AwsCredentialsProviderBase): Credentials provider
+        credentials_provider (AwsCredentialsProvider): Credentials provider
             to fetch signing credentials with.
 
         region (str): The region to sign against.
@@ -470,7 +483,7 @@ def __init__(self,
 
         assert isinstance(algorithm, AwsSigningAlgorithm)
         assert isinstance(signature_type, AwsSignatureType)
-        assert isinstance(credentials_provider, AwsCredentialsProviderBase)
+        assert isinstance(credentials_provider, AwsCredentialsProvider)
         assert isinstance(region, str)
         assert isinstance(service, str)
         assert callable(should_sign_header) or should_sign_header is None
@@ -533,7 +546,7 @@ def signature_type(self):
 
     @property
     def credentials_provider(self):
-        """AwsCredentialsProviderBase: Credentials provider to fetch signing credentials with"""
+        """AwsCredentialsProvider: Credentials provider to fetch signing credentials with"""
         return _awscrt.signing_config_get_credentials_provider(self._binding)
 
     @property

awscrt/awsiot_mqtt_connection_builder.py

@@ -226,7 +226,7 @@ def websockets_with_default_aws_signing(region, credentials_provider, websocket_
     Arguments:
         region (str): AWS region to use when signing.
 
-        credentials_provider (awscrt.auth.AwsCredentialsProviderBase): Source of AWS credentials to use when signing.
+        credentials_provider (awscrt.auth.AwsCredentialsProvider): Source of AWS credentials to use when signing.
 
         websocket_proxy_options (awscrt.http.HttpProxyOptions): If specified, a proxy is used when connecting.
 

crt/aws-c-auth

@@ -21,6 +21,7 @@ PyObject *aws_py_credentials_provider_new_profile(PyObject *self, PyObject *args
 PyObject *aws_py_credentials_provider_new_process(PyObject *self, PyObject *args);
 PyObject *aws_py_credentials_provider_new_environment(PyObject *self, PyObject *args);
 PyObject *aws_py_credentials_provider_new_chain(PyObject *self, PyObject *args);
+PyObject *aws_py_credentials_provider_new_delegate(PyObject *self, PyObject *args);
 
 PyObject *aws_py_signing_config_new(PyObject *self, PyObject *args);
 PyObject *aws_py_signing_config_get_algorithm(PyObject *self, PyObject *args);

source/auth.h

@@ -135,10 +135,16 @@ PyObject *aws_py_credentials_expiration_timestamp_seconds(PyObject *self, PyObje
  */
 struct credentials_provider_binding {
     struct aws_credentials_provider *native;
+
+    /* Python get_credentials() callable.
+     * Only used by "delegate" provider type */
+    PyObject *py_delegate;
 };
 
 /* Finally clean up binding (after capsule destructor runs and credentials provider shutdown completes) */
 static void s_credentials_provider_binding_clean_up(struct credentials_provider_binding *binding) {
+    Py_XDECREF(binding->py_delegate);
+
     aws_mem_release(aws_py_get_allocator(), binding);
 }
 
@@ -166,7 +172,7 @@ struct aws_credentials_provider *aws_py_get_credentials_provider(PyObject *crede
     AWS_PY_RETURN_NATIVE_FROM_BINDING(
         credentials_provider,
         s_capsule_name_credentials_provider,
-        "AwsCredentialsProviderBase",
+        "AwsCredentialsProvider",
         credentials_provider_binding);
 }
 
@@ -559,3 +565,103 @@ PyObject *aws_py_credentials_provider_new_chain(PyObject *self, PyObject *args)
     Py_XDECREF(capsule);
     return NULL;
 }
+
+static int s_credentials_provider_delegate_get_credentials(
+    void *delegate_user_data,
+    aws_on_get_credentials_callback_fn callback,
+    void *callback_user_data) {
+
+    struct credentials_provider_binding *binding = delegate_user_data;
+
+    PyGILState_STATE state;
+    if (aws_py_gilstate_ensure(&state)) {
+        /* Python has shut down. Nothing matters anymore, but don't crash */
+        return aws_raise_error(AWS_ERROR_INVALID_STATE);
+    }
+
+    struct aws_credentials *native_credentials = NULL;
+
+    PyObject *py_result = PyObject_CallFunction(binding->py_delegate, "()");
+    if (!py_result) {
+        AWS_LOGF_ERROR(
+            AWS_LS_AUTH_CREDENTIALS_PROVIDER,
+            "(id=%p) Exception in get_credentials() delegate callback",
+            (void *)binding->native);
+
+        PyErr_WriteUnraisable(binding->py_delegate);
+        goto done;
+    }
+
+    /* Expect py_result to be AwsCredentials (which wraps native aws_credentials). */
+    native_credentials = aws_py_get_credentials(py_result);
+    if (!native_credentials) {
+        AWS_LOGF_ERROR(
+            AWS_LS_AUTH_CREDENTIALS_PROVIDER,
+            "(id=%p) get_credentials() delegate callback failed to return AwsCredentials",
+            (void *)binding->native);
+
+        PyErr_WriteUnraisable(binding->py_delegate);
+        goto done;
+    }
+
+    /* Keep native aws_credentials alive until we pass them to callback. */
+    aws_credentials_acquire(native_credentials);
+
+done:
+    /* Decref the python AwsCredentials (or whatever else was returned) before releasing the GIL */
+    Py_XDECREF(py_result);
+
+    PyGILState_Release(state);
+
+    if (!native_credentials) {
+        return aws_raise_error(AWS_ERROR_CRT_CALLBACK_EXCEPTION);
+    }
+
+    callback(native_credentials, AWS_ERROR_SUCCESS, callback_user_data);
+    aws_credentials_release(native_credentials);
+    return AWS_OP_SUCCESS;
+}
+
+PyObject *aws_py_credentials_provider_new_delegate(PyObject *self, PyObject *args) {
+    (void)self;
+    struct aws_allocator *allocator = aws_py_get_allocator();
+
+    PyObject *py_delegate;
+
+    if (!PyArg_ParseTuple(args, "O", &py_delegate)) {
+        return NULL;
+    }
+
+    struct credentials_provider_binding *binding;
+    PyObject *capsule = s_new_credentials_provider_binding_and_capsule(&binding);
+    if (!capsule) {
+        return NULL;
+    }
+
+    binding->py_delegate = py_delegate;
+    Py_INCREF(py_delegate);
+
+    /* From hereon, we need to clean up if errors occur.
+     * Fortunately, the capsule destructor will clean up anything stored inside the binding */
+
+    struct aws_credentials_provider_delegate_options options = {
+        .get_credentials = s_credentials_provider_delegate_get_credentials,
+        .delegate_user_data = binding,
+        .shutdown_options =
+            {
+                .shutdown_callback = s_credentials_provider_shutdown_complete,
+                .shutdown_user_data = binding,
+            },
+    };
+
+    binding->native = aws_credentials_provider_new_delegate(allocator, &options);
+    if (!binding->native) {
+        PyErr_SetAwsLastError();
+        goto error;
+    }
+
+    return capsule;
+error:
+    Py_DECREF(capsule);
+    return NULL;
+}

source/auth_credentials.c

@@ -382,7 +382,7 @@ int aws_py_gilstate_ensure(PyGILState_STATE *out_state) {
 
 void *aws_py_get_binding(PyObject *obj, const char *capsule_name, const char *class_name) {
     if (!obj || obj == Py_None) {
-        return PyErr_Format(PyExc_TypeError, "Excepted '%s', received 'NoneType'", class_name);
+        return PyErr_Format(PyExc_TypeError, "Expected '%s', received 'NoneType'", class_name);
     }
 
     PyObject *py_binding = PyObject_GetAttrString(obj, "_binding"); /* new reference */
@@ -554,6 +554,7 @@ static PyMethodDef s_module_methods[] = {
     AWS_PY_METHOD_DEF(credentials_provider_new_process, METH_VARARGS),
     AWS_PY_METHOD_DEF(credentials_provider_new_environment, METH_VARARGS),
     AWS_PY_METHOD_DEF(credentials_provider_new_chain, METH_VARARGS),
+    AWS_PY_METHOD_DEF(credentials_provider_new_delegate, METH_VARARGS),
     AWS_PY_METHOD_DEF(signing_config_new, METH_VARARGS),
     AWS_PY_METHOD_DEF(signing_config_get_algorithm, METH_VARARGS),
     AWS_PY_METHOD_DEF(signing_config_get_signature_type, METH_VARARGS),

source/module.c

@@ -181,6 +181,40 @@ def test_process_provider(self):
             self.assertTrue('process_secret_access_key' == credentials.secret_access_key)
             self.assertTrue(credentials.session_token is None)
 
+    def test_delegate_provider(self):
+        def delegate_get_credentials():
+            return awscrt.auth.AwsCredentials("accesskey", "secretAccessKey", "sessionToken")
+
+        provider = awscrt.auth.AwsCredentialsProvider.new_delegate(delegate_get_credentials)
+        credentials = provider.get_credentials().result(TIMEOUT)
+
+        # Don't use assertEqual(), which could log actual credentials if test fails.
+        self.assertTrue('accesskey' == credentials.access_key_id)
+        self.assertTrue('secretAccessKey' == credentials.secret_access_key)
+        self.assertTrue('sessionToken' == credentials.session_token)
+
+    def test_delegate_provider_exception(self):
+        # delegate that raises exception should result in exception
+        def delegate_get_credentials():
+            raise Exception("purposefully thrown exception")
+
+        provider = awscrt.auth.AwsCredentialsProvider.new_delegate(delegate_get_credentials)
+
+        with self.assertRaises(Exception):
+            credentials_future = provider.get_credentials()
+            credentials = credentials_future.result(TIMEOUT)
+
+    def test_delegate_provider_exception_from_bad_return_type(self):
+        # delegate that returns wrong type should result in exception
+        def delegate_get_credentials():
+            return "purposefully return wrong type"
+
+        provider = awscrt.auth.AwsCredentialsProvider.new_delegate(delegate_get_credentials)
+
+        with self.assertRaises(Exception):
+            credentials_future = provider.get_credentials()
+            credentials = credentials_future.result(TIMEOUT)
+
 
 class TestSigningConfig(NativeResourceTest):
     def test_create(self):