| 1 |
|
| 2 |
|
| 3 |
use super::*; |
| 4 |
use axum::body::Body; |
| 5 |
use axum::http::Request as HttpRequest; |
| 6 |
use tower::ServiceExt; |
| 7 |
|
| 8 |
fn test_config(api_token: Option<&str>) -> Config { |
| 9 |
let mut config = Config { |
| 10 |
serve: crate::config::ServeConfig::default(), |
| 11 |
instance: crate::config::InstanceConfig::default(), |
| 12 |
targets: HashMap::new(), |
| 13 |
peers: HashMap::new(), |
| 14 |
storage: crate::config::StorageConfig::default(), |
| 15 |
alerts: None, |
| 16 |
}; |
| 17 |
config.serve.api_token = api_token.map(std::string::ToString::to_string); |
| 18 |
config |
| 19 |
} |
| 20 |
|
| 21 |
#[tokio::test] |
| 22 |
async fn no_token_configured_allows_all_requests() { |
| 23 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 24 |
let app = router(pool, test_config(None), None); |
| 25 |
|
| 26 |
let resp = app |
| 27 |
.oneshot(with_connect_info("/api/status", None)) |
| 28 |
.await |
| 29 |
.unwrap(); |
| 30 |
assert_eq!(resp.status(), StatusCode::OK); |
| 31 |
} |
| 32 |
|
| 33 |
#[tokio::test] |
| 34 |
async fn valid_token_allows_request() { |
| 35 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 36 |
let app = router(pool, test_config(Some("secret123")), None); |
| 37 |
|
| 38 |
let resp = app |
| 39 |
.oneshot(with_connect_info("/api/status", Some("Bearer secret123"))) |
| 40 |
.await |
| 41 |
.unwrap(); |
| 42 |
assert_eq!(resp.status(), StatusCode::OK); |
| 43 |
} |
| 44 |
|
| 45 |
#[tokio::test] |
| 46 |
async fn dashboard_uses_cookie_not_embedded_token() { |
| 47 |
|
| 48 |
|
| 49 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 50 |
let mut config = test_config(Some("supersecret-token")); |
| 51 |
config.serve.dashboard = true; |
| 52 |
let app = router(pool, config, None); |
| 53 |
|
| 54 |
let req = HttpRequest::builder().uri("/").body(Body::empty()).unwrap(); |
| 55 |
let resp = app.clone().oneshot(req).await.unwrap(); |
| 56 |
assert_eq!(resp.status(), StatusCode::OK); |
| 57 |
|
| 58 |
let cookie_hdr = resp |
| 59 |
.headers() |
| 60 |
.get(axum::http::header::SET_COOKIE) |
| 61 |
.expect("dashboard must set a session cookie") |
| 62 |
.to_str() |
| 63 |
.unwrap() |
| 64 |
.to_string(); |
| 65 |
assert!(cookie_hdr.contains("pom_dash=")); |
| 66 |
assert!(cookie_hdr.contains("HttpOnly")); |
| 67 |
assert!( |
| 68 |
!cookie_hdr.contains("supersecret-token"), |
| 69 |
"cookie must not be the api_token" |
| 70 |
); |
| 71 |
|
| 72 |
let body = axum::body::to_bytes(resp.into_body(), usize::MAX) |
| 73 |
.await |
| 74 |
.unwrap(); |
| 75 |
let html = String::from_utf8_lossy(&body); |
| 76 |
assert!( |
| 77 |
!html.contains("supersecret-token"), |
| 78 |
"the api_token must never appear in served HTML" |
| 79 |
); |
| 80 |
|
| 81 |
|
| 82 |
let dash = cookie_hdr.split(';').next().unwrap().trim().to_string(); |
| 83 |
let mut api_req = HttpRequest::builder() |
| 84 |
.uri("/api/status") |
| 85 |
.header(axum::http::header::COOKIE, dash) |
| 86 |
.body(Body::empty()) |
| 87 |
.unwrap(); |
| 88 |
api_req |
| 89 |
.extensions_mut() |
| 90 |
.insert(axum::extract::ConnectInfo(std::net::SocketAddr::from(( |
| 91 |
[127, 0, 0, 1], |
| 92 |
40001, |
| 93 |
)))); |
| 94 |
let api_resp = app.oneshot(api_req).await.unwrap(); |
| 95 |
assert_eq!( |
| 96 |
api_resp.status(), |
| 97 |
StatusCode::OK, |
| 98 |
"dashboard cookie must authenticate /api/*" |
| 99 |
); |
| 100 |
} |
| 101 |
|
| 102 |
#[tokio::test] |
| 103 |
async fn wrong_token_returns_401() { |
| 104 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 105 |
let app = router(pool, test_config(Some("secret123")), None); |
| 106 |
|
| 107 |
let req = HttpRequest::builder() |
| 108 |
.uri("/api/status") |
| 109 |
.header("authorization", "Bearer wrong-token") |
| 110 |
.body(Body::empty()) |
| 111 |
.unwrap(); |
| 112 |
let resp = app.oneshot(req).await.unwrap(); |
| 113 |
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); |
| 114 |
} |
| 115 |
|
| 116 |
#[tokio::test] |
| 117 |
async fn missing_header_returns_401() { |
| 118 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 119 |
let app = router(pool, test_config(Some("secret123")), None); |
| 120 |
|
| 121 |
let req = HttpRequest::builder() |
| 122 |
.uri("/api/status") |
| 123 |
.body(Body::empty()) |
| 124 |
.unwrap(); |
| 125 |
let resp = app.oneshot(req).await.unwrap(); |
| 126 |
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); |
| 127 |
} |
| 128 |
|
| 129 |
#[tokio::test] |
| 130 |
async fn malformed_header_returns_401() { |
| 131 |
let pool = crate::db::connect_in_memory().await.unwrap(); |
| 132 |
let app = router(pool, test_config(Some("secret123")), None); |
| 133 |
|
| 134 |
let req = HttpRequest::builder() |
| 135 |
.uri("/api/status") |
| 136 |
.header("authorization", "Basic dXNlcjpwYXNz") |
| 137 |
.body(Body::empty()) |
| 138 |
.unwrap(); |
| 139 |
let resp = app.oneshot(req).await.unwrap(); |
| 140 |
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); |
| 141 |
} |
| 142 |
|
| 143 |
fn ip(n: u8) -> std::net::IpAddr { |
| 144 |
std::net::IpAddr::V4(std::net::Ipv4Addr::new(10, 0, 0, n)) |
| 145 |
} |
| 146 |
|
| 147 |
#[test] |
| 148 |
fn reproject_drops_injected_structure() { |
| 149 |
|
| 150 |
let hostile = serde_json::json!({ |
| 151 |
"instance": { "id": "abc", "version": "9.9", "evil_field": {"x": 1} }, |
| 152 |
"targets": { "mnw": { "status": "operational", "response_time_ms": 5, "evil": "inject" } }, |
| 153 |
"peers": { "p2": { "status": "up", "latency_ms": 3 } }, |
| 154 |
"top_level_injection": [1, 2, 3] |
| 155 |
}); |
| 156 |
let clean = reproject_peer_status(&hostile); |
| 157 |
|
| 158 |
|
| 159 |
let obj = clean.as_object().unwrap(); |
| 160 |
let mut keys: Vec<&String> = obj.keys().collect(); |
| 161 |
keys.sort(); |
| 162 |
assert_eq!(keys, vec!["instance", "peers", "targets"]); |
| 163 |
assert!(clean.get("top_level_injection").is_none()); |
| 164 |
|
| 165 |
|
| 166 |
assert_eq!(clean["instance"]["id"], "abc"); |
| 167 |
assert_eq!(clean["instance"]["version"], "9.9"); |
| 168 |
assert!(clean["instance"].get("evil_field").is_none()); |
| 169 |
assert_eq!(clean["targets"]["mnw"]["status"], "operational"); |
| 170 |
assert!(clean["targets"]["mnw"].get("evil").is_none()); |
| 171 |
assert_eq!(clean["peers"]["p2"]["latency_ms"], 3); |
| 172 |
} |
| 173 |
|
| 174 |
|
| 175 |
|
| 176 |
|
| 177 |
fn with_connect_info(uri: &str, bearer: Option<&str>) -> HttpRequest<Body> { |
| 178 |
let mut b = HttpRequest::builder().uri(uri); |
| 179 |
if let Some(h) = bearer { |
| 180 |
b = b.header("authorization", h); |
| 181 |
} |
| 182 |
let mut req = b.body(Body::empty()).unwrap(); |
| 183 |
req.extensions_mut() |
| 184 |
.insert(axum::extract::ConnectInfo(std::net::SocketAddr::from(( |
| 185 |
[127, 0, 0, 1], |
| 186 |
40000, |
| 187 |
)))); |
| 188 |
req |
| 189 |
} |
| 190 |
|
| 191 |
#[test] |
| 192 |
fn rate_limiter_allows_within_limit() { |
| 193 |
let limiter = PerIpRateLimiter::new(3, std::time::Duration::from_mins(1)); |
| 194 |
assert!(limiter.try_acquire(ip(1))); |
| 195 |
assert!(limiter.try_acquire(ip(1))); |
| 196 |
assert!(limiter.try_acquire(ip(1))); |
| 197 |
} |
| 198 |
|
| 199 |
#[test] |
| 200 |
fn rate_limiter_blocks_over_limit() { |
| 201 |
let limiter = PerIpRateLimiter::new(2, std::time::Duration::from_mins(1)); |
| 202 |
assert!(limiter.try_acquire(ip(1))); |
| 203 |
assert!(limiter.try_acquire(ip(1))); |
| 204 |
assert!(!limiter.try_acquire(ip(1))); |
| 205 |
} |
| 206 |
|
| 207 |
#[test] |
| 208 |
fn rate_limiter_isolates_clients_by_ip() { |
| 209 |
|
| 210 |
let limiter = PerIpRateLimiter::new(1, std::time::Duration::from_mins(1)); |
| 211 |
assert!(limiter.try_acquire(ip(1))); |
| 212 |
assert!( |
| 213 |
!limiter.try_acquire(ip(1)), |
| 214 |
"ip(1) is now over its own limit" |
| 215 |
); |
| 216 |
assert!(limiter.try_acquire(ip(2)), "ip(2) has its own fresh bucket"); |
| 217 |
} |
| 218 |
|
| 219 |
#[tokio::test] |
| 220 |
async fn rate_limiter_resets_after_window() { |
| 221 |
let limiter = PerIpRateLimiter::new(1, std::time::Duration::from_millis(10)); |
| 222 |
assert!(limiter.try_acquire(ip(1))); |
| 223 |
assert!(!limiter.try_acquire(ip(1))); |
| 224 |
tokio::time::sleep(std::time::Duration::from_millis(15)).await; |
| 225 |
assert!(limiter.try_acquire(ip(1))); |
| 226 |
} |
| 227 |
|