483 lines
16 KiB
Rust
483 lines
16 KiB
Rust
//! Agent binary download endpoints
|
|
//!
|
|
//! Provides endpoints for downloading the attune-agent binary for injection
|
|
//! into arbitrary containers. This supports deployments where shared Docker
|
|
//! volumes are impractical (Kubernetes, ECS, remote Docker hosts).
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::{Query, State},
|
|
http::{header, HeaderMap, StatusCode},
|
|
response::IntoResponse,
|
|
routing::get,
|
|
Json, Router,
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::sync::Arc;
|
|
use subtle::ConstantTimeEq;
|
|
use tokio::fs;
|
|
use tokio_util::io::ReaderStream;
|
|
use utoipa::{IntoParams, ToSchema};
|
|
|
|
use crate::state::AppState;
|
|
|
|
/// Query parameters for the binary download endpoint
|
|
#[derive(Debug, Deserialize, IntoParams)]
|
|
pub struct BinaryDownloadParams {
|
|
/// Target architecture (x86_64, aarch64). Defaults to x86_64.
|
|
#[param(example = "x86_64")]
|
|
pub arch: Option<String>,
|
|
/// Optional bootstrap token for authentication
|
|
pub token: Option<String>,
|
|
}
|
|
|
|
/// Agent binary metadata
|
|
#[derive(Debug, Serialize, ToSchema)]
|
|
pub struct AgentBinaryInfo {
|
|
/// Available architectures
|
|
pub architectures: Vec<AgentArchInfo>,
|
|
/// Agent version (from build)
|
|
pub version: String,
|
|
}
|
|
|
|
/// Per-architecture binary info
|
|
#[derive(Debug, Serialize, ToSchema)]
|
|
pub struct AgentArchInfo {
|
|
/// Architecture name
|
|
pub arch: String,
|
|
/// Binary size in bytes
|
|
pub size_bytes: u64,
|
|
/// Whether this binary is available
|
|
pub available: bool,
|
|
}
|
|
|
|
/// Validate that the architecture name is safe (no path traversal) and normalize it.
|
|
fn validate_arch(arch: &str) -> Result<&str, (StatusCode, Json<serde_json::Value>)> {
|
|
match arch {
|
|
"x86_64" | "aarch64" => Ok(arch),
|
|
// Accept arm64 as an alias for aarch64
|
|
"arm64" => Ok("aarch64"),
|
|
_ => Err((
|
|
StatusCode::BAD_REQUEST,
|
|
Json(serde_json::json!({
|
|
"error": "Invalid architecture",
|
|
"message": format!("Unsupported architecture '{}'. Supported: x86_64, aarch64", arch),
|
|
})),
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Validate bootstrap token if configured.
|
|
///
|
|
/// If the agent config has a `bootstrap_token` set, the request must provide it
|
|
/// via the `X-Agent-Token` header or the `token` query parameter. If no token
|
|
/// is configured, access is unrestricted.
|
|
fn validate_token(
|
|
config: &attune_common::config::Config,
|
|
headers: &HeaderMap,
|
|
query_token: &Option<String>,
|
|
) -> Result<(), (StatusCode, Json<serde_json::Value>)> {
|
|
let expected_token = config
|
|
.agent
|
|
.as_ref()
|
|
.and_then(|ac| ac.bootstrap_token.as_ref());
|
|
|
|
let expected_token = match expected_token {
|
|
Some(t) => t,
|
|
None => {
|
|
use std::sync::Once;
|
|
static WARN_ONCE: Once = Once::new();
|
|
WARN_ONCE.call_once(|| {
|
|
tracing::warn!(
|
|
"Agent binary download endpoint has no bootstrap_token configured. \
|
|
Anyone with network access to the API can download the agent binary. \
|
|
Set agent.bootstrap_token in config to restrict access."
|
|
);
|
|
});
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
// Check X-Agent-Token header first, then query param
|
|
let provided_token = headers
|
|
.get("x-agent-token")
|
|
.and_then(|v| v.to_str().ok())
|
|
.map(|s| s.to_string())
|
|
.or_else(|| query_token.clone());
|
|
|
|
match provided_token {
|
|
Some(ref t) if bool::from(t.as_bytes().ct_eq(expected_token.as_bytes())) => Ok(()),
|
|
Some(_) => Err((
|
|
StatusCode::UNAUTHORIZED,
|
|
Json(serde_json::json!({
|
|
"error": "Invalid token",
|
|
"message": "The provided bootstrap token is invalid",
|
|
})),
|
|
)),
|
|
None => Err((
|
|
StatusCode::UNAUTHORIZED,
|
|
Json(serde_json::json!({
|
|
"error": "Token required",
|
|
"message": "A bootstrap token is required. Provide via X-Agent-Token header or token query parameter.",
|
|
})),
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Download the agent binary
|
|
///
|
|
/// Returns the statically-linked attune-agent binary for the requested architecture.
|
|
/// The binary can be injected into any container to turn it into an Attune worker.
|
|
#[utoipa::path(
|
|
get,
|
|
path = "/api/v1/agent/binary",
|
|
params(BinaryDownloadParams),
|
|
responses(
|
|
(status = 200, description = "Agent binary", content_type = "application/octet-stream"),
|
|
(status = 400, description = "Invalid architecture"),
|
|
(status = 401, description = "Invalid or missing bootstrap token"),
|
|
(status = 404, description = "Agent binary not found"),
|
|
(status = 503, description = "Agent binary distribution not configured"),
|
|
),
|
|
tag = "agent"
|
|
)]
|
|
pub async fn download_agent_binary(
|
|
State(state): State<Arc<AppState>>,
|
|
headers: HeaderMap,
|
|
Query(params): Query<BinaryDownloadParams>,
|
|
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
|
// Validate bootstrap token if configured
|
|
validate_token(&state.config, &headers, ¶ms.token)?;
|
|
|
|
let agent_config = state.config.agent.as_ref().ok_or_else(|| {
|
|
(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
Json(serde_json::json!({
|
|
"error": "Not configured",
|
|
"message": "Agent binary distribution is not configured. Set agent.binary_dir in config.",
|
|
})),
|
|
)
|
|
})?;
|
|
|
|
let arch = params.arch.as_deref().unwrap_or("x86_64");
|
|
let arch = validate_arch(arch)?;
|
|
|
|
let binary_dir = std::path::Path::new(&agent_config.binary_dir);
|
|
|
|
// Try arch-specific binary first, then fall back to generic name.
|
|
// IMPORTANT: The generic `attune-agent` binary is only safe to serve for
|
|
// x86_64 requests, because the current build pipeline produces an
|
|
// x86_64-unknown-linux-musl binary. Serving it for aarch64/arm64 would
|
|
// give the caller an incompatible executable (exec format error).
|
|
let arch_specific = binary_dir.join(format!("attune-agent-{}", arch));
|
|
let generic = binary_dir.join("attune-agent");
|
|
|
|
let binary_path = if arch_specific.exists() {
|
|
arch_specific
|
|
} else if arch == "x86_64" && generic.exists() {
|
|
tracing::debug!(
|
|
"Arch-specific binary not found at {:?}, falling back to generic {:?} (safe for x86_64)",
|
|
arch_specific,
|
|
generic
|
|
);
|
|
generic
|
|
} else {
|
|
tracing::warn!(
|
|
"Agent binary not found. Checked: {:?} and {:?}",
|
|
arch_specific,
|
|
generic
|
|
);
|
|
return Err((
|
|
StatusCode::NOT_FOUND,
|
|
Json(serde_json::json!({
|
|
"error": "Not found",
|
|
"message": format!(
|
|
"Agent binary not found for architecture '{}'. Ensure the agent binary is built and placed in '{}'.",
|
|
arch,
|
|
agent_config.binary_dir
|
|
),
|
|
})),
|
|
));
|
|
};
|
|
|
|
// Get file metadata for Content-Length
|
|
let metadata = fs::metadata(&binary_path).await.map_err(|e| {
|
|
tracing::error!(
|
|
"Failed to read agent binary metadata at {:?}: {}",
|
|
binary_path,
|
|
e
|
|
);
|
|
(
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
Json(serde_json::json!({
|
|
"error": "Internal error",
|
|
"message": "Failed to read agent binary",
|
|
})),
|
|
)
|
|
})?;
|
|
|
|
// Open file for streaming
|
|
let file = fs::File::open(&binary_path).await.map_err(|e| {
|
|
tracing::error!("Failed to open agent binary at {:?}: {}", binary_path, e);
|
|
(
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
Json(serde_json::json!({
|
|
"error": "Internal error",
|
|
"message": "Failed to open agent binary",
|
|
})),
|
|
)
|
|
})?;
|
|
|
|
let stream = ReaderStream::new(file);
|
|
let body = Body::from_stream(stream);
|
|
|
|
let headers_response = [
|
|
(header::CONTENT_TYPE, "application/octet-stream".to_string()),
|
|
(
|
|
header::CONTENT_DISPOSITION,
|
|
"attachment; filename=\"attune-agent\"".to_string(),
|
|
),
|
|
(header::CONTENT_LENGTH, metadata.len().to_string()),
|
|
(header::CACHE_CONTROL, "public, max-age=3600".to_string()),
|
|
];
|
|
|
|
tracing::info!(
|
|
arch = arch,
|
|
size_bytes = metadata.len(),
|
|
path = ?binary_path,
|
|
"Serving agent binary download"
|
|
);
|
|
|
|
Ok((headers_response, body))
|
|
}
|
|
|
|
/// Get agent binary metadata
|
|
///
|
|
/// Returns information about available agent binaries, including
|
|
/// supported architectures and binary sizes.
|
|
#[utoipa::path(
|
|
get,
|
|
path = "/api/v1/agent/info",
|
|
responses(
|
|
(status = 200, description = "Agent binary info", body = AgentBinaryInfo),
|
|
(status = 503, description = "Agent binary distribution not configured"),
|
|
),
|
|
tag = "agent"
|
|
)]
|
|
pub async fn agent_info(
|
|
State(state): State<Arc<AppState>>,
|
|
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
|
|
let agent_config = state.config.agent.as_ref().ok_or_else(|| {
|
|
(
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
Json(serde_json::json!({
|
|
"error": "Not configured",
|
|
"message": "Agent binary distribution is not configured.",
|
|
})),
|
|
)
|
|
})?;
|
|
|
|
let binary_dir = std::path::Path::new(&agent_config.binary_dir);
|
|
let architectures = ["x86_64", "aarch64"];
|
|
|
|
let mut arch_infos = Vec::new();
|
|
for arch in &architectures {
|
|
let arch_specific = binary_dir.join(format!("attune-agent-{}", arch));
|
|
let generic = binary_dir.join("attune-agent");
|
|
|
|
// Only fall back to the generic binary for x86_64, since the build
|
|
// pipeline currently produces x86_64-only generic binaries.
|
|
let (available, size_bytes) = if arch_specific.exists() {
|
|
match fs::metadata(&arch_specific).await {
|
|
Ok(m) => (true, m.len()),
|
|
Err(_) => (false, 0),
|
|
}
|
|
} else if *arch == "x86_64" && generic.exists() {
|
|
match fs::metadata(&generic).await {
|
|
Ok(m) => (true, m.len()),
|
|
Err(_) => (false, 0),
|
|
}
|
|
} else {
|
|
(false, 0)
|
|
};
|
|
|
|
arch_infos.push(AgentArchInfo {
|
|
arch: arch.to_string(),
|
|
size_bytes,
|
|
available,
|
|
});
|
|
}
|
|
|
|
Ok(Json(AgentBinaryInfo {
|
|
architectures: arch_infos,
|
|
version: env!("CARGO_PKG_VERSION").to_string(),
|
|
}))
|
|
}
|
|
|
|
/// Create agent routes
|
|
pub fn routes() -> Router<Arc<AppState>> {
|
|
Router::new()
|
|
.route("/agent/binary", get(download_agent_binary))
|
|
.route("/agent/info", get(agent_info))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use attune_common::config::AgentConfig;
|
|
use axum::http::{HeaderMap, HeaderValue};
|
|
|
|
// ── validate_arch tests ─────────────────────────────────────────
|
|
|
|
#[test]
|
|
fn test_validate_arch_valid_x86_64() {
|
|
let result = validate_arch("x86_64");
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), "x86_64");
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_arch_valid_aarch64() {
|
|
let result = validate_arch("aarch64");
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), "aarch64");
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_arch_arm64_alias() {
|
|
// "arm64" is an alias for "aarch64"
|
|
let result = validate_arch("arm64");
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), "aarch64");
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_arch_invalid() {
|
|
let result = validate_arch("mips");
|
|
assert!(result.is_err());
|
|
let (status, body) = result.unwrap_err();
|
|
assert_eq!(status, StatusCode::BAD_REQUEST);
|
|
assert_eq!(body.0["error"], "Invalid architecture");
|
|
}
|
|
|
|
// ── validate_token tests ────────────────────────────────────────
|
|
|
|
/// Helper: build a minimal Config with the given agent config.
|
|
/// Only the `agent` field is relevant for `validate_token`.
|
|
fn test_config(agent: Option<AgentConfig>) -> attune_common::config::Config {
|
|
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| ".".to_string());
|
|
let config_path = format!("{}/../../config.test.yaml", manifest_dir);
|
|
let mut config = attune_common::config::Config::load_from_file(&config_path)
|
|
.expect("Failed to load test config");
|
|
config.agent = agent;
|
|
config
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_no_config() {
|
|
// When no agent config is set at all, no token is required.
|
|
let config = test_config(None);
|
|
let headers = HeaderMap::new();
|
|
let query_token = None;
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_no_bootstrap_token_configured() {
|
|
// Agent config exists but bootstrap_token is None → no token required.
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: None,
|
|
}));
|
|
let headers = HeaderMap::new();
|
|
let query_token = None;
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_valid_from_header() {
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: Some("s3cret-bootstrap".to_string()),
|
|
}));
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert(
|
|
"x-agent-token",
|
|
HeaderValue::from_static("s3cret-bootstrap"),
|
|
);
|
|
let query_token = None;
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_valid_from_query() {
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: Some("s3cret-bootstrap".to_string()),
|
|
}));
|
|
let headers = HeaderMap::new();
|
|
let query_token = Some("s3cret-bootstrap".to_string());
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_invalid() {
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: Some("correct-token".to_string()),
|
|
}));
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert("x-agent-token", HeaderValue::from_static("wrong-token"));
|
|
let query_token = None;
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_err());
|
|
let (status, body) = result.unwrap_err();
|
|
assert_eq!(status, StatusCode::UNAUTHORIZED);
|
|
assert_eq!(body.0["error"], "Invalid token");
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_missing_when_required() {
|
|
// bootstrap_token is configured but caller provides nothing.
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: Some("required-token".to_string()),
|
|
}));
|
|
let headers = HeaderMap::new();
|
|
let query_token = None;
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_err());
|
|
let (status, body) = result.unwrap_err();
|
|
assert_eq!(status, StatusCode::UNAUTHORIZED);
|
|
assert_eq!(body.0["error"], "Token required");
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_token_header_takes_precedence_over_query() {
|
|
// When both header and query provide a token, the header value is
|
|
// checked first (it appears first in the or_else chain). Provide a
|
|
// valid token in the header and an invalid one in the query — should
|
|
// succeed because the header matches.
|
|
let config = test_config(Some(AgentConfig {
|
|
binary_dir: "/tmp/test".to_string(),
|
|
bootstrap_token: Some("the-real-token".to_string()),
|
|
}));
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert("x-agent-token", HeaderValue::from_static("the-real-token"));
|
|
let query_token = Some("wrong-token".to_string());
|
|
|
|
let result = validate_token(&config, &headers, &query_token);
|
|
assert!(result.is_ok());
|
|
}
|
|
}
|