truffle-core 0.3.19

Truffle mesh networking core (clean architecture)
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
//! File transfer subsystem for truffle-core.
//!
//! Provides a first-class API for sending and receiving files between peers
//! in the truffle mesh. The [`FileTransfer`] manager handles:
//!
//! - **Sending**: Push a local file to a remote peer.
//! - **Receiving**: Accept/reject incoming file offers via an offer channel.
//! - **Pulling**: Request a file from a remote peer.
//! - **Events**: Subscribe to transfer lifecycle events (progress, completed, failed).
//!
//! # Architecture
//!
//! The file transfer module sits on top of the Node API, using:
//! - WS namespace `"ft"` for signaling (offers, accepts, rejects, pull requests)
//! - Raw TCP streams for bulk data transfer
//! - SHA-256 for integrity verification
//!
//! # Example
//!
//! ```ignore
//! let ft = node.file_transfer();
//!
//! // Subscribe to events
//! let mut events = ft.subscribe();
//!
//! // Send a file
//! let result = ft.send_file("peer-id", "/path/to/file.txt", "file.txt").await?;
//!
//! // Receive with manual accept/reject
//! let mut offers = ft.offer_channel().await;
//! while let Some((offer, responder)) = offers.recv().await {
//!     responder.accept("/tmp/received-file.txt");
//! }
//! ```

pub mod types;
pub mod sender;
pub mod receiver;

pub use types::*;

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use tokio::sync::{broadcast, mpsc, Mutex};
use tokio::task::JoinHandle;
use tracing::info;

use crate::network::NetworkProvider;
use crate::node::Node;

/// Default maximum transfer size: 1 GB.
const DEFAULT_MAX_TRANSFER_SIZE: u64 = 1_073_741_824;

/// File transfer manager (internal state).
///
/// This struct holds the channels and configuration for the file transfer
/// subsystem. It is stored inside [`Node`] and accessed via
/// [`Node::file_transfer()`], which returns a [`FileTransferHandle`]
/// that has access to both this state and the Node's networking capabilities.
pub(crate) struct FileTransferState {
    /// Broadcast sender for file transfer events.
    pub(crate) event_tx: broadcast::Sender<FileTransferEvent>,
    /// Sender side of the offer channel (held to clone for new receivers).
    pub(crate) offer_tx: Mutex<Option<mpsc::UnboundedSender<(FileOffer, OfferResponder)>>>,
    /// Maximum allowed transfer size in bytes.
    pub(crate) max_transfer_size: AtomicU64,
    /// Handle to the background receiver task (lazy-started).
    pub(crate) receiver_handle: Mutex<Option<JoinHandle<()>>>,
}

impl FileTransferState {
    /// Create a new file transfer state.
    pub(crate) fn new() -> Self {
        let (event_tx, _) = broadcast::channel(256);
        Self {
            event_tx,
            offer_tx: Mutex::new(None),
            max_transfer_size: AtomicU64::new(DEFAULT_MAX_TRANSFER_SIZE),
            receiver_handle: Mutex::new(None),
        }
    }
}

/// File transfer handle — the public API for file transfers.
///
/// This is a borrowed reference to a [`Node`] that provides file transfer
/// operations. Obtained via [`Node::file_transfer()`].
///
/// Generic over `N: NetworkProvider` to match the [`Node`] it wraps.
pub struct FileTransfer<'a, N: NetworkProvider + 'static> {
    /// Reference to the owning node.
    node: &'a Node<N>,
}

