Skip to main content

truffle_core/file_transfer/
mod.rs

1//! File transfer subsystem for truffle-core.
2//!
3//! Provides a first-class API for sending and receiving files between peers
4//! in the truffle mesh. The [`FileTransfer`] manager handles:
5//!
6//! - **Sending**: Push a local file to a remote peer.
7//! - **Receiving**: Accept/reject incoming file offers via an offer channel.
8//! - **Pulling**: Request a file from a remote peer.
9//! - **Events**: Subscribe to transfer lifecycle events (progress, completed, failed).
10//!
11//! # Architecture
12//!
13//! The file transfer module sits on top of the Node API, using:
14//! - WS namespace `"ft"` for signaling (offers, accepts, rejects, pull requests)
15//! - Raw TCP streams for bulk data transfer
16//! - SHA-256 for integrity verification
17//!
18//! # Example
19//!
20//! ```ignore
21//! let ft = node.file_transfer();
22//!
23//! // Subscribe to events
24//! let mut events = ft.subscribe();
25//!
26//! // Send a file
27//! let result = ft.send_file("peer-id", "/path/to/file.txt", "file.txt").await?;
28//!
29//! // Receive with manual accept/reject
30//! let mut offers = ft.offer_channel().await;
31//! while let Some((offer, responder)) = offers.recv().await {
32//!     responder.accept("/tmp/received-file.txt");
33//! }
34//! ```
35
36pub 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
52/// Default maximum transfer size: 1 GB.
53const DEFAULT_MAX_TRANSFER_SIZE: u64 = 1_073_741_824;
54
55/// File transfer manager (internal state).
56///
57/// This struct holds the channels and configuration for the file transfer
58/// subsystem. It is stored inside [`Node`] and accessed via
59/// [`Node::file_transfer()`], which returns a [`FileTransferHandle`]
60/// that has access to both this state and the Node's networking capabilities.
61pub(crate) struct FileTransferState {
62    /// Broadcast sender for file transfer events.
63    pub(crate) event_tx: broadcast::Sender<FileTransferEvent>,
64    /// Sender side of the offer channel (held to clone for new receivers).
65    pub(crate) offer_tx: Mutex<Option<mpsc::UnboundedSender<(FileOffer, OfferResponder)>>>,
66    /// Maximum allowed transfer size in bytes.
67    pub(crate) max_transfer_size: AtomicU64,
68    /// Handle to the background receiver task (lazy-started).
69    pub(crate) receiver_handle: Mutex<Option<JoinHandle<()>>>,
70}
71
72impl FileTransferState {
73    /// Create a new file transfer state.
74    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
85/// File transfer handle — the public API for file transfers.
86///
87/// This is a borrowed reference to a [`Node`] that provides file transfer
88/// operations. Obtained via [`Node::file_transfer()`].
89///
90/// Generic over `N: NetworkProvider` to match the [`Node`] it wraps.
91pub struct FileTransfer<'a, N: NetworkProvider + 'static> {
92    /// Reference to the owning node.
93    node: &'a Node<N>,
94}
95
96impl<'a, N: NetworkProvider + 'static> FileTransfer<'a, N> {
97    /// Create a new file transfer handle for the given node.
98    pub(crate) fn new(node: &'a Node<N>) -> Self {
99        Self { node }
100    }
101
102    /// Access the internal file transfer state.
103    fn state(&self) -> &FileTransferState {
104        &self.node.file_transfer_state
105    }
106
107    /// Subscribe to file transfer events.
108    ///
109    /// Returns a broadcast receiver that yields [`FileTransferEvent`]s.
110    /// Multiple subscribers are supported.
111    pub fn subscribe(&self) -> broadcast::Receiver<FileTransferEvent> {
112        self.state().event_tx.subscribe()
113    }
114
115    /// Set the maximum allowed transfer size in bytes.
116    ///
117    /// Transfers exceeding this size will be rejected. Default is 1 GB.
118    pub fn set_max_transfer_size(&self, bytes: u64) {
119        self.state()
120            .max_transfer_size
121            .store(bytes, Ordering::Relaxed);
122    }
123
124    /// Get the current maximum transfer size in bytes.
125    pub fn max_transfer_size(&self) -> u64 {
126        self.state().max_transfer_size.load(Ordering::Relaxed)
127    }
128
129    /// Get a channel for receiving incoming file offers.
130    ///
131    /// Each offer comes with an [`OfferResponder`] that must be used to
132    /// accept or reject the transfer. Calling this starts the background
133    /// receiver listener (lazy start, idempotent).
134    ///
135    /// Only one offer channel can be active at a time. Calling this again
136    /// replaces the previous channel (the old receiver will stop getting offers).
137    ///
138    /// **Important**: The returned receiver must be polled. The `node` reference
139    /// used here must remain valid — store the `Node` in an `Arc` and pass it
140    /// to spawned tasks as needed.
141    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    /// Convenience: auto-accept all incoming offers, saving files to
155    /// `output_dir`.
156    ///
157    /// This starts the receiver listener (lazy start) and spawns a task
158    /// that automatically accepts every offer with the save path set to
159    /// `{output_dir}/{file_name}`.
160    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    /// Convenience: auto-reject all incoming offers.
182    ///
183    /// This starts the receiver listener (lazy start) and spawns a task
184    /// that automatically rejects every offer.
185    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    /// Send a file to a peer.
196    ///
197    /// Hashes the local file, sends an OFFER via the `"ft"` namespace,
198    /// waits for ACCEPT, then streams the file over TCP.
199    ///
200    /// # Errors
201    ///
202    /// Returns [`TransferError`] on I/O failure, rejection, timeout, or
203    /// integrity check failure.
204    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    /// Pull a file from a remote peer.
223    ///
224    /// Sends a PULL_REQUEST, waits for the peer's OFFER, accepts it, then
225    /// receives the file over TCP.
226    ///
227    /// # Errors
228    ///
229    /// Returns [`TransferError`] on I/O failure, rejection, timeout, or
230    /// integrity check failure.
231    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    /// Ensure the background receiver task is running.
248    ///
249    /// The receiver is lazy-started on the first call to `offer_channel()`,
250    /// `auto_accept()`, or `auto_reject()`.
251    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 there's an existing handle, abort it (we're replacing the offer channel)
258        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
270// ---------------------------------------------------------------------------
271// Pull file — download from a remote peer
272// ---------------------------------------------------------------------------
273
274/// Pull (download) a file from a remote peer.
275///
276/// 1. Send PULL_REQUEST via WS
277/// 2. Wait for OFFER from peer
278/// 3. Send ACCEPT
279/// 4. Listen for TCP connection from peer
280/// 5. Read [size][sha256][file_bytes], streaming to disk
281/// 6. Verify SHA-256, rename temp file
282/// 7. Send ACK
283async 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    // 0. Resolve peer_id to the canonical Tailscale node ID.
299    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    // 1. Send PULL_REQUEST and wait for OFFER from peer
306    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    // 3. Start TCP listener, then send ACCEPT with our port
359    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    // 4. Accept incoming TCP connection with timeout
382    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    // 5. Read [8-byte size][64-byte sha256_hex][file_bytes]
394    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]; // hex-encoded SHA-256
399    stream.read_exact(&mut sha_buf).await?;
400    let received_sha = String::from_utf8_lossy(&sha_buf).to_string();
401
402    // Verify the metadata matches the OFFER
403    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    // Stream file data to disk (instead of memory buffer)
417    // Create parent directories if needed
418    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        // Emit progress event (best-effort)
451        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    // Flush temp file
462    AsyncWriteExt::flush(&mut temp_file).await?;
463
464    // 6. Verify SHA-256
465    let actual_sha = hex::encode(hasher.finalize());
466
467    if actual_sha != offer_sha256 {
468        // Send NACK
469        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    // Rename temp to final
478    tokio::fs::rename(&temp_path, local_path).await?;
479
480    // 7. Send ACK
481    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    // Emit completed event
492    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}