Skip to content

Commit

Permalink
fix(file_upload): validate mimetype as configured
Browse files Browse the repository at this point in the history
  • Loading branch information
qvalentin committed Oct 22, 2024
1 parent 22d198d commit 2e4689b
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 6 deletions.
13 changes: 13 additions & 0 deletions backend/chainlit/server.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
import fnmatch
import glob
import json
import mimetypes
Expand Down Expand Up @@ -870,6 +871,18 @@ async def upload_file(
assert file.filename, "No filename for uploaded file"
assert file.content_type, "No content type for uploaded file"

if (
config.features.spontaneous_file_upload is not None
and config.features.spontaneous_file_upload.accept is not None
):
for pattern in config.features.spontaneous_file_upload.accept:
if fnmatch.fnmatch(file.content_type, pattern):
return True
raise HTTPException(
status_code=400,
detail="File type not allowed",
)

file_response = await session.persist_file(
name=file.filename, content=content, mime=file.content_type
)
Expand Down
71 changes: 65 additions & 6 deletions backend/tests/test_server.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,24 @@
import datetime # Added import for datetime
import os
from pathlib import Path
import pathlib
import tempfile
from pathlib import Path
from typing import Callable
from unittest.mock import AsyncMock, Mock, create_autospec, mock_open
import datetime # Added import for datetime

import pytest
import tempfile
from chainlit.session import WebsocketSession
from chainlit.auth import get_current_user
from chainlit.config import APP_ROOT, ChainlitConfig, load_config
from chainlit.config import (
APP_ROOT,
ChainlitConfig,
SpontaneousFileUploadFeature,
load_config,
)
from chainlit.server import app
from fastapi.testclient import TestClient
from chainlit.session import WebsocketSession
from chainlit.types import FileReference
from chainlit.user import PersistedUser # Added import for PersistedUser
from fastapi.testclient import TestClient


@pytest.fixture
Expand Down Expand Up @@ -509,6 +514,60 @@ def test_upload_file_unauthorized(
assert response.status_code == 422


@pytest.mark.parametrize(
"accept_pattern, mime_type, expected_status",
[
(["image/*"], "text/plain", 400),
(["image/*", "application/*"], "text/plain", 400),
(["image/png", "application/pdf"], "image/jpeg", 400),
(["text/*"], "text/plain", 200),
(["application/*"], "application/pdf", 200),
(["image/*"], "image/jpeg", 200),
(["image/*", "text/*"], "text/plain", 200),
(["*/*"], "text/plain", 200),
(["*/*"], "image/jpeg", 200),
(["*/*"], "application/pdf", 200),
(["image/*", "application/*"], "application/pdf", 200),
(["image/*", "application/*"], "image/jpeg", 200),
(["image/png", "application/pdf"], "image/png", 200),
(["image/png", "application/pdf"], "application/pdf", 200),
],
)
def test_upload_file_invalid_mime_type(
test_client: TestClient,
test_config: ChainlitConfig,
mock_session_get_by_id_patched: Mock,
monkeypatch: pytest.MonkeyPatch,
accept_pattern: list[str],
mime_type: str,
expected_status: int,
):
"""Test successful file upload."""

# Set accept in config
monkeypatch.setattr(
test_config.features,
"spontaneous_file_upload",
SpontaneousFileUploadFeature(enabled=True, accept=accept_pattern),
)

# Prepare the files to upload
file_content = b"Sample file content"
files = {
"file": ("test_upload.txt", file_content, mime_type),
}

# Make the POST request to upload the file
response = test_client.post(
"/project/file",
files=files,
params={"session_id": mock_session_get_by_id_patched.id},
)

# Verify the response
assert response.status_code == expected_status


def test_project_translations_file_path_traversal(
test_client: TestClient, monkeypatch: pytest.MonkeyPatch
):
Expand Down

0 comments on commit 2e4689b

Please sign in to comment.