shanxia2.0/backend/tests/test_doclib_api.py

"""使用真实路由和 SQLite；对象存储替身仅隔离外部 MinIO。"""
import io

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from test_doclib import db, svc, seed_file
from app.api.doclib import router
from app.core.auth import get_current_user
from app.core.config import settings
from app.core.database import get_db
from app.services import doclib_storage


class StoredResponse(io.BytesIO):
    def release_conn(self):
        self.released = True

    def stream(self, amt=65536):
        while True:
            chunk = self.read(amt)
            if not chunk:
                break
            yield chunk


class MemoryStorage:
    data = {}
    response = None

    def upload(self, key, file, length):
        self.data[key] = file.read()

    def open(self, key):
        self.response = StoredResponse(self.data[key])
        return self.response

    def discard(self, key):
        self.data.pop(key, None)


@pytest.fixture
def client(db, monkeypatch):
    storage = MemoryStorage()
    storage.data = {}
    monkeypatch.setattr(doclib_storage, 'private_storage', lambda: storage)
    app = FastAPI()
    app.include_router(router, prefix='/api')
    app.dependency_overrides[get_db] = lambda: db
    identity = {'user_id': 'owner', 'tenant_no': 'a', 'role': 'admin'}
    app.dependency_overrides[get_current_user] = lambda: identity
    with TestClient(app, headers={'Authorization': 'Bearer test'}) as client:
        yield client, identity, storage


def test_upload_download_and_revoked_access(client):
    c, identity, storage = client
    result = c.post('/api/web/doclib/upload', files={'file': ('中文.json', b'{"hello":true}', 'application/json')})
    assert result.status_code == 200, result.text
    doc = result.json()['data']
    assert doc['name'] == '中文.json' and doc['size'] == 14
    assert c.put('/api/web/doclib/files/'+doc['id']+'/shares', json={'userIds': ['reader']}).status_code == 200
    identity['user_id'] = 'reader'
    downloaded = c.get('/api/web/doclib/files/'+doc['id']+'/download')
    assert downloaded.content == b'{"hello":true}'
    assert downloaded.headers['content-type'] == 'application/octet-stream'
    assert 'attachment;' in downloaded.headers['content-disposition']
    assert downloaded.headers['cache-control'] == 'no-store'
    assert storage.response.closed and storage.response.released
    assert c.post('/api/web/doclib/files/batch-delete', json={'ids': [doc['id']]}).status_code == 403
    identity['user_id'] = 'owner'
    c.put('/api/web/doclib/files/'+doc['id']+'/shares', json={'userIds': []})
    identity['user_id'] = 'reader'
    assert c.get('/api/web/doclib/files/'+doc['id']+'/download').status_code == 404


def test_upload_limits_folder_access_and_malicious_body(client, db, monkeypatch):
    c, identity, storage = client
    monkeypatch.setattr(settings, 'DOCLIB_MAX_UPLOAD_BYTES', 5)
    assert c.post('/api/web/doclib/upload', files={'file': ('large.txt', b'123456')}).status_code == 413
    assert not storage.data
    foreign = svc(db, 'reader').create_folder('other', '')
    assert c.post('/api/web/doclib/upload', data={'folderId': foreign['folderId']}, files={'file': ('a', b'1')}).status_code == 403
    assert not storage.data
    assert c.post('/api/web/doclib/folders', json={'name': 'bad', 'tenant_no': 'b'}).status_code == 422
    assert c.get('/api/web/doclib/files?page=0').status_code == 422
    assert c.get('/api/web/doclib/files?pageSize=999').status_code == 422


def test_download_requires_authentication(db):
    app = FastAPI()
    app.include_router(router, prefix='/api')
    app.dependency_overrides[get_db] = lambda: db
    with TestClient(app) as c:
        assert c.get('/api/web/doclib/files/any/download').status_code == 401


def test_audit_does_not_consume_document_stream_and_still_logs(monkeypatch):
    import asyncio
    from starlette.requests import Request
    from starlette.responses import StreamingResponse
    from app.core.api_request_log_middleware import ApiRequestLogMiddleware
    from app.core import api_request_log_middleware as audit
    logs, consumed = [], []
    monkeypatch.setattr(settings, 'API_REQUEST_LOG_ENABLED', True)
    monkeypatch.setattr(audit, '_schedule_persist_log', lambda **data: logs.append(data))
    async def content():
        consumed.append(True)
        yield b'large binary'
    response = StreamingResponse(content(), media_type='application/octet-stream', headers={'Content-Disposition': 'attachment; filename=test'})
    async def receive():
        return {'type': 'http.request', 'body': b'', 'more_body': False}
    request = Request({'type': 'http', 'method': 'GET', 'path': '/api/web/doclib/files/x/download', 'headers': [], 'query_string': b''}, receive)
    async def call_next(req):
        return response
    result = asyncio.run(ApiRequestLogMiddleware(FastAPI()).dispatch(request, call_next))
    assert result is response
    assert not consumed
    assert logs and logs[0]['status_code'] == 200
    assert 'omitted' in logs[0]['response_body']


