feat(cloud): harden multi-tenant runtime resources

src/langbot/libs/dify_service_api/v1/client.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import httpx
 import typing
 import json
@@ -8,6 +9,75 @@ from .errors import DifyAPIError
 from pathlib import Path
 import os
 
+_MAX_DIFY_RESPONSE_BYTES = 1024 * 1024
+_MAX_DIFY_SSE_LINE_BYTES = 1024 * 1024
+_MAX_DIFY_STREAM_BYTES = 16 * 1024 * 1024
+_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
+
+
+async def _read_limited_response(
+    response: httpx.Response,
+    *,
+    max_bytes: int = _MAX_DIFY_RESPONSE_BYTES,
+) -> bytes:
+    content_length = response.headers.get('Content-Length')
+    if content_length is not None:
+        try:
+            if int(content_length) > max_bytes:
+                raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
+        except (TypeError, ValueError):
+            pass
+
+    body = bytearray()
+    async for chunk in response.aiter_bytes(chunk_size=8192):
+        body.extend(chunk)
+        if len(body) > max_bytes:
+            raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
+    return bytes(body)
+
+
+async def _iter_sse_json(
+    response: httpx.Response,
+) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
+    """Parse Dify's one-JSON-per-data-line SSE without unbounded line buffering."""
+
+    buffer = bytearray()
+    total = 0
+    async for chunk in response.aiter_bytes(chunk_size=8192):
+        total += len(chunk)
+        if total > _MAX_DIFY_STREAM_BYTES:
+            raise DifyAPIError('Dify SSE stream exceeds the runtime limit')
+        buffer.extend(chunk)
+        while b'\n' in buffer:
+            raw_line, _, remainder = buffer.partition(b'\n')
+            buffer = bytearray(remainder)
+            if len(raw_line) > _MAX_DIFY_SSE_LINE_BYTES:
+                raise DifyAPIError('Dify SSE event exceeds the runtime limit')
+            line = raw_line.rstrip(b'\r').strip()
+            if not line or not line.startswith(b'data:'):
+                continue
+            payload = json.loads(line[5:].decode('utf-8', errors='replace'))
+            if isinstance(payload, dict):
+                yield payload
+        if len(buffer) > _MAX_DIFY_SSE_LINE_BYTES:
+            raise DifyAPIError('Dify SSE event exceeds the runtime limit')
+
+    line = bytes(buffer).rstrip(b'\r').strip()
+    if line.startswith(b'data:'):
+        payload = json.loads(line[5:].decode('utf-8', errors='replace'))
+        if isinstance(payload, dict):
+            yield payload
+
+
+def _read_local_file_limited(path: Path) -> bytes:
+    if path.stat().st_size > _MAX_DIFY_UPLOAD_BYTES:
+        raise ValueError('Dify upload exceeds the size limit')
+    with path.open('rb') as handle:
+        body = handle.read(_MAX_DIFY_UPLOAD_BYTES + 1)
+    if len(body) > _MAX_DIFY_UPLOAD_BYTES:
+        raise ValueError('Dify upload exceeds the size limit')
+    return body
+
 
 class AsyncDifyServiceClient:
     """Dify Service API 客户端"""
@@ -22,6 +92,21 @@ class AsyncDifyServiceClient:
     ) -> None:
         self.api_key = api_key
         self.base_url = base_url
