使用 Flama JWT 身份驗證保護 ML API

您可能已經聽說過最近發布的 Flama 1.7，它帶來了一些令人興奮的新功能，可以幫助您開發和生產 ML API。這篇文章專門討論該版本的主要亮點之一：支援 JWT 身份驗證。但是，在我們透過實際範例深入了解細節之前，我們建議您記住以下資源（如果您還沒有熟悉它們，請先熟悉一下）：

• Flama 官方文件：Flama 文檔
• 介紹 Flama for ML API 的貼文：用於穩健機器學習 API 的 Flama 簡介

現在，讓我們開始使用新功能，看看如何透過標頭或 Cookie 進行基於令牌的身份驗證來保護 API 端點。

目錄

這篇文章的架構如下：

• 什麼是 JSON Web Token (JWT)？
• 使用 Flama 進行 JWT 身份驗證
  • 設定開發環境
  • 基礎應用
  • Flama JWT 組件
  • 受保護的端點
  • 運行應用程式
  • 登入端點
• 結論
• 支持我們的工作
• 參考文獻
• 關於作者

什麼是 JSON Web 令牌？

如果您已經熟悉 JSON Web Token (JWT) 的概念及其工作原理，請隨意跳過本節並直接跳至使用 Flama 實現 JWT 身份驗證。否則，我們將嘗試向您簡要解釋 JWT 是什麼以及為什麼它對授權目的如此有用。

簡介

我們可以從RFC 7519文件中給出的官方定義開始：

JSON Web Token (JWT) 是一種緊湊、URL 安全的表示方式
債權在兩方之間轉讓。 JWT 中的聲明
被編碼為 JSON 對象，用作 JSON
的負載 Web 簽章 (JWS) 結構或作為 JSON Web 的明文
加密（JWE）結構，使索賠能夠數位化
使用訊息驗證程式碼進行簽署或完整性保護
(MAC) 和/或加密。

所以，簡單來說，JWT 是一個開放標準，它定義了一種以 JSON 物件的形式在兩方之間傳輸訊息的方式。所傳輸的資訊經過數位簽名，以確保其完整性和真實性。這就是為什麼 JWT 的主要用例之一是用於授權目的。

基於 JWT 的原型驗證流程如下所示：

• 使用者登入系統並收到 JWT 令牌。
• 使用者隨每個後續請求發送此令牌到伺服器。
• 伺服器驗證令牌，如果令牌有效，則授予對所請求資源的存取權限。
• 如果令牌無效，伺服器將拒絕存取該資源。

但是，JWT 令牌不僅可用於識別用戶，而且還可用於透過令牌的有效負載以安全的方式在不同服務之間共享資訊。此外，鑑於令牌的簽章是使用標頭和有效負載計算的，接收者可以驗證令牌的完整性並確保它在傳輸過程中沒有被更改。

智威湯遜結構

JWT 表示為一系列由點 (.) 分隔的 URL 安全部分，每個部分包含一個 base64url 編碼的 JSON 物件。 JWT 中的部分數量可能會有所不同，具體取決於所使用的表示形式，JWS（JSON Web 簽章）RFC-7515 或 JWE（JSON Web 加密）RFC-7516。
到目前為止，我們可以假設我們將使用 JWS，這是 JWT 令牌最常見的表示形式。在這種情況下，JWT 令牌由三個部分組成：

• 標頭：包含有關令牌以及應用於其的加密操作的元資料。
• Payload：包含令牌攜帶的聲明（資訊）。
• 簽章：確保令牌的完整性，用來驗證令牌的發送者是否是其聲稱的身分。

因此，使用扁平化 JWS JSON 序列化的原型 JWS JSON 語法如下（有關更多信息，請參閱 RFC-7515 第 7.2.2 節）：

{
   "payload":"<payload contents>",
   "header":<header contents>,
   "signature":"<signature contents>"
}
</signature></header></payload>

在 JWS 緊湊序列化中，JTW 表示為串聯：

