truffle_core/file_transfer/
mod.rs1pub mod types;
37pub mod sender;
38pub mod receiver;
39
40pub use types::*;
41
42use std::sync::atomic::{AtomicU64, Ordering};
43use std::sync::Arc;
44
45use tokio::sync::{broadcast, mpsc, Mutex};
46use tokio::task::JoinHandle;
47use tracing::info;
48
49use crate::network::NetworkProvider;
50use crate::node::Node;
51
52const DEFAULT_MAX_TRANSFER_SIZE: u64 = 1_073_741_824;
54
55pub(crate) struct FileTransferState {
62 pub(crate) event_tx: broadcast::Sender<FileTransferEvent>,
64 pub(crate) offer_tx: Mutex<Option<mpsc::UnboundedSender<(FileOffer, OfferResponder)>>>,
66 pub(crate) max_transfer_size: AtomicU64,
68 pub(crate) receiver_handle: Mutex<Option<JoinHandle<()>>>,
70}
71
72impl FileTransferState {
73 pub(crate) fn new() -> Self {
75 let (event_tx, _) = broadcast::channel(256);
76 Self {
77 event_tx,
78 offer_tx: Mutex::new(None),
79 max_transfer_size: AtomicU64::new(DEFAULT_MAX_TRANSFER_SIZE),
80 receiver_handle: Mutex::new(None),
81 }
82 }
83}
84
85pub struct FileTransfer<'a, N: NetworkProvider + 'static> {
92 node: &'a Node<N>,
94}
95
96impl<'a, N: NetworkProvider + 'static> FileTransfer<'a, N> {
97 pub(crate) fn new(node: &'a Node<N>) -> Self {
99 Self { node }
100 }
101
102 fn state(&self) -> &FileTransferState {
104 &self.node.file_transfer_state
105 }
106
107 pub fn subscribe(&self) -> broadcast::Receiver<FileTransferEvent> {
112 self.state().event_tx.subscribe()
113 }
114
115 pub fn set_max_transfer_size(&self, bytes: u64) {
119 self.state()
120 .max_transfer_size
121 .store(bytes, Ordering::Relaxed);
122 }
123
124 pub fn max_transfer_size(&self) -> u64 {
126 self.state().max_transfer_size.load(Ordering::Relaxed)
127 }
128
129 pub async fn offer_channel(
142 &self,
143 node: Arc<Node<N>>,
144 ) -> mpsc::UnboundedReceiver<(FileOffer, OfferResponder)> {
145 let (tx, rx) = mpsc::unbounded_channel();
146 {
147 let mut offer_tx = self.state().offer_tx.lock().await;
148 *offer_tx = Some(tx.clone());
149 }
150 self.ensure_receiver_started(node, tx).await;
151 rx
152 }
153
154 pub async fn auto_accept(&self, node: Arc<Node<N>>, output_dir: &str) {
161 let mut rx = self.offer_channel(node).await;
162 let output_dir = output_dir.to_string();
163
164 tokio::spawn(async move {
165 while let Some((offer, responder)) = rx.recv().await {
166 let dest = if offer.suggested_path.is_empty() || offer.suggested_path == "." {
167 format!("{}/{}", output_dir, offer.file_name)
168 } else {
169 offer.suggested_path.clone()
170 };
171 info!(
172 file = offer.file_name.as_str(),
173 dest = dest.as_str(),
174 "Auto-accepting file offer"
175 );
176 responder.accept(&dest);
177 }
178 });
179 }
180
181 pub async fn auto_reject(&self, node: Arc<Node<N>>) {
186 let mut rx = self.offer_channel(node).await;
187
188 tokio::spawn(async move {
189 while let Some((_offer, responder)) = rx.recv().await {
190 responder.reject("auto-rejected");
191 }
192 });
193 }
194
195 pub async fn send_file(
205 &self,
206 peer_id: &str,
207 local_path: &str,
208 remote_path: &str,
209 ) -> Result<TransferResult, TransferError> {
210 let max_size = self.state().max_transfer_size.load(Ordering::Relaxed);
211 sender::send_file(
212 self.node,
213 peer_id,
214 local_path,
215 remote_path,
216 max_size,
217 &self.state().event_tx,
218 )
219 .await
220 }
221
222 pub async fn pull_file(
232 &self,
233 peer_id: &str,
234 remote_path: &str,
235 local_path: &str,
236 ) -> Result<TransferResult, TransferError> {
237 pull_file(
238 self.node,
239 peer_id,
240 remote_path,
241 local_path,
242 &self.state().event_tx,
243 )
244 .await
245 }
246
247 async fn ensure_receiver_started(
252 &self,
253 node: Arc<Node<N>>,
254 offer_tx: mpsc::UnboundedSender<(FileOffer, OfferResponder)>,
255 ) {
256 let mut handle = self.state().receiver_handle.lock().await;
257 if let Some(h) = handle.take() {
259 h.abort();
260 }
261 let h = receiver::spawn_receive_handler(
262 node,
263 offer_tx,
264 self.state().event_tx.clone(),
265 );
266 *handle = Some(h);
267 }
268}
269
270async fn pull_file<N: NetworkProvider + 'static>(
284 node: &Node<N>,
285 peer_id: &str,
286 remote_path: &str,
287 local_path: &str,
288 event_tx: &broadcast::Sender<FileTransferEvent>,
289) -> Result<TransferResult, TransferError> {
290 use sha2::{Digest, Sha256};
291 use std::time::Instant;
292 use tokio::io::{AsyncReadExt, AsyncWriteExt};
293
294 let start = Instant::now();
295 let token = uuid::Uuid::new_v4().to_string();
296 let requester_id = node.local_info().id;
297
298 let peer_id = node
300 .resolve_peer_id(peer_id)
301 .await
302 .map_err(|e| TransferError::Node(e.to_string()))?;
303 let peer_id = peer_id.as_str();
304
305 let pull_req = FtMessage::PullRequest {
307 path: remote_path.to_string(),
308 requester_id,
309 token: token.clone(),
310 };
311
312 let pull_bytes = serde_json::to_vec(&pull_req)
313 .map_err(|e| TransferError::Protocol(format!("Failed to serialize pull request: {e}")))?;
314
315 node.send(peer_id, "ft", &pull_bytes)
316 .await
317 .map_err(|e| TransferError::Node(e.to_string()))?;
318
319 tracing::info!(peer = peer_id, path = remote_path, "Sent PULL_REQUEST");
320
321 let mut rx = node.subscribe("ft");
323 let offer_deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(30);
324
325 let (offer_sha256, offer_size, offer_token, file_name) = loop {
326 let msg = tokio::time::timeout(
327 offer_deadline.duration_since(tokio::time::Instant::now()),
328 rx.recv(),
329 )
330 .await
331 .map_err(|_| TransferError::Timeout)?
332 .map_err(|e| TransferError::Protocol(format!("Channel error: {e}")))?;
333
334 if msg.from != peer_id {
335 continue;
336 }
337
338 let ft_msg: FtMessage = serde_json::from_value(msg.payload.clone())
339 .map_err(|e| TransferError::Protocol(format!("Bad FT message: {e}")))?;
340
341 match ft_msg {
342 FtMessage::Offer {
343 file_name,
344 sha256,
345 size,
346 token: offer_token,
347 ..
348 } => {
349 tracing::info!(size = size, sha256 = sha256.as_str(), "Received OFFER from peer");
350 break (sha256, size, offer_token, file_name);
351 }
352 FtMessage::Reject { reason, .. } => {
353 return Err(TransferError::Rejected(reason));
354 }
355 _ => continue,
356 }
357 };
358
359 let mut listener = node
361 .listen_tcp(0)
362 .await
363 .map_err(|e| TransferError::Node(format!("Failed to listen TCP: {e}")))?;
364
365 let accept = FtMessage::Accept {
366 token: offer_token,
367 tcp_port: listener.port,
368 };
369
370 let accept_bytes = serde_json::to_vec(&accept)
371 .map_err(|e| TransferError::Protocol(format!("Failed to serialize accept: {e}")))?;
372
373 node.send(peer_id, "ft", &accept_bytes)
374 .await
375 .map_err(|e| TransferError::Node(e.to_string()))?;
376
377 tracing::info!(
378 port = listener.port,
379 "Sent ACCEPT, listening for TCP connection"
380 );
381
382 let incoming = tokio::time::timeout(
384 tokio::time::Duration::from_secs(30),
385 listener.accept(),
386 )
387 .await
388 .map_err(|_| TransferError::Timeout)?
389 .ok_or_else(|| TransferError::Protocol("Listener closed before accepting".to_string()))?;
390
391 let mut stream = incoming.stream;
392 tracing::info!("TCP connection accepted from peer");
393
394 let mut size_buf = [0u8; 8];
396 stream.read_exact(&mut size_buf).await?;
397 let file_size = u64::from_be_bytes(size_buf);
398
399 let mut sha_buf = [0u8; 64]; stream.read_exact(&mut sha_buf).await?;
401 let received_sha = String::from_utf8_lossy(&sha_buf).to_string();
402
403 if received_sha != offer_sha256 {
405 return Err(TransferError::IntegrityError {
406 expected: offer_sha256,
407 actual: received_sha,
408 });
409 }
410
411 if file_size != offer_size {
412 return Err(TransferError::Protocol(format!(
413 "Size mismatch: offer said {offer_size}, stream header says {file_size}"
414 )));
415 }
416
417 if let Some(parent) = std::path::Path::new(local_path).parent() {
420 tokio::fs::create_dir_all(parent).await?;
421 }
422
423 let temp_path = format!("{local_path}.truffle-tmp");
424 let mut temp_file = tokio::fs::File::create(&temp_path).await?;
425 let mut hasher = Sha256::new();
426 let mut bytes_received: u64 = 0;
427 let progress_start = Instant::now();
428 let mut buf = vec![0u8; 64 * 1024];
429
430 while bytes_received < file_size {
431 let to_read = ((file_size - bytes_received) as usize).min(buf.len());
432 let n = stream.read(&mut buf[..to_read]).await?;
433 if n == 0 {
434 tokio::fs::remove_file(&temp_path).await.ok();
435 return Err(TransferError::Io(std::io::Error::new(
436 std::io::ErrorKind::UnexpectedEof,
437 format!("Connection closed after {bytes_received}/{file_size} bytes"),
438 )));
439 }
440 hasher.update(&buf[..n]);
441 AsyncWriteExt::write_all(&mut temp_file, &buf[..n]).await?;
442 bytes_received += n as u64;
443
444 let elapsed = progress_start.elapsed().as_secs_f64();
445 let speed = if elapsed > 0.0 {
446 bytes_received as f64 / elapsed
447 } else {
448 0.0
449 };
450
451 let _ = event_tx.send(FileTransferEvent::Progress(TransferProgress {
453 token: token.clone(),
454 direction: TransferDirection::Receive,
455 file_name: file_name.clone(),
456 bytes_transferred: bytes_received,
457 total_bytes: file_size,
458 speed_bps: speed,
459 }));
460 }
461
462 AsyncWriteExt::flush(&mut temp_file).await?;
464
465 let actual_sha = hex::encode(hasher.finalize());
467
468 if actual_sha != offer_sha256 {
469 stream.write_all(&[0x00]).await?;
471 tokio::fs::remove_file(&temp_path).await.ok();
472 return Err(TransferError::IntegrityError {
473 expected: offer_sha256,
474 actual: actual_sha,
475 });
476 }
477
478 tokio::fs::rename(&temp_path, local_path).await?;
480
481 stream.write_all(&[0x01]).await?;
483
484 let elapsed = start.elapsed().as_secs_f64();
485 tracing::info!(
486 bytes = file_size,
487 elapsed_ms = (elapsed * 1000.0) as u64,
488 path = local_path,
489 "Download complete"
490 );
491
492 let _ = event_tx.send(FileTransferEvent::Completed {
494 token,
495 direction: TransferDirection::Receive,
496 file_name,
497 bytes_transferred: file_size,
498 sha256: actual_sha.clone(),
499 elapsed_secs: elapsed,
500 });
501
502 Ok(TransferResult {
503 bytes_transferred: file_size,
504 sha256: actual_sha,
505 elapsed_secs: elapsed,
506 })
507}