From ed5866f639511ea7783cf9db6515466514d79f40 Mon Sep 17 00:00:00 2001
From: CK Ng <ngchoonkhon@gmail.com>
Date: Sun, 26 Feb 2023 11:10:06 +0800
Subject: [PATCH] Add  client_id validation for access token, and get_user_info
 method.

The client_id from access token will be validated whether it matches the one defined in the `settings.COGNITO_AUDIENCE`.
When using access token, the backend will call the userinfo endpoint from Amazon Cognito, to obtain the user info based on the access key, and pass into the payload inside the Django user_model's get_or_create_for_cognito method.
---
 README.rst                          | 10 +++++++++-
 src/django_cognito_jwt/backend.py   | 18 +++++++++++++++++-
 src/django_cognito_jwt/validator.py |  5 +++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 1b9bd08..896f4e0 100644
--- a/README.rst
+++ b/README.rst
@@ -77,4 +77,12 @@ The library by default uses id token. To use access token, add the following lin
 
 .. code-block:: python
 
-	COGNITO_TOKEN_TYPE = "access"  # '{'id', 'access'} Default: 'id'
+	COGNITO_TOKEN_TYPE = "access"  # {'id', 'access'}, default 'id'
+
+
+As the payload of access token only contains basic user info, we could obtain further info from the `UserInfo endpoint`.
+You need to specify the Cognito domain in the ``settings.py`` file to obtain the user info from the endpoint, as follows:
+
+.. code-block:: python
+
+	COGNITO_DOMAIN = "your-user-pool-domain"  # eg, exampledomain.auth.ap-southeast-1.amazoncognito.com
diff --git a/src/django_cognito_jwt/backend.py b/src/django_cognito_jwt/backend.py
index c66fd54..97173ae 100644
--- a/src/django_cognito_jwt/backend.py
+++ b/src/django_cognito_jwt/backend.py
@@ -1,4 +1,6 @@
 import logging
+import requests
+import json
 
 from django.apps import apps as django_apps
 from django.conf import settings
@@ -29,12 +31,26 @@ def authenticate(self, request):
             raise exceptions.AuthenticationFailed()
 
         USER_MODEL = self.get_user_model()
-        user = USER_MODEL.objects.get_or_create_for_cognito(jwt_payload)
+        if settings.COGNITO_TOKEN_TYPE == "access":
+            user_info = self.get_user_info(jwt_token.decode("UTF-8"))
+            user_info = json.loads(user_info.decode("UTF-8"))
+            user = USER_MODEL.objects.get_or_create_for_cognito(user_info)
+        else:
+            user = USER_MODEL.objects.get_or_create_for_cognito(jwt_payload)
         return (user, jwt_token)
 
     def get_user_model(self):
         user_model = getattr(settings, "COGNITO_USER_MODEL", settings.AUTH_USER_MODEL)
         return django_apps.get_model(user_model, require_ready=False)
+    
+    def get_user_info(self, access_token):
+        if settings.COGNITO_TOKEN_TYPE == "access":
+            url = f"https://{settings.COGNITO_DOMAIN}/oauth2/userInfo"
+
+            headers = {'Authorization': f'Bearer {access_token}'}
+
+            res = requests.get(url, headers=headers)
+            return res.content
 
     def get_jwt_token(self, request):
         auth = get_authorization_header(request).split()
diff --git a/src/django_cognito_jwt/validator.py b/src/django_cognito_jwt/validator.py
index 9d4fc79..df3e120 100644
--- a/src/django_cognito_jwt/validator.py
+++ b/src/django_cognito_jwt/validator.py
@@ -73,6 +73,11 @@ def validate(self, token):
                 params.update({"audience": self.audience})
 
             jwt_data = jwt.decode(**params)
+            if self.token_type == "access":
+                if "access" not in jwt_data["token_use"]:
+                    raise TokenError("Incorrect token use")
+                if jwt_data["client_id"] not in self.audience:
+                    raise TokenError("Incorrect client_id")
         except (
             jwt.InvalidTokenError,
             jwt.ExpiredSignatureError,