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_payload = serde_json::to_value(&pull_req)
313 .map_err(|e| TransferError::Protocol(format!("Failed to serialize pull request: {e}")))?;
314
315 tracing::info!(peer = peer_id, path = remote_path, "Sending PULL_REQUEST");
316
317 let (offer_sha256, offer_size, offer_token, file_name) =
318 crate::request_reply::send_and_wait(
319 node,
320 peer_id,
321 "ft",
322 "pull_request",
323 &pull_payload,
324 std::time::Duration::from_secs(30),
325 |msg| {
326 if msg.from != peer_id {
327 return None;
328 }
329 let ft_msg: FtMessage = serde_json::from_value(msg.payload.clone()).ok()?;
330 match ft_msg {
331 FtMessage::Offer {
332 file_name,
333 sha256,
334 size,
335 token: offer_token,
336 ..
337 } => {
338 tracing::info!(size = size, "Received OFFER from peer");
339 Some(Ok((sha256, size, offer_token, file_name)))
340 }
341 FtMessage::Reject { reason, .. } => {
342 Some(Err(TransferError::Rejected(reason)))
343 }
344 _ => None,
345 }
346 },
347 )
348 .await
349 .map_err(|e| match e {
350 crate::request_reply::RequestError::Timeout => TransferError::Timeout,
351 crate::request_reply::RequestError::Send(e) => TransferError::Node(e.to_string()),
352 crate::request_reply::RequestError::ChannelClosed => {
353 TransferError::Protocol("Channel closed".into())
354 }
355 })?
356 .map_err(|e| e)?;
357
358 let mut listener = node
360 .listen_tcp(0)
361 .await
362 .map_err(|e| TransferError::Node(format!("Failed to listen TCP: {e}")))?;
363
364 let accept = FtMessage::Accept {
365 token: offer_token,
366 tcp_port: listener.port,
367 };
368
369 let accept_payload = serde_json::to_value(&accept)
370 .map_err(|e| TransferError::Protocol(format!("Failed to serialize accept: {e}")))?;
371
372 node.send_typed(peer_id, "ft", "accept", &accept_payload)
373 .await
374 .map_err(|e| TransferError::Node(e.to_string()))?;
375
376 tracing::info!(
377 port = listener.port,
378 "Sent ACCEPT, listening for TCP connection"
379 );
380
381 let incoming = tokio::time::timeout(
383 tokio::time::Duration::from_secs(30),
384 listener.accept(),
385 )
386 .await
387 .map_err(|_| TransferError::Timeout)?
388 .ok_or_else(|| TransferError::Protocol("Listener closed before accepting".to_string()))?;
389
390 let mut stream = incoming.stream;
391 tracing::info!("TCP connection accepted from peer");
392
393 let mut size_buf = [0u8; 8];
395 stream.read_exact(&mut size_buf).await?;
396 let file_size = u64::from_be_bytes(size_buf);
397
398 let mut sha_buf = [0u8; 64]; stream.read_exact(&mut sha_buf).await?;
400 let received_sha = String::from_utf8_lossy(&sha_buf).to_string();
401
402 if received_sha != offer_sha256 {
404 return Err(TransferError::IntegrityError {
405 expected: offer_sha256,
406 actual: received_sha,
407 });
408 }
409
410 if file_size != offer_size {
411 return Err(TransferError::Protocol(format!(
412 "Size mismatch: offer said {offer_size}, stream header says {file_size}"
413 )));
414 }
415
416 if let Some(parent) = std::path::Path::new(local_path).parent() {
419 tokio::fs::create_dir_all(parent).await?;
420 }
421
422 let temp_path = format!("{local_path}.truffle-tmp");
423 let mut temp_file = tokio::fs::File::create(&temp_path).await?;
424 let mut hasher = Sha256::new();
425 let mut bytes_received: u64 = 0;
426 let progress_start = Instant::now();
427 let mut buf = vec![0u8; 64 * 1024];
428
429 while bytes_received < file_size {
430 let to_read = ((file_size - bytes_received) as usize).min(buf.len());
431 let n = stream.read(&mut buf[..to_read]).await?;
432 if n == 0 {
433 tokio::fs::remove_file(&temp_path).await.ok();
434 return Err(TransferError::Io(std::io::Error::new(
435 std::io::ErrorKind::UnexpectedEof,
436 format!("Connection closed after {bytes_received}/{file_size} bytes"),
437 )));
438 }
439 hasher.update(&buf[..n]);
440 AsyncWriteExt::write_all(&mut temp_file, &buf[..n]).await?;
441 bytes_received += n as u64;
442
443 let elapsed = progress_start.elapsed().as_secs_f64();
444 let speed = if elapsed > 0.0 {
445 bytes_received as f64 / elapsed
446 } else {
447 0.0
448 };
449
450 let _ = event_tx.send(FileTransferEvent::Progress(TransferProgress {
452 token: token.clone(),
453 direction: TransferDirection::Receive,
454 file_name: file_name.clone(),
455 bytes_transferred: bytes_received,
456 total_bytes: file_size,
457 speed_bps: speed,
458 }));
459 }
460
461 AsyncWriteExt::flush(&mut temp_file).await?;
463
464 let actual_sha = hex::encode(hasher.finalize());
466
467 if actual_sha != offer_sha256 {
468 stream.write_all(&[0x00]).await?;
470 tokio::fs::remove_file(&temp_path).await.ok();
471 return Err(TransferError::IntegrityError {
472 expected: offer_sha256,
473 actual: actual_sha,
474 });
475 }
476
477 tokio::fs::rename(&temp_path, local_path).await?;
479
480 stream.write_all(&[0x01]).await?;
482
483 let elapsed = start.elapsed().as_secs_f64();
484 tracing::info!(
485 bytes = file_size,
486 elapsed_ms = (elapsed * 1000.0) as u64,
487 path = local_path,
488 "Download complete"
489 );
490
491 let _ = event_tx.send(FileTransferEvent::Completed {
493 token,
494 direction: TransferDirection::Receive,
495 file_name,
496 bytes_transferred: file_size,
497 sha256: actual_sha.clone(),
498 elapsed_secs: elapsed,
499 });
500
501 Ok(TransferResult {
502 bytes_transferred: file_size,
503 sha256: actual_sha,
504 elapsed_secs: elapsed,
505 })
506}