531 lines
15 KiB
Rust
531 lines
15 KiB
Rust
use axum::{
|
|
body::Body,
|
|
http::{Request, StatusCode},
|
|
routing::{get, post},
|
|
Router,
|
|
};
|
|
use http_body_util::BodyExt;
|
|
use serde_json::{json, Value};
|
|
use tower::ServiceExt; // for oneshot
|
|
|
|
mod common;
|
|
use common::test_state;
|
|
|
|
// ─── Restricted SVG Server Tests ────────────────────────────────────────
|
|
|
|
fn svg_app(state: server18004::state::AppState) -> Router {
|
|
Router::new()
|
|
.fallback(server18004::handlers::restricted::handle_svg)
|
|
.with_state(state)
|
|
}
|
|
|
|
fn png_app(state: server18004::state::AppState) -> Router {
|
|
Router::new()
|
|
.fallback(server18004::handlers::restricted::handle_png)
|
|
.with_state(state)
|
|
}
|
|
|
|
fn api_app(state: server18004::state::AppState) -> Router {
|
|
Router::new()
|
|
.route("/generate", post(server18004::handlers::api::generate_qr))
|
|
.route("/domains/add", post(server18004::handlers::api::add_domain))
|
|
.route("/domains/remove", post(server18004::handlers::api::remove_domain))
|
|
.route("/domains", get(server18004::handlers::api::list_domains))
|
|
.route("/health", get(server18004::handlers::api::health))
|
|
.with_state(state)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_svg_allowed_domain() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/some/path")
|
|
.header("host", "qr.example.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert_eq!(
|
|
response.headers().get("content-type").unwrap(),
|
|
"image/svg+xml"
|
|
);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let svg = String::from_utf8(body.to_vec()).unwrap();
|
|
assert!(svg.contains("<svg"), "Response should contain SVG content");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_svg_forbidden_domain() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.header("host", "qr.forbidden.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_svg_missing_host() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_svg_x_forwarded_host() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/path")
|
|
.header("x-forwarded-host", "qr.example.com")
|
|
.header("host", "internal:8080")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_svg_default_domain() {
|
|
// The default domain "test-default.com" should always be allowed
|
|
let state = test_state(&[]); // empty allowlist
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.header("host", "qr.test-default.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
// ─── Restricted PNG Server Tests ────────────────────────────────────────
|
|
|
|
#[tokio::test]
|
|
async fn test_png_allowed_domain() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = png_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.header("host", "qr.example.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert_eq!(
|
|
response.headers().get("content-type").unwrap(),
|
|
"image/png"
|
|
);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
assert!(body.len() > 8, "PNG should have content");
|
|
assert_eq!(&body[0..4], &[0x89, 0x50, 0x4E, 0x47], "Should have PNG magic bytes");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_png_forbidden_domain() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = png_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.header("host", "qr.evil.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
// ─── API Server Tests ───────────────────────────────────────────────────
|
|
|
|
#[tokio::test]
|
|
async fn test_health_endpoint() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/health")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_svg() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"text": "https://example.com",
|
|
"format": "svg"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert_eq!(
|
|
response.headers().get("content-type").unwrap(),
|
|
"image/svg+xml"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_png() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"text": "https://example.com",
|
|
"format": "png",
|
|
"ecl": "H"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert_eq!(
|
|
response.headers().get("content-type").unwrap(),
|
|
"image/png"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_base64() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"text": "https://example.com",
|
|
"format": "base64",
|
|
"base64_source": "png"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
|
assert_eq!(json["format"], "base64");
|
|
assert_eq!(json["source"], "png");
|
|
assert!(!json["data"].as_str().unwrap().is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_base64url_svg_source() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"text": "hello world",
|
|
"format": "base64url",
|
|
"base64_source": "svg"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
|
assert_eq!(json["format"], "base64url");
|
|
assert_eq!(json["source"], "svg");
|
|
// base64url should not contain + or / or =
|
|
let data = json["data"].as_str().unwrap();
|
|
assert!(!data.contains('+'), "base64url should not contain +");
|
|
assert!(!data.contains('/'), "base64url should not contain /");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_generate_invalid_format() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"text": "test",
|
|
"format": "bmp"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ─── Domain Management Tests ────────────────────────────────────────────
|
|
|
|
#[tokio::test]
|
|
async fn test_list_domains_empty() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/domains")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
|
assert!(json["success"].as_bool().unwrap());
|
|
assert!(json["domains"].as_array().unwrap().is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_add_domain() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/domains/add")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"domain": "newdomain.com"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
|
assert!(json["success"].as_bool().unwrap());
|
|
assert!(json["domains"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.any(|d| d == "newdomain.com"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_add_empty_domain() {
|
|
let state = test_state(&[]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/domains/add")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"domain": " "
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_remove_domain() {
|
|
let state = test_state(&["example.com", "test.com"]);
|
|
let app = api_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/domains/remove")
|
|
.header("content-type", "application/json")
|
|
.body(Body::from(
|
|
serde_json::to_string(&json!({
|
|
"domain": "example.com"
|
|
}))
|
|
.unwrap(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let json: Value = serde_json::from_slice(&body).unwrap();
|
|
assert!(json["success"].as_bool().unwrap());
|
|
assert!(!json["domains"]
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.any(|d| d == "example.com"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_cache_control_headers() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/test")
|
|
.header("host", "qr.example.com")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
assert_eq!(
|
|
response.headers().get("cache-control").unwrap(),
|
|
"public, max-age=86400"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_x_forwarded_proto() {
|
|
let state = test_state(&["example.com"]);
|
|
let app = svg_app(state);
|
|
|
|
let response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.uri("/mypage?q=1")
|
|
.header("host", "qr.example.com")
|
|
.header("x-forwarded-proto", "http")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// The QR code should encode http://example.com/mypage?q=1
|
|
let body = response.into_body().collect().await.unwrap().to_bytes();
|
|
let svg = String::from_utf8(body.to_vec()).unwrap();
|
|
assert!(svg.contains("<svg"), "Response should be SVG");
|
|
}
|