Skip to main content

protest_extras/generators/
network.rs

1//! Network-related generators
2//!
3//! This module provides generators for:
4//! - IPv4 and IPv6 addresses
5//! - Email addresses
6//! - URLs
7//!
8//! All generators use std library only.
9
10use protest::{Generator, GeneratorConfig};
11use rand::Rng;
12use std::net::{Ipv4Addr, Ipv6Addr};
13
14// ============================================================================
15// IP Address Generator
16// ============================================================================
17
18/// Generator for IP addresses (IPv4 or IPv6)
19#[derive(Debug, Clone)]
20pub enum IpAddressGenerator {
21    /// Generate IPv4 addresses
22    V4,
23    /// Generate IPv6 addresses
24    V6,
25    /// Generate both IPv4 and IPv6 addresses
26    Both,
27}
28
29impl IpAddressGenerator {
30    /// Create a generator for IPv4 addresses only
31    pub fn ipv4() -> Self {
32        Self::V4
33    }
34
35    /// Create a generator for IPv6 addresses only
36    pub fn ipv6() -> Self {
37        Self::V6
38    }
39
40    /// Create a generator for both IPv4 and IPv6 addresses
41    pub fn both() -> Self {
42        Self::Both
43    }
44}
45
46impl Generator<String> for IpAddressGenerator {
47    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
48        match self {
49            IpAddressGenerator::V4 => {
50                let octets: [u8; 4] = [rng.r#gen(), rng.r#gen(), rng.r#gen(), rng.r#gen()];
51                Ipv4Addr::from(octets).to_string()
52            }
53            IpAddressGenerator::V6 => {
54                let segments: [u16; 8] = [
55                    rng.r#gen(),
56                    rng.r#gen(),
57                    rng.r#gen(),
58                    rng.r#gen(),
59                    rng.r#gen(),
60                    rng.r#gen(),
61                    rng.r#gen(),
62                    rng.r#gen(),
63                ];
64                Ipv6Addr::from(segments).to_string()
65            }
66            IpAddressGenerator::Both => {
67                if rng.r#gen::<bool>() {
68                    IpAddressGenerator::V4.generate(rng, _config)
69                } else {
70                    IpAddressGenerator::V6.generate(rng, _config)
71                }
72            }
73        }
74    }
75
76    fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
77        // Try to shrink to simpler IP addresses
78        let mut shrinks = Vec::new();
79
80        // Try localhost
81        if value != "127.0.0.1" && value != "::1" {
82            shrinks.push("127.0.0.1".to_string());
83        }
84
85        // Try zero address
86        if value != "0.0.0.0" && value != "::" {
87            shrinks.push("0.0.0.0".to_string());
88        }
89
90        Box::new(shrinks.into_iter())
91    }
92}
93
94// ============================================================================
95// Email Generator
96// ============================================================================
97
98/// Generator for RFC-compliant email addresses
99#[derive(Debug, Clone)]
100pub struct EmailGenerator {
101    min_local_len: usize,
102    max_local_len: usize,
103    min_domain_len: usize,
104    max_domain_len: usize,
105}
106
107impl EmailGenerator {
108    /// Create a new email generator with default lengths
109    pub fn new() -> Self {
110        Self {
111            min_local_len: 3,
112            max_local_len: 20,
113            min_domain_len: 5,
114            max_domain_len: 30,
115        }
116    }
117
118    /// Create an email generator with custom length constraints
119    pub fn with_lengths(
120        min_local_len: usize,
121        max_local_len: usize,
122        min_domain_len: usize,
123        max_domain_len: usize,
124    ) -> Self {
125        Self {
126            min_local_len,
127            max_local_len,
128            min_domain_len,
129            max_domain_len,
130        }
131    }
132
133    fn generate_local_part(&self, rng: &mut dyn rand::RngCore) -> String {
134        let len = rng.r#gen_range(self.min_local_len..=self.max_local_len);
135        let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789._-";
136
137        (0..len)
138            .map(|i| {
139                if i == 0 {
140                    // First char must be alphanumeric
141                    let alphanumeric = b"abcdefghijklmnopqrstuvwxyz0123456789";
142                    alphanumeric[rng.r#gen_range(0..alphanumeric.len())] as char
143                } else {
144                    valid_chars[rng.r#gen_range(0..valid_chars.len())] as char
145                }
146            })
147            .collect()
148    }
149
150    fn generate_domain(&self, rng: &mut dyn rand::RngCore) -> String {
151        let len = rng.r#gen_range(self.min_domain_len..=self.max_domain_len);
152        let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789-";
153
154        let domain_name: String = (0..len)
155            .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
156            .collect();
157
158        // Add TLD
159        let tlds = ["com", "org", "net", "edu", "gov", "io", "dev"];
160        let tld = tlds[rng.r#gen_range(0..tlds.len())];
161
162        format!("{}.{}", domain_name, tld)
163    }
164}
165
166impl Default for EmailGenerator {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172impl Generator<String> for EmailGenerator {
173    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
174        let local = self.generate_local_part(rng);
175        let domain = self.generate_domain(rng);
176        format!("{}@{}", local, domain)
177    }
178
179    fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
180        let mut shrinks = Vec::new();
181
182        // Try minimal email
183        if value != "a@b.com" {
184            shrinks.push("a@b.com".to_string());
185        }
186
187        // Try to shorten local part
188        if let Some((local, domain)) = value.split_once('@') {
189            if local.len() > 1 {
190                shrinks.push(format!("{}@{}", &local[..local.len() - 1], domain));
191            }
192
193            // Try to shorten domain
194            if let Some((domain_name, tld)) = domain.rsplit_once('.')
195                && domain_name.len() > 1
196            {
197                shrinks.push(format!(
198                    "{}@{}.{}",
199                    local,
200                    &domain_name[..domain_name.len() - 1],
201                    tld
202                ));
203            }
204        }
205
206        Box::new(shrinks.into_iter())
207    }
208}
209
210// ============================================================================
211// URL Generator
212// ============================================================================
213
214/// Generator for URLs with various schemes
215#[derive(Debug, Clone)]
216pub struct UrlGenerator {
217    schemes: Vec<&'static str>,
218    include_port: bool,
219    include_path: bool,
220    include_query: bool,
221}
222
223impl UrlGenerator {
224    /// Create a new URL generator with http/https schemes
225    pub fn new() -> Self {
226        Self {
227            schemes: vec!["http", "https"],
228            include_port: true,
229            include_path: true,
230            include_query: true,
231        }
232    }
233
234    /// Create a URL generator with custom schemes
235    pub fn with_schemes(schemes: Vec<&'static str>) -> Self {
236        Self {
237            schemes,
238            include_port: true,
239            include_path: true,
240            include_query: true,
241        }
242    }
243
244    /// Set whether to include port numbers
245    pub fn with_port(mut self, include: bool) -> Self {
246        self.include_port = include;
247        self
248    }
249
250    /// Set whether to include paths
251    pub fn with_path(mut self, include: bool) -> Self {
252        self.include_path = include;
253        self
254    }
255
256    /// Set whether to include query parameters
257    pub fn with_query(mut self, include: bool) -> Self {
258        self.include_query = include;
259        self
260    }
261
262    fn generate_host(&self, rng: &mut dyn rand::RngCore) -> String {
263        let len = rng.r#gen_range(5..20);
264        let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789-";
265
266        let domain: String = (0..len)
267            .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
268            .collect();
269
270        let tlds = ["com", "org", "net", "io", "dev"];
271        let tld = tlds[rng.r#gen_range(0..tlds.len())];
272
273        format!("{}.{}", domain, tld)
274    }
275
276    fn generate_path(&self, rng: &mut dyn rand::RngCore) -> String {
277        let num_segments = rng.r#gen_range(1..=4);
278        let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789_-";
279
280        let segments: Vec<String> = (0..num_segments)
281            .map(|_| {
282                let len = rng.r#gen_range(3..10);
283                (0..len)
284                    .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
285                    .collect()
286            })
287            .collect();
288
289        format!("/{}", segments.join("/"))
290    }
291
292    fn generate_query(&self, rng: &mut dyn rand::RngCore) -> String {
293        let num_params = rng.r#gen_range(1..=3);
294        let valid_chars = b"abcdefghijklmnopqrstuvwxyz0123456789";
295
296        let params: Vec<String> = (0..num_params)
297            .map(|_| {
298                let key_len = rng.r#gen_range(3..8);
299                let val_len = rng.r#gen_range(3..10);
300
301                let key: String = (0..key_len)
302                    .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
303                    .collect();
304
305                let val: String = (0..val_len)
306                    .map(|_| valid_chars[rng.r#gen_range(0..valid_chars.len())] as char)
307                    .collect();
308
309                format!("{}={}", key, val)
310            })
311            .collect();
312
313        format!("?{}", params.join("&"))
314    }
315}
316
317impl Default for UrlGenerator {
318    fn default() -> Self {
319        Self::new()
320    }
321}
322
323impl Generator<String> for UrlGenerator {
324    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> String {
325        let scheme = self.schemes[rng.r#gen_range(0..self.schemes.len())];
326        let host = self.generate_host(rng);
327
328        let mut url = format!("{}://{}", scheme, host);
329
330        if self.include_port && rng.r#gen::<bool>() {
331            let port: u16 = rng.r#gen_range(1024..65535);
332            url.push_str(&format!(":{}", port));
333        }
334
335        if self.include_path && rng.r#gen::<bool>() {
336            url.push_str(&self.generate_path(rng));
337        }
338
339        if self.include_query && rng.r#gen::<bool>() {
340            url.push_str(&self.generate_query(rng));
341        }
342
343        url
344    }
345
346    fn shrink(&self, value: &String) -> Box<dyn Iterator<Item = String>> {
347        let mut shrinks = Vec::new();
348
349        // Try minimal URL
350        if value != "http://a.com" {
351            shrinks.push("http://a.com".to_string());
352        }
353
354        // Try removing query parameters
355        if let Some(base) = value.split('?').next()
356            && base != value
357        {
358            shrinks.push(base.to_string());
359        }
360
361        // Try removing path
362        if let Some(scheme_and_host) = value
363            .split('/')
364            .take(3)
365            .collect::<Vec<_>>()
366            .join("/")
367            .into()
368            && scheme_and_host != *value
369        {
370            shrinks.push(scheme_and_host);
371        }
372
373        Box::new(shrinks.into_iter())
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380    use rand::thread_rng;
381
382    #[test]
383    fn test_ipv4_generator() {
384        let generator = IpAddressGenerator::ipv4();
385        let mut rng = thread_rng();
386        let config = GeneratorConfig::default();
387
388        for _ in 0..10 {
389            let ip = generator.generate(&mut rng, &config);
390            // Should parse as valid IPv4
391            let parts: Vec<&str> = ip.split('.').collect();
392            assert_eq!(parts.len(), 4);
393            for part in parts {
394                // Parse as u8 to verify each octet is valid (0-255)
395                let _num: u8 = part.parse().unwrap();
396            }
397        }
398    }
399
400    #[test]
401    fn test_email_generator() {
402        let generator = EmailGenerator::new();
403        let mut rng = thread_rng();
404        let config = GeneratorConfig::default();
405
406        for _ in 0..10 {
407            let email = generator.generate(&mut rng, &config);
408            assert!(email.contains('@'));
409            assert_eq!(email.matches('@').count(), 1);
410
411            let parts: Vec<&str> = email.split('@').collect();
412            assert_eq!(parts.len(), 2);
413            assert!(!parts[0].is_empty());
414            assert!(!parts[1].is_empty());
415            assert!(parts[1].contains('.'));
416        }
417    }
418
419    #[test]
420    fn test_url_generator() {
421        let generator = UrlGenerator::new();
422        let mut rng = thread_rng();
423        let config = GeneratorConfig::default();
424
425        for _ in 0..10 {
426            let url = generator.generate(&mut rng, &config);
427            assert!(url.starts_with("http://") || url.starts_with("https://"));
428        }
429    }
430}