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
306    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    // 2. Wait for OFFER from peer
322    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    // 3. Start TCP listener, then send ACCEPT with our port
360    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    // 4. Accept incoming TCP connection with timeout
383    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    // 5. Read [8-byte size][64-byte sha256_hex][file_bytes]
395    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]; // hex-encoded SHA-256
400    stream.read_exact(&mut sha_buf).await?;
401    let received_sha = String::from_utf8_lossy(&sha_buf).to_string();
402
403    // Verify the metadata matches the OFFER
404    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    // Stream file data to disk (instead of memory buffer)
418    // Create parent directories if needed
419    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        // Emit progress event (best-effort)
452        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    // Flush temp file
463    AsyncWriteExt::flush(&mut temp_file).await?;
464
465    // 6. Verify SHA-256
466    let actual_sha = hex::encode(hasher.finalize());
467
468    if actual_sha != offer_sha256 {
469        // Send NACK
470        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    // Rename temp to final
479    tokio::fs::rename(&temp_path, local_path).await?;
480
481    // 7. Send ACK
482    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    // Emit completed event
493    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}