Skip to content

Commit 887f72b

Browse files
authored
feat(auth): require API key alongside whitelisted hotkey (#3)
* feat(auth): require API key alongside hotkey for authentication Add a WORKER_API_KEY environment variable (required) that whitelisted hotkeys must provide via the X-Api-Key HTTP header to access protected endpoints. Authentication now requires both a valid whitelisted hotkey (X-Hotkey) and a matching API key (X-Api-Key) for submitting batches via POST /submit. This adds a second layer of request validation controlled by the worker operator. Changes: - auth.rs: Add api_key field to AuthHeaders, extract X-Api-Key header in extract_auth_headers, add InvalidApiKey variant to AuthError, and validate the API key in verify_request against the expected value. Added tests for invalid API key rejection and missing header detection. - config.rs: Add worker_api_key field to Config, loaded from WORKER_API_KEY env var (panics if unset). Log confirmation at startup. - handlers.rs: Pass worker_api_key from config to verify_request, update error message to list X-Api-Key as a required header. - AGENTS.md: Document WORKER_API_KEY env var and updated auth flow. - Cargo.lock: Version bump to 1.1.0. * ci: trigger CI run
1 parent dc8d8d4 commit 887f72b

File tree

5 files changed

+73
-18
lines changed

5 files changed

+73
-18
lines changed

AGENTS.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,8 @@ Both hooks are activated via `git config core.hooksPath .githooks`.
164164
| `MAX_ARCHIVE_BYTES` | `524288000` | Max uploaded archive size (500MB) |
165165
| `MAX_OUTPUT_BYTES` | `1048576` | Max captured output per command (1MB) |
166166
| `WORKSPACE_BASE` | `/tmp/sessions` | Base directory for session workspaces |
167+
| `WORKER_API_KEY` | *(required)* | API key that whitelisted hotkeys must provide via `X-Api-Key` header |
167168

168169
## Authentication
169170

170-
Authentication uses SS58 hotkey validation via the `X-Hotkey` HTTP header. The authorized hotkey is hardcoded as `AUTHORIZED_HOTKEY` in `src/config.rs`. Only requests with a matching hotkey can submit batches via `POST /submit`. All other endpoints are open.
171+
Authentication uses SS58 hotkey validation via the `X-Hotkey` HTTP header combined with an API key via the `X-Api-Key` header. The authorized hotkey is hardcoded as `AUTHORIZED_HOTKEY` in `src/config.rs`. The API key is configured via the `WORKER_API_KEY` environment variable (required). Only requests with both a matching hotkey and a valid API key can submit batches via `POST /submit`. All other endpoints are open.

Cargo.lock

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

src/auth.rs

Lines changed: 58 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ pub struct AuthHeaders {
4141
pub hotkey: String,
4242
pub nonce: String,
4343
pub signature: String,
44+
pub api_key: String,
4445
}
4546

4647
pub fn extract_auth_headers(headers: &axum::http::HeaderMap) -> Option<AuthHeaders> {
@@ -62,18 +63,33 @@ pub fn extract_auth_headers(headers: &axum::http::HeaderMap) -> Option<AuthHeade
6263
.and_then(|v| v.to_str().ok())
6364
.map(|s| s.to_string())?;
6465

66+
let api_key = headers
67+
.get("X-Api-Key")
68+
.or_else(|| headers.get("x-api-key"))
69+
.and_then(|v| v.to_str().ok())
70+
.map(|s| s.to_string())?;
71+
6572
Some(AuthHeaders {
6673
hotkey,
6774
nonce,
6875
signature,
76+
api_key,
6977
})
7078
}
7179

72-
pub fn verify_request(auth: &AuthHeaders, nonce_store: &NonceStore) -> Result<(), AuthError> {
80+
pub fn verify_request(
81+
auth: &AuthHeaders,
82+
nonce_store: &NonceStore,
83+
expected_api_key: &str,
84+
) -> Result<(), AuthError> {
7385
if auth.hotkey != AUTHORIZED_HOTKEY {
7486
return Err(AuthError::UnauthorizedHotkey);
7587
}
7688

89+
if auth.api_key != expected_api_key {
90+
return Err(AuthError::InvalidApiKey);
91+
}
92+
7793
if !validate_ss58(&auth.hotkey) {
7894
return Err(AuthError::InvalidHotkey);
7995
}
@@ -94,6 +110,7 @@ pub fn verify_request(auth: &AuthHeaders, nonce_store: &NonceStore) -> Result<()
94110
pub enum AuthError {
95111
UnauthorizedHotkey,
96112
InvalidHotkey,
113+
InvalidApiKey,
97114
NonceReused,
98115
InvalidSignature,
99116
}
@@ -103,6 +120,7 @@ impl AuthError {
103120
match self {
104121
AuthError::UnauthorizedHotkey => "Hotkey is not authorized",
105122
AuthError::InvalidHotkey => "Invalid SS58 hotkey format",
123+
AuthError::InvalidApiKey => "Invalid API key",
106124
AuthError::NonceReused => "Nonce has already been used",
107125
AuthError::InvalidSignature => "Signature verification failed",
108126
}
@@ -112,6 +130,7 @@ impl AuthError {
112130
match self {
113131
AuthError::UnauthorizedHotkey => "unauthorized_hotkey",
114132
AuthError::InvalidHotkey => "invalid_hotkey",
133+
AuthError::InvalidApiKey => "invalid_api_key",
115134
AuthError::NonceReused => "nonce_reused",
116135
AuthError::InvalidSignature => "invalid_signature",
117136
}
@@ -182,6 +201,18 @@ pub fn validate_ss58(address: &str) -> bool {
182201
ss58_to_public_key_bytes(address).is_some()
183202
}
184203

204+
#[cfg(test)]
205+
fn sp_ss58_checksum(data: &[u8]) -> [u8; 64] {
206+
use sha2::{Digest, Sha512};
207+
let mut hasher = Sha512::new();
208+
hasher.update(b"SS58PRE");
209+
hasher.update(data);
210+
let result = hasher.finalize();
211+
let mut out = [0u8; 64];
212+
out.copy_from_slice(&result);
213+
out
214+
}
215+
185216
#[cfg(test)]
186217
mod tests {
187218
use super::*;
@@ -219,12 +250,14 @@ mod tests {
219250
headers.insert("X-Hotkey", AUTHORIZED_HOTKEY.parse().unwrap());
220251
headers.insert("X-Nonce", "test-nonce-123".parse().unwrap());
221252
headers.insert("X-Signature", "0xdeadbeef".parse().unwrap());
253+
headers.insert("X-Api-Key", "my-secret-key".parse().unwrap());
222254
let auth = extract_auth_headers(&headers);
223255
assert!(auth.is_some());
224256
let auth = auth.unwrap();
225257
assert_eq!(auth.hotkey, AUTHORIZED_HOTKEY);
226258
assert_eq!(auth.nonce, "test-nonce-123");
227259
assert_eq!(auth.signature, "0xdeadbeef");
260+
assert_eq!(auth.api_key, "my-secret-key");
228261
}
229262

230263
#[test]
@@ -240,11 +273,34 @@ mod tests {
240273
hotkey: "5InvalidHotkey".to_string(),
241274
nonce: "nonce-1".to_string(),
242275
signature: "0x00".to_string(),
276+
api_key: "test-key".to_string(),
243277
};
244-
let err = verify_request(&auth, &store).unwrap_err();
278+
let err = verify_request(&auth, &store, "test-key").unwrap_err();
245279
assert!(matches!(err, AuthError::UnauthorizedHotkey));
246280
}
247281

282+
#[test]
283+
fn test_verify_request_invalid_api_key() {
284+
let store = NonceStore::new();
285+
let auth = AuthHeaders {
286+
hotkey: AUTHORIZED_HOTKEY.to_string(),
287+
nonce: "nonce-1".to_string(),
288+
signature: "0x00".to_string(),
289+
api_key: "wrong-key".to_string(),
290+
};
291+
let err = verify_request(&auth, &store, "correct-key").unwrap_err();
292+
assert!(matches!(err, AuthError::InvalidApiKey));
293+
}
294+
295+
#[test]
296+
fn test_extract_auth_headers_missing_api_key() {
297+
let mut headers = axum::http::HeaderMap::new();
298+
headers.insert("X-Hotkey", AUTHORIZED_HOTKEY.parse().unwrap());
299+
headers.insert("X-Nonce", "test-nonce-123".parse().unwrap());
300+
headers.insert("X-Signature", "0xdeadbeef".parse().unwrap());
301+
assert!(extract_auth_headers(&headers).is_none());
302+
}
303+
248304
#[test]
249305
fn test_verify_sr25519_roundtrip() {
250306
use schnorrkel::{Keypair, MiniSecretKey};
@@ -272,15 +328,3 @@ mod tests {
272328
assert!(!verify_sr25519_signature(&ss58, "wrong-message", &sig_hex));
273329
}
274330
}
275-
276-
#[cfg(test)]
277-
fn sp_ss58_checksum(data: &[u8]) -> [u8; 64] {
278-
use sha2::{Digest, Sha512};
279-
let mut hasher = Sha512::new();
280-
hasher.update(b"SS58PRE");
281-
hasher.update(data);
282-
let result = hasher.finalize();
283-
let mut out = [0u8; 64];
284-
out.copy_from_slice(&result);
285-
out
286-
}

src/config.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ pub struct Config {
2525
#[allow(dead_code)]
2626
pub max_output_bytes: usize,
2727
pub workspace_base: PathBuf,
28+
pub worker_api_key: String,
2829
}
2930

3031
impl Config {
@@ -41,6 +42,8 @@ impl Config {
4142
workspace_base: PathBuf::from(
4243
std::env::var("WORKSPACE_BASE").unwrap_or_else(|_| DEFAULT_WORKSPACE_BASE.into()),
4344
),
45+
worker_api_key: std::env::var("WORKER_API_KEY")
46+
.expect("WORKER_API_KEY environment variable must be set"),
4447
}
4548
}
4649

@@ -62,6 +65,7 @@ impl Config {
6265
"║ Workspace: {:<28}║",
6366
self.workspace_base.display()
6467
);
68+
tracing::info!("║ API key: {:<28}║", "configured");
6569
tracing::info!("╚══════════════════════════════════════════════════╝");
6670
}
6771
}
@@ -79,9 +83,11 @@ mod tests {
7983

8084
#[test]
8185
fn test_config_defaults() {
86+
std::env::set_var("WORKER_API_KEY", "test-api-key-123");
8287
let cfg = Config::from_env();
8388
assert_eq!(cfg.port, DEFAULT_PORT);
8489
assert_eq!(cfg.max_concurrent_tasks, DEFAULT_MAX_CONCURRENT);
90+
assert_eq!(cfg.worker_api_key, "test-api-key-123");
8591
}
8692

8793
#[test]

src/handlers.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,12 +103,16 @@ async fn submit_batch(
103103
StatusCode::UNAUTHORIZED,
104104
Json(serde_json::json!({
105105
"error": "missing_auth",
106-
"message": "Missing required headers: X-Hotkey, X-Nonce, X-Signature"
106+
"message": "Missing required headers: X-Hotkey, X-Nonce, X-Signature, X-Api-Key"
107107
})),
108108
)
109109
})?;
110110

111-
if let Err(e) = auth::verify_request(&auth_headers, &state.nonce_store) {
111+
if let Err(e) = auth::verify_request(
112+
&auth_headers,
113+
&state.nonce_store,
114+
&state.config.worker_api_key,
115+
) {
112116
return Err((
113117
StatusCode::UNAUTHORIZED,
114118
Json(serde_json::json!({

0 commit comments

Comments
 (0)