Skip to main content

utils/md5/
bruteforce.rs

1use crate::{md5, multithreading, multiversion};
2use core::fmt::NumBuffer;
3use std::array;
4use std::num::NonZeroUsize;
5use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
6
7/// Brute force hashes of a prefix followed by an increasing integer.
8///
9/// This function calls the predicate repeatedly until it returns true from a pool of worker threads
10/// each using the [`FASTEST`](super::FASTEST) supported vectorized MD5 implementation to hash
11/// multiple inputs at once.
12///
13/// When `additional_hashes` is zero, the predicate will be called with:
14/// ```ignore
15/// predicate(i, hash(prefix + i.to_string()))
16/// ```
17///
18/// When `additional_hashes` is more than zero, key stretching is used. For example, passing 2 will
19/// cause the predicate to be called with:
20/// ```ignore
21/// predicate(i, hash(to_hex(hash(to_hex(hash(prefix + i.to_string()))))))
22/// ```
23pub fn find_hash_with_appended_count(
24    prefix: &str,
25    additional_hashes: u32,
26    predicate: impl Fn(u32, [u32; 4]) -> bool + Copy + Sync,
27) {
28    let counter = AtomicU32::new(0);
29    let done = AtomicBool::new(false);
30    multithreading::worker_pool(|| {
31        worker(
32            prefix.as_bytes(),
33            additional_hashes,
34            &predicate,
35            &counter,
36            &done,
37        );
38    });
39}
40
41multiversion! {
42    use {crate::simd::*, crate::md5::*};
43
44    #[dyn_dispatch = md5::FASTEST]
45    #[expect(clippy::cast_possible_truncation)]
46    fn worker(
47        prefix: &[u8],
48        additional_hashes: u32,
49        predicate: impl Fn(u32, [u32; 4]) -> bool + Copy + Send,
50        counter: &AtomicU32,
51        done: &AtomicBool,
52    ) {
53        let lane_size = prefix.len() + 10; // u32::MAX is 10 digits long
54
55        let mut buf = vec![0u8; lane_size * U32Vector::LANES];
56        for i in 0..prefix.len() {
57            buf[i * U32Vector::LANES..(i + 1) * U32Vector::LANES].fill(prefix[i]);
58        }
59
60        let mut single = vec![0u8; lane_size];
61        single[..prefix.len()].copy_from_slice(prefix);
62
63        let batch_size = if additional_hashes > 0 {
64            U32Vector::LANES as u32
65        }  else {
66            1000u32.next_multiple_of(U32Vector::LANES as u32)
67        };
68
69        while !done.load(Ordering::Acquire) {
70            let batch_start = counter.fetch_add(batch_size, Ordering::AcqRel);
71            for base in (batch_start..batch_start + batch_size).step_by(U32Vector::LANES) {
72                let mut hashes = match u32_to_ascii_multi(&mut buf[U32Vector::LANES * prefix.len()..], base) {
73                    Some(length) => hash(&buf[..U32Vector::LANES * (prefix.len() + length.get())]),
74                    None => {
75                        // Lengths are different
76                        array::from_fn(|i| {
77                            let mut num_buf = NumBuffer::new();
78                            let digits = (base + i as u32).format_into(&mut num_buf).as_bytes();
79                            let end = prefix.len() + digits.len();
80                            single[prefix.len()..end].copy_from_slice(digits);
81                            md5::hash(&single[..end])
82                        })
83                    }
84                };
85
86                let mut hex_buf = [0u8; 32 * U32Vector::LANES];
87                for _ in 0..additional_hashes {
88                    for i in 0..U32Vector::LANES {
89                        let hex = md5::to_hex(hashes[i]);
90                        for h in 0..32 {
91                            hex_buf[h * U32Vector::LANES + i] = hex[h];
92                        }
93                    }
94                    hashes = hash(&hex_buf);
95                }
96
97                for (i, &hash) in hashes.iter().enumerate() {
98                    if predicate(base + i as u32, hash) {
99                        // Don't return early. For example, in 2016 day 5, this block of a thousand
100                        // could include more than one password letter. If we break early after
101                        // completing the password with the first letter, we won't process the
102                        // second letter which may have a lower count than the letter stored at that
103                        // position.
104                        done.store(true, Ordering::Release);
105                    }
106                }
107            }
108        }
109    }
110
111    #[inline]
112    #[expect(clippy::cast_possible_truncation)]
113    pub fn u32_to_ascii_multi(buf: &mut [u8], base: u32) -> Option<NonZeroUsize> {
114        assert!(buf.len() >= U32Vector::LANES * 10);
115
116        let length = 1 + base.checked_ilog10().unwrap_or(0) as usize;
117        assert!(length <= 10);
118
119        let mut values: [u32; U32Vector::LANES] = array::from_fn(|i| base + i as u32);
120        for d in (0..length).rev() {
121            let digits: &mut [u8; U32Vector::LANES] =
122                buf[d * U32Vector::LANES..(d + 1) * U32Vector::LANES]
123                    .as_mut_array()
124                    .unwrap();
125            for i in 0..U32Vector::LANES {
126                digits[i] = (values[i] % 10) as u8 + b'0';
127                values[i] /= 10;
128            }
129        }
130
131        if values.iter().any(|&x| x > 0) {
132            // At least one number has an extra digit, fallback to scalar code
133            return None;
134        }
135
136        Some(NonZeroUsize::new(length).unwrap())
137    }
138}