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_hex: &str,
116 ) -> Result<Vec<u8>> {
117 if let Some(data) = store.get_chunk(hash_hex)? {
119 return Ok(data);
120 }
121
122 let data = self.fetch_chunk(webrtc_state, hash_hex).await?;
124 store.put_blob(&data)?;
125 Ok(data)
126 }
127
128 pub async fn fetch_tree(
131 &self,
132 store: &HashtreeStore,
133 webrtc_state: Option<&Arc<WebRTCState>>,
134 root_hash: &[u8; 32],
135 ) -> Result<(usize, u64)> {
136 let mut chunks_fetched = 0usize;
137 let mut bytes_fetched = 0u64;
138
139 let root_hex = to_hex(root_hash);
140
141 if store.blob_exists(&root_hex)? {
143 return Ok((0, 0));
144 }
145
146 let mut queue: VecDeque<[u8; 32]> = VecDeque::new();
148 queue.push_back(*root_hash);
149
150 while let Some(hash) = queue.pop_front() {
151 let hash_hex = to_hex(&hash);
152
153 if store.blob_exists(&hash_hex)? {
155 continue;
156 }
157
158 let data = self.fetch_chunk(webrtc_state, &hash_hex).await?;
160
161 store.put_blob(&data)?;
163 chunks_fetched += 1;
164 bytes_fetched += data.len() as u64;
165
166 if let Ok(node) = decode_tree_node(&data) {
168 for link in node.links {
169 queue.push_back(link.hash);
170 }
171 }
172 }
173
174 Ok((chunks_fetched, bytes_fetched))
175 }
176
177 pub async fn fetch_file(
180 &self,
181 store: &HashtreeStore,
182 webrtc_state: Option<&Arc<WebRTCState>>,
183 hash_hex: &str,
184 ) -> Result<Option<Vec<u8>>> {
185 if let Some(content) = store.get_file(hash_hex)? {
187 return Ok(Some(content));
188 }
189
190 let hash = hashtree_core::from_hex(hash_hex)
192 .map_err(|e| anyhow::anyhow!("Invalid hash: {}", e))?;
193
194 self.fetch_tree(store, webrtc_state, &hash).await?;
196
197 store.get_file(hash_hex)
199 }
200
201 pub async fn fetch_directory(
203 &self,
204 store: &HashtreeStore,
205 webrtc_state: Option<&Arc<WebRTCState>>,
206 hash_hex: &str,
207 ) -> Result<Option<crate::storage::DirectoryListing>> {
208 if let Ok(Some(listing)) = store.get_directory_listing(hash_hex) {
210 return Ok(Some(listing));
211 }
212
213 let hash = hashtree_core::from_hex(hash_hex)
215 .map_err(|e| anyhow::anyhow!("Invalid hash: {}", e))?;
216
217 self.fetch_tree(store, webrtc_state, &hash).await?;
219
220 store.get_directory_listing(hash_hex)
222 }
223
224 pub async fn upload(&self, data: &[u8]) -> Result<String> {
226 self.blossom
227 .upload(data)
228 .await
229 .map_err(|e| anyhow::anyhow!("Blossom upload failed: {}", e))
230 }
231
232 pub async fn upload_if_missing(&self, data: &[u8]) -> Result<(String, bool)> {
234 self.blossom
235 .upload_if_missing(data)
236 .await
237 .map_err(|e| anyhow::anyhow!("Blossom upload failed: {}", e))
238 }
239}