BASE64URL(UTF8(JWS Header)) || '.' || BASE64URL(JWS Payload) || '.' || BASE64URL(JWS Signature)  

JWT 令牌的範例如下所示（取自此處的 Flama 測試之一）：

eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9.eyJkYXRhIjogeyJmb28iOiAiYmFyIn0sICJpYXQiOiAwfQ==.J3zdedMZSFNOimstjJat0V28rM_b1UU62XCp9dg_5kg=

標題

對於一般 JWT 對象，標頭可以包含各種欄位（例如，描述應用於令牌的加密操作），具體取決於 JWT 是 JWS 還是 JWE。然而，有些欄位是兩種情況共有的，也是最常用的：

• alg: The algorithm used to sign the token, e.g. HS256, HS384, HS512. To see the list of supported algorithms in flama, have a look at it here.
• typ: The type parameter is used to declare the media type of the JWT. If present, it is recommended to have the value JWT to indicate that the object is a JWT. For more information on this field, see RFC-7519 Sec. 5.1.
• cty: The content type is used to communicate structural information about the JWT. For example, if the JWT is a Nested JWT, the value of this field would be JWT. For more information on this field, see RFC-7519 Sec. 5.2.

The Payload

The payload contains the claims (information) that the token is carrying. The claims are represented as a JSON object, and they can be divided into three categories, according to the standard RFC-7519 Sec. 4:

• Registered claims: These are a set of predefined claims that are not mandatory but recommended. Some of the most common ones are iss (issuer), sub (subject), aud (audience), exp (expiration time), nbf (not before), iat (issued at), and jti (JWT ID). For a detailed explanation of each of these claims, see RFC-7519 Sec. 4.1.
• Public claims: These are claims that are defined by the user and can be used to share information between different services.
• Private claims: These are claims that are defined by the user and are intended to be shared between the parties that are involved in the communication.

The Signature

The signature is used to ensure the integrity of the token and to verify that the sender of the token is who it claims to be. The signature is calculated using the header and the payload, and it is generated using the algorithm specified in the header. The signature is then appended to the token to create the final JWT.

Implementing JWT Authentication with Flama

Having introduced the concept of JWT, its potential applications, besides a prototypical authentication flow, we can now move on to the implementation of a JWT-authenticated API using Flama. But, before we start, if you need to review the basics on how to create a simple API with flama, or how to run the API once you already have the code ready, then you might want to check out the quick start guide. There, you'll find the fundamental concepts and steps required to follow this post. Now, without further ado, let's get started with the implementation.

Setting up the development environment

Our first step is to create our development environment, and install all required dependencies for this project. The good thing is that for this example we only need to install flama to have all the necessary tools to implement JWT authentication. We'll be using poetry to manage our dependencies, but you can also use pip if you prefer:

poetry add flama[full]

If you want to know how we typically organise our projects, have a look at our previous post here, where we explain in detail how to set up a python project with poetry, and the project folder structure we usually follow.

Base application

Let's start with a simple application that has a single public endpoint. This endpoint will return a brief description of the API.

from flama import Flama

app = Flama(
    title="JWT protected API",
    version="1.0.0",
    description="JWT Authentication with Flama ?",
    docs="/docs/",
)

@app.get("/public/info/", name="public-info")
def info():
    """
    tags:
        - Public
    summary:
        Ping
    description:
        Returns a brief description of the API
    responses:
        200:
            description:
                Successful ping.
    """
    return {"title": app.schema.title, "description": app.schema.description, "public": True}

If you want to run this application, you can save the code above in a file called main.py under the src folder, and then run the following command (remember to have the poetry environment activated, otherwise you'll need to prefix the command with poetry run):

flama run --server-reload src.main:app

INFO:     Started server process [3267]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)

where the --server-reload flag is optional and is used to reload the server automatically when the code changes. This is very useful during development, but you can remove it if you don't need it. For a full list of the available options, you can run flama run --help, or check the documentation.

