stackforge_core/layer/tls/
builder.rs1use super::types::TlsContentType;
9
10#[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, length: None, fragment: Vec::new(),
26 }
27 }
28}
29
30impl TlsRecordBuilder {
31 pub fn new() -> Self {
33 Self::default()
34 }
35
36 pub fn content_type(mut self, ct: TlsContentType) -> Self {
38 self.content_type = ct.as_u8();
39 self
40 }
41
42 pub fn content_type_raw(mut self, ct: u8) -> Self {
44 self.content_type = ct;
45 self
46 }
47
48 pub fn version(mut self, version: u16) -> Self {
50 self.version = version;
51 self
52 }
53
54 pub fn length(mut self, length: Option<u16>) -> Self {
57 self.length = length;
58 self
59 }
60
61 pub fn fragment(mut self, data: Vec<u8>) -> Self {
63 self.fragment = data;
64 self
65 }
66
67 pub fn fragment_from_slice(mut self, data: &[u8]) -> Self {
69 self.fragment = data.to_vec();
70 self
71 }
72
73 pub fn record_size(&self) -> usize {
75 5 + self.fragment.len()
76 }
77
78 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#[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, description: 0, 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#[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); assert_eq!(record[1..3], [0x03, 0x03]); assert_eq!(u16::from_be_bytes([record[3], record[4]]), 9); assert_eq!(record[5], 0x01); }
189
190 #[test]
191 fn test_build_with_explicit_length() {
192 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); }
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) .version(0x0303)
231 .build();
232
233 assert_eq!(alert[0], 0x15); 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); assert_eq!(alert[6], 40); }
239
240 #[test]
241 fn test_build_ccs() {
242 let ccs = TlsCcsBuilder::new().version(0x0303).build();
243
244 assert_eq!(ccs[0], 0x14); 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}