impl<'a, N: NetworkProvider + 'static> FileTransfer<'a, N> {
    /// Create a new file transfer handle for the given node.
    pub(crate) fn new(node: &'a Node<N>) -> Self {
        Self { node }
    }

    /// Access the internal file transfer state.
    fn state(&self) -> &FileTransferState {
        &self.node.file_transfer_state
    }

    /// Subscribe to file transfer events.
    ///
    /// Returns a broadcast receiver that yields [`FileTransferEvent`]s.
    /// Multiple subscribers are supported.
    pub fn subscribe(&self) -> broadcast::Receiver<FileTransferEvent> {
        self.state().event_tx.subscribe()
    }

    /// Set the maximum allowed transfer size in bytes.
    ///
    /// Transfers exceeding this size will be rejected. Default is 1 GB.
    pub fn set_max_transfer_size(&self, bytes: u64) {
        self.state()
            .max_transfer_size
            .store(bytes, Ordering::Relaxed);
    }

    /// Get the current maximum transfer size in bytes.
    pub fn max_transfer_size(&self) -> u64 {
        self.state().max_transfer_size.load(Ordering::Relaxed)
    }

    /// Get a channel for receiving incoming file offers.
    ///
    /// Each offer comes with an [`OfferResponder`] that must be used to
    /// accept or reject the transfer. Calling this starts the background
    /// receiver listener (lazy start, idempotent).
    ///
    /// Only one offer channel can be active at a time. Calling this again
    /// replaces the previous channel (the old receiver will stop getting offers).
    ///
    /// **Important**: The returned receiver must be polled. The `node` reference
    /// used here must remain valid — store the `Node` in an `Arc` and pass it
    /// to spawned tasks as needed.
    pub async fn offer_channel(
        &self,
        node: Arc<Node<N>>,
    ) -> mpsc::UnboundedReceiver<(FileOffer, OfferResponder)> {
        let (tx, rx) = mpsc::unbounded_channel();
        {
            let mut offer_tx = self.state().offer_tx.lock().await;
            *offer_tx = Some(tx.clone());
        }
        self.ensure_receiver_started(node, tx).await;
        rx
    }

    /// Convenience: auto-accept all incoming offers, saving files to
    /// `output_dir`.
    ///
    /// This starts the receiver listener (lazy start) and spawns a task
    /// that automatically accepts every offer with the save path set to
    /// `{output_dir}/{file_name}`.
    pub async fn auto_accept(&self, node: Arc<Node<N>>, output_dir: &str) {
        let mut rx = self.offer_channel(node).await;
        let output_dir = output_dir.to_string();

        tokio::spawn(async move {
            while let Some((offer, responder)) = rx.recv().await {
                let dest = if offer.suggested_path.is_empty() || offer.suggested_path == "." {
                    format!("{}/{}", output_dir, offer.file_name)
                } else {
                    offer.suggested_path.clone()
                };
                info!(
                    file = offer.file_name.as_str(),
                    dest = dest.as_str(),
                    "Auto-accepting file offer"
                );
                responder.accept(&dest);
            }
        });
    }

    /// Convenience: auto-reject all incoming offers.
    ///
    /// This starts the receiver listener (lazy start) and spawns a task
    /// that automatically rejects every offer.
    pub async fn auto_reject(&self, node: Arc<Node<N>>) {
        let mut rx = self.offer_channel(node).await;

        tokio::spawn(async move {
            while let Some((_offer, responder)) = rx.recv().await {
                responder.reject("auto-rejected");
            }
        });
    }

    /// Send a file to a peer.
    ///
    /// Hashes the local file, sends an OFFER via the `"ft"` namespace,
    /// waits for ACCEPT, then streams the file over TCP.
    ///
    /// # Errors
    ///
    /// Returns [`TransferError`] on I/O failure, rejection, timeout, or
    /// integrity check failure.
    pub async fn send_file(
        &self,
        peer_id: &str,
        local_path: &str,
        remote_path: &str,
    ) -> Result<TransferResult, TransferError> {
        let max_size = self.state().max_transfer_size.load(Ordering::Relaxed);
        sender::send_file(
            self.node,
            peer_id,
            local_path,
            remote_path,
            max_size,
            &self.state().event_tx,
        )
        .await
    }

    /// Pull a file from a remote peer.
    ///
    /// Sends a PULL_REQUEST, waits for the peer's OFFER, accepts it, then
    /// receives the file over TCP.
    ///
    /// # Errors
    ///
    /// Returns [`TransferError`] on I/O failure, rejection, timeout, or
    /// integrity check failure.
    pub async fn pull_file(
        &self,
        peer_id: &str,
        remote_path: &str,
        local_path: &str,
    ) -> Result<TransferResult, TransferError> {
        pull_file(
            self.node,
            peer_id,
            remote_path,
            local_path,
            &self.state().event_tx,
        )
        .await
    }

    /// Ensure the background receiver task is running.
    ///
    /// The receiver is lazy-started on the first call to `offer_channel()`,
    /// `auto_accept()`, or `auto_reject()`.
    async fn ensure_receiver_started(
        &self,
        node: Arc<Node<N>>,
        offer_tx: mpsc::UnboundedSender<(FileOffer, OfferResponder)>,
    ) {
        let mut handle = self.state().receiver_handle.lock().await;
        // If there's an existing handle, abort it (we're replacing the offer channel)
        if let Some(h) = handle.take() {
            h.abort();
        }
        let h = receiver::spawn_receive_handler(
            node,
            offer_tx,
            self.state().event_tx.clone(),
        );
        *handle = Some(h);
    }
}