Authentication: Flama JWT Component

Ok, now that we have our base application running, let's add a new endpoint that requires authentication. To do this, we'll need to use the following flama functionality:

• Components (specifically JWTComponent): They're the building blocks for dependency injection in flama.
• Middlewares (specifically AuthenticationMiddleware): They're used as a processing layer between the incoming requests from clients and the responses sent by the server.

Let's proceed and modify our base application to include the JWT authentication as intended, and then we'll explain the code in more detail.

import uuid

from flama import Flama
from flama.authentication import AuthenticationMiddleware, JWTComponent
from flama.middleware import Middleware

JWT_SECRET_KEY = uuid.UUID(int=0).bytes  # The secret key used to signed the token
JWT_HEADER_KEY = "Authorization"  # Authorization header identifie
JWT_HEADER_PREFIX = "Bearer"  # Bearer prefix
JWT_ALGORITHM = "HS256"  # Algorithm used to sign the token
JWT_TOKEN_EXPIRATION = 300  # 5 minutes in seconds
JWT_REFRESH_EXPIRATION = 2592000  # 30 days in seconds
JWT_ACCESS_COOKIE_KEY = "access_token"
JWT_REFRESH_COOKIE_KEY = "refresh_token"

app = Flama(
    title="JWT-protected API",
    version="1.0.0",
    description="JWT Authentication with Flama ?",
    docs="/docs/",
)

app.add_component(
    JWTComponent(
        secret=JWT_SECRET_KEY,
        header_key=JWT_HEADER_KEY,
        header_prefix=JWT_HEADER_PREFIX,
        cookie_key=JWT_ACCESS_COOKIE_KEY,
    )
)

app.add_middleware(
    Middleware(AuthenticationMiddleware),
)

# The same code as before here
# ...

Although we've decided to add the component and middleware after the initialisation of the app, you can also add them directly to the Flama constructor by passing the components and middlewares arguments:

app = Flama(
    title="JWT-protected API",
    version="1.0.0",
    description="JWT Authentication with Flama ?",
    docs="/docs/",
    components=[
        JWTComponent(
            secret=JWT_SECRET_KEY,
            header_key=JWT_HEADER_KEY,
            header_prefix=JWT_HEADER_PREFIX,
            cookie_key=JWT_ACCESS_COOKIE_KEY,
        )
    ],
    middlewares=[
        Middleware(AuthenticationMiddleware),
    ]
)

This is just a matter of preference, and both ways are completely valid.

With the modifications introduced above, we can proceed to add a new (and JWT protected) endpoint. However, before we do that, let's briefly explain in more detail the functionality we've just introduced, namely components and middlewares.

Flama Components

As you might've already noticed, whenever we create a new application we're instantiating a Flama object. As the application grows, as is the case right now, also grows the need to add more dependencies to it. Without dependency injection (DI), this would mean that the Flama class would have to create and manage all its dependencies internally. This would make the class tightly coupled to specific implementations and harder to test or modify. With DI the dependencies are provided to the class from the outside, which decouples the class from specific implementations, making it more flexible and easier to test. And, this is where components come into play in flama.

In flama, a Component is a building block for dependency injection. It encapsulates logic that determines how specific dependencies should be resolved when the application runs. Components can be thought of as self-contained units responsible for providing the necessary resources that the application's constituents need. Each Component in flama has a unique identity, making it easy to look up and inject the correct dependency during execution. Components are highly flexible, allowing you to handle various types of dependencies, whether they're simple data types, complex objects, or asynchronous operations. There are several built-in components in flama, although in this post we're going to exclusively focus on the JWTComponent, which (as all others) derives from the Component class.

The JWTComponent contains all the information and logic necessary for extracting a JWT from either the headers or cookies of an incoming request, decoding it, and then validating its authenticity. The component is initialised with the following parameters:

