1use anyhow::Result;
9use hashtree_blossom::BlossomClient;
10use hashtree_core::{decode_tree_node, to_hex};
11use nostr::Keys;
12use std::collections::VecDeque;
13use std::sync::Arc;
14use std::time::Duration;
15use tracing::debug;
16
17use crate::storage::HashtreeStore;
18use crate::webrtc::WebRTCState;
19
20#[derive(Clone)]
22pub struct FetchConfig {
23 pub webrtc_timeout: Duration,
25 pub blossom_timeout: Duration,
27}
28
29impl Default for FetchConfig {
30 fn default() -> Self {
31 Self {
32 webrtc_timeout: Duration::from_millis(2000),
33 blossom_timeout: Duration::from_millis(10000),
34 }
35 }
36}
37
38pub struct Fetcher {
40 config: FetchConfig,
41 blossom: BlossomClient,
42}
43
44impl Fetcher {
45 pub fn new(config: FetchConfig) -> Self {
48 let keys = Keys::generate();
50 let blossom = BlossomClient::new(keys)
51 .with_timeout(config.blossom_timeout);
52
53 Self { config, blossom }
54 }
55
56 pub fn with_keys(config: FetchConfig, keys: Keys) -> Self {
58 let blossom = BlossomClient::new(keys)
59 .with_timeout(config.blossom_timeout);
60
61 Self { config, blossom }
62 }
63
64 pub fn blossom(&self) -> &BlossomClient {
66 &self.blossom
67 }
68
69 pub async fn fetch_chunk(
71 &self,
72 webrtc_state: Option<&Arc<WebRTCState>>,
73 hash_hex: &str,
74 ) -> Result<Vec<u8>> {
75 let short_hash = if hash_hex.len() >= 12 {
76 &hash_hex[..12]
77 } else {
78 hash_hex
79 };
80
81 if let Some(state) = webrtc_state {
83 debug!("Trying WebRTC for {}", short_hash);
84 let webrtc_result = tokio::time::timeout(
85 self.config.webrtc_timeout,
86 state.request_from_peers(hash_hex),
87 )
88 .await;
89
90 if let Ok(Some(data)) = webrtc_result {
91 debug!("Got {} from WebRTC ({} bytes)", short_hash, data.len());
92 return Ok(data);
93 }
94 }
95
96 debug!("Trying Blossom for {}", short_hash);
98 match self.blossom.download(hash_hex).await {
99 Ok(data) => {
100 debug!("Got {} from Blossom ({} bytes)", short_hash, data.len());
101 Ok(data)
102 }
103 Err(e) => {
104 debug!("Blossom download failed for {}: {}", short_hash, e);
105 Err(anyhow::anyhow!("Failed to fetch {} from any source: {}", short_hash, e))
106 }
107 }
108 }
109
110 pub async fn fetch_chunk_with_store(
112 &self,
113 store: &HashtreeStore,
114 webrtc_state: Option<&Arc<WebRTCState>>,
115 hash: &[u8; 32],
116 ) -> Result<Vec<u8>> {
117 if let Some(data) = store.get_chunk(hash)? {
119 return Ok(data);
120 }
121
122 let hash_hex = to_hex(hash);
124 let data = self.fetch_chunk(webrtc_state, &hash_hex).await?;
125 store.put_blob(&data)?;
126 Ok(data)
127 }
128
129 pub async fn fetch_tree(
132 &self,
133 store: &HashtreeStore,
134 webrtc_state: Option<&Arc<WebRTCState>>,
135 root_hash: &[u8; 32],
136 ) -> Result<(usize, u64)> {
137 let mut chunks_fetched = 0usize;
138 let mut bytes_fetched = 0u64;
139
140 if store.blob_exists(root_hash)? {
142 return Ok((0, 0));
143 }
144
145 let mut queue: VecDeque<[u8; 32]> = VecDeque::new();
147 queue.push_back(*root_hash);
148
149 while let Some(hash) = queue.pop_front() {
150 if store.blob_exists(&hash)? {
152 continue;
153 }
154
155 let hash_hex = to_hex(&hash);
156
157 let data = self.fetch_chunk(webrtc_state, &hash_hex).await?;
159
160 store.put_blob(&data)?;
162 chunks_fetched += 1;
163 bytes_fetched += data.len() as u64;
164
165 if let Ok(node) = decode_tree_node(&data) {
167 for link in node.links {
168 queue.push_back(link.hash);
169 }
170 }
171 }
172
173 Ok((chunks_fetched, bytes_fetched))
174 }
175
176 pub async fn fetch_file(
179 &self,
180 store: &HashtreeStore,
181 webrtc_state: Option<&Arc<WebRTCState>>,
182 hash: &[u8; 32],
183 ) -> Result<Option<Vec<u8>>> {
184 if let Some(content) = store.get_file(hash)? {
186 return Ok(Some(content));
187 }
188
189 self.fetch_tree(store, webrtc_state, hash).await?;
191
192 store.get_file(hash)
194 }
195
196 pub async fn fetch_directory(
198 &self,
199 store: &HashtreeStore,
200 webrtc_state: Option<&Arc<WebRTCState>>,
201 hash: &[u8; 32],
202 ) -> Result<Option<crate::storage::DirectoryListing>> {
203 if let Ok(Some(listing)) = store.get_directory_listing(hash) {
205 return Ok(Some(listing));
206 }
207
208 self.fetch_tree(store, webrtc_state, hash).await?;
210
211 store.get_directory_listing(hash)
213 }
214
215 pub async fn upload(&self, data: &[u8]) -> Result<String> {
217 self.blossom
218 .upload(data)
219 .await
220 .map_err(|e| anyhow::anyhow!("Blossom upload failed: {}", e))
221 }
222
223 pub async fn upload_if_missing(&self, data: &[u8]) -> Result<(String, bool)> {
225 self.blossom
226 .upload_if_missing(data)
227 .await
228 .map_err(|e| anyhow::anyhow!("Blossom upload failed: {}", e))
229 }
230}