// ---------------------------------------------------------------------------
// Pull file — download from a remote peer
// ---------------------------------------------------------------------------

/// Pull (download) a file from a remote peer.
///
/// 1. Send PULL_REQUEST via WS
/// 2. Wait for OFFER from peer
/// 3. Send ACCEPT
/// 4. Listen for TCP connection from peer
/// 5. Read [size][sha256][file_bytes], streaming to disk
/// 6. Verify SHA-256, rename temp file
/// 7. Send ACK
async fn pull_file<N: NetworkProvider + 'static>(
    node: &Node<N>,
    peer_id: &str,
    remote_path: &str,
    local_path: &str,
    event_tx: &broadcast::Sender<FileTransferEvent>,
) -> Result<TransferResult, TransferError> {
    use sha2::{Digest, Sha256};
    use std::time::Instant;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    let start = Instant::now();
    let token = uuid::Uuid::new_v4().to_string();
    let requester_id = node.local_info().id;

    // 0. Resolve peer_id to the canonical Tailscale node ID.
    let peer_id = node
        .resolve_peer_id(peer_id)
        .await
        .map_err(|e| TransferError::Node(e.to_string()))?;
    let peer_id = peer_id.as_str();

    // 1. Send PULL_REQUEST and wait for OFFER from peer
    let pull_req = FtMessage::PullRequest {
        path: remote_path.to_string(),
        requester_id,
        token: token.clone(),
    };

    let pull_payload = serde_json::to_value(&pull_req)
        .map_err(|e| TransferError::Protocol(format!("Failed to serialize pull request: {e}")))?;

    tracing::info!(peer = peer_id, path = remote_path, "Sending PULL_REQUEST");

    let (offer_sha256, offer_size, offer_token, file_name) =
        crate::request_reply::send_and_wait(
            node,
            peer_id,
            "ft",
            "pull_request",
            &pull_payload,
            std::time::Duration::from_secs(30),
            |msg| {
                if msg.from != peer_id {
                    return None;
                }
                let ft_msg: FtMessage = serde_json::from_value(msg.payload.clone()).ok()?;
                match ft_msg {
                    FtMessage::Offer {
                        file_name,
                        sha256,
                        size,
                        token: offer_token,
                        ..
                    } => {
                        tracing::info!(size = size, "Received OFFER from peer");
                        Some(Ok((sha256, size, offer_token, file_name)))
                    }
                    FtMessage::Reject { reason, .. } => {
                        Some(Err(TransferError::Rejected(reason)))
                    }
                    _ => None,
                }
            },
        )
        .await
        .map_err(|e| match e {
            crate::request_reply::RequestError::Timeout => TransferError::Timeout,
            crate::request_reply::RequestError::Send(e) => TransferError::Node(e.to_string()),
            crate::request_reply::RequestError::ChannelClosed => {
                TransferError::Protocol("Channel closed".into())
            }
        })?
        .map_err(|e| e)?;

    // 3. Start TCP listener, then send ACCEPT with our port
    let mut listener = node
        .listen_tcp(0)
        .await
        .map_err(|e| TransferError::Node(format!("Failed to listen TCP: {e}")))?;

    let accept = FtMessage::Accept {
        token: offer_token,
        tcp_port: listener.port,
    };

    let accept_payload = serde_json::to_value(&accept)
        .map_err(|e| TransferError::Protocol(format!("Failed to serialize accept: {e}")))?;

    node.send_typed(peer_id, "ft", "accept", &accept_payload)
        .await
        .map_err(|e| TransferError::Node(e.to_string()))?;

    tracing::info!(
        port = listener.port,
        "Sent ACCEPT, listening for TCP connection"
    );

    // 4. Accept incoming TCP connection with timeout
    let incoming = tokio::time::timeout(
        tokio::time::Duration::from_secs(30),
        listener.accept(),
    )
    .await
    .map_err(|_| TransferError::Timeout)?
    .ok_or_else(|| TransferError::Protocol("Listener closed before accepting".to_string()))?;

    let mut stream = incoming.stream;
    tracing::info!("TCP connection accepted from peer");

    // 5. Read [8-byte size][64-byte sha256_hex][file_bytes]
    let mut size_buf = [0u8; 8];
    stream.read_exact(&mut size_buf).await?;
    let file_size = u64::from_be_bytes(size_buf);

    let mut sha_buf = [0u8; 64]; // hex-encoded SHA-256
    stream.read_exact(&mut sha_buf).await?;
    let received_sha = String::from_utf8_lossy(&sha_buf).to_string();

    // Verify the metadata matches the OFFER
    if received_sha != offer_sha256 {
        return Err(TransferError::IntegrityError {
            expected: offer_sha256,
            actual: received_sha,
        });
    }

    if file_size != offer_size {
        return Err(TransferError::Protocol(format!(
            "Size mismatch: offer said {offer_size}, stream header says {file_size}"
        )));
    }

    // Stream file data to disk (instead of memory buffer)
    // Create parent directories if needed
    if let Some(parent) = std::path::Path::new(local_path).parent() {
        tokio::fs::create_dir_all(parent).await?;
    }

    let temp_path = format!("{local_path}.truffle-tmp");
    let mut temp_file = tokio::fs::File::create(&temp_path).await?;
    let mut hasher = Sha256::new();
    let mut bytes_received: u64 = 0;
    let progress_start = Instant::now();
    let mut buf = vec![0u8; 64 * 1024];

    while bytes_received < file_size {
        let to_read = ((file_size - bytes_received) as usize).min(buf.len());
        let n = stream.read(&mut buf[..to_read]).await?;
        if n == 0 {
            tokio::fs::remove_file(&temp_path).await.ok();
            return Err(TransferError::Io(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                format!("Connection closed after {bytes_received}/{file_size} bytes"),
            )));
        }
        hasher.update(&buf[..n]);
        AsyncWriteExt::write_all(&mut temp_file, &buf[..n]).await?;
        bytes_received += n as u64;

        let elapsed = progress_start.elapsed().as_secs_f64();
        let speed = if elapsed > 0.0 {
            bytes_received as f64 / elapsed
        } else {
            0.0
        };

        // Emit progress event (best-effort)
        let _ = event_tx.send(FileTransferEvent::Progress(TransferProgress {
            token: token.clone(),
            direction: TransferDirection::Receive,
            file_name: file_name.clone(),
            bytes_transferred: bytes_received,
            total_bytes: file_size,
            speed_bps: speed,
        }));
    }

    // Flush temp file
    AsyncWriteExt::flush(&mut temp_file).await?;

    // 6. Verify SHA-256
    let actual_sha = hex::encode(hasher.finalize());

    if actual_sha != offer_sha256 {
        // Send NACK
        stream.write_all(&[0x00]).await?;
        tokio::fs::remove_file(&temp_path).await.ok();
        return Err(TransferError::IntegrityError {
            expected: offer_sha256,
            actual: actual_sha,
        });
    }

    // Rename temp to final
    tokio::fs::rename(&temp_path, local_path).await?;

    // 7. Send ACK
    stream.write_all(&[0x01]).await?;

    let elapsed = start.elapsed().as_secs_f64();
    tracing::info!(
        bytes = file_size,
        elapsed_ms = (elapsed * 1000.0) as u64,
        path = local_path,
        "Download complete"
    );

    // Emit completed event
    let _ = event_tx.send(FileTransferEvent::Completed {
        token,
        direction: TransferDirection::Receive,
        file_name,
        bytes_transferred: file_size,
        sha256: actual_sha.clone(),
        elapsed_secs: elapsed,
    });

    Ok(TransferResult {
        bytes_transferred: file_size,
        sha256: actual_sha,
        elapsed_secs: elapsed,
    })
}