• secret: The secret key used to decode the JWT.
• header_key: The key used to identify the JWT in the request headers.
• header_prefix: The prefix used to identify the JWT in the request headers.
• cookie_key: The key used to identify the JWT in the request cookies.

In the code above, we've initialised the JWTComponent with some dummy values for the secret key, header key, header prefix, and cookie key. In a real-world scenario, you should replace these values with your own secret key and identifiers.

Flama Middlewares

Middleware is a crucial concept that acts as a processing layer between the incoming requests from clients and the responses sent by the server. In simpler terms, middleware functions as a gatekeeper or intermediary that can inspect, modify, or act on requests before they reach the core logic of your application, and also on the responses before they are sent back to the client.

In flama, middleware is used to handle various tasks that need to occur before a request is processed or after a response is generated. In this particular case, the task we want to handle is the authentication of incoming requests using JWT. To achieve this, we're going to use the built-in class AuthenticationMiddleware. This middleware is designed to ensure that only authorised users can access certain parts of your application. It works by intercepting incoming requests, checking for the necessary credentials (such as a valid JWT token), and then allowing or denying access based on the user's permissions which are encoded in the token.

Here’s how it works:

• Initialization: The middleware is initialized with the Flama application instance. This ties the middleware to your app, so it can interact with incoming requests and outgoing responses.
• Handling Requests: The __call__ method of the middleware is called every time a request is made to your application. If the request type is http or websocket, the middleware checks if the route (the path the request is trying to access) requires any specific permissions.
• Permission Check: The middleware extracts the required permissions for the route from the route's tags. If no permissions are required, the request is passed on to the next layer of the application. If permissions are required, the middleware attempts to resolve a JWT token from the request. This is done using the previously discussed JWTComponent.
• Validating Permissions: Once the JWT token is resolved, the middleware checks the permissions associated with the token against those required by the route. If the user’s permissions match or exceed the required permissions, the request is allowed to proceed. Otherwise, the middleware returns a 403 Forbidden response, indicating that the user does not have sufficient permissions.
• Handling Errors: If the JWT token is invalid or cannot be resolved, the middleware catches the exception and returns an appropriate error response, such as 401 Unauthorized.

Adding a protected endpoint

By now, we should have a pretty solid understanding on what our code does, and how it does it. Nevertheless, we still need to see what we have to do to add a protected endpoint to our application. Let's do it:

@app.get("/private/info/", name="private-info", tags={"permissions": ["my-permission-name"]})
def protected_info():
    """
    tags:
        - Private
    summary:
        Ping
    description:
        Returns a brief description of the API
    responses:
        200:
            description:
                Successful ping.
    """
    return {"title": app.schema.title, "description": app.schema.description, "public": False}

And that's it! We've added a new endpoint that requires authentication to access. The functionality is exactly the same as the previous endpoint, but this time we've added the tags parameter to the @app.get decorator. The tag parameter can be used to specify additional metadata for an endpoint. But, if we use the special key permissions whilst using the AuthenticationMiddleware, we can specify the permissions required to access the endpoint. In this case, we've set the permissions to ["my-permission-name"], which means that only users with the permission my-permission-name will be able to access this endpoint. If a user tries to access the endpoint without the required permission, they will receive a 403 Forbidden response.

Running the application

If we put all the pieces together, we should have a fully functional application that has a public endpoint and a private endpoint that requires authentication. The full code should look something like this:

import uuid

from flama import Flama
from flama.authentication import AuthenticationMiddleware, JWTComponent
from flama.middleware import Middleware

JWT_SECRET_KEY = uuid.UUID(int=0).bytes  # The secret key used to signed the token
JWT_HEADER_KEY = "Authorization"  # Authorization header identifie
JWT_HEADER_PREFIX = "Bearer"  # Bearer prefix
JWT_ALGORITHM = "HS256"  # Algorithm used to sign the token
JWT_TOKEN_EXPIRATION = 300  # 5 minutes in seconds
JWT_REFRESH_EXPIRATION = 2592000  # 30 days in seconds
JWT_ACCESS_COOKIE_KEY = "access_token"
JWT_REFRESH_COOKIE_KEY = "refresh_token"

