Skip to main content

truffle_core/file_transfer/
receiver.rs

1//! Receive handler — background task that handles incoming file transfer
2//! requests from other peers.
3//!
4//! Instead of auto-accepting, offers are forwarded through an offer channel
5//! so the application can decide whether to accept or reject each transfer.
6//! Pull requests are auto-served (same as the CLI version).
7
8use 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
23/// Spawn a background task that listens for incoming file transfer messages.
24///
25/// - **OFFER**: Creates a [`FileOffer`] + [`OfferResponder`] pair, sends them
26///   on the `offer_tx` channel, and waits up to 60 seconds for a decision.
27/// - **PULL_REQUEST**: Auto-serves the requested file (same as CLI).
28/// - **ACCEPT / REJECT**: Ignored (handled by send/pull initiators).
29pub 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                    // Handle incoming OFFER (someone wants to push a file to us)
74                    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                            // Distinguish rejection from actual failures
89                            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                    // Handle incoming PULL_REQUEST (someone wants to download from us)
125                    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                    // ACCEPT / REJECT are handled by the upload/download initiators
139                }
140            }
141        }
142    })
143}
144
145/// Handle an incoming OFFER: forward to offer channel, wait for decision,
146/// then accept/reject accordingly.
147async 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    // Build the FileOffer
166    let offer = FileOffer {
167        from_peer: from.to_string(),
168        from_name: from.to_string(), // Best we have — the peer ID
169        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    // Emit OfferReceived event (informational)
177    let _ = event_tx.send(FileTransferEvent::OfferReceived(offer.clone()));
178
179    // Create oneshot channel for the decision
180    let (decision_tx, decision_rx) = oneshot::channel::<OfferDecision>();
181    let responder = OfferResponder::new(decision_tx);
182
183    // Send offer + responder to the offer channel
184    offer_tx
185        .send((offer, responder))
186        .map_err(|_| TransferError::Protocol("Offer channel closed".to_string()))?;
187
188    // Wait for decision with 60s timeout
189    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            // Send REJECT message to sender
203            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
225/// Accept an incoming offer: open TCP listener, receive file streaming to
226/// disk, verify SHA-256, and send ACK.
227async 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    // Create parent directories
240    if let Some(parent) = std::path::Path::new(save_path).parent() {
241        tokio::fs::create_dir_all(parent).await?;
242    }
243
244    // Start TCP listener
245    let mut listener = node
246        .listen_tcp(0)
247        .await
248        .map_err(|e| TransferError::Node(format!("Failed to listen TCP: {e}")))?;
249
250    // Send ACCEPT
251    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    // Wait for TCP connection with 30s timeout
267    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    // Read header: [8-byte size][64-byte sha256_hex]
278    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    // Verify the metadata matches the OFFER
287    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    // Stream file data to disk (instead of memory buffer)
301    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        // Throttle progress events to max 4/sec
324        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    // Flush temp file
344    tokio::io::AsyncWriteExt::flush(&mut temp_file).await?;
345
346    // Verify SHA-256
347    let actual_sha = hex::encode(hasher.finalize());
348
349    if actual_sha != sha256 {
350        // Send NACK
351        stream.write_all(&[0x00]).await?;
352        // Clean up temp file
353        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    // Resolve final path: if save_path is a directory, append the file name
361    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    // Create parent directories for the final path
371    if let Some(parent) = std::path::Path::new(&final_path).parent() {
372        tokio::fs::create_dir_all(parent).await?;
373    }
374
375    // Move temp file to final path. Use rename first (fast, atomic on same
376    // filesystem), fall back to copy+delete for cross-device moves.
377    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    // Send ACK and flush
397    stream.write_all(&[0x01]).await?;
398    tokio::io::AsyncWriteExt::flush(&mut stream).await?;
399
400    // Wait briefly for the ACK to propagate through the Go bridge.
401    // The bridge uses bidirectional io.Copy — when we drop the stream,
402    // the bridge may close both directions before the ACK byte has been
403    // forwarded to the sender. This small delay ensures the ACK is delivered.
404    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    // Emit completed event
415    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
427/// Handle a PULL_REQUEST: read file, send OFFER, wait for ACCEPT, stream via TCP.
428async 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    // Read and hash the file
438    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    // Send OFFER and wait for ACCEPT
456    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    // Open TCP to peer and stream the file
507    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    // Write [size][sha256_hex][file_data]
519    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        // Emit progress event (best-effort)
540        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    // Read ACK
553    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    // Emit completed event
567    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}