iris-chat 0.1.47

use async_trait::async_trait;
use hashtree_core::{sha256, to_hex, Hash, Store, StoreError};
use reqwest::StatusCode;
use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync::RwLock;

/// Read-only Blossom store whose owned response buffer cannot exceed its
/// per-blob cap. Every received body chunk also shares one traversal budget,
/// including chunks from failed servers and snapshot fallback attempts.
pub(super) struct BoundedBlossomStore {
    client: reqwest::Client,
    servers: Vec<String>,
    max_blob_bytes: usize,
    budget: ReadBudget,
    cache: Arc<RwLock<BlobCache>>,
}

const MAX_CACHED_BYTES: usize = 32 * 1024 * 1024;
static SHARED_CACHE: OnceLock<Arc<RwLock<BlobCache>>> = OnceLock::new();

#[derive(Default)]
struct BlobCache {
    entries: HashMap<Hash, Vec<u8>>,
    order: VecDeque<Hash>,
    bytes: usize,
}

impl BlobCache {
    fn insert(&mut self, hash: Hash, data: Vec<u8>) {
        if self.entries.contains_key(&hash) || data.len() > MAX_CACHED_BYTES {
            return;
        }
        while self.bytes + data.len() > MAX_CACHED_BYTES {
            let Some(oldest) = self.order.pop_front() else {
                break;
            };
            if let Some(removed) = self.entries.remove(&oldest) {
                self.bytes -= removed.len();
            }
        }
        self.bytes += data.len();
        self.order.push_back(hash);
        self.entries.insert(hash, data);
    }
}

struct ReadBudget {
    reads_left: AtomicUsize,
    bytes_left: AtomicUsize,
}

impl ReadBudget {
    fn new(reads: usize, bytes: usize) -> Self {
        Self {
            reads_left: AtomicUsize::new(reads),
            bytes_left: AtomicUsize::new(bytes),
        }
    }

    fn consume(counter: &AtomicUsize, amount: usize) -> bool {
        #[allow(
            deprecated,
            reason = "The Linux release toolchain uses Rust 1.93, before Atomic::try_update stabilized."
        )]
        let previous = counter
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |left| {
                Some(left.saturating_sub(amount))
            })
            .unwrap_or_else(|left| left);
        previous >= amount
    }

    fn consume_read(&self) -> bool {
        Self::consume(&self.reads_left, 1)
    }

    fn consume_bytes(&self, bytes: usize) -> bool {
        Self::consume(&self.bytes_left, bytes)
    }
}

#[derive(Debug, PartialEq, Eq)]
enum BlobFetchError {
    Request(String),
    Status(u16),
    TooLarge { buffered_bytes: usize, limit: usize },
    HashMismatch,
    BudgetExceeded,
}

impl std::fmt::Display for BlobFetchError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Request(error) => formatter.write_str(error),
            Self::Status(status) => write!(formatter, "HTTP {status}"),
            Self::TooLarge {
                buffered_bytes,
                limit,
            } => write!(
                formatter,
                "blob exceeds {limit}-byte limit after buffering {buffered_bytes} bytes"
            ),
            Self::HashMismatch => formatter.write_str("downloaded blob hash does not match CID"),
            Self::BudgetExceeded => formatter.write_str("profile search byte budget exceeded"),
        }
    }
}

impl BoundedBlossomStore {
    pub(super) fn new(
        servers: Vec<String>,
        timeout: Duration,
        max_blob_bytes: usize,
        max_reads: usize,
        max_total_bytes: usize,
    ) -> Result<Self, String> {
        let client = reqwest::Client::builder()
            .timeout(timeout)
            .build()
            .map_err(|error| format!("could not start profile blob client: {error}"))?;
        Ok(Self {
            client,
            servers,
            max_blob_bytes,
            budget: ReadBudget::new(max_reads, max_total_bytes),
            cache: Arc::new(RwLock::new(BlobCache::default())),
        })
    }

    pub(super) fn with_shared_cache(mut self) -> Self {
        self.cache = SHARED_CACHE
            .get_or_init(|| Arc::new(RwLock::new(BlobCache::default())))
            .clone();
        self
    }