app = Flama(
    title="JWT-protected API",
    version="1.0.0",
    description="JWT Authentication with Flama ?",
    docs="/docs/",
)

app.add_component(
    JWTComponent(
        secret=JWT_SECRET_KEY,
        header_key=JWT_HEADER_KEY,
        header_prefix=JWT_HEADER_PREFIX,
        cookie_key=JWT_ACCESS_COOKIE_KEY,
    )
)

app.add_middleware(
    Middleware(AuthenticationMiddleware),
)

@app.get("/public/info/", name="public-info")
def info():
    """
    tags:
        - Public
    summary:
        Info
    description:
        Returns a brief description of the API
    responses:
        200:
            description:
                Successful ping.
    """
    return {"title": app.schema.title, "description": app.schema.description, "public": True}


@app.get("/private/info/", name="private-info", tags={"permissions": ["my-permission-name"]})
def protected_info():
    """
    tags:
        - Private
    summary:
        Info
    description:
        Returns a brief description of the API
    responses:
        200:
            description:
                Successful ping.
    """
    return {"title": app.schema.title, "description": app.schema.description, "public": False}

Running the application as before, we should see the following output:

flama run --server-reload src.main:app

INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     Started reloader process [48145] using WatchFiles
INFO:     Started server process [48149]
INFO:     Waiting for application startup.
INFO:     Application startup complete.

If we navigate with our favourite browser to http://127.0.0.1:8000/docs/ we should see the documentation of our API as shown below:

As we can already expect, if we send a request to the public endpoint /public/info/ we should receive a successful response:

curl --request GET \
  --url http://localhost:8000/public/info/ \
  --header 'Accept: application/json'
{"title":"JWT protected API","description":"JWT Authentication with Flama ?","public":true}

However, if we try to access the private endpoint /private/info/ without providing a valid JWT token, we should receive a 401 Unauthorized response:

curl --request GET \
  --url http://localhost:8000/private/info/ \
  --header 'Accept: application/json'
{"status_code":401,"detail":"Unauthorized","error":null}% 

Thus, we can say that the JWT authentication is working as expected, and only users with a valid JWT token will be able to access the private endpoint.

Adding the login endpoint

As we've just seen, we have a private endpoint which requires a valid JWT token to be accessed. But, how do we get this token? The answer is simple: we need to create a login endpoint that will authenticate the user and return a valid JWT token. To do this, we're going to define the schemas for the input and output of the login endpoint (feel free to use your own schemas if you prefer, for instance, adding more fields to the input schema such as username or email):

import uuid
import typing
import pydantic

class User(pydantic.BaseModel):
    id: uuid.UUID
    password: typing.Optional[str] = None


class UserToken(pydantic.BaseModel):
    id: uuid.UUID
    token: str

Now, let's add the login endpoint to our application:

import http

from flama import types
from flama.authentication.jwt import JWT
from flama.http import APIResponse

@app.post("/auth/", name="signin")
def signin(user: types.Schema[User]) -> types.Schema[UserToken]:
    """
    tags:
        - Public
    summary:
        Authenticate
    description:
        Returns a user token to access protected endpoints
    responses:
        200:
            description:
                Successful ping.
    """
    token = (
        JWT(
            header={"alg": JWT_ALGORITHM},
            payload={
                "iss": "vortico",
                "data": {"id": str(user["id"]), "permissions": ["my-permission-name"]},
            },
        )
        .encode(JWT_SECRET_KEY)
        .decode()
    )

    return APIResponse(
        status_code=http.HTTPStatus.OK, schema=types.Schema[UserToken], content={"id": str(user["id"]), "token": token}
    )

