1use std::sync::Arc;
9
10use sha2::{Digest, Sha256};
11use tokio::io::{AsyncReadExt, AsyncWriteExt};
12use tokio::sync::{broadcast, mpsc, oneshot};
13use tracing::{error, info, warn};
14
15use crate::network::NetworkProvider;
16use crate::node::Node;
17
18use super::types::{
19 FileOffer, FileTransferEvent, FtMessage, OfferDecision, OfferResponder, TransferDirection,
20 TransferError, TransferProgress,
21};
22
23pub fn spawn_receive_handler<N: NetworkProvider + 'static>(
30 node: Arc<Node<N>>,
31 offer_tx: mpsc::UnboundedSender<(FileOffer, OfferResponder)>,
32 event_tx: broadcast::Sender<FileTransferEvent>,
33) -> tokio::task::JoinHandle<()> {
34 tokio::spawn(async move {
35 let mut rx = node.subscribe("ft");
36 info!("File transfer receive handler started");
37
38 loop {
39 let msg = match rx.recv().await {
40 Ok(m) => m,
41 Err(broadcast::error::RecvError::Lagged(n)) => {
42 warn!("FT receive handler lagged, missed {n} messages");
43 continue;
44 }
45 Err(broadcast::error::RecvError::Closed) => {
46 info!("FT receive handler: channel closed, exiting");
47 break;
48 }
49 };
50
51 let ft_msg: FtMessage = match serde_json::from_value(msg.payload.clone()) {
52 Ok(m) => m,
53 Err(e) => {
54 warn!(from = msg.from.as_str(), "Bad FT message: {e}");
55 continue;
56 }
57 };
58
59 let node = node.clone();
60 let from = msg.from.clone();
61 let offer_tx = offer_tx.clone();
62 let event_tx = event_tx.clone();
63
64 match ft_msg {
65 FtMessage::Offer {
66 file_name,
67 size,
68 sha256,
69 save_path,
70 token,
71 tcp_port: _,
72 } => {
73 tokio::spawn(async move {
75 if let Err(e) = handle_incoming_offer(
76 &node,
77 &from,
78 &file_name,
79 size,
80 &sha256,
81 &save_path,
82 &token,
83 &offer_tx,
84 &event_tx,
85 )
86 .await
87 {
88 match &e {
90 TransferError::Rejected(reason) => {
91 info!(
92 from = from.as_str(),
93 file = file_name.as_str(),
94 "File offer rejected: {reason}"
95 );
96 let _ = event_tx.send(FileTransferEvent::Rejected {
97 token,
98 file_name,
99 reason: reason.clone(),
100 });
101 }
102 _ => {
103 error!(
104 from = from.as_str(),
105 file = file_name.as_str(),
106 "Failed to receive file: {e}"
107 );
108 let _ = event_tx.send(FileTransferEvent::Failed {
109 token,
110 direction: TransferDirection::Receive,
111 file_name,
112 reason: e.to_string(),
113 });
114 }
115 }
116 }
117 });
118 }
119 FtMessage::PullRequest {
120 path,
121 requester_id: _,
122 token,
123 } => {
124 tokio::spawn(async move {
126 if let Err(e) =
127 handle_pull_request(&node, &from, &path, &token, &event_tx).await
128 {
129 error!(
130 from = from.as_str(),
131 path = path.as_str(),
132 "Failed to serve file: {e}"
133 );
134 }
135 });
136 }
137 _ => {
138 }
140 }
141 }
142 })
143}
144
145async fn handle_incoming_offer<N: NetworkProvider + 'static>(
148 node: &Node<N>,
149 from: &str,
150 file_name: &str,
151 size: u64,
152 sha256: &str,
153 save_path: &str,
154 token: &str,
155 offer_tx: &mpsc::UnboundedSender<(FileOffer, OfferResponder)>,
156 event_tx: &broadcast::Sender<FileTransferEvent>,
157) -> Result<(), TransferError> {
158 info!(
159 from = from,
160 file = file_name,
161 size = size,
162 "Received incoming file offer"
163 );
164
165 let offer = FileOffer {
167 from_peer: from.to_string(),
168 from_name: from.to_string(), file_name: file_name.to_string(),
170 size,
171 sha256: sha256.to_string(),
172 suggested_path: save_path.to_string(),
173 token: token.to_string(),
174 };
175
176 let _ = event_tx.send(FileTransferEvent::OfferReceived(offer.clone()));
178
179 let (decision_tx, decision_rx) = oneshot::channel::<OfferDecision>();
181 let responder = OfferResponder::new(decision_tx);
182
183 offer_tx
185 .send((offer, responder))
186 .map_err(|_| TransferError::Protocol("Offer channel closed".to_string()))?;
187
188 let decision = tokio::time::timeout(
190 tokio::time::Duration::from_secs(60),
191 decision_rx,
192 )
193 .await
194 .map_err(|_| TransferError::Timeout)?
195 .map_err(|_| TransferError::Protocol("Offer responder dropped without decision".to_string()))?;
196
197 match decision {
198 OfferDecision::Accept { save_path: dest } => {
199 accept_and_receive(node, from, file_name, size, sha256, token, &dest, event_tx).await
200 }
201 OfferDecision::Reject { reason } => {
202 let reject = FtMessage::Reject {
204 token: token.to_string(),
205 reason: reason.clone(),
206 };
207 let reject_payload = serde_json::to_value(&reject)
208 .map_err(|e| TransferError::Protocol(format!("Serialize error: {e}")))?;
209 node.send_typed(from, "ft", "reject", &reject_payload)
210 .await
211 .map_err(|e| TransferError::Node(format!("Failed to send REJECT: {e}")))?;
212
213 info!(
214 from = from,
215 file = file_name,
216 reason = reason.as_str(),
217 "Rejected file offer"
218 );
219
220 Err(TransferError::Rejected(reason))
221 }
222 }
223}
224
225async fn accept_and_receive<N: NetworkProvider + 'static>(
228 node: &Node<N>,
229 from: &str,
230 file_name: &str,
231 size: u64,
232 sha256: &str,
233 token: &str,
234 save_path: &str,
235 event_tx: &broadcast::Sender<FileTransferEvent>,
236) -> Result<(), TransferError> {
237 let start = std::time::Instant::now();
238
239 if let Some(parent) = std::path::Path::new(save_path).parent() {
241 tokio::fs::create_dir_all(parent).await?;
242 }
243
244 let mut listener = node
246 .listen_tcp(0)
247 .await
248 .map_err(|e| TransferError::Node(format!("Failed to listen TCP: {e}")))?;
249
250 let accept = FtMessage::Accept {
252 token: token.to_string(),
253 tcp_port: listener.port,
254 };
255 let accept_payload = serde_json::to_value(&accept)
256 .map_err(|e| TransferError::Protocol(format!("Serialize error: {e}")))?;
257 node.send_typed(from, "ft", "accept", &accept_payload)
258 .await
259 .map_err(|e| TransferError::Node(format!("Failed to send ACCEPT: {e}")))?;
260
261 info!(
262 port = listener.port,
263 "Sent ACCEPT, listening for TCP connection"
264 );
265
266 let incoming = tokio::time::timeout(
268 tokio::time::Duration::from_secs(30),
269 listener.accept(),
270 )
271 .await
272 .map_err(|_| TransferError::Timeout)?
273 .ok_or_else(|| TransferError::Protocol("Listener closed before accepting".to_string()))?;
274
275 let mut stream = incoming.stream;
276
277 let mut size_buf = [0u8; 8];
279 stream.read_exact(&mut size_buf).await?;
280 let file_size = u64::from_be_bytes(size_buf);
281
282 let mut sha_buf = [0u8; 64];
283 stream.read_exact(&mut sha_buf).await?;
284 let received_sha = String::from_utf8_lossy(&sha_buf).to_string();
285
286 if received_sha != sha256 {
288 return Err(TransferError::IntegrityError {
289 expected: sha256.to_string(),
290 actual: received_sha,
291 });
292 }
293
294 if file_size != size {
295 return Err(TransferError::Protocol(format!(
296 "Size mismatch: offer said {size}, stream header says {file_size}"
297 )));
298 }
299
300 let temp_path = format!("{save_path}.truffle-tmp");
302 let mut temp_file = tokio::fs::File::create(&temp_path).await?;
303 let mut hasher = Sha256::new();
304 let mut bytes_received: u64 = 0;
305 let progress_start = std::time::Instant::now();
306 let mut last_progress = std::time::Instant::now();
307 let mut buf = vec![0u8; 64 * 1024];
308
309 while bytes_received < file_size {
310 let to_read = ((file_size - bytes_received) as usize).min(buf.len());
311 let n = stream.read(&mut buf[..to_read]).await?;
312 if n == 0 {
313 tokio::fs::remove_file(&temp_path).await.ok();
314 return Err(TransferError::Io(std::io::Error::new(
315 std::io::ErrorKind::UnexpectedEof,
316 format!("Connection closed after {bytes_received}/{file_size} bytes"),
317 )));
318 }
319 hasher.update(&buf[..n]);
320 tokio::io::AsyncWriteExt::write_all(&mut temp_file, &buf[..n]).await?;
321 bytes_received += n as u64;
322
323 if last_progress.elapsed() >= std::time::Duration::from_millis(250) {
325 let elapsed = progress_start.elapsed().as_secs_f64();
326 let speed = if elapsed > 0.0 {
327 bytes_received as f64 / elapsed
328 } else {
329 0.0
330 };
331 let _ = event_tx.send(FileTransferEvent::Progress(TransferProgress {
332 token: token.to_string(),
333 direction: TransferDirection::Receive,
334 file_name: file_name.to_string(),
335 bytes_transferred: bytes_received,
336 total_bytes: file_size,
337 speed_bps: speed,
338 }));
339 last_progress = std::time::Instant::now();
340 }
341 }
342
343 tokio::io::AsyncWriteExt::flush(&mut temp_file).await?;
345
346 let actual_sha = hex::encode(hasher.finalize());
348
349 if actual_sha != sha256 {
350 stream.write_all(&[0x00]).await?;
352 tokio::fs::remove_file(&temp_path).await.ok();
354 return Err(TransferError::IntegrityError {
355 expected: sha256.to_string(),
356 actual: actual_sha,
357 });
358 }
359
360 let final_path = {
362 let p = std::path::Path::new(save_path);
363 if p.is_dir() || save_path.ends_with('/') || save_path.ends_with('\\') {
364 format!("{}/{}", save_path.trim_end_matches(['/', '\\']), file_name)
365 } else {
366 save_path.to_string()
367 }
368 };
369
370 if let Some(parent) = std::path::Path::new(&final_path).parent() {
372 tokio::fs::create_dir_all(parent).await?;
373 }
374
375 info!(
378 temp = temp_path.as_str(),
379 final_path = final_path.as_str(),
380 "Moving temp file to final destination"
381 );
382 if let Err(rename_err) = tokio::fs::rename(&temp_path, &final_path).await {
383 info!(
384 err = %rename_err,
385 "Rename failed, trying copy+delete fallback"
386 );
387 tokio::fs::copy(&temp_path, &final_path).await?;
388 tokio::fs::remove_file(&temp_path).await.ok();
389 }
390 info!(
391 final_path = final_path.as_str(),
392 exists = std::path::Path::new(&final_path).exists(),
393 "File save completed"
394 );
395
396 stream.write_all(&[0x01]).await?;
398 tokio::io::AsyncWriteExt::flush(&mut stream).await?;
399
400 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
405
406 let elapsed = start.elapsed().as_secs_f64();
407 info!(
408 file = final_path.as_str(),
409 bytes = file_size,
410 elapsed_ms = (elapsed * 1000.0) as u64,
411 "File received and verified"
412 );
413
414 let _ = event_tx.send(FileTransferEvent::Completed {
416 token: token.to_string(),
417 direction: TransferDirection::Receive,
418 file_name: file_name.to_string(),
419 bytes_transferred: file_size,
420 sha256: actual_sha,
421 elapsed_secs: elapsed,
422 });
423
424 Ok(())
425}
426
427async fn handle_pull_request<N: NetworkProvider + 'static>(
429 node: &Node<N>,
430 from: &str,
431 path: &str,
432 _token: &str,
433 event_tx: &broadcast::Sender<FileTransferEvent>,
434) -> Result<(), TransferError> {
435 info!(from = from, path = path, "Processing PULL_REQUEST");
436
437 let data = tokio::fs::read(path)
439 .await
440 .map_err(|e| TransferError::Io(e))?;
441 let size = data.len() as u64;
442
443 let mut hasher = Sha256::new();
444 hasher.update(&data);
445 let sha256 = hex::encode(hasher.finalize());
446
447 let file_name = std::path::Path::new(path)
448 .file_name()
449 .and_then(|n| n.to_str())
450 .unwrap_or("file")
451 .to_string();
452
453 let offer_token = uuid::Uuid::new_v4().to_string();
454
455 let offer = FtMessage::Offer {
457 file_name: file_name.clone(),
458 size,
459 sha256: sha256.clone(),
460 save_path: String::new(),
461 token: offer_token.clone(),
462 tcp_port: 0,
463 };
464 let offer_payload = serde_json::to_value(&offer)
465 .map_err(|e| TransferError::Protocol(format!("Serialize error: {e}")))?;
466
467 let _accept_port = crate::request_reply::send_and_wait(
468 node,
469 from,
470 "ft",
471 "offer",
472 &offer_payload,
473 std::time::Duration::from_secs(30),
474 |msg| {
475 if msg.from != from {
476 return None;
477 }
478 let ft_msg: FtMessage = serde_json::from_value(msg.payload.clone()).ok()?;
479 match ft_msg {
480 FtMessage::Accept {
481 token: ref t,
482 tcp_port,
483 } if *t == offer_token => Some(Ok(tcp_port)),
484 FtMessage::Reject {
485 token: ref t,
486 reason,
487 } if *t == offer_token => {
488 Some(Err(TransferError::Rejected(format!("Peer rejected: {reason}"))))
489 }
490 _ => None,
491 }
492 },
493 )
494 .await
495 .map_err(|e| match e {
496 crate::request_reply::RequestError::Timeout => TransferError::Timeout,
497 crate::request_reply::RequestError::Send(e) => {
498 TransferError::Node(format!("Failed to send OFFER: {e}"))
499 }
500 crate::request_reply::RequestError::ChannelClosed => {
501 TransferError::Protocol("Channel closed".into())
502 }
503 })?
504 .map_err(|e| e)?;
505
506 let mut stream = node
508 .open_tcp(from, _accept_port)
509 .await
510 .map_err(|e| {
511 TransferError::Node(format!(
512 "Failed to open TCP to {from}:{_accept_port}: {e}"
513 ))
514 })?;
515
516 let start = std::time::Instant::now();
517
518 stream.write_all(&size.to_be_bytes()).await?;
520 stream.write_all(sha256.as_bytes()).await?;
521
522 let chunk_size = 64 * 1024;
523 let mut offset = 0;
524 let mut bytes_sent: u64 = 0;
525
526 while offset < data.len() {
527 let end = (offset + chunk_size).min(data.len());
528 stream.write_all(&data[offset..end]).await?;
529 bytes_sent += (end - offset) as u64;
530 offset = end;
531
532 let elapsed = start.elapsed().as_secs_f64();
533 let speed = if elapsed > 0.0 {
534 bytes_sent as f64 / elapsed
535 } else {
536 0.0
537 };
538
539 let _ = event_tx.send(FileTransferEvent::Progress(TransferProgress {
541 token: offer_token.clone(),
542 direction: TransferDirection::Send,
543 file_name: file_name.clone(),
544 bytes_transferred: bytes_sent,
545 total_bytes: size,
546 speed_bps: speed,
547 }));
548 }
549
550 stream.flush().await?;
551
552 let mut ack = [0u8; 1];
554 stream.read_exact(&mut ack).await?;
555
556 if ack[0] != 0x01 {
557 return Err(TransferError::IntegrityError {
558 expected: sha256,
559 actual: "peer reported integrity failure".to_string(),
560 });
561 }
562
563 let elapsed = start.elapsed().as_secs_f64();
564 info!(path = path, bytes = size, "File served successfully");
565
566 let _ = event_tx.send(FileTransferEvent::Completed {
568 token: offer_token,
569 direction: TransferDirection::Send,
570 file_name,
571 bytes_transferred: size,
572 sha256,
573 elapsed_secs: elapsed,
574 });
575
576 Ok(())
577}