TDD Approach to Create an Authentication System With FastAPI Part 5

Introduction

We have been using FastAPI along with the Test-Driven Development process to come up with an authentication system. Till now we can register to own tiny application using a username and password. We are also able to generate a JWT token to be classified as a logged-in user. But there is no point in logging in until the logged-in user can do something privileged. In this iteration of the series, we’d see how to access a protected endpoint using FastAPI and JWT. Given that we already have a JWT token.

Series Index

We have already seen:

Project Setup and FastAPI introduction
Database Setup and User Registration
Fix the failing test with mocking and dependency injection
Authentication Premier and Login Endpoint
Token Validation and Accessing Protected Endpoints (you are here)

What happens when you access protected routes?

I would like to start this section by saying that HTTP is stateless, meaning that applications written on top of HTTP won’t be able to identify a client if they have authorized. Meaning that if you have sent the server the correct username and password, they’d send you an authentication token. This token has to be sent with each request made to be identified as logged in.

What happens when you have successfully logged in?

Let us start where we left off in the previous post, take a look at the below workflow.

On successful login, you are given a Bearer token. In the above gif, you can see that server responds with a JSON with access_token and token_type. The former contains the encoded JWT token.

Now when you don’t send this token with every request, you’d be considered logged out. Here is what I’m trying to say in an image.

Trying to access protected route with and without Bearer token

Points to remember:

On successful login, the server sends a JWT token to the client.
This token needs to be sent to the server in every request when trying to access a protected route.

Access Protected Endpoints

For sake of contrast, we are going to have two endpoints, we will reuse our /ping endpoint, which we can access without being signed up or logged in. And then we’ll have a dummy /protected endpoint which will return a dummy JSON { "message": "protected resource" }.

Let’s first write down the requirements.

Requirement for /protected endpoint

GET request to /protected without JWT token returns 403.
GET request to /protected with JWT token returns 200.

Requirements turned into tests

I am going to put this test inside tests/test_users.py along with other user-related tests.

tests/test_users.py

 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
diff --git a/tests/test_users.py b/tests/test_users.py
index 2239b5b..dbd5750 100644
--- a/tests/test_users.py
+++ b/tests/test_users.py
@@ -84,3 +84,31 @@ class TestUserLogin:
         )
         assert response.status_code == 200
         assert len(response.json()) == 2
+
+
+class TestProtectedEndpoints:
+    def test_get_request_without_JWT_token_returns_403(self):
+        """sends GET request to /protected without JWT token"""
+
+        response = client.get("/protected")
+        assert response.status_code == 403
+
+    def test_get_request_without_JWT_token_returns_200_and_body(self):
+        """sends GET request to /protected"""
+
+        response = client.post(
+            "/users/auth",
+            json={"username": "santosh", "password": "sntsh"}
+        )
+
+        login_response = response.json()
+        access_token = login_response['access_token']
+        token_type = login_response['token_type']
+
+        response = client.get(
+            "/protected",
+            headers={'Authorization': '{} {}'.format(token_type, access_token)}
+        )
+
+        assert response.status_code == 200
+        assert len(response.json()) == 1

Before we try to implement or even run the test, I would like to go through the definition of the tests to tell what they are doing.

The test_get_request_without_JWT_token_returns_403 route tries to access a protected route (route which we are yet to define). It returns a status code of 403 in return. This happens because we don’t have the Bearer token in the Authorization header.

If you don’t know about the Bearer token in Authorization, I’d recommend you read the previous post in this series.

The test_get_request_without_JWT_token_returns_200_and_body does exactly what the first test does. But this time, it first fetches the JWT token by sending a login request. Thereafter, it uses the token and token type to put inside the Authorization header. And when we send the GET request again, we see 200 in response.

So without further ado, let’s run the tests.

$ pytest -k protected
========================== test session starts ==========================
platform linux -- Python 3.8.10, pytest-6.2.5, py-1.11.0, pluggy-1.0.0
rootdir: /efs/repos/fastauth
plugins: anyio-3.4.0
collected 11 items / 9 deselected / 2 selected                          

tests/test_users.py FF                                            [100%]

=============================== FAILURES ================================
_ TestProtectedEndpoints.test_get_request_without_JWT_token_returns_403 _

self = <tests.test_users.TestProtectedEndpoints object at 0x7f1401a20880>

    def test_get_request_without_JWT_token_returns_403(self):
        """sends GET request to /protected without JWT token"""
    
        response = client.get("/protected")
>       assert response.status_code == 403
E       assert 404 == 403
E        +  where 404 = <Response [404]>.status_code
tests/test_users.py:94: AssertionError
_ TestProtectedEndpoints.test_get_request_without_JWT_token_returns_200_and_body _

self = <tests.test_users.TestProtectedEndpoints object at 0x7f14019e6460>

    def test_get_request_without_JWT_token_returns_200_and_body(self):
        """sends GET request to /protected"""
    
        response = client.post(
            "/users/auth",
            json={"username": "santosh", "password": "sntsh"}
        )
    
        login_response = response.json()
        access_token = login_response['access_token']
        token_type = login_response['token_type']
    
        response = client.get(
            "/protected",            headers={'Authorization': '{} {}'.format(token_type, access_token)}
        )
    
>       assert response.status_code == 200
E       assert 404 == 200
E        +  where 404 = <Response [404]>.status_code

