aprender-serve 0.64.0

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! Falsifiers for the `RouterConfig` hardening flags behind
//! `apr serve run --no-metrics` / `--no-cors`.
//!
//! These assert observable HTTP behaviour (status code, response headers), not
//! the presence of a config field: an operator who passes `--no-metrics` must
//! get a 404 from `/metrics`, and one who passes `--no-cors` must see no
//! `access-control-*` header on the wire.

use axum::{
    body::Body,
    http::{Request, StatusCode},
};
use tower::util::ServiceExt;

use crate::api::*;

fn router_with(cors: bool, metrics: bool) -> axum::Router {
    let state = AppState::with_cache(10);
    create_router_with_config(
        state,
        RouterConfig {
            openai_api: true,
            cors,
            metrics,
        },
    )
}

async fn get(app: axum::Router, uri: &str) -> axum::http::Response<Body> {
    app.oneshot(
        Request::builder()
            .method("GET")
            .uri(uri)
            .header("origin", "http://example.invalid")
            .body(Body::empty())
            .expect("request builds"),
    )
    .await
    .expect("router responds")
}

async fn body_string(response: axum::http::Response<Body>) -> String {
    let bytes = axum::body::to_bytes(response.into_body(), 1024 * 1024)
        .await
        .expect("body collects");
    String::from_utf8_lossy(&bytes).into_owned()
}

/// Baseline: with metrics enabled the endpoint really does serve Prometheus
/// text. Without this, a "fix" that deleted the route would look correct.
#[tokio::test]
async fn metrics_enabled_serves_prometheus_exposition() {
    let response = get(router_with(true, true), "/metrics").await;
    assert_eq!(response.status(), StatusCode::OK);

    let body = body_string(response).await;
    assert!(
        body.contains("realizar_requests_total"),
        "enabled /metrics must expose the counter set, got: {body}"
    );
}

/// `--no-metrics` must withhold telemetry, not merely hide the banner line.
#[tokio::test]
async fn metrics_disabled_makes_metrics_endpoints_404() {
    for uri in ["/metrics", "/metrics/dispatch"] {
        let response = get(router_with(true, false), uri).await;
        assert_eq!(
            response.status(),
            StatusCode::NOT_FOUND,
            "{uri} must 404 when metrics are disabled"
        );

        let body = body_string(response).await;
        assert!(
            !body.contains("realizar_requests_total"),
            "{uri} leaked metrics while disabled: {body}"
        );
    }
}

/// Disabling metrics must not take the rest of the server down with it.
#[tokio::test]
async fn metrics_disabled_keeps_health_serving() {
    let response = get(router_with(true, false), "/health").await;
    assert_eq!(response.status(), StatusCode::OK);
}

/// Baseline: CORS enabled really does advertise a permissive origin.
#[tokio::test]
async fn cors_enabled_sends_allow_origin_header() {
    let response = get(router_with(true, true), "/health").await;
    assert_eq!(
        response
            .headers()
            .get("access-control-allow-origin")
            .map(|v| v.to_str().unwrap_or_default().to_owned()),
        Some("*".to_owned()),
        "CORS-enabled server must advertise a permissive origin"
    );
}

/// `--no-cors` must strip every `access-control-*` header from the response.
#[tokio::test]
async fn cors_disabled_sends_no_access_control_headers() {
    let response = get(router_with(false, true), "/health").await;
    assert_eq!(response.status(), StatusCode::OK);

    let leaked: Vec<String> = response
        .headers()
        .keys()
        .map(|k| k.as_str().to_owned())
        .filter(|k| k.starts_with("access-control"))
        .collect();
    assert!(
        leaked.is_empty(),
        "CORS-disabled server still sent: {leaked:?}"
    );
}

/// A browser preflight against a CORS-disabled server must not be answered
/// with permissive method/header allowances either.
#[tokio::test]
async fn cors_disabled_does_not_answer_preflight_permissively() {
    let response = router_with(false, true)
        .oneshot(
            Request::builder()
                .method("OPTIONS")
                .uri("/health")
                .header("origin", "http://example.invalid")
                .header("access-control-request-method", "POST")
                .body(Body::empty())
                .expect("request builds"),
        )
        .await
        .expect("router responds");

    let leaked: Vec<String> = response
        .headers()
        .keys()
        .map(|k| k.as_str().to_owned())
        .filter(|k| k.starts_with("access-control"))
        .collect();
    assert!(
        leaked.is_empty(),
        "CORS-disabled preflight still sent: {leaked:?}"
    );
}