+        self._client: httpx.AsyncClient | None = None
+
+    def _get_client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                base_url=self.base_url,
+                trust_env=True,
+            )
+        return self._client
+
+    async def aclose(self) -> None:
+        client = self._client
+        self._client = None
+        if client is not None:
+            await client.aclose()
 
     async def chat_messages(
         self,
@@ -38,37 +123,32 @@ class AsyncDifyServiceClient:
         if response_mode != 'streaming':
             raise DifyAPIError('当前仅支持 streaming 模式')
 
-        async with httpx.AsyncClient(
-            base_url=self.base_url,
-            trust_env=True,
-            timeout=timeout,
-        ) as client:
-            payload = {
-                'inputs': inputs,
-                'query': query,
-                'user': user,
-                'response_mode': response_mode,
-                'conversation_id': conversation_id,
-                'files': files,
-                'model_config': model_config or {},
-            }
+        client = self._get_client()
+        payload = {
+            'inputs': inputs,
+            'query': query,
+            'user': user,
+            'response_mode': response_mode,
+            'conversation_id': conversation_id,
+            'files': files,
+            'model_config': model_config or {},
+        }
 
-            async with client.stream(
-                'POST',
-                '/chat-messages',
-                headers={
-                    'Authorization': f'Bearer {self.api_key}',
-                    'Content-Type': 'application/json',
-                },
-                json=payload,
-            ) as r:
-                async for chunk in r.aiter_lines():
-                    if r.status_code != 200:
-                        raise DifyAPIError(f'{r.status_code} {chunk}')
-                    if chunk.strip() == '':
-                        continue
-                    if chunk.startswith('data:'):
-                        yield json.loads(chunk[5:])
+        async with client.stream(
+            'POST',
+            '/chat-messages',
+            headers={
+                'Authorization': f'Bearer {self.api_key}',
+                'Content-Type': 'application/json',
+            },
+            json=payload,
+            timeout=timeout,
+        ) as r:
+            if r.status_code != 200:
+                body = await _read_limited_response(r)
+                raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
+            async for event in _iter_sse_json(r):
+                yield event
 
     async def workflow_run(
         self,
@@ -82,32 +162,27 @@ class AsyncDifyServiceClient:
         if response_mode != 'streaming':
             raise DifyAPIError('当前仅支持 streaming 模式')
 
-        async with httpx.AsyncClient(
-            base_url=self.base_url,
-            trust_env=True,
+        client = self._get_client()
+        async with client.stream(
+            'POST',
+            '/workflows/run',
+            headers={
+                'Authorization': f'Bearer {self.api_key}',
+                'Content-Type': 'application/json',
+            },
+            json={
+                'inputs': inputs,
+                'user': user,
+                'response_mode': response_mode,
+                'files': files,
+            },
             timeout=timeout,
-        ) as client:
-            async with client.stream(
-                'POST',
-                '/workflows/run',
-                headers={
-                    'Authorization': f'Bearer {self.api_key}',
-                    'Content-Type': 'application/json',
-                },
-                json={
-                    'inputs': inputs,
-                    'user': user,
-                    'response_mode': response_mode,
-                    'files': files,
-                },
-            ) as r:
-                async for chunk in r.aiter_lines():
-                    if r.status_code != 200:
-                        raise DifyAPIError(f'{r.status_code} {chunk}')
-                    if chunk.strip() == '':
-                        continue
-                    if chunk.startswith('data:'):
-                        yield json.loads(chunk[5:])
+        ) as r:
+            if r.status_code != 200:
+                body = await _read_limited_response(r)
+                raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
+            async for event in _iter_sse_json(r):
+                yield event
 
     async def workflow_submit(
         self,
@@ -129,41 +204,38 @@ class AsyncDifyServiceClient:
             'Content-Type': 'application/json',
         }
 
-        async with httpx.AsyncClient(
-            base_url=self.base_url,
-            trust_env=True,
+        client = self._get_client()
+        # Step 1: Submit the form
+        payload: dict[str, typing.Any] = {
+            'inputs': inputs if isinstance(inputs, dict) else {},
+            'user': user,
+            'action': action,
+        }
+
+        async with client.stream(
+            'POST',
+            f'/form/human_input/{form_token}',
+            headers=headers,
+            json=payload,
             timeout=timeout,
-        ) as client:
-            # Step 1: Submit the form
-            payload: dict[str, typing.Any] = {
-                'inputs': inputs if isinstance(inputs, dict) else {},
-                'user': user,
-                'action': action,
-            }
-
-            submit_resp = await client.post(
-                f'/form/human_input/{form_token}',
-                headers=headers,
-                json=payload,
-            )
+        ) as submit_resp:
+            submit_body = await _read_limited_response(submit_resp)
             if submit_resp.status_code != 200:
-                raise DifyAPIError(f'{submit_resp.status_code} {submit_resp.text}')
+                raise DifyAPIError(f'{submit_resp.status_code} {submit_body.decode(errors="replace")}')
 
-            # Step 2: Stream resumed workflow events
-            async with client.stream(
-                'GET',
-                f'/workflow/{workflow_run_id}/events',
-                headers={'Authorization': f'Bearer {self.api_key}'},
-                params={'user': user},
-            ) as r:
-                if r.status_code != 200:
-                    body = (await r.aread()).decode(errors='replace')
-                    raise DifyAPIError(f'{r.status_code} {body}')
-                async for chunk in r.aiter_lines():
-                    if chunk.strip() == '':
-                        continue
-                    if chunk.startswith('data:'):
-                        yield json.loads(chunk[5:])
+        # Step 2: Stream resumed workflow events
+        async with client.stream(
+            'GET',
+            f'/workflow/{workflow_run_id}/events',
+            headers={'Authorization': f'Bearer {self.api_key}'},
+            params={'user': user},
+            timeout=timeout,
+        ) as r:
+            if r.status_code != 200:
+                body = await _read_limited_response(r)
+                raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
+            async for event in _iter_sse_json(r):
+                yield event
 
     async def upload_file(
         self,
@@ -175,37 +247,30 @@ class AsyncDifyServiceClient:
         if isinstance(file, Path):
             if not file.exists():
                 raise ValueError(f'File not found: {file}')
-            with open(file, 'rb') as f:
-                file = f.read()
+            file = await asyncio.to_thread(_read_local_file_limited, file)
 
         # 处理文件路径字符串
         elif isinstance(file, str):
             if not os.path.isfile(file):
                 raise ValueError(f'File not found: {file}')
-            with open(file, 'rb') as f:
-                file = f.read()
+            file = await asyncio.to_thread(_read_local_file_limited, Path(file))
 
         # 处理文件对象
         elif hasattr(file, 'read'):
-            file = file.read()
-        async with httpx.AsyncClient(
-            base_url=self.base_url,
-            trust_env=True,
+            file = await asyncio.to_thread(file.read, _MAX_DIFY_UPLOAD_BYTES + 1)
+            if len(file) > _MAX_DIFY_UPLOAD_BYTES:
+                raise ValueError('Dify upload exceeds the size limit')
+        client = self._get_client()
+        # multipart/form-data
+        async with client.stream(
+            'POST',
+            '/files/upload',
+            headers={'Authorization': f'Bearer {self.api_key}'},
+            files={'file': file},
+            data={'user': user},
             timeout=timeout,
-        ) as client:
-            # multipart/form-data
-            response = await client.post(
-                '/files/upload',
-                headers={'Authorization': f'Bearer {self.api_key}'},
-                files={
-                    'file': file,
-                },
-                data={
-                    'user': user,
-                },
-            )
-
+        ) as response:
+            body = await _read_limited_response(response)
             if response.status_code != 201:
-                raise DifyAPIError(f'{response.status_code} {response.text}')
-
-            return response.json()
+                raise DifyAPIError(f'{response.status_code} {body.decode(errors="replace")}')
+        return json.loads(body)