tests/test_users.py:113: AssertionError
======================== short test summary info ========================
FAILED tests/test_users.py::TestProtectedEndpoints::test_get_request_without_JWT_token_returns_403
FAILED tests/test_users.py::TestProtectedEndpoints::test_get_request_without_JWT_token_returns_200_and_body
==================== 2 failed, 9 deselected in 1.44s ====================

As we can see, we are getting 404 as the status code usual instead of 403 and 200 respectively. This is obvious because as usual, we haven’t implemented the route.

Write code to fix failing test

Before we write the actual route, we need to have some helper functions. In the last post, we saw how to encode some data, along with the expiry time of a token. Now we are going to do the reverse of that.

fastauth/auth.py

34
35
36
37
def decode_jwt_token(token: str):
    decoded_token = jwt.decode(token, SECRET_KEY, algorithms=JWT_ALGORITHM)
    print(decoded_token)
    return decoded_token if decoded_token['exp'] >= time.time() else None

This function takes in the token and returns decoded version if the expiry is remaining, otherwise returns None.

Now in crude words, this decoded token needs to be shown to FastAPI with every request to identify the request as authorized; then only FastAPI is letting the client access the protected route.

We could have written that part manually, but we have a certain structure we have to follow with FastAPI. So we are going to write a class that inherits from HTTPBearer.

Below I’m going to post the entire diff along with the function I wrote above.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
diff --git a/fastauth/auth.py b/fastauth/auth.py
index ead9785..6363f3c 100644
--- a/fastauth/auth.py
+++ b/fastauth/auth.py
@@ -1,6 +1,9 @@
 import datetime
+import time
  
 import bcrypt
+from fastapi import Request, HTTPException
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials import jwt
 from sqlalchemy.orm import Session
  
@@ -26,3 +29,37 @@ def encode_jwt_token(*, data: dict, expires_delta: datetime.timedelta = None):
  
     encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=JWT_ALGORITHM)
     return encoded_jwt
+
+
+def decode_jwt_token(token: str):
+    decoded_token = jwt.decode(token, SECRET_KEY, algorithms=JWT_ALGORITHM)
+    print(decoded_token)
+    return decoded_token if decoded_token['exp'] >= time.time() else None
+
+
+class JWTBearer(HTTPBearer):
+    def __init__(self, auto_error: bool = True):
+        super(JWTBearer, self).__init__(auto_error=auto_error)+
+    async def __call__(self, request: Request):
+        credentials : HTTPAuthorizationCredentials = await super(JWTBearer, self).__call__(request)
+        if credentials:+            if not credentials.scheme == "Bearer":
+                raise HTTPException(status_code=401, detail="Invalid authentication scheme.")
+            if not self.verify_jwt(credentials.credentials):
+                raise HTTPException(status_code=401, detail="Invalid token or expired token.")
+            return credentials.credentials
+        else:
+            raise HTTPException(status_code=401, detail="Invalid authentication code.")
+
+    def verify_jwt(self, jwtToken: str) -> bool:
+        isTokenValid: bool = False
+
+        try:
+            payload = decode_jwt_token(jwtToken)
+        except:
+            payload = None
+
+        if payload:
+            isTokenValid = True
+        return isTokenValid

main.py

Let us come to the file with our route definition.

Remember I told you there is a certain way we handle token validation in FastAPI. Yes. The class we wrote, we can use that class as a dependency to the route in the following way.

49
50
51
52

@app.get("/protected", dependencies=[Depends(auth.JWTBearer())])
def get_protected_resource():
    return { "message": "protected resource" }

The thing to note in the above code is that the dependencies=[Depends(auth.JWTBearer())]. dependencies array takes multiple dependencies of the route. We have our single JWTBearer class marked as a dependency to the /protected route.

You can assume dependencies similar to a middleware. What a middleware does is adds or removes more data into the request and then pass the data to the next middleware.

I would ask you to read fastapi dependency vs middleware

Confirm if tests are passing

Let us run only the required tests.

$ pytest -k protected
=================================== test session starts ====================================
platform linux -- Python 3.8.10, pytest-6.2.5, py-1.11.0, pluggy-1.0.0
rootdir: /efs/repos/fastauth
plugins: anyio-3.4.0
collected 11 items / 9 deselected / 2 selected                                             

tests/test_users.py ..                                                               [100%]

============================= 2 passed, 9 deselected in 1.66s ==============================

Test manually

If you wish to, you are free to test it manually by not passing JWT token while accessing the protected route.

Then in the next iteration, obtain a token, and put it inside the Authorization header with the authorization type as Bearer.

Bonus

I would like to give out the complete code we have written till now in a git repository.

https://github.com/santosh/fastauth

I have also gone ahead and included a Postman collection of endpoints to test. Feel free to play around with data and values to get a better understanding of what’s going on.

Conclusion

This post is a full stop to the 5 post series on this topic. I tried to cover the basics of Test Driven Development, Dependency Injection, and Mocking. I also had this itch to dig into information security. I did that by writing my first authentication system. I learned concepts related to authorization and authentication concerning HTTP while writing this series. I plan to keep digging into infosec space.

Thank you all of you if you have been following this post till now. I appreciate it and respect your time. I hope you have learned something new here. You can connect with on on LinkedIn, or say Hi on Twitter. You can also put your email into the below box to stay tuned with similar posts around the tech and programming space.

Have a nice time ahead.