def test_encrypted_file_list_and_download_keep_existing_middleware(client, db, monkeypatch):
    import time
    from jose import jwt
    from app.core import api_crypto, api_crypto_middleware as crypto, api_request_log_middleware as audit
    c, identity, storage = client
    app = FastAPI()
    app.include_router(router, prefix="/api")
    app.dependency_overrides = c.app.dependency_overrides.copy()
    app.add_middleware(crypto.ApiCryptoMiddleware)
    app.add_middleware(audit.ApiRequestLogMiddleware)
    c = TestClient(app, headers={'Authorization': 'Bearer test'})
    logs = []
    monkeypatch.setattr(settings, 'API_ENCRYPT_ENABLED', True)
    monkeypatch.setattr(settings, 'API_ENCRYPT_FIXED_SECRET', 'doclib-test-key')
    monkeypatch.setattr(settings, 'API_REQUEST_LOG_ENABLED', True)
    monkeypatch.setattr(audit, '_schedule_persist_log', lambda **kwargs: logs.append(kwargs))
    async def nonce_ok(*args):
        return True
    monkeypatch.setattr(crypto, '_mark_nonce_once', nonce_ok)
    token = jwt.encode({'sub': 'owner'}, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
    uploaded = c.post('/api/web/doclib/upload', files={'file': ('file.txt', b'streamed')}).json()['data']
    assert any(log['request_body'] == '[multipart/form-data body omitted]' for log in logs)
    def encrypted_get(path, params):
        ts, nonce = int(time.time()*1000), api_crypto.gen_nonce_hex()
        key = api_crypto.derive_key(fixed_secret='doclib-test-key', user_id='owner', ts=ts, nonce=nonce)
        return c.post(path, json={'enc': 1, 'ts': ts, 'nonce': nonce, 'edata': api_crypto.encrypt_json({'data': params}, key)},
                      headers={'Authorization': 'Bearer '+token, 'X-HTTP-Method-Override': 'GET'})
    listed = encrypted_get('/api/web/doclib/files', {'scope': 'mine', 'page': 1, 'pageSize': 20})
    body = listed.json()
    assert body['encrypted'] is True
    key = api_crypto.derive_key(fixed_secret='doclib-test-key', user_id='owner', ts=body['ts'], nonce=body['nonce'])
    assert api_crypto.decrypt_json(body['edata'], key)['data']['total'] == 1
    downloaded = encrypted_get('/api/web/doclib/files/'+uploaded['id']+'/download', {})
    assert downloaded.content == b'streamed'
    assert any(log['response_body'] == '[document attachment body omitted]' for log in logs)

    c.close()


def test_failed_upload_discards_ambiguous_object(client):
    c, identity, storage = client
    def upload_then_timeout(key, file, length):
        storage.data[key] = file.read()
        raise TimeoutError('reply lost after MinIO stored object')
    storage.upload = upload_then_timeout
    result = c.post('/api/web/doclib/upload', files={'file': ('file.txt', b'abc')})
    assert result.status_code == 503
    assert storage.data == {}


def test_chunked_upload_stops_before_consuming_full_body(client, monkeypatch):
    import asyncio
    c, _, _ = client
    monkeypatch.setattr(settings, 'DOCLIB_MAX_UPLOAD_BYTES', 5)
    prefix = b'--test\r\nContent-Disposition: form-data; name="file"; filename="large.txt"\r\n\r\n'
    chunks = [prefix] + [b'x' * 16384] * 128 + [b'\r\n--test--\r\n']
    consumed, output = [], []
    async def receive():
        index = len(consumed)
        consumed.append(index)
        return {'type': 'http.request', 'body': chunks[index], 'more_body': index < len(chunks)-1}
    async def send(message):
        output.append(message)
    scope = {'type': 'http', 'asgi': {'version': '3.0'}, 'http_version': '1.1', 'method': 'POST',
             'scheme': 'http', 'path': '/api/web/doclib/upload', 'raw_path': b'/api/web/doclib/upload',
             'query_string': b'', 'headers': [(b'content-type', b'multipart/form-data; boundary=test'), (b'authorization', b'Bearer test')],
             'server': ('testserver', 80), 'client': ('testclient', 1234)}
    asyncio.run(c.app(scope, receive, send))
    assert next(m for m in output if m['type'] == 'http.response.start')['status'] == 413
    assert len(consumed) < 10


@pytest.mark.parametrize('encrypted', [True, False])
def test_upload_rejects_nonmultipart_without_reading_body(client, monkeypatch, encrypted):
    from app.core import api_crypto_middleware as crypto, api_request_log_middleware as audit
    c, _, _ = client
    app = FastAPI()
    app.include_router(router, prefix='/api')
    app.dependency_overrides = c.app.dependency_overrides.copy()
    app.add_middleware(audit.ApiRequestLogMiddleware)
    app.add_middleware(crypto.ApiCryptoMiddleware)
    monkeypatch.setattr(settings, 'API_ENCRYPT_ENABLED', encrypted)
    monkeypatch.setattr(settings, 'API_ENCRYPT_FIXED_SECRET', 'test-key')
    monkeypatch.setattr(settings, 'API_ENCRYPT_EXCLUDE_PATHS', '')
    monkeypatch.setattr(settings, 'API_REQUEST_LOG_ENABLED', True)
    monkeypatch.setattr(settings, 'DOCLIB_MAX_UPLOAD_BYTES', 5)
    logs, body_reads = [], []
    monkeypatch.setattr(audit, '_schedule_persist_log', lambda **data: logs.append(data))
    from starlette.requests import Request
    original = Request.body
    async def tracked_body(request):
        body_reads.append(True)
        return await original(request)
    monkeypatch.setattr(Request, 'body', tracked_body)
    with TestClient(app) as test_client:
        response = test_client.post('/api/web/doclib/upload', content=b'x' * (2*1024*1024),
                                    headers={'Content-Type': 'application/octet-stream', 'Authorization': 'Bearer test', 'X-HTTP-Method-Override': 'GET'})
    assert response.status_code == 415
    assert not body_reads
    assert logs and logs[0]['status_code'] == 415