@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d841386d27
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Rust Frontend] Support API key authentication (#44321)

Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Chao-Ju Chen committed Jun 9, 2026 at 18:15 UTC d841386d272200dd381a5791833771db9a47adf7
14 files changed +384 -15
rust/Cargo.lock
+3
@@ -5799,6 +5799,7 @@ dependencies = [
5799 "axum",
5800 "bytes",
5801 "clap",
5802 + "educe",
5803 "expect-test",
5804 "futures",
5805 "http-body",
@@ -5813,7 +5814,9 @@ dependencies = [
5814 "serde_json",
5815 "serde_with",
5816 "serial_test",
5817 + "sha2",
5818 "socket2",
5819 + "subtle",
5820 "thiserror-ext",
5821 "tokio",
5822 "tokio-stream",
rust/Cargo.toml
+2
@@ -75,8 +75,10 @@ serde_repr = "0.1.20"
75 serde_tuple = "1.1.3"
76 serde_with = "3.18.0"
77 serial_test = { version = "3.2.0", features = ["file_locks"] }
78 +sha2 = "0.10.9"
79 socket2 = "0.6.3"
80 subenum = "1.1.3"
81 +subtle = "2.6"
82 task-local = "0.1.1"
83 tekken = { package = "tekken-rs", version = "0.1.1", default-features = false }
84 tempfile = "3.23.0"
rust/src/cmd/src/cli.rs
+25 -1
@@ -16,6 +16,7 @@ use educe::Educe;
16 use serde::Deserialize;
17 use serde::de::DeserializeOwned;
18 use serde_json::Value;
19 +use serde_with::{DefaultOnNull, OneOrMany, serde_as};
20 use thiserror_ext::AsReport as _;
21 use uuid::Uuid;
22 use vllm_engine_core_client::TransportMode;
@@ -84,6 +85,7 @@ pub enum Command {
85 }
86
87 /// Runtime arguments shared by the external-engine and managed-engine paths.
88 +#[serde_as]
89 #[derive(Educe, Clone, Args, PartialEq, Eq, Deserialize)]
90 #[educe(Debug)]
91 pub struct SharedRuntimeArgs {
@@ -178,6 +180,14 @@ pub struct SharedRuntimeArgs {
180 #[serde(default)]
181 pub enable_request_id_headers: bool,
182
183 + /// If provided, the server will require one of these keys to be presented
184 + /// in the Authorization header.
185 + #[educe(Debug(ignore))]
186 + #[arg(long, env = "VLLM_API_KEY", value_delimiter = ' ')]
187 + #[serde_as(as = "DefaultOnNull<OneOrMany<_>>")]
188 + #[serde(default)]
189 + pub api_key: Vec<String>,
190 +
191 /// Disable periodic logging of engine statistics (throughput, queue depth,
192 /// cache usage).
193 #[arg(long)]
@@ -215,6 +225,15 @@ impl SharedRuntimeArgs {
225 Duration::from_secs(self.shutdown_timeout)
226 }
227
228 + /// Apply fallback logic for API key configuration from env variables.
229 + fn apply_env_api_key_fallback(&mut self) {
230 + if self.api_key.is_empty()
231 + && let Ok(api_key) = std::env::var("VLLM_API_KEY")
232 + {
233 + self.api_key.push(api_key);
234 + }
235 + }
236 +
237 /// Build the OpenAI-server config for the Python-bootstrap worker contract.
238 ///
239 /// The resulting config binds the Python-supplied transport addresses and
@@ -253,6 +272,7 @@ impl SharedRuntimeArgs {
272 chat_template_content_format: self.chat_template_content_format,
273 enable_log_requests: self.enable_log_requests,
274 enable_request_id_headers: self.enable_request_id_headers,
275 + api_keys: self.api_key,
276 disable_log_stats: self.disable_log_stats,
277 grpc_port: self.grpc_port,
278 shutdown_timeout,
@@ -295,6 +315,7 @@ impl SharedRuntimeArgs {
315 chat_template_content_format: self.chat_template_content_format,
316 enable_log_requests: self.enable_log_requests,
317 enable_request_id_headers: self.enable_request_id_headers,
318 + api_keys: self.api_key,
319 disable_log_stats: self.disable_log_stats,
320 grpc_port: self.grpc_port,
321 shutdown_timeout,
@@ -311,8 +332,11 @@ fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
332 }
333
334 fn parse_runtime_args_json(value: &str) -> Result<SharedRuntimeArgs, String> {
314 - let args: SharedRuntimeArgs = serde_json::from_str(value)
335 + let mut args: SharedRuntimeArgs = serde_json::from_str(value)
336 .map_err(|e| format!("invalid JSON arguments: {}", e.as_report()))?;
337 + // --args-json is parsed with serde, so clap's env support does not run for
338 + // the Python-supervised frontend path.
339 + args.apply_env_api_key_fallback();
340 args.unsupported.check()?;
341 Ok(args)
342 }
rust/src/cmd/src/cli/tests.rs
+78 -4
@@ -169,6 +169,76 @@ fn frontend_args_json_passes_enable_request_id_headers_into_config() {
169 assert!(config.enable_request_id_headers);
170 }
171
172 +#[test]
173 +fn serve_passes_api_keys_into_config() {
174 + let cli = Cli::try_parse_from([
175 + "vllm-rs",
176 + "serve",
177 + "Qwen/Qwen3-0.6B",
178 + "--api-key",
179 + "secret-a",
180 + "--api-key",
181 + "secret-b",
182 + ])
183 + .unwrap();
184 +
185 + let Command::Serve(args) = cli.command else {
186 + panic!("expected serve args");
187 + };
188 + let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
189 + assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
190 + let debug = format!("{config:#?}");
191 + assert!(debug.contains("api_keys: [<redacted>; 2]"));
192 + assert!(!debug.contains("secret-a"));
193 + assert!(!debug.contains("secret-b"));
194 +}
195 +
196 +#[test]
197 +fn frontend_args_json_accepts_api_key_string() {
198 + let cli = Cli::try_parse_from([
199 + "vllm-rs",
200 + "frontend",
201 + "--listen-fd",
202 + "3",
203 + "--input-address",
204 + "ipc:///tmp/input.sock",
205 + "--output-address",
206 + "ipc:///tmp/output.sock",
207 + "--args-json",
208 + r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":"secret"}"#,
209 + ])
210 + .unwrap();
211 +
212 + let Command::Frontend(args) = cli.command else {
213 + panic!("expected frontend args");
214 + };
215 + let config = args.into_config();
216 + assert_eq!(config.api_keys, vec!["secret"]);
217 +}
218 +
219 +#[test]
220 +fn frontend_args_json_accepts_api_key_list() {
221 + let cli = Cli::try_parse_from([
222 + "vllm-rs",
223 + "frontend",
224 + "--listen-fd",
225 + "3",
226 + "--input-address",
227 + "ipc:///tmp/input.sock",
228 + "--output-address",
229 + "ipc:///tmp/output.sock",
230 + "--args-json",
231 + r#"{"model_tag":"Qwen/Qwen3-0.6B","api_key":["secret-a","secret-b"]}"#,
232 + ])
233 + .unwrap();
234 +
235 + let Command::Frontend(args) = cli.command else {
236 + panic!("expected frontend args");
237 + };
238 + let config = args.into_config();
239 + assert_eq!(config.api_keys, vec!["secret-a", "secret-b"]);
240 +}
241 +
242 #[test]
243 fn serve_args_reject_unknown_renderer_value() {
244 let error = Cli::try_parse_from([
@@ -446,20 +516,21 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() {
516 "--output-address",
517 "ipc:///tmp/output.sock",
518 "--args-json",
449 - r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}"#,
519 + r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}"#,
520 ])
521 .unwrap_err();
522
523 + let actual = error.to_string().replace(": \n", ":\n");
524 expect![[r#"
454 - error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"api_key":"secret"}' for '--args-json <JSON>':
525 + error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>':
526 The following arguments are not implemented in Rust frontend yet:
527 - allow_credentials
457 - - api_key
528 + - ssl_keyfile
529
530 Remove these arguments to continue.
531
532 For more information, try '--help'.
462 - "#]].assert_eq(&error.to_string());
533 + "#]].assert_eq(&actual);
534 }
535
536 #[test]
@@ -793,6 +864,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
864 chat_template_content_format: Auto,
865 enable_log_requests: false,
866 enable_request_id_headers: false,
867 + api_keys: [],
868 disable_log_stats: false,
869 grpc_port: None,
870 shutdown_timeout: 0ns,
@@ -857,6 +929,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
929 chat_template_content_format: Auto,
930 enable_log_requests: false,
931 enable_request_id_headers: false,
932 + api_keys: [],
933 disable_log_stats: false,
934 grpc_port: None,
935 shutdown_timeout: 0ns,
@@ -936,6 +1009,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
1009 chat_template_content_format: Auto,
1010 enable_log_requests: false,
1011 enable_request_id_headers: false,
1012 + api_keys: [],
1013 disable_log_stats: false,
1014 grpc_port: None,
1015 shutdown_timeout: 0ns,
rust/src/cmd/src/cli/unsupported.rs
-5
@@ -564,11 +564,6 @@ pub struct ServerUnsupportedArgs {
564 #[arg(long)]
565 pub allowed_headers: Option<Unsupported>,
566
567 - /// If provided, the server will require one of these keys to be presented
568 - /// in the header.
569 - #[arg(long)]
570 - pub api_key: Option<Unsupported>,
571 -
567 /// The file path to the SSL key file.
568 #[arg(long)]
569 pub ssl_keyfile: Option<Unsupported>,
rust/src/server/Cargo.toml
+3
@@ -8,6 +8,7 @@ license.workspace = true
8 anyhow.workspace = true
9 asynk-strim-attr.workspace = true
10 axum.workspace = true
11 +educe.workspace = true
12 futures.workspace = true
13 http-body.workspace = true
14 itertools.workspace = true
@@ -19,7 +20,9 @@ rmpv.workspace = true
20 serde.workspace = true
21 serde_json.workspace = true
22 serde_with.workspace = true
23 +sha2.workspace = true
24 socket2.workspace = true
25 +subtle.workspace = true
26 thiserror-ext.workspace = true
27 tokio.workspace = true
28 tokio-stream.workspace = true
rust/src/server/examples/external_engine_openai_qwen.rs
+1
@@ -70,6 +70,7 @@ async fn main() -> Result<()> {
70 chat_template_content_format: ChatTemplateContentFormatOption::Auto,
71 enable_log_requests: false,
72 enable_request_id_headers: false,
73 + api_keys: Vec::new(),
74 disable_log_stats: false,
75 grpc_port: None,
76 shutdown_timeout: Duration::ZERO,
rust/src/server/src/config.rs
+24 -1
@@ -1,7 +1,9 @@
1 use std::collections::HashMap;
2 +use std::fmt;
3 use std::time::Duration;
4
5 use anyhow::Result;
6 +use educe::Educe;
7 use serde::Serialize;
8 use serde_json::Value;
9 use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
@@ -33,7 +35,8 @@ pub enum CoordinatorMode {
35 }
36
37 /// Normalized runtime configuration for the minimal OpenAI-compatible server.
36 -#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
38 +#[derive(Educe, Clone, PartialEq, Eq, Serialize)]
39 +#[educe(Debug)]
40 pub struct Config {
41 /// Frontend-to-engine transport setup.
42 pub transport_mode: TransportMode,
@@ -67,6 +70,10 @@ pub struct Config {
70 pub enable_log_requests: bool,
71 /// When `true`, set `X-Request-Id` on every HTTP response.
72 pub enable_request_id_headers: bool,
73 + /// API keys accepted as bearer tokens for guarded routes.
74 + #[serde(skip_serializing)]
75 + #[educe(Debug(method(fmt_redacted_api_keys)))]
76 + pub api_keys: Vec<String>,
77 /// When `true`, suppress periodic stats logging (throughput, queue depth,
78 /// cache usage).
79 pub disable_log_stats: bool,
@@ -114,3 +121,19 @@ impl Config {
121 }
122 }
123 }
124 +
125 +struct RedactedApiKeys<'a>(&'a [String]);
126 +
127 +impl fmt::Debug for RedactedApiKeys<'_> {
128 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129 + if self.0.is_empty() {
130 + f.debug_list().finish()
131 + } else {
132 + write!(f, "[<redacted>; {}]", self.0.len())
133 + }
134 + }
135 +}
136 +
137 +fn fmt_redacted_api_keys(api_keys: &[String], f: &mut fmt::Formatter<'_>) -> fmt::Result {
138 + fmt::Debug::fmt(&RedactedApiKeys(api_keys), f)
139 +}
rust/src/server/src/lib.rs
+2 -1
@@ -92,7 +92,8 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
92 AppState::new(served_model_names, chat)
93 .with_log_requests(config.enable_log_requests)
94 .with_request_id_headers(config.enable_request_id_headers)
95 - .with_server_info(ServerInfoSnapshot::from_config(config)),
95 + .with_server_info(ServerInfoSnapshot::from_config(config))
96 + .with_api_keys(config.api_keys.clone()),
97 ))
98 }
99
rust/src/server/src/middleware/auth.rs new
+91
@@ -0,0 +1,91 @@
1 +use std::sync::Arc;
2 +
3 +use axum::Json;
4 +use axum::extract::{Request, State};
5 +use axum::http::header::AUTHORIZATION;
6 +use axum::http::{HeaderValue, Method, StatusCode};
7 +use axum::middleware::Next;
8 +use axum::response::{IntoResponse, Response};
9 +use serde_json::json;
10 +
11 +use crate::state::{ApiKeyHash, AppState, hash_api_key};
12 +
13 +const GUARDED_PREFIXES: &[&str] = &["/v1", "/v2", "/inference"];
14 +
15 +/// Authenticate guarded HTTP routes with an OpenAI-compatible bearer token.
16 +///
17 +/// Mirrors Python `AuthenticationMiddleware`: OPTIONS requests and non-guarded
18 +/// helper endpoints such as `/health` are allowed through without a token.
19 +pub async fn authenticate_api_key(
20 + State(state): State<Arc<AppState>>,
21 + req: Request,
22 + next: Next,
23 +) -> Response {
24 + if req.method() == Method::OPTIONS || !requires_auth(req.uri().path()) {
25 + return next.run(req).await;
26 + }
27 +
28 + if verify_token(req.headers().get(AUTHORIZATION), state.api_key_hashes()) {
29 + return next.run(req).await;
30 + }
31 +
32 + (
33 + StatusCode::UNAUTHORIZED,
34 + Json(json!({ "error": "Unauthorized" })),
35 + )
36 + .into_response()
37 +}
38 +
39 +fn requires_auth(path: &str) -> bool {
40 + GUARDED_PREFIXES.iter().any(|prefix| path.starts_with(prefix))
41 +}
42 +
43 +fn verify_token(authorization: Option<&HeaderValue>, api_key_hashes: &[ApiKeyHash]) -> bool {
44 + let Some(authorization) = authorization else {
45 + return false;
46 + };
47 + let Ok(authorization) = authorization.to_str() else {
48 + return false;
49 + };
50 + let Some((scheme, token)) = authorization.split_once(' ') else {
51 + return false;
52 + };
53 + if !scheme.eq_ignore_ascii_case("bearer") {
54 + return false;
55 + }
56 +
57 + let token_hash = hash_api_key(token);
58 + let mut token_match = false;
59 + for api_key_hash in api_key_hashes {
60 + token_match |= constant_time_eq(&token_hash, api_key_hash);
61 + }
62 + token_match
63 +}
64 +
65 +fn constant_time_eq(left: &ApiKeyHash, right: &ApiKeyHash) -> bool {
66 + use subtle::ConstantTimeEq;
67 +
68 + bool::from(left.ct_eq(right))
69 +}
70 +
71 +#[cfg(test)]
72 +mod tests {
73 + use super::constant_time_eq;
74 + use crate::state::hash_api_key;
75 +
76 + #[test]
77 + fn constant_time_eq_checks_sha256_digests() {
78 + assert!(constant_time_eq(
79 + &hash_api_key("secret"),
80 + &hash_api_key("secret")
81 + ));
82 + assert!(!constant_time_eq(
83 + &hash_api_key("secret"),
84 + &hash_api_key("secrex")
85 + ));
86 + assert!(!constant_time_eq(
87 + &hash_api_key("secret"),
88 + &hash_api_key("secret-more")
89 + ));
90 + }
91 +}
rust/src/server/src/middleware/mod.rs
+2
@@ -1,7 +1,9 @@
1 +mod auth;
2 mod load;
3 mod metrics;
4 mod request_id;
5
6 +pub use auth::authenticate_api_key;
7 pub use load::track_server_load;
8 pub use metrics::track_http_metrics;
9 pub use request_id::set_request_id_header;
rust/src/server/src/routes.rs
+17 -3
@@ -98,11 +98,25 @@ fn build_router_with_options(
98 }
99
100 let enable_request_id_headers = state.enable_request_id_headers;
101 + let enable_api_key_auth = state.has_api_keys();
102 let mut router = router
103 .with_state(state.clone())
103 - .layer(from_fn_with_state(state, middleware::track_server_load))
104 - .layer(from_fn(middleware::track_http_metrics))
105 - .layer(TraceLayer::new_for_http());
104 + .layer(from_fn_with_state(
105 + state.clone(),
106 + middleware::track_server_load,
107 + ))
108 + .layer(from_fn(middleware::track_http_metrics));
109 +
110 + if enable_api_key_auth {
111 + router = router.layer(from_fn_with_state(
112 + state.clone(),
113 + middleware::authenticate_api_key,
114 + ));
115 + }
116 +
117 + // Later layers wrap earlier ones. Keep tracing outside auth so rejected
118 + // requests are visible, while metrics/load only see authenticated traffic.
119 + router = router.layer(TraceLayer::new_for_http());
120
121 if enable_request_id_headers {
122 router = router.layer(from_fn(middleware::set_request_id_header));
rust/src/server/src/routes/tests.rs
+108
@@ -775,6 +775,19 @@ async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
775 (app, engine_task)
776 }
777
778 +async fn test_app_with_api_keys(api_keys: Vec<String>) -> (axum::Router, MockEngineTask) {
779 + let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
780 + b"engine-openai-api-key",
781 + default_stream_output_specs(),
782 + Arc::new(FakeChatBackend::new()),
783 + )
784 + .await;
785 + let app = build_router(Arc::new(
786 + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat).with_api_keys(api_keys),
787 + ));
788 + (app, engine_task)
789 +}
790 +
791 async fn test_health_app_with_engine_script<F>(
792 script: F,
793 ) -> (axum::Router, Arc<AppState>, MockEngineTask)
@@ -1073,6 +1086,101 @@ async fn request_id_header_echoes_incoming_header_when_enabled() {
1086 assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
1087 }
1088
1089 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1090 +#[serial]
1091 +async fn api_key_auth_rejects_missing_token_on_guarded_route() {
1092 + let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
1093 + let response = app
1094 + .call(
1095 + Request::builder()
1096 + .method("GET")
1097 + .uri("/v1/models")
1098 + .body(Body::empty())
1099 + .expect("build request"),
1100 + )
1101 + .await
1102 + .expect("call app");
1103 +
1104 + assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
1105 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
1106 + let json: serde_json::Value = serde_json::from_slice(&body).expect("json body");
1107 + assert_eq!(json, json!({ "error": "Unauthorized" }));
1108 +}
1109 +
1110 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1111 +#[serial]
1112 +async fn api_key_auth_rejects_wrong_token_on_guarded_route() {
1113 + let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
1114 + let response = app
1115 + .call(
1116 + Request::builder()
1117 + .method("GET")
1118 + .uri("/v1/models")
1119 + .header("authorization", "Bearer wrong")
1120 + .body(Body::empty())
1121 + .expect("build request"),
1122 + )
1123 + .await
1124 + .expect("call app");
1125 +
1126 + assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
1127 +}
1128 +
1129 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1130 +#[serial]
1131 +async fn api_key_auth_accepts_matching_bearer_token_on_guarded_route() {
1132 + let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
1133 + let response = app
1134 + .call(
1135 + Request::builder()
1136 + .method("GET")
1137 + .uri("/v1/models")
1138 + .header("authorization", "Bearer secret")
1139 + .body(Body::empty())
1140 + .expect("build request"),
1141 + )
1142 + .await
1143 + .expect("call app");
1144 +
1145 + assert_eq!(response.status(), StatusCode::OK);
1146 +}
1147 +
1148 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1149 +#[serial]
1150 +async fn api_key_auth_allows_options_without_token() {
1151 + let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
1152 + let response = app
1153 + .call(
1154 + Request::builder()
1155 + .method("OPTIONS")
1156 + .uri("/v1/models")
1157 + .body(Body::empty())
1158 + .expect("build request"),
1159 + )
1160 + .await
1161 + .expect("call app");
1162 +
1163 + assert_ne!(response.status(), StatusCode::UNAUTHORIZED);
1164 +}
1165 +
1166 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1167 +#[serial]
1168 +async fn api_key_auth_allows_unguarded_route_without_token() {
1169 + let (mut app, _engine_task) = test_app_with_api_keys(vec!["secret".to_string()]).await;
1170 + let response = app
1171 + .call(
1172 + Request::builder()
1173 + .method("GET")
1174 + .uri("/health")
1175 + .body(Body::empty())
1176 + .expect("build request"),
1177 + )
1178 + .await
1179 + .expect("call app");
1180 +
1181 + assert_eq!(response.status(), StatusCode::OK);
1182 +}
1183 +
1184 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1185 #[serial]
1186 async fn version_returns_engine_vllm_version() {
rust/src/server/src/state.rs
+28
@@ -2,6 +2,7 @@ use std::sync::Arc;
2 use std::sync::atomic::{AtomicU64, Ordering};
3
4 use serde_json::Value;
5 +use sha2::{Digest, Sha256};
6 use tokio::time::{Duration, Instant, sleep_until};
7 use tracing::warn;
8 use vllm_chat::ChatLlm;
@@ -14,6 +15,12 @@ use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};
15
16 const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
17
18 +pub(crate) type ApiKeyHash = [u8; 32];
19 +
20 +pub(crate) fn hash_api_key(api_key: &str) -> ApiKeyHash {
21 + Sha256::digest(api_key.as_bytes()).into()
22 +}
23 +
24 /// Shared router state for the minimal single-model OpenAI server.
25 pub struct AppState {
26 /// All public model IDs served by this frontend. The first entry is the
@@ -27,6 +34,8 @@ pub struct AppState {
34 pub enable_request_id_headers: bool,
35 /// Runtime server information returned by `/server_info`, when available.
36 server_info: Option<ServerInfoSnapshot>,
37 + /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes.
38 + api_key_hashes: Vec<ApiKeyHash>,
39 /// Number of in-flight inference requests currently owned by this frontend.
40 server_load: AtomicU64,
41 /// Dynamic LoRA adapter registry.
@@ -53,6 +62,7 @@ impl AppState {
62 enable_log_requests: false,
63 enable_request_id_headers: false,
64 server_info: None,
65 + api_key_hashes: Vec::new(),
66 server_load: AtomicU64::new(0),
67 lora_manager: LoraManager::new(),
68 }
@@ -84,6 +94,24 @@ impl AppState {
94 self.server_info.as_ref().map(|server_info| server_info.response(config_format))
95 }
96
97 + /// Configure API keys accepted by guarded HTTP routes.
98 + pub fn with_api_keys(mut self, api_keys: Vec<String>) -> Self {
99 + self.api_key_hashes = api_keys
100 + .into_iter()
101 + .filter(|key| !key.is_empty())
102 + .map(|key| hash_api_key(&key))
103 + .collect();
104 + self
105 + }
106 +
107 + pub(crate) fn has_api_keys(&self) -> bool {
108 + !self.api_key_hashes.is_empty()
109 + }
110 +
111 + pub(crate) fn api_key_hashes(&self) -> &[ApiKeyHash] {
112 + &self.api_key_hashes
113 + }
114 +
115 /// The primary model name echoed back in API responses (the first served
116 /// name).
117 pub fn primary_model_name(&self) -> &str {