add prometheus metrics for failed/successful connections

.gitignore

@@ -6,3 +6,5 @@ __pycache__/
 *.egg-info/
 build/
 dist/
+*.yml
+*.yaml

Dockerfile

@@ -3,4 +3,4 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
-CMD ["python", "server.py"]
+ENTRYPOINT ["python", "server.py"]

README.md

@@ -33,15 +33,26 @@ The server runs at `http://127.0.0.1:5000`.
 
 In one terminal, connect to the WebSocket tunnel:
 
-```powershell
-websocat ws://127.0.0.1:5000/tunnel
+```sh
+websocat ws://127.0.0.1:5000/tunnel/asdf
+
+# with an api key
+websocat \
+  --header="X-API-Key:hello" \
+  - ws://127.0.0.1:5000/tunnel/asdf
 ```
 
 In another terminal, send a webhook payload:
 
 ### bash
 ```sh
-curl -X POST http://127.0.0.1:5000/webhook \
+curl -X POST http://127.0.0.1:5000/webhook/asdf \
+  -H "Content-Type: application/json" \
+  -d '{"message":"hello from webhook"}'
+
+# with an api key
+curl -X POST http://127.0.0.1:5000/webhook/asdf \
+  -H "X-API-Key: hello" \
   -H "Content-Type: application/json" \
   -d '{"message":"hello from webhook"}'
 ```
@@ -52,3 +63,30 @@ Invoke-RestMethod -Method Post -Uri http://127.0.0.1:5000/webhook -ContentType "
 ```
 
 The connected WebSocket client should receive the JSON payload.
+
+## on the real website
+create an file called `config.yml` with content like
+```yml
+api_key: thesecretofalltime
+```
+run the server with
+```sh
+docker-compose up --build -d
+```
+
+if you wanna see logs
+```sh
+docker logs smee2-smee2-1 --tail 300 -f
+```
+
+to test that its working, use the same api key like:
+
+```sh
+curl -X POST https://sce.sjsu.edu/webhook/asdf \
+  -H "Content-Type: application/json" \
+  -d '{"message":"hello from webhook"}'
+
+websocat \
+  --header="X-API-Key:hello" \
+  - wss://sce.sjsu.edu/tunnel/asdf
+```

args.py

@@ -8,7 +8,8 @@ def get_args():
         "-v",
         action="count",
         default=0,
-        help="increase logging verbosity; can be used multiple times"
+        help="increase logging verbosity; can be used multiple times like -vvv"
     )
+    parser.add_argument("--config", help="path to yaml file for api key, see readme")
 
     return parser.parse_args()

docker-compose.yml

@@ -3,6 +3,12 @@ services:
     build:
       context: .
       dockerfile: ./Dockerfile
-    ports:
-      - "5000:5000"
+    volumes:
+      - ./config.yml:/app/config.yml
+    command:
+        - --config=/app/config.yml
 
+networks:
+  default:
+    external:
+      name: sce

requirements.txt

@@ -18,4 +18,3 @@ typing_extensions==4.15.0
 uvicorn==0.49.0
 watchfiles==1.2.0
 websockets==16.0
-

server.py

@@ -1,9 +1,11 @@
 import collections
 import logging
+import collections
 
 from fastapi import FastAPI, WebSocket, Request, HTTPException, Response
 import prometheus_client
 import uvicorn
+import yaml
 
 from args import get_args
 from metrics import MetricsHandler
@@ -19,6 +21,8 @@
 )
 logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
 logging.getLogger("uvicorn.error").setLevel(logging.WARNING)
+logging.getLogger("asyncio").setLevel(logging.WARNING)
+logging.getLogger("watchfiles").setLevel(logging.WARNING)
 
 logger = logging.getLogger(__name__)
 
@@ -27,35 +31,45 @@
 clients = collections.defaultdict(list)
 
 subscribers = {}
+# see if __name__ == '__server__' section for setting this value
+API_KEY = None
 
 @app.post("/webhook/{subscription_id}")
 async def webhook(subscription_id: str, request: Request):
