stackforge_core/layer/tftp/
builder.rs

1//! TFTP packet builder.
2//!
3//! Provides a fluent API for constructing TFTP packets (RFC 1350).
4//!
5//! # Examples
6//!
7//! ```rust
8//! use stackforge_core::layer::tftp::builder::TftpBuilder;
9//!
10//! // Build a Read Request
11//! let pkt = TftpBuilder::new().rrq("file.txt", "octet").build();
12//!
13//! // Build a DATA packet (block 1)
14//! let pkt = TftpBuilder::new().data(1, b"hello world").build();
15//!
16//! // Build an ACK for block 1
17//! let pkt = TftpBuilder::new().ack(1).build();
18//! ```
19
20use super::{OPCODE_ACK, OPCODE_DATA, OPCODE_ERROR, OPCODE_RRQ, OPCODE_WRQ};
21
22/// Builder for TFTP packets.
23#[must_use]
24#[derive(Debug, Clone)]
25pub struct TftpBuilder {
26    opcode: u16,
27    filename: Vec<u8>,
28    mode: Vec<u8>,
29    block_num: u16,
30    payload: Vec<u8>,
31    error_code: u16,
32    error_msg: Vec<u8>,
33}
34
35impl Default for TftpBuilder {
36    fn default() -> Self {
37        Self {
38            opcode: OPCODE_RRQ,
39            filename: b"file.bin".to_vec(),
40            mode: b"octet".to_vec(),
41            block_num: 0,
42            payload: Vec::new(),
43            error_code: 0,
44            error_msg: Vec::new(),
45        }
46    }
47}
48
49impl TftpBuilder {
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    // ========================================================================
55    // Request builders
56    // ========================================================================
57
58    /// Build a Read Request (RRQ) packet.
59    ///
60    /// Mode is typically "netascii", "octet", or "mail".
61    pub fn rrq(mut self, filename: impl Into<Vec<u8>>, mode: impl Into<Vec<u8>>) -> Self {
62        self.opcode = OPCODE_RRQ;
63        self.filename = filename.into();
64        self.mode = mode.into();
65        self
66    }
67
68    /// Build a Write Request (WRQ) packet.
69    pub fn wrq(mut self, filename: impl Into<Vec<u8>>, mode: impl Into<Vec<u8>>) -> Self {
70        self.opcode = OPCODE_WRQ;
71        self.filename = filename.into();
72        self.mode = mode.into();
73        self
74    }
75
76    // ========================================================================
77    // Data / ACK builders
78    // ========================================================================
79
80    /// Build a DATA packet with the given block number and payload.
81    pub fn data(mut self, block_num: u16, payload: impl Into<Vec<u8>>) -> Self {
82        self.opcode = OPCODE_DATA;
83        self.block_num = block_num;
84        self.payload = payload.into();
85        self
86    }
87
88    /// Build an ACK packet for the given block number.
89    pub fn ack(mut self, block_num: u16) -> Self {
90        self.opcode = OPCODE_ACK;
91        self.block_num = block_num;
92        self
93    }
94
95    // ========================================================================
96    // Error builder
97    // ========================================================================
98
99    /// Build an ERROR packet.
100    pub fn error(mut self, error_code: u16, msg: impl Into<Vec<u8>>) -> Self {
101        self.opcode = OPCODE_ERROR;
102        self.error_code = error_code;
103        self.error_msg = msg.into();
104        self
105    }
106
107    /// Build "File not found" error (code 1).
108    pub fn error_file_not_found(self) -> Self {
109        self.error(1, b"File not found".as_ref())
110    }
111
112    /// Build "Access violation" error (code 2).
113    pub fn error_access_violation(self) -> Self {
114        self.error(2, b"Access violation".as_ref())
115    }
116
117    /// Build "Disk full" error (code 3).
118    pub fn error_disk_full(self) -> Self {
119        self.error(3, b"Disk full or allocation exceeded".as_ref())
120    }
121
122    /// Build "Illegal operation" error (code 4).
123    pub fn error_illegal_op(self) -> Self {
124        self.error(4, b"Illegal TFTP operation".as_ref())
125    }
126
127    /// Build "File already exists" error (code 6).
128    pub fn error_file_exists(self) -> Self {
129        self.error(6, b"File already exists".as_ref())
130    }
131
132    // ========================================================================
133    // Build
134    // ========================================================================
135
136    /// Serialize this TFTP packet to bytes.
137    #[must_use]
138    pub fn build(&self) -> Vec<u8> {
139        match self.opcode {
140            OPCODE_RRQ | OPCODE_WRQ => self.build_request(),
141            OPCODE_DATA => self.build_data(),
142            OPCODE_ACK => self.build_ack(),
143            OPCODE_ERROR => self.build_error(),
144            _ => vec![],
145        }
146    }
147
148    fn build_request(&self) -> Vec<u8> {
149        let mut out = Vec::new();
150        out.extend_from_slice(&self.opcode.to_be_bytes());
151        out.extend_from_slice(&self.filename);
152        out.push(0); // null terminator
153        out.extend_from_slice(&self.mode);
154        out.push(0); // null terminator
155        out
156    }
157
158    fn build_data(&self) -> Vec<u8> {
159        let mut out = Vec::with_capacity(4 + self.payload.len());
160        out.extend_from_slice(&self.opcode.to_be_bytes());
161        out.extend_from_slice(&self.block_num.to_be_bytes());
162        out.extend_from_slice(&self.payload);
163        out
164    }
165
166    fn build_ack(&self) -> Vec<u8> {
167        let mut out = Vec::with_capacity(4);
168        out.extend_from_slice(&self.opcode.to_be_bytes());
169        out.extend_from_slice(&self.block_num.to_be_bytes());
170        out
171    }
172
173    fn build_error(&self) -> Vec<u8> {
174        let mut out = Vec::new();
175        out.extend_from_slice(&self.opcode.to_be_bytes());
176        out.extend_from_slice(&self.error_code.to_be_bytes());
177        out.extend_from_slice(&self.error_msg);
178        out.push(0); // null terminator
179        out
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use crate::layer::LayerIndex;
187    use crate::layer::LayerKind;
188    use crate::layer::tftp::{
189        OPCODE_ACK, OPCODE_DATA, OPCODE_ERROR, OPCODE_RRQ, OPCODE_WRQ, TftpLayer,
190    };
191
192    fn parse(data: Vec<u8>) -> (TftpLayer, Vec<u8>) {
193        let len = data.len();
194        let layer = TftpLayer::new(LayerIndex::new(LayerKind::Tftp, 0, len));
195        (layer, data)
196    }
197
198    #[test]
199    fn test_build_rrq() {
200        let pkt = TftpBuilder::new().rrq("test.txt", "octet").build();
201        let (layer, buf) = parse(pkt);
202        assert_eq!(layer.opcode(&buf).unwrap(), OPCODE_RRQ);
203        assert_eq!(layer.filename(&buf).unwrap(), "test.txt");
204        assert_eq!(layer.mode(&buf).unwrap(), "octet");
205    }
206
207    #[test]
208    fn test_build_wrq() {
209        let pkt = TftpBuilder::new().wrq("upload.txt", "netascii").build();
210        let (layer, buf) = parse(pkt);
211        assert_eq!(layer.opcode(&buf).unwrap(), OPCODE_WRQ);
212        assert_eq!(layer.filename(&buf).unwrap(), "upload.txt");
213        assert_eq!(layer.mode(&buf).unwrap(), "netascii");
214    }
215
216    #[test]
217    fn test_build_data() {
218        let payload = b"Hello TFTP world!";
219        let pkt = TftpBuilder::new().data(1, payload.as_ref()).build();
220        let (layer, buf) = parse(pkt);
221        assert_eq!(layer.opcode(&buf).unwrap(), OPCODE_DATA);
222        assert_eq!(layer.block_num(&buf).unwrap(), 1);
223        assert_eq!(layer.data(&buf).unwrap(), payload);
224    }
225
226    #[test]
227    fn test_build_ack() {
228        let pkt = TftpBuilder::new().ack(3).build();
229        let (layer, buf) = parse(pkt);
230        assert_eq!(layer.opcode(&buf).unwrap(), OPCODE_ACK);
231        assert_eq!(layer.block_num(&buf).unwrap(), 3);
232        assert_eq!(buf.len(), 4);
233    }
234
235    #[test]
236    fn test_build_error() {
237        let pkt = TftpBuilder::new().error_file_not_found().build();
238        let (layer, buf) = parse(pkt);
239        assert_eq!(layer.opcode(&buf).unwrap(), OPCODE_ERROR);
240        assert_eq!(layer.error_code(&buf).unwrap(), 1);
241        assert_eq!(layer.error_msg(&buf).unwrap(), "File not found");
242    }
243
244    #[test]
245    fn test_build_custom_error() {
246        let pkt = TftpBuilder::new()
247            .error(0, b"Custom error message".as_ref())
248            .build();
249        let (layer, buf) = parse(pkt);
250        assert_eq!(layer.error_code(&buf).unwrap(), 0);
251        assert_eq!(layer.error_msg(&buf).unwrap(), "Custom error message");
252    }
253
254    #[test]
255    fn test_data_roundtrip() {
256        let large_data: Vec<u8> = (0u8..=255u8).collect();
257        let pkt = TftpBuilder::new().data(42, large_data.clone()).build();
258        let (layer, buf) = parse(pkt);
259        assert_eq!(layer.block_num(&buf).unwrap(), 42);
260        assert_eq!(layer.data(&buf).unwrap(), large_data);
261    }
262}
stackforge_core/layer/tftp/builder.rs

stackforge_core/layer/tftp/
builder.rs
