使用 fastapi UploadFile 验证文件类型和扩展名

Question

我目前正在从事一个小项目，该项目涉及创建一个允许用户上传jar文件的fastapi服务器。

基本上我有这条路线：

@app.post("/upload")
async def upload(jar_file: UploadFile = File(...)):

我真的很想检查并验证该文件是否真的是jar文件。

我可以自己实现它，但我很好奇fastapi或任何其他 package 是否提供此功能。

Answer 1

您可以检查 MIME 类型 (https://fastapi.tiangolo.com/tutorial/request-files/#uploadfile )。

@app.post("/upload")
async def upload(jar_file: UploadFile = File(...)):
    if jar_file.content_type != "application/java-archive":
        raise HTTPException(400, detail="Invalid document type")
    return {"filename": "jar_file.filename"}

Answer 2

我有同样的需求，并且由于我的文件相对较大，我希望能够在将文件上传到后端（至少不是整个文件）之前获得错误消息，如rezan21所述。

这是使它起作用的方法。 请注意，由于 Stalette 的一些限制，有多种解决方法，例如1. 这个用于读取请求主体异步生成器和2. 这个问题处理这个确切的需要。

首先，我直接从 SwaggerUI选择文件输入读取文件，因此没有传递额外的标头来指示前端或 api 消费者可以读取的文件扩展名或 MIME 类型。

然后，我想直接在路由定义中设置文件，就像任何其他依赖项一样。 单独的依赖项在这里不起作用，因为它仅在整个文件上传后才被调用。

因此，我的 csv 和 excel 文件的当前工作解决方案是使用自定义BaseHTTPMiddleware ，异步读取请求正文并从文件本身获取“标头”。

根据我的推断，这获取了 body 异步生成器的第一个块，并且它具有正在上传的文件的信息。 为了防止程序停顿，get_body function按照1.

import re

from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from dependencies import ContentTypeChecker


def get_content_type_from_body(body):
    content_type_match = re.search(rb'Content-Type: ([^\r\n]+)', body)

    if content_type_match:
        content_type = content_type_match.group(1).decode("utf-8")
    return content_type


async def set_body(request: Request, body: bytes):
    async def receive():
        return {"type": "http.request", "body": body}
    request._receive = receive


async def get_body(request: Request) -> bytes:
    body = await request.body()
    await set_body(request, body)
    return body


class ValidateContentTypeMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next):
        content_type = request.headers.get("Content-Type", "")
        file_content_type = ''

        if content_type.startswith("multipart/form-data"):
            bd = await get_body(request)
            file_content_type = get_content_type_from_body(bd)

        if file_content_type:
            for route in request.app.routes:
                try:
                    for dependency in route.dependant.dependencies:
                        if not isinstance(dependency.cache_key[0], ContentTypeChecker):
                            continue

                        valid_content_type = dependency.call(
                            content_type=file_content_type)

                        if not valid_content_type:
                            exc = HTTPException(
                                detail=f'File of type {file_content_type} not in {dependency.cache_key[0].content_types}',
                                status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE)

                            return JSONResponse(status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, content={'message': exc.detail})

                except AttributeError as e:
                    if e.name == 'dependant':
                        pass

        response = await call_next(request)
        return response

然后，为了使其工作，内容类型检查器是一个简单的 class，它使用允许的内容类型列表和一个__call__方法实例化，该方法在中间件中接收内容类型

class ContentTypeChecker:
    def __init__(self, content_types: List[str]) -> None:
        self.content_types = content_types

    def __call__(self, content_type: str = ''):
        if content_type and content_type not in self.content_types:
            return False
        return True

这种方法的一个警告是，如果内容类型与允许的类型匹配并且中间件转发请求，FastAPI 将再次调用它。 因此， __call__方法上 content_type 的默认值为''并在 FastAPI 自行进行检查时返回 True 。

最后，这是路由定义：

@router.post('/upload',
             dependencies=[Depends(ContentTypeChecker(['text/csv']))]
             )
async def upload(file: UploadFile = File(...)):
    ...

我不确定是否有更好的方法来调用依赖于验证过程。