added await

README.md

@@ -4,8 +4,9 @@
 
 ```json
 {
-    "host": "127.0.0.1",
-    "port": 9250,
+    "binds": [
+        "127.0.0.1:9250"
+    ],
     "log": "-",
     "buffer-size": "4KiB",
     "max-size": "2GB",
@@ -13,6 +14,6 @@
         "TESTKEY"
     ],
     "max-data": "10GB",
-    "database": "databse.json"
+    "database": "database.json"
 }
 ```

src/main.py

@@ -4,7 +4,7 @@ import argparse
 import json
 from os.path import exists
 
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request, HTTPException, Query
 from fastapi.responses import StreamingResponse
 from fastapi import status
 from hypercorn.config import Config
@@ -32,7 +32,7 @@ if not exists(DATABASE):
     save_database(DATABASE, {'data-used': 0})
 
 
-api = FastAPI()
+api = FastAPI(docs_url=None, redoc_url=None)
 
 
 class MaxSizePerRequestError(Exception):
@@ -44,7 +44,7 @@ class MinSizePerRequestError(Exception):
 
 
 @api.get('/')
-def test_data(api_key: str, size: str, request: Request) -> StreamingResponse:
+async def test_data(api_key: str, request: Request, size: str) -> StreamingResponse:
     try:
         if api_key not in AUTHORIZED_KEYS:
             raise HTTPException(
@@ -52,7 +52,13 @@ def test_data(api_key: str, size: str, request: Request) -> StreamingResponse:
                 detail='Invalid API Key.'
             )
 
+        try:
             size = convert_to_bytes(size)
+        except ValueError:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail='Invalid format format for size.'
+            )
 
         if size < 0:
             raise MinSizePerRequestError
@@ -71,7 +77,8 @@ def test_data(api_key: str, size: str, request: Request) -> StreamingResponse:
 
         return StreamingResponse(
             status_code=status.HTTP_200_OK,
-            content=generate_data(size, BUFFER_SIZE)
+            content=generate_data(size, BUFFER_SIZE),
+            media_type='application/octet-stream'
         )
 
     except MinSizePerRequestError:
@@ -90,7 +97,7 @@ def main():
     asyncio.run(serve(
         api,
         Config().from_mapping(
-            bind=[f"{CONFIG['host']}:{CONFIG['port']}"],
+            bind=CONFIG['binds'],
             accesslog='-'
         )
     ))

src/utils.py

@@ -1,14 +1,15 @@
 import json
+import asyncio
 
 def convert_to_bytes(size: int | str) -> int:
     try:
         return int(size)
     except ValueError:  # treat as string
         units = {
-            'TB': 10 ** 12, 'TiB': 2 ** 40,
-            'GB': 10 ** 9, 'GiB': 2 ** 30,
-            'MB': 10 ** 6, 'MiB': 2 ** 20,
-            'KB': 10 ** 3, 'KiB': 2 ** 10,
+            'TB': 1000 ** 4, 'TiB': 1024 ** 4,
+            'GB': 1000 ** 3, 'GiB': 1024 ** 3,
+            'MB': 1000 ** 2, 'MiB': 1024 ** 2,
+            'KB': 1000, 'KiB': 1024,
             'B': 1
         }
 
@@ -17,15 +18,25 @@ def convert_to_bytes(size: int | str) -> int:
                 return int(float(size.removesuffix(unit)) * units[unit])
                 break
 
+        raise ValueError
+
 
 async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
     size_left = size
 
+    # https://github.com/tiangolo/fastapi/issues/5183
+    # https://github.com/encode/starlette/discussions/1776#discussioncomment-3207518
+
+    try:
         while size_left > buffer_size:
             size_left -= buffer_size
             yield b'\0' * buffer_size
+            await asyncio.sleep(0)
         else:
             yield b'\0' * size_left
+            await asyncio.sleep(0)
+    except asyncio.CancelledError:
+        raise GeneratorExit
 
 
 def check_policies(ip: str) -> None: