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