use crate::crypto::{Keypair, now_ts, random_nonce, signing_bytes, verify_envelope}; use anyhow::{Context, Result, anyhow}; use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use serde_json::Value; pub const TYPE_QUERY: &str = "query"; pub const TYPE_RESPONSE: &str = "response"; pub const TYPE_AGGREGATE: &str = "aggregate"; pub const EXPOSURE_METADATA: &str = "metadata"; pub const EXPOSURE_FULL: &str = "full"; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Envelope { #[serde(rename = "type")] pub msg_type: String, pub from: String, pub ts: u64, pub nonce: String, pub body: Value, pub sig: String, } impl Envelope { pub fn new(key: &Keypair, msg_type: &str, body: Value) -> Self { let mut envelope = Self { msg_type: msg_type.to_string(), from: key.public_hex(), ts: now_ts(), nonce: random_nonce(), body, sig: String::new(), }; envelope.sig = key.sign(&signing_bytes(&envelope)); envelope } pub fn verify(&self) -> Result<()> { verify_envelope(self) } pub fn parse_body(&self) -> Result { serde_json::from_value(self.body.clone()).context("malformed body") } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Budget { pub max_results: usize, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueryBody { pub qid: String, pub text: String, #[serde(default)] pub entities: Vec, pub budget: Budget, } impl QueryBody { pub fn new(text: &str, max_results: usize) -> Self { Self { qid: crate::crypto::random_id(), text: text.to_string(), entities: Vec::new(), budget: Budget { max_results: max_results.clamp(1, 1000), }, } } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponseItem { pub url: String, pub title: String, pub summary: String, pub published: String, pub exposure: String, pub content: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponseBody { pub qid: String, pub results: Vec, pub truncated: bool, pub more_available: u64, pub cursor: Option, } pub fn build_response( qid: &str, hits: Vec, total: u64, max_results: usize, ) -> ResponseBody { let max_results = max_results.max(1); let mut results = hits; results.truncate(max_results); let more_available = total.saturating_sub(results.len() as u64); ResponseBody { qid: qid.to_string(), results, truncated: more_available > 0, more_available, cursor: None, } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AggregateBody { pub period: String, pub sent: u64, pub passed: u64, } pub fn require_type(envelope: &Envelope, expected: &str) -> Result<()> { if envelope.msg_type != expected { return Err(anyhow!( "expected message type {expected}, got {}", envelope.msg_type )); } Ok(()) } #[cfg(test)] mod tests { use super::*; use serde_json::json; fn item(n: usize) -> ResponseItem { ResponseItem { url: format!("file:///doc{n}.txt"), title: format!("doc {n}"), summary: String::new(), published: String::new(), exposure: EXPOSURE_METADATA.to_string(), content: None, } } #[test] fn max_results_is_respected() { let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2); assert_eq!(response.results.len(), 2); } #[test] fn truncation_is_honest() { let response = build_response("q1", vec![item(1), item(2), item(3)], 10, 2); assert!(response.truncated); assert_eq!(response.more_available, 8); let response = build_response("q2", vec![item(1), item(2)], 2, 5); assert!(!response.truncated); assert_eq!(response.more_available, 0); } #[test] fn response_carries_no_scores() { let response = build_response("q1", vec![item(1)], 1, 5); let json = serde_json::to_string(&response).unwrap(); assert!(!json.contains("score")); assert!(!json.contains("relevance")); } #[test] fn zero_budget_is_clamped() { let query = QueryBody::new("anything", 0); assert_eq!(query.budget.max_results, 1); } #[test] fn envelope_rejects_wrong_type() { let key = Keypair::generate(); let envelope = Envelope::new(&key, TYPE_RESPONSE, json!({})); assert!(require_type(&envelope, TYPE_QUERY).is_err()); } }