add api key auth

dunctk · dunctk · commit 6d28f2e129dc · 2025-06-11T12:42:08.000Z
diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml
@@ -28,4 +28,6 @@ jobs:
         uses: docker/setup-buildx-action@v2
 
       - name: Run Deployment Script
-        run: ./deploy-working.sh 
+        run: ./deploy-working.sh 
+        env:
+          API_KEY: ${{ secrets.SCREENSHOT_API_KEY }} 
diff --git a/README.md b/README.md
@@ -164,6 +164,44 @@ aws apigatewayv2 create-api \
 The function works out of the box but you can configure:
 
 - `RUST_LOG`: Set logging level (e.g., "info", "debug")
+- `API_KEY`: **Optional.** If set, every request must provide this value either in the `x-api-key` HTTP header **or** as the `key` query-string parameter. Omit the variable to disable authentication (useful for local testing).
+
+### Setting `API_KEY`
+
+Once the Lambda function exists you can set / change the key at any time:
+
+```bash
+# one-off update (keeps existing image & config)
+aws lambda update-function-configuration \
+  --function-name screenshotapi \
+  --environment "Variables={API_KEY=your-secret-value}" \
+  --region eu-central-1
+```
+
+The key persists for the lifetime of the function. **However:** the sample `deploy-working.sh` script currently deletes and recreates the function on each deploy. If you use that script you have two options:
+
+1. Add `--environment "Variables={API_KEY=your-secret-value}"` to the `aws lambda create-function` call inside the script so the key is applied on every deploy.
+2. Stop deleting the function (remove the `aws lambda delete-function` call). Updating an existing function keeps its environment variables intact.
+
+### Calling the API with the key
+
+```bash
+# Header style
+curl -H "x-api-key: your-secret-value" "${FUNCTION_URL}?url=https://example.com" | jq .
+
+# Query-string style (handy for quick browser tests)
+curl "${FUNCTION_URL}?url=https://example.com&key=your-secret-value" | jq .
+```
+
+Requests without the correct key receive:
+
+```json
+HTTP/1.1 401 Unauthorized
+{
+  "error": "UNAUTHORIZED",
+  "message": "Invalid or missing API key"
+}
+```
 
 ### Lambda Configuration Recommendations
 
diff --git a/deploy-working.sh b/deploy-working.sh
@@ -86,6 +86,7 @@ aws lambda create-function \
     --region "${REGION}" \
     --memory-size 2048 \
     --timeout 90 \
+    ${API_KEY:+--environment "Variables={API_KEY=$API_KEY}"} \
     --architectures x86_64 > /dev/null
 
 echo "Waiting for function to become active..."
diff --git a/src/http_handler.rs b/src/http_handler.rs
@@ -1,6 +1,7 @@
 use lambda_http::{Body, Error, Request, RequestExt, Response};
 use serde::Serialize;
 use crate::screenshot::ScreenshotService;
+use std::env;
 
 #[derive(Serialize)]
 struct ErrorResponse {
@@ -17,6 +18,41 @@ struct SuccessResponse {
 
 /// Main function handler for the screenshot API
 pub(crate) async fn function_handler(event: Request) -> Result<Response<Body>, Error> {
+    // --- Simple API key check ---
+    // Expected key (set via Lambda environment variable `API_KEY`)
+    let expected_key = env::var("API_KEY").unwrap_or_default();
+
+    if !expected_key.is_empty() {
+        // Try to read from `x-api-key` header first
+        let provided_key_header = event
+            .headers()
+            .get("x-api-key")
+            .and_then(|v| v.to_str().ok());
+
+        // Fallback: query parameter `key`
+        let provided_key_query = event
+            .query_string_parameters()
+            .first("key");
+
+        let provided_key = provided_key_header.or(provided_key_query);
+
+        if provided_key != Some(expected_key.as_str()) {
+            // Unauthorized
+            let error_response = ErrorResponse {
+                error: "UNAUTHORIZED".to_string(),
+                message: "Invalid or missing API key".to_string(),
+            };
+
+            let resp = Response::builder()
+                .status(401)
+                .header("content-type", "application/json")
+                .body(serde_json::to_string(&error_response)?.into())
+                .map_err(Box::new)?;
+
+            return Ok(resp);
+        }
+    }
+
     match handle_screenshot_request(event).await {
         Ok(response) => Ok(response),
         Err(e) => {