-    if subscription_id is not None:
-        header_val = request.headers.get("X-API-Key")
-        if (header_val != "hello"):
-            raise HTTPException(status_code=403, detail="API key is not valid ")
-
-        data = await request.json()
-        logger.debug("Data pushed to webhook %s, received: %s", subscription_id, data)
-
-        subscribers[subscription_id] = data
-
-        for client in clients.get(subscription_id, []):
-            await client.send_json(data)
-
-        logger.error("Data sent to websocket client")
-        return {"message":"received"}  
-
-    else:   
-        logger.error("Invalid subscription '%s', connection not accepted", subscription_id)
-        return
-
-
+    if subscription_id is None:
+        logger.error(
+            "Invalid subscription '%s', connection not accepted", subscription_id
+        )
+        MetricsHandler.failed_connections.labels(
+            subscription_id=subscription_id,
+            reason="no_subscription_id",
+        ).inc()
+        raise HTTPException(
+            status_code=404,
+            detail="subscripiton id was empty. expecting soemthing like /webhook/1234",
+        )
+    header_val = request.headers.get("X-API-Key")
+
+    data = await request.json()
+    logger.debug("Data pushed to webhook %s, received: %s", subscription_id, data)
+
+    subscribers[subscription_id] = data
+
+    for client in clients.get(subscription_id, []):
+        await client.send_json(data)
+
+    number_of_clients = len(clients.get(subscription_id, []))
+    message = f"forwarded to {number_of_clients} client(s) for subscription {subscription_id}"
+    logger.debug(message)
+    return {"message": message}
+
+
 @app.websocket("/tunnel/{subscription_id}")
 async def websocket_endpoint(subscription_id: str, websocket: WebSocket):
-    api_key = websocket.headers.get("X-API-Key")
-
-    if (api_key != "hello"):
+    api_key_from_header = websocket.headers.get("X-API-Key")
+
+    if API_KEY is not None and api_key_from_header != API_KEY:
+        logger.error('/tunnel recieved invalid X-API-Key of "%s"', api_key_from_header)
         MetricsHandler.failed_connections.labels(
             subscription_id=subscription_id,
             reason="bad_api_key",
@@ -64,29 +78,33 @@ async def websocket_endpoint(subscription_id: str, websocket: WebSocket):
         return
 
     await websocket.accept()
-    
+
     MetricsHandler.connected_clients.labels(subscription_id).inc()
-    
+
     clients[subscription_id].append(websocket)
 
-    logger.debug(f"Websocket connection successfully established at id: {subscription_id}")
-
+    logger.debug(
+        f"Websocket connection successfully established at id: {subscription_id}"
+    )
+
     try:
         while True:
             data = await websocket.receive_text()
             await websocket.send_text("Message received")
-    except Exception as e:
+    except Exception:
+        logger.exception("ok")
         MetricsHandler.failed_connections.labels(
             subscription_id=subscription_id,
             reason="websocket_receive_failed",
         ).inc()
         MetricsHandler.connected_clients.labels(subscription_id).dec()
+    finally:
         clients[subscription_id].remove(websocket)
         if not clients[subscription_id]:
             clients.pop(subscription_id, None)
 
         logger.debug(f"Websocket connection disconnected at id: {subscription_id}")
-            
+
 
 @app.get("/metrics")
 def get_metrics():
@@ -95,6 +113,7 @@ def get_metrics():
         media_type="text/plain",
     )
 
+
 # we have a separate __name__ check here due to how FastAPI starts
 # a server. the file is first ran (where __name__ == "__main__")
 # and then calls `uvicorn.run`. the call to run() reruns the file,
@@ -105,6 +124,19 @@ def get_metrics():
 # server uses
 if __name__ == "server":
     MetricsHandler.init()
+    try:
+        with open(args.config, "r") as stream:
+            data = yaml.safe_load(stream)
+            API_KEY = data.get("api_key", None)
+        logger.info(f'loaded api key from {args.config}')
+    except Exception:
+        logging.warning("unable to open yaml file, smee2 is not checking for api keys")
 
 if __name__ == "__main__":
-    uvicorn.run("server:app", host="0.0.0.0", port=5000)
+    uvicorn.run(
+        "server:app",
+        host="0.0.0.0",
+        port=5000,
+        reload=True,
+        timeout_graceful_shutdown=1,
+    )