    async fn fetch_from_server(
        &self,
        server: &str,
        hash: &Hash,
    ) -> Result<Option<Vec<u8>>, BlobFetchError> {
        let url = format!("{}/{}.bin", server.trim_end_matches('/'), to_hex(hash));
        let mut response = self
            .client
            .get(url)
            .send()
            .await
            .map_err(|error| BlobFetchError::Request(error.to_string()))?;
        if response.status() == StatusCode::NOT_FOUND {
            return Ok(None);
        }
        if !response.status().is_success() {
            return Err(BlobFetchError::Status(response.status().as_u16()));
        }
        if response
            .content_length()
            .is_some_and(|length| length > self.max_blob_bytes as u64)
        {
            return Err(BlobFetchError::TooLarge {
                buffered_bytes: 0,
                limit: self.max_blob_bytes,
            });
        }

        let mut data = Vec::new();
        while let Some(chunk) = response
            .chunk()
            .await
            .map_err(|error| BlobFetchError::Request(error.to_string()))?
        {
            if !self.budget.consume_bytes(chunk.len()) {
                return Err(BlobFetchError::BudgetExceeded);
            }
            if chunk.len() > self.max_blob_bytes.saturating_sub(data.len()) {
                return Err(BlobFetchError::TooLarge {
                    buffered_bytes: data.len(),
                    limit: self.max_blob_bytes,
                });
            }
            data.extend_from_slice(&chunk);
        }
        if sha256(&data) != *hash {
            return Err(BlobFetchError::HashMismatch);
        }
        Ok(Some(data))
    }
}

#[async_trait]
impl Store for BoundedBlossomStore {
    async fn put(&self, _hash: Hash, _data: Vec<u8>) -> Result<bool, StoreError> {
        Err(StoreError::Other(
            "profile search store is read-only".to_string(),
        ))
    }

    async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
        if !self.budget.consume_read() {
            return Err(StoreError::Other(
                "profile search read budget exceeded".to_string(),
            ));
        }
        {
            let cache = self.cache.read().await;
            if let Some(data) = cache.entries.get(hash) {
                if data.len() > self.max_blob_bytes {
                    return Err(StoreError::Other(
                        "cached profile blob exceeds byte limit".to_string(),
                    ));
                }
                if !self.budget.consume_bytes(data.len()) {
                    return Err(StoreError::Other(
                        "profile search byte budget exceeded".to_string(),
                    ));
                }
                return Ok(Some(data.clone()));
            }
        }