In this code snippet, we've defined a new endpoint /auth/ that receives a User object as input and returns a UserToken object as output. The User object contains the id of the user and the password (which is optional for this example, since we're not going to be comparing it with any stored password). The UserToken object contains the id of the user and the generated token that will be used to authenticate the user in the private endpoints. The token is generated using the JWT class, and contains the permissions granted to the user, in this case, the permission my-permission-name, which will allow the user to access the private endpoint /private/info/.

Now, let's test the login endpoint to see if it's working as expected. For this, we can proceed via the /docs/ page:

Or, we can use curl to send a request to the login endpoint:

curl --request POST \
  --url http://localhost:8000/auth/ \
  --header 'Accept: application/json' \
  --header 'Content-Type: application/json' \
  --data '{
  "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
  "password": "string"
}'

which should return a response similar to this:

{"id":"497f6eca-6276-4993-bfeb-53cbbbba6f08","token":"eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9.eyJkYXRhIjogeyJpZCI6ICI0OTdmNmVjYS02Mjc2LTQ5OTMtYmZlYi01M2NiYmJiYTZmMDgiLCAicGVybWlzc2lvbnMiOiBbIm15LXBlcm1pc3Npb24tbmFtZSJdfSwgImlhdCI6IDE3MjU5ODk5NzQsICJpc3MiOiAidm9ydGljbyJ9.vwwgqahgtALckMAzQHWpNDNwhS8E4KAGwNiFcqEZ_04="}

That string is the JWT token that we can use to authenticate the user in the private endpoint. Let's try to access the private endpoint using the token:

curl --request GET \
  --url http://localhost:8000/private/info/ \
  --header 'Accept: application/json' \
  --header 'Authorization: Bearer eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9.eyJkYXRhIjogeyJpZCI6ICI0OTdmNmVjYS02Mjc2LTQ5OTMtYmZlYi01M2NiYmJiYTZmMDgiLCAicGVybWlzc2lvbnMiOiBbIm15LXBlcm1pc3Npb24tbmFtZSJdfSwgImlhdCI6IDE3MjU5ODk5NzQsICJpc3MiOiAidm9ydGljbyJ9.vwwgqahgtALckMAzQHWpNDNwhS8E4KAGwNiFcqEZ_04='

which now returns a successful response:

{"title":"JWT protected API","description":"JWT Authentication with Flama ?","public":false}

And that's it! We've successfully implemented JWT authentication in our Flama application. We now have a public endpoint that can be accessed without authentication, a private endpoint that requires a valid JWT token to be accessed, and a login endpoint that generates a valid JWT token for the user.

Conclusions

The privatisation of some endpoints (even all at some instances) is a common requirement in many applications, even more so when dealing with sensitive data as is often the case in Machine Learning APIs which process personal or financial information, to name some examples. In this post, we've covered the fundamentals of token-based authentication for APIs, and how this can be implemented without much of a hassle using the new features introduced in the latest release of flama (introduced in a previous post). Thanks to the JWTComponent and AuthenticationMiddleware, we can secure our API endpoints and control the access to them based on the permissions granted to the user, and all this with just a few modifications to our base unprotected application.

We hope you've found this post useful, and that you're now ready to implement JWT authentication in your own flama applications. If you have any questions or comments, feel free to reach out to us. We're always happy to help!

Stay tuned for more posts on flama and other exciting topics in the world of AI and software development. Until next time!

Support our work

If you like what we do, there is a free and easy way to support our work. Gift us a ⭐ at Flama.

GitHub ⭐'s mean a world to us, and give us the sweetest fuel to keep working on it to help others on its journey to build robust Machine Learning APIs.

You can also follow us on ?, where we share our latest news and updates, besides interesting threads on AI, software development, and much more.

References

• Flama documentation
• Flama GitHub repository
• Flama PyPI package

About the authors

• Vortico: We're specialised in software development to help businesses enhance and expand their AI and technology capabilities.

以上是使用 Flama JWT 身份驗證保護 ML API的詳細內容。更多資訊請關注PHP中文網其他相關文章！