Skip to main content

stackforge_core/layer/tls/
builder.rs

1//! TLS Record builder.
2//!
3//! Provides a builder for constructing TLS record layer packets.
4//! Supports the "Permissive Parser, Explicit Builder" pattern:
5//! if `length` is set to `Some(val)`, that value is used directly
6//! (even if incorrect), enabling fuzzing and security testing.
7
8use super::types::TlsContentType;
9
10/// Builder for TLS record layer packets.
11#[derive(Debug, Clone)]
12pub struct TlsRecordBuilder {
13    content_type: u8,
14    version: u16,
15    length: Option<u16>,
16    fragment: Vec<u8>,
17}
18
19impl Default for TlsRecordBuilder {
20    fn default() -> Self {
21        Self {
22            content_type: TlsContentType::Handshake.as_u8(),
23            version: 0x0303, // TLS 1.2
24            length: None,    // auto-calculate
25            fragment: Vec::new(),
26        }
27    }
28}
29
30impl TlsRecordBuilder {
31    /// Create a new TLS record builder with defaults.
32    #[must_use]
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Set the content type.
38    #[must_use]
39    pub fn content_type(mut self, ct: TlsContentType) -> Self {
40        self.content_type = ct.as_u8();
41        self
42    }
43
44    /// Set the content type as raw u8 (for fuzzing/testing).
45    #[must_use]
46    pub fn content_type_raw(mut self, ct: u8) -> Self {
47        self.content_type = ct;
48        self
49    }
50
51    /// Set the protocol version.
52    #[must_use]
53    pub fn version(mut self, version: u16) -> Self {
54        self.version = version;
55        self
56    }
57
58    /// Set the length field explicitly (overrides auto-calculation).
59    /// Pass `None` to auto-calculate from fragment size.
60    #[must_use]
61    pub fn length(mut self, length: Option<u16>) -> Self {
62        self.length = length;
63        self
64    }
65
66    /// Set the fragment data.
67    #[must_use]
68    pub fn fragment(mut self, data: Vec<u8>) -> Self {
69        self.fragment = data;
70        self
71    }
72
73    /// Set the fragment data from a slice.
74    #[must_use]
75    pub fn fragment_from_slice(mut self, data: &[u8]) -> Self {
76        self.fragment = data.to_vec();
77        self
78    }
79
80    /// Get the total size of the built record.
81    #[must_use]
82    pub fn record_size(&self) -> usize {
83        5 + self.fragment.len()
84    }
85
86    /// Build the TLS record into bytes.
87    #[must_use]
88    pub fn build(&self) -> Vec<u8> {
89        let frag_len = self.length.unwrap_or(self.fragment.len() as u16);
90        let mut out = Vec::with_capacity(5 + self.fragment.len());
91
92        out.push(self.content_type);
93        out.extend_from_slice(&self.version.to_be_bytes());
94        out.extend_from_slice(&frag_len.to_be_bytes());
95        out.extend_from_slice(&self.fragment);
96
97        out
98    }
99}
100
101/// Builder for TLS Alert messages.
102#[derive(Debug, Clone)]
103pub struct TlsAlertBuilder {
104    level: u8,
105    description: u8,
106    version: u16,
107}
108
109impl Default for TlsAlertBuilder {
110    fn default() -> Self {
111        Self {
112            level: 2,       // fatal
113            description: 0, // close_notify
114            version: 0x0303,
115        }
116    }
117}
118
119impl TlsAlertBuilder {
120    #[must_use]
121    pub fn new() -> Self {
122        Self::default()
123    }
124
125    #[must_use]
126    pub fn level(mut self, level: u8) -> Self {
127        self.level = level;
128        self
129    }
130
131    #[must_use]
132    pub fn description(mut self, desc: u8) -> Self {
133        self.description = desc;
134        self
135    }
136
137    #[must_use]
138    pub fn version(mut self, version: u16) -> Self {
139        self.version = version;
140        self
141    }
142
143    #[must_use]
144    pub fn build(&self) -> Vec<u8> {
145        TlsRecordBuilder::new()
146            .content_type(TlsContentType::Alert)
147            .version(self.version)
148            .fragment(vec![self.level, self.description])
149            .build()
150    }
151}
152
153/// Builder for TLS `ChangeCipherSpec` messages.
154#[derive(Debug, Clone)]
155pub struct TlsCcsBuilder {
156    version: u16,
157}
158
159impl Default for TlsCcsBuilder {
160    fn default() -> Self {
161        Self { version: 0x0303 }
162    }
163}
164
165impl TlsCcsBuilder {
166    #[must_use]
167    pub fn new() -> Self {
168        Self::default()
169    }
170
171    #[must_use]
172    pub fn version(mut self, version: u16) -> Self {
173        self.version = version;
174        self
175    }
176
177    #[must_use]
178    pub fn build(&self) -> Vec<u8> {
179        TlsRecordBuilder::new()
180            .content_type(TlsContentType::ChangeCipherSpec)
181            .version(self.version)
182            .fragment(vec![0x01])
183            .build()
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use crate::layer::tls::record::TlsLayer;
191    use crate::layer::{LayerIndex, LayerKind};
192
193    #[test]
194    fn test_build_handshake_record() {
195        let record = TlsRecordBuilder::new()
196            .content_type(TlsContentType::Handshake)
197            .version(0x0303)
198            .fragment(vec![0x01, 0x00, 0x00, 0x05, 0x03, 0x03, 0x00, 0x00, 0x00])
199            .build();
200
201        assert_eq!(record[0], 0x16); // Handshake
202        assert_eq!(record[1..3], [0x03, 0x03]); // TLS 1.2
203        assert_eq!(u16::from_be_bytes([record[3], record[4]]), 9); // length
204        assert_eq!(record[5], 0x01); // ClientHello
205    }
206
207    #[test]
208    fn test_build_with_explicit_length() {
209        // Set length to 100 even though fragment is only 5 bytes (for fuzzing)
210        let record = TlsRecordBuilder::new()
211            .content_type(TlsContentType::Handshake)
212            .length(Some(100))
213            .fragment(vec![0x01, 0x02, 0x03, 0x04, 0x05])
214            .build();
215
216        assert_eq!(u16::from_be_bytes([record[3], record[4]]), 100);
217        assert_eq!(record.len(), 10); // 5 header + 5 fragment
218    }
219
220    #[test]
221    fn test_roundtrip() {
222        let fragment = vec![0x01, 0x00, 0x00, 0x01, 0x00];
223        let record = TlsRecordBuilder::new()
224            .content_type(TlsContentType::Handshake)
225            .version(0x0301)
226            .fragment(fragment.clone())
227            .build();
228
229        let layer = TlsLayer {
230            index: LayerIndex::new(LayerKind::Tls, 0, record.len()),
231        };
232
233        assert_eq!(
234            layer.content_type(&record).unwrap(),
235            TlsContentType::Handshake
236        );
237        assert_eq!(layer.version(&record).unwrap(), 0x0301);
238        assert_eq!(layer.length(&record).unwrap(), 5);
239        assert_eq!(layer.fragment(&record), &fragment);
240    }
241
242    #[test]
243    fn test_build_alert() {
244        let alert = TlsAlertBuilder::new()
245            .level(2)
246            .description(40) // handshake_failure
247            .version(0x0303)
248            .build();
249
250        assert_eq!(alert[0], 0x15); // Alert
251        assert_eq!(alert[1..3], [0x03, 0x03]);
252        assert_eq!(u16::from_be_bytes([alert[3], alert[4]]), 2);
253        assert_eq!(alert[5], 2); // fatal
254        assert_eq!(alert[6], 40); // handshake_failure
255    }
256
257    #[test]
258    fn test_build_ccs() {
259        let ccs = TlsCcsBuilder::new().version(0x0303).build();
260
261        assert_eq!(ccs[0], 0x14); // ChangeCipherSpec
262        assert_eq!(ccs[1..3], [0x03, 0x03]);
263        assert_eq!(u16::from_be_bytes([ccs[3], ccs[4]]), 1);
264        assert_eq!(ccs[5], 0x01);
265    }
266
267    #[test]
268    fn test_record_size() {
269        let builder = TlsRecordBuilder::new().fragment(vec![0; 100]);
270        assert_eq!(builder.record_size(), 105);
271    }
272}