        let mut last_error = None;
        for server in &self.servers {
            match self.fetch_from_server(server, hash).await {
                Ok(Some(data)) => {
                    self.cache.write().await.insert(*hash, data.clone());
                    return Ok(Some(data));
                }
                Ok(None) => {}
                Err(BlobFetchError::BudgetExceeded) => {
                    return Err(StoreError::Other(
                        "profile search byte budget exceeded".to_string(),
                    ));
                }
                Err(error) => last_error = Some(format!("{server}: {error}")),
            }
        }
        match last_error {
            Some(error) => Err(StoreError::Other(format!(
                "profile search blob download failed: {error}"
            ))),
            None => Ok(None),
        }
    }

    async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
        if !self.budget.consume_read() {
            return Err(StoreError::Other(
                "profile search read budget exceeded".to_string(),
            ));
        }
        if self.cache.read().await.entries.contains_key(hash) {
            return Ok(true);
        }
        let hash = to_hex(hash);
        for server in &self.servers {
            let url = format!("{}/{}.bin", server.trim_end_matches('/'), hash);
            if self
                .client
                .head(url)
                .send()
                .await
                .is_ok_and(|response| response.status().is_success())
            {
                return Ok(true);
            }
        }
        Ok(false)
    }

    async fn delete(&self, _hash: &Hash) -> Result<bool, StoreError> {
        Err(StoreError::Other(
            "profile search store is read-only".to_string(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpListener;
    use std::thread;

    fn serve_chunked_blob(chunks: Vec<Vec<u8>>) -> (String, thread::JoinHandle<String>) {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        let server = thread::spawn(move || {
            let (mut stream, _) = listener.accept().unwrap();
            let mut request = Vec::new();
            let mut buffer = [0; 1024];
            while !request.windows(4).any(|window| window == b"\r\n\r\n") {
                let count = stream.read(&mut buffer).unwrap();
                if count == 0 {
                    break;
                }
                request.extend_from_slice(&buffer[..count]);
            }
            stream
                .write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
                .unwrap();
            for chunk in chunks {
                if write!(stream, "{:x}\r\n", chunk.len()).is_err()
                    || stream.write_all(&chunk).is_err()
                    || stream.write_all(b"\r\n").is_err()
                    || stream.flush().is_err()
                {
                    break;
                }
                thread::sleep(Duration::from_millis(5));
            }
            let _ = stream.write_all(b"0\r\n\r\n");
            String::from_utf8_lossy(&request).into_owned()
        });
        (format!("http://{address}"), server)
    }

    fn test_store(
        servers: Vec<String>,
        blob: usize,
        reads: usize,
        total: usize,
    ) -> BoundedBlossomStore {
        BoundedBlossomStore::new(servers, Duration::from_secs(1), blob, reads, total).unwrap()
    }

    fn assert_requested(server: thread::JoinHandle<String>, hash: &Hash) {
        assert!(server
            .join()
            .unwrap()
            .contains(&format!("/{}.bin", to_hex(hash))));
    }

    #[tokio::test]
    async fn chunked_response_never_grows_owned_buffer_past_blob_limit() {
        let body = vec![7; 48];
        let hash = sha256(&body);
        let (url, server) = serve_chunked_blob(vec![body[..16].to_vec(); 3]);
        let store = test_store(vec![url.clone()], 24, 1, 100);

        let error = store.fetch_from_server(&url, &hash).await.unwrap_err();
        assert!(matches!(
            error,
            BlobFetchError::TooLarge {
                buffered_bytes: 0..=24,
                limit: 24
            }
        ));
        assert_requested(server, &hash);
    }

    #[tokio::test]
    async fn oversized_server_falls_back_to_valid_hash_matching_blob() {
        let valid = b"bounded profile blob";
        let hash = sha256(valid);
        let (large_url, large_server) = serve_chunked_blob(vec![vec![9; 16], vec![9; 16]]);
        let (valid_url, valid_server) = serve_chunked_blob(vec![valid.to_vec()]);
        let store = test_store(vec![large_url, valid_url], valid.len(), 1, 100);

        assert_eq!(
            store.get(&hash).await.unwrap().as_deref(),
            Some(valid.as_slice())
        );
        assert_requested(large_server, &hash);
        assert_requested(valid_server, &hash);
    }

    #[tokio::test]
    async fn hash_mismatch_falls_back_to_valid_server() {
        let valid = b"bounded profile index blob";
        let hash = sha256(valid);
        let (bad_url, bad_server) = serve_chunked_blob(vec![b"wrong body".to_vec()]);
        let (valid_url, valid_server) = serve_chunked_blob(vec![valid.to_vec()]);
        let store = test_store(vec![bad_url, valid_url], valid.len(), 1, 100);

        assert_eq!(
            store.get(&hash).await.unwrap().as_deref(),
            Some(valid.as_slice())
        );
        assert_requested(bad_server, &hash);
        assert_requested(valid_server, &hash);
    }

    #[tokio::test]
    async fn failed_servers_share_and_exhaust_aggregate_byte_budget() {
        let hash = sha256(b"expected body");
        let (first_url, first_server) = serve_chunked_blob(vec![vec![1; 16]]);
        let (second_url, second_server) = serve_chunked_blob(vec![vec![2; 16]]);
        let store = test_store(vec![first_url, second_url], 32, 1, 24);

        let error = store.get(&hash).await.unwrap_err();
        assert!(error.to_string().contains("byte budget exceeded"));
        assert_requested(first_server, &hash);
        assert_requested(second_server, &hash);
        assert_eq!(store.budget.bytes_left.load(Ordering::Relaxed), 0);
    }

    #[tokio::test]
    async fn cache_hits_consume_read_and_byte_budgets() {
        let body = b"cached profile blob";
        let hash = sha256(body);
        let (url, server) = serve_chunked_blob(vec![body.to_vec()]);
        let store = test_store(vec![url], body.len(), 3, body.len() * 2);

        assert!(store.get(&hash).await.unwrap().is_some());
        assert!(store.get(&hash).await.unwrap().is_some());
        assert_eq!(store.budget.reads_left.load(Ordering::Relaxed), 1);
        assert_eq!(store.budget.bytes_left.load(Ordering::Relaxed), 0);
        assert!(store
            .get(&hash)
            .await
            .unwrap_err()
            .to_string()
            .contains("byte budget exceeded"));
        assert_requested(server, &hash);
    }

    #[tokio::test]
    async fn separate_searches_reuse_verified_blobs_with_independent_budgets() {
        let body = b"shared cross-query profile blob";
        let hash = sha256(body);
        let (url, server) = serve_chunked_blob(vec![body.to_vec()]);
        let first = test_store(vec![url], body.len(), 1, body.len()).with_shared_cache();
        assert_eq!(
            first.get(&hash).await.unwrap().as_deref(),
            Some(body.as_slice())
        );
        assert_requested(server, &hash);
        let second = test_store(Vec::new(), body.len(), 1, body.len()).with_shared_cache();
        assert_eq!(
            second.get(&hash).await.unwrap().as_deref(),
            Some(body.as_slice())
        );
        assert!(second.get(&hash).await.is_err());
        let limited = test_store(Vec::new(), body.len() - 1, 1, body.len()).with_shared_cache();
        assert!(limited.get(&hash).await.is_err());
    }
}