File size: 1,506 Bytes
32e1e2e
 
 
a415c67
32e1e2e
fdc8fdb
 
32e1e2e
a415c67
fdc8fdb
65526b0
 
 
 
a415c67
 
32e1e2e
a415c67
 
 
fdc8fdb
a415c67
32e1e2e
a415c67
fdc8fdb
 
32e1e2e
a415c67
fdc8fdb
65526b0
 
 
 
a415c67
 
fdc8fdb
a415c67
fdc8fdb
a415c67
32e1e2e
a415c67
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import os
import json
import bittensor as bt
from aiohttp.web import Request, Response, middleware

EXPECTED_ACCESS_KEY = os.environ.get("EXPECTED_ACCESS_KEY")


@middleware
async def api_key_middleware(request: Request, handler):
    if request.path.startswith("/docs") or request.path.startswith("/static/swagger"):
        # Skip checks when accessing OpenAPI documentation.
        return await handler(request)

    # Logging the request
    bt.logging.info(f"Handling {request.method} request to {request.path}")

    # Check access key
    access_key = request.headers.get("api_key")
    if EXPECTED_ACCESS_KEY is not None and access_key != EXPECTED_ACCESS_KEY:
        bt.logging.error(f"Invalid access key: {access_key}")
        return Response(status=401, reason="Invalid access key")

    # Continue to the next handler if the API key is valid
    return await handler(request)


@middleware
async def json_parsing_middleware(request: Request, handler):
    if request.path.startswith("/docs") or request.path.startswith("/static/swagger"):
        # Skip checks when accessing OpenAPI documentation.
        return await handler(request)

    try:
        # Parsing JSON data from the request
        request["data"] = await request.json()
    except json.JSONDecodeError as e:
        bt.logging.error(f"Invalid JSON data: {str(e)}")
        return Response(status=400, text="Invalid JSON")

    # Continue to the next handler if JSON is successfully parsed
    return await handler(request)