use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use serde::{Deserialize, Serialize};
use crate::config::{KibbleConfig, load_config};
use crate::net::{build_client, resolve_proxy, url_is_allowed};
#[derive(Serialize, Clone)]
#[serde(rename_all = "lowercase")]
pub enum JobStatus { Queued, Running, Done, Error }
#[derive(Serialize, Clone)]
pub struct Job {
pub id: u64,
pub url: String,
pub status: JobStatus,
pub handler: Option<String>,
pub files: usize,
pub bytes: u64,
pub error: Option<String>,
}
#[derive(Clone)]
pub struct AppState {
pub cfg: Arc<KibbleConfig>,
pub client: reqwest::Client,
pub repo_root: PathBuf,
pub jobs: Arc<Mutex<HashMap<u64, Job>>>,
pub counter: Arc<AtomicU64>,
pub token: String,
pub allow_hosts: Vec<String>,
}
#[derive(Deserialize)]
struct IngestReq { url: String, name: Option<String> }
fn check_auth(headers: &HeaderMap, token: &str) -> Result<(), StatusCode> {
let got = headers.get("authorization").and_then(|v| v.to_str().ok()).unwrap_or("");
if got.strip_prefix("Bearer ") == Some(token) { Ok(()) } else { Err(StatusCode::UNAUTHORIZED) }
}
async fn health() -> &'static str { "ok" }
async fn ingest(State(state): State<AppState>, headers: HeaderMap, Json(req): Json<IngestReq>) -> Response {
if let Err(code) = check_auth(&headers, &state.token) {
return code.into_response();
}
if let Err(msg) = url_is_allowed(&req.url, &state.allow_hosts) {
return (StatusCode::BAD_REQUEST, msg).into_response();
}
let id = state.counter.fetch_add(1, Ordering::SeqCst) + 1;
state.jobs.lock().unwrap().insert(id, Job {
id, url: req.url.clone(), status: JobStatus::Queued,
handler: None, files: 0, bytes: 0, error: None,
});
let st = state.clone();
let url = req.url.clone();
let name = req.name.clone();
tokio::spawn(async move {
if let Some(j) = st.jobs.lock().unwrap().get_mut(&id) { j.status = JobStatus::Running; }
let res = crate::fetch::stage_url(&st.repo_root, &st.cfg, &st.client, &url, name.as_deref()).await;
let mut jobs = st.jobs.lock().unwrap();
if let Some(j) = jobs.get_mut(&id) {
match res {
Ok(r) => { j.status = JobStatus::Done; j.handler = Some(r.handler); j.files = r.files; j.bytes = r.bytes; }
Err(e) => { j.status = JobStatus::Error; j.error = Some(e.to_string()); }
}
}
});
(StatusCode::ACCEPTED, Json(serde_json::json!({ "job_id": id }))).into_response()
}
async fn job_status(State(state): State<AppState>, headers: HeaderMap, Path(id): Path<u64>) -> Response {
if let Err(code) = check_auth(&headers, &state.token) {
return code.into_response();
}
match state.jobs.lock().unwrap().get(&id) {
Some(job) => Json(job.clone()).into_response(),
None => StatusCode::NOT_FOUND.into_response(),
}
}
pub fn build_router(state: AppState) -> Router {
Router::new()
.route("/health", get(health))
.route("/ingest", post(ingest))
.route("/jobs/{id}", get(job_status))
.with_state(state)
}
pub fn state_from_config(repo_root: &std::path::Path, token: String) -> std::io::Result<AppState> {
let cfg = load_config(&repo_root.join(crate::config::CONFIG_FILE));
let proxy = resolve_proxy(cfg.network.proxy.as_deref(), |k| std::env::var(k).ok());
let client = build_client(proxy.as_deref())
.map_err(|e| std::io::Error::other(format!("http client: {e}")))?;
let allow_hosts = cfg.serve.allow_hosts.clone();
Ok(AppState {
cfg: Arc::new(cfg),
client,
repo_root: repo_root.to_path_buf(),
jobs: Arc::new(Mutex::new(HashMap::new())),
counter: Arc::new(AtomicU64::new(0)),
token,
allow_hosts,
})
}
pub async fn run_serve(repo_root: &std::path::Path) -> std::io::Result<()> {
let token = std::env::var("KIBBLE_API_TOKEN").unwrap_or_default();
if token.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"KIBBLE_API_TOKEN must be set (refusing to start an unauthenticated daemon)",
));
}
let cfg = load_config(&repo_root.join(crate::config::CONFIG_FILE));
let bind = cfg.serve.bind.clone();
let state = state_from_config(repo_root, token)?;
let router = build_router(state);
let listener = tokio::net::TcpListener::bind(&bind).await?;
println!("🦴 kibble serve — ready on {bind}");
axum::serve(listener, router).await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::KibbleConfig;
use std::sync::{Arc, Mutex};
use std::sync::atomic::AtomicU64;
fn test_state(repo_root: std::path::PathBuf, allow: Vec<String>) -> AppState {
AppState {
cfg: Arc::new(KibbleConfig::default()),
client: crate::net::build_client(None).unwrap(),
repo_root,
jobs: Arc::new(Mutex::new(std::collections::HashMap::new())),
counter: Arc::new(AtomicU64::new(0)),
token: "secret".to_string(),
allow_hosts: allow,
}
}
async fn body_string(resp: axum::response::Response) -> String {
let b = axum::body::to_bytes(resp.into_body(), usize::MAX).await.unwrap();
String::from_utf8(b.to_vec()).unwrap()
}
#[tokio::test]
async fn health_is_open() {
use tower::ServiceExt;
let router = build_router(test_state(std::env::temp_dir(), vec![]));
let req = axum::http::Request::builder().uri("/health").body(axum::body::Body::empty()).unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), 200);
}
#[tokio::test]
async fn ingest_requires_token() {
use tower::ServiceExt;
let router = build_router(test_state(std::env::temp_dir(), vec![]));
let req = axum::http::Request::builder()
.method("POST").uri("/ingest")
.header("content-type", "application/json")
.body(axum::body::Body::from(r#"{"url":"http://example.com/"}"#)).unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), 401);
}
#[tokio::test]
async fn ingest_rejects_ssrf_url() {
use tower::ServiceExt;
let router = build_router(test_state(std::env::temp_dir(), vec![]));
let req = axum::http::Request::builder()
.method("POST").uri("/ingest")
.header("authorization", "Bearer secret")
.header("content-type", "application/json")
.body(axum::body::Body::from(r#"{"url":"http://127.0.0.1/x"}"#)).unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), 400);
}
#[tokio::test]
async fn unknown_job_is_404() {
use tower::ServiceExt;
let router = build_router(test_state(std::env::temp_dir(), vec![]));
let req = axum::http::Request::builder().uri("/jobs/999")
.header("authorization", "Bearer secret")
.body(axum::body::Body::empty()).unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), 404);
}
#[tokio::test]
async fn ingest_then_job_completes_with_allow_hosts() {
use tower::ServiceExt;
use std::io::{Read, Write};
use std::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
std::thread::spawn(move || {
if let Ok((mut s, _)) = listener.accept() {
let mut b = [0u8; 1024];
let _ = s.read(&mut b);
let body = b"served bytes";
let head = format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len());
let _ = s.write_all(head.as_bytes());
let _ = s.write_all(body);
}
});
let root = std::env::temp_dir().join(format!("kibble_serve_{}", std::process::id()));
let _ = std::fs::remove_dir_all(&root);
std::fs::create_dir_all(&root).unwrap();
let state = test_state(root.clone(), vec!["127.0.0.1".to_string()]);
let body = format!("{{\"url\":\"http://127.0.0.1:{port}/f.txt\",\"name\":\"j\"}}");
let req = axum::http::Request::builder().method("POST").uri("/ingest")
.header("authorization", "Bearer secret").header("content-type", "application/json")
.body(axum::body::Body::from(body)).unwrap();
let resp = build_router(state.clone()).oneshot(req).await.unwrap();
assert_eq!(resp.status(), 202);
let txt = body_string(resp).await;
let v: serde_json::Value = serde_json::from_str(&txt).unwrap();
let id = v.get("job_id").unwrap().as_u64().unwrap();
let mut done = false;
for _ in 0..50 {
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
let req = axum::http::Request::builder().uri(format!("/jobs/{id}"))
.header("authorization", "Bearer secret").body(axum::body::Body::empty()).unwrap();
let resp = build_router(state.clone()).oneshot(req).await.unwrap();
let j: serde_json::Value = serde_json::from_str(&body_string(resp).await).unwrap();
if j.get("status").and_then(|s| s.as_str()) == Some("done") { done = true; break; }
}
assert!(done, "job did not reach done");
assert_eq!(std::fs::read_to_string(root.join("data/ingest/j/f.txt")).unwrap(), "served bytes");
}
#[test]
fn state_from_config_requires_token_at_callsite() {
let st = state_from_config(&std::env::temp_dir(), "tok".to_string()).unwrap();
assert_eq!(st.token, "tok");
}
}