pg-api 0.1.0

A high-performance PostgreSQL REST API driver with rate limiting, connection pooling, and observability
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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

use axum::{
    extract::Request,
    http::{HeaderMap, StatusCode},
    middleware::Next,
    response::Response,
};
use uuid::Uuid;

#[allow(unused_imports)]
use crate::models::{Account, AppState};

pub async fn auth_middleware(
    axum::extract::State(state): axum::extract::State<AppState>,
    mut request: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    // Skip auth for health check
    if request.uri().path() == "/health" {
        return Ok(next.run(request).await);
    }
    
    let api_key = extract_api_key(request.headers()).map_err(|e| {
        eprintln!("[AUTH] Failed to extract API key: {}", e);
        StatusCode::UNAUTHORIZED
    })?;
    
    eprintln!("[AUTH] Checking API key: {}...", &api_key[..20.min(api_key.len())]);
    
    // Validate API key
    let account = {
        let mut accounts = state.accounts.write().await;
        eprintln!("[AUTH] Total accounts loaded: {}", accounts.len());
        if let Some(account) = accounts.get_mut(&api_key) {
            account.last_used = chrono::Utc::now();
            Some(account.clone())
        } else {
            eprintln!("[AUTH] API key not found in {} accounts", accounts.len());
            None
        }
    }; // Release the lock here
    
    if let Some(account) = account {
        // Add account info to request extensions
        request.extensions_mut().insert(account);
        Ok(next.run(request).await)
    } else {
        Err(StatusCode::UNAUTHORIZED)
    }
}

pub async fn request_id_middleware(
    mut request: Request,
    next: Next,
) -> Response {
    let request_id = Uuid::new_v4().to_string();
    request.headers_mut().insert(
        "x-request-id",
        request_id.parse().unwrap(),
    );
    
    next.run(request).await
}

pub fn extract_api_key(headers: &HeaderMap) -> Result<String, String> {
    headers
        .get("x-api-key")
        .or_else(|| headers.get("authorization"))
        .and_then(|h| h.to_str().ok())
        .map(|h| {
            if h.starts_with("Bearer ") {
                h.strip_prefix("Bearer ").unwrap().to_string()
            } else {
                h.to_string()
            }
        })
        .ok_or_else(|| "Missing API key".to_string())
}