Skip to main content

stenoxide_core/image_io/
phash.rs

1//! Perceptual hash margin filter with a hard stability limit of `k <= 1`.
2//!
3//! The Argon2id salt is not stored anywhere. It is recomputed from the image
4//! itself on both sides, which is what lets extraction work without any
5//! metadata travelling next to the payload. That only holds if the hash the
6//! receiver computes over the *stego* image is bit-for-bit the one the sender
7//! computed over the *cover* image.
8//!
9//! Embedding perturbs the least significant bits of a few thousand pixels. A
10//! 32x32 DCT coefficient is an average over the whole image, so those flips
11//! move it by a vanishing amount — unless the coefficient happened to sit right
12//! on the median, in which case an arbitrarily small perturbation flips its
13//! bit and the salt, the master key and the whole payload are lost.
14//!
15//! This module therefore refuses to hash images whose coefficients sit too
16//! close to the median. Each bit carries a margin, and:
17//!
18//! - `k == 0` unstable bits: the hash is reproducible, embedding proceeds.
19//! - `k == 1`: the image is still usable, but the receiver has to try both
20//!   values of the uncertain bit; `recover_phash_salt` does exactly that.
21//! - `k >= 2`: rejected. Two uncertain bits would mean four hypotheses, each
22//!   costing a full Argon2id derivation, and the number doubles from there.
23//!
24//! The margin filter also absorbs a second, subtler source of divergence: the
25//! DCT is built on `cos`, which is not specified bit-exactly by IEEE 754 and
26//! may differ in the last place between platforms. A bit whose margin exceeds
27//! `DELTA_MIN` cannot be flipped by an error of that size.
28
29use std::f32::consts::PI;
30use std::fmt;
31
32use image::{imageops, imageops::FilterType, GrayImage, Luma};
33use sha3::{Digest, Sha3_256};
34use zeroize::{ZeroizeOnDrop, Zeroizing};
35
36use crate::crypto::expand::{expand_master_key, DerivedKeys};
37use crate::crypto::kdf::KeyDeriver;
38use crate::image_io::buffer::{ColorSpace, CoverSource, ImageBuffer};
39
40/// Side length of the square thumbnail the DCT is computed over.
41const PHASH_THUMBNAIL_SIZE: usize = 32;
42
43/// Number of AC coefficients turned into hash bits.
44const N_HASH_BITS: usize = 64;
45
46/// Smallest distance from the median a coefficient may have and still be
47/// considered stable, in coefficient units.
48///
49/// The number is only meaningful against a fixed transform scale; see
50/// [`dct_2d`], which is deliberately left unnormalised for that reason.
51///
52/// At that scale the threshold is a texture requirement in disguise. A smooth
53/// container concentrates almost all of its energy in a handful of low
54/// coefficients and leaves the rest piled up around a near-zero median, so
55/// dozens of bits fall inside the margin and the image is rejected. A textured
56/// one spreads its energy across the spectrum and clears the threshold on
57/// every bit. That is the same property the cost layer wants, arrived at
58/// independently.
59const DELTA_MIN: f32 = 5.0;
60
61/// Largest number of unstable bits an image may have and still be accepted.
62///
63/// One uncertain bit means two hypotheses on the receiving side. Two would
64/// mean four Argon2id derivations, which at 128 MiB and four passes each is
65/// already an unreasonable price to pay for a container the user can simply
66/// replace.
67const MAX_UNSTABLE_BITS: usize = 1;
68
69/// Domain separator mixed into the salt so that the 64 hash bits cannot be
70/// reused as an input to any other construction in this crate.
71const PHASH_SALT_DOMAIN: &[u8] = b"STENOXIDE-v1-phash-salt";
72
73/// Length of the derived salt, in bytes.
74const PHASH_SALT_LEN: usize = 32;
75
76/// Magic number of a Zstandard frame, little-endian, as it appears on the wire.
77///
78/// See [`prefix_matches_key`] for why a compression magic number, of all
79/// things, is what distinguishes the two hypotheses.
80const ZSTD_FRAME_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD];
81
82/// Number of keystream bytes XChaCha20-Poly1305 spends on the Poly1305 key
83/// before it starts encrypting the message.
84///
85/// The AEAD construction of RFC 8439 reserves block zero for the one-time MAC
86/// key and encrypts the plaintext from block one onwards, so a raw XChaCha20
87/// instance has to be seeked past those 64 bytes to line its keystream up with
88/// the ciphertext.
89const AEAD_KEYSTREAM_OFFSET: u64 = 64;
90
91/// Perceptual hash of the container image, used as the Argon2id salt.
92///
93/// The bytes are wiped when the value is dropped. Like [`MasterKey`], the type
94/// implements neither [`Clone`], [`Copy`] nor [`Debug`]: the salt is not a
95/// secret in the sense a key is, but it is a deterministic function of the
96/// container, and leaking it into a log would let an attacker confirm which
97/// image a given payload belongs to.
98///
99/// [`MasterKey`]: crate::crypto::kdf::MasterKey
100#[derive(ZeroizeOnDrop)]
101pub struct PHashSalt([u8; PHASH_SALT_LEN]);
102
103impl PHashSalt {
104    /// Wraps the 32 bytes of a perceptual hash as a salt.
105    ///
106    /// Restricted to the crate: outside code must obtain a salt through
107    /// [`compute_stable_phash`], the only path that applies the stability
108    /// filter.
109    pub(crate) fn new(bytes: [u8; PHASH_SALT_LEN]) -> Self {
110        Self(bytes)
111    }
112
113    /// Borrows the salt bytes, for use as the Argon2id salt.
114    pub(crate) fn as_bytes(&self) -> &[u8] {
115        &self.0
116    }
117}
118
119/// Every way perceptual hashing can fail.
120#[derive(Debug)]
121pub enum PHashError {
122    /// Too many coefficients sit within `DELTA_MIN` of the median for the
123    /// hash to be reproducible after embedding.
124    InsufficientStability {
125        /// How many of the 64 bits are uncertain.
126        unstable_bits: usize,
127        /// The margin threshold that was applied.
128        threshold: f32,
129    },
130    /// Neither hypothesis for the uncertain bit produced a key that matches the
131    /// extracted payload. The image, the password or the payload is wrong.
132    RecoveryFailed,
133    /// Key derivation failed while testing a hypothesis.
134    KdfError(String),
135}
136
137impl fmt::Display for PHashError {
138    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139        match self {
140            PHashError::InsufficientStability {
141                unstable_bits,
142                threshold,
143            } => write!(
144                f,
145                "image is perceptually unstable: {unstable_bits} hash bits sit within {threshold} \
146                 of the median; choose a container with more texture"
147            ),
148            PHashError::RecoveryFailed => write!(
149                f,
150                "could not recover the perceptual hash salt from the stego image"
151            ),
152            PHashError::KdfError(message) => {
153                write!(f, "key derivation failed during salt recovery: {message}")
154            }
155        }
156    }
157}
158
159impl std::error::Error for PHashError {}
160
161/// The 64 hash bits together with how far each one is from the decision
162/// boundary.
163///
164/// Kept as one value because the two arrays are always produced and consumed
165/// together: a bit without its margin cannot be judged reproducible, and a
166/// margin without its bit says nothing about the hash.
167struct HashBits {
168    /// `true` where the AC coefficient is above the median.
169    bits: [bool; N_HASH_BITS],
170    /// `|coefficient - median|` for each bit.
171    margins: [f32; N_HASH_BITS],
172}
173
174impl HashBits {
175    /// Indices of the bits whose margin is below [`DELTA_MIN`].
176    fn unstable_indices(&self) -> Vec<usize> {
177        self.margins
178            .iter()
179            .enumerate()
180            .filter(|(_, &margin)| margin < DELTA_MIN)
181            .map(|(index, _)| index)
182            .collect()
183    }
184}
185
186/// BT.601 luma of one pixel, from the raw sample bytes of that pixel.
187///
188/// `sample` is exactly one pixel as laid out by `color_space`, so every index
189/// below is in bounds by construction: the callers slice the buffer with
190/// [`ColorSpace::bytes_per_pixel`].
191///
192/// Shared with [`crate::image_io::jpeg_detect`], which measures its block
193/// energies on the same luma plane. The two analyses must agree on what "the
194/// brightness of a pixel" means, so there is exactly one implementation of it.
195pub(super) fn luminance(sample: &[u8], color_space: ColorSpace) -> u8 {
196    let (red, green, blue) = match color_space {
197        // Grayscale is already luma; running it through the coefficients would
198        // only add a rounding error.
199        ColorSpace::Luma8 => return sample[0],
200        ColorSpace::Rgb8 | ColorSpace::Rgba8 => {
201            (sample[0] as f32, sample[1] as f32, sample[2] as f32)
202        }
203        // Samples were stored as explicit little-endian pairs by the validation
204        // stage, so they are re-read the same way here rather than as native
205        // `u16`. Normalised to [0, 255] before the coefficients are applied, so
206        // that DELTA_MIN means the same thing at both bit depths.
207        ColorSpace::Rgb16 => {
208            const SCALE: f32 = 255.0 / 65535.0;
209            let red = u16::from_le_bytes([sample[0], sample[1]]) as f32 * SCALE;
210            let green = u16::from_le_bytes([sample[2], sample[3]]) as f32 * SCALE;
211            let blue = u16::from_le_bytes([sample[4], sample[5]]) as f32 * SCALE;
212            (red, green, blue)
213        }
214    };
215
216    let luma = 0.299 * red + 0.587 * green + 0.114 * blue;
217    luma.clamp(0.0, 255.0) as u8
218}
219
220/// Step 1 and 2 — converts the image to luma and resizes it to 32x32.
221///
222/// The resize is bilinear ([`FilterType::Triangle`]): the filter kernel is
223/// scaled to the source resolution, so every input pixel contributes to the
224/// thumbnail. That is what makes the result insensitive to the handful of
225/// least significant bits embedding will flip.
226fn luminance_thumbnail(img: &ImageBuffer) -> GrayImage {
227    let (width, height) = img.dimensions();
228    let color_space = img.color_space();
229    let bytes_per_pixel = color_space.bytes_per_pixel();
230    let pixels = img.pixels();
231
232    let full = GrayImage::from_fn(width, height, |x, y| {
233        let offset = img.pixel_offset(x, y);
234        // In range for every coordinate the closure is called with, by the
235        // `CoverSource` length contract. Falling back to black instead of
236        // indexing keeps the function total.
237        let sample = pixels.get(offset..offset + bytes_per_pixel);
238        Luma([sample.map_or(0, |sample| luminance(sample, color_space))])
239    });
240
241    let side = PHASH_THUMBNAIL_SIZE as u32;
242    imageops::resize(&full, side, side, FilterType::Triangle)
243}
244
245/// Step 3 — separable DCT-II of the 32x32 thumbnail.
246///
247/// Rows first, then columns; the output is row-major, and the receiver
248/// reproduces exactly this order.
249///
250/// # Why the transform is left unnormalised
251///
252/// No orthonormal scale factor is applied, following the convention every
253/// mainstream perceptual hash uses. The choice is not cosmetic: it fixes the
254/// scale that [`DELTA_MIN`] is measured against. Orthonormalising would divide
255/// every AC coefficient by about sixteen, and a threshold of `5.0` would then
256/// reject even densely textured containers, which is the opposite of what the
257/// filter is for. Only the ratio between the threshold and the transform scale
258/// has any meaning; this function pins one half of it.
259fn dct_2d(thumbnail: &GrayImage) -> [f32; PHASH_THUMBNAIL_SIZE * PHASH_THUMBNAIL_SIZE] {
260    const N: usize = PHASH_THUMBNAIL_SIZE;
261
262    // basis[k][n] = cos(pi/N * (n + 1/2) * k). Constant for a fixed N, but
263    // rebuilt per call because `cos` is not available in a const context on
264    // stable Rust.
265    let mut basis = [[0.0f32; N]; N];
266    for (k, row) in basis.iter_mut().enumerate() {
267        for (n, value) in row.iter_mut().enumerate() {
268            *value = (PI / N as f32 * (n as f32 + 0.5) * k as f32).cos();
269        }
270    }
271
272    let mut rows = [0.0f32; N * N];
273    for y in 0..N {
274        for (k, basis_row) in basis.iter().enumerate() {
275            let mut acc = 0.0f32;
276            for (n, weight) in basis_row.iter().enumerate() {
277                acc += thumbnail.get_pixel(n as u32, y as u32).0[0] as f32 * weight;
278            }
279            rows[y * N + k] = acc;
280        }
281    }
282
283    let mut coefficients = [0.0f32; N * N];
284    for x in 0..N {
285        for (k, basis_row) in basis.iter().enumerate() {
286            let mut acc = 0.0f32;
287            for (n, weight) in basis_row.iter().enumerate() {
288                acc += rows[n * N + x] * weight;
289            }
290            coefficients[k * N + x] = acc;
291        }
292    }
293
294    coefficients
295}
296
297/// Median of the 64 AC coefficients.
298///
299/// An even number of samples, so the two central values are averaged. Sorted
300/// with [`f32::total_cmp`] rather than `partial_cmp`: the latter is fallible on
301/// NaN, and a comparator that has to decide what to do about NaN is a
302/// comparator that can panic.
303fn median(coefficients: &[f32; N_HASH_BITS]) -> f32 {
304    let mut sorted = *coefficients;
305    sorted.sort_by(f32::total_cmp);
306    (sorted[N_HASH_BITS / 2 - 1] + sorted[N_HASH_BITS / 2]) / 2.0
307}
308
309/// Steps 1 to 4 — everything up to, but not including, the stability verdict.
310///
311/// Split out because [`recover_phash_salt`] needs the margins as well as the
312/// bits, and recomputing a DCT over a multi-megapixel image to get them back
313/// would be wasteful.
314fn compute_hash_bits(img: &ImageBuffer) -> HashBits {
315    let thumbnail = luminance_thumbnail(img);
316    let coefficients = dct_2d(&thumbnail);
317
318    // The DC term at [0, 0] is the mean brightness of the whole image. It
319    // carries no structure, dwarfs every other coefficient and would drag the
320    // median with it, so it is dropped and the next 64 coefficients are taken
321    // in row-major order.
322    let mut ac = [0.0f32; N_HASH_BITS];
323    ac.copy_from_slice(&coefficients[1..=N_HASH_BITS]);
324
325    let median = median(&ac);
326
327    let mut bits = [false; N_HASH_BITS];
328    let mut margins = [0.0f32; N_HASH_BITS];
329    for index in 0..N_HASH_BITS {
330        bits[index] = ac[index] > median;
331        margins[index] = (ac[index] - median).abs();
332    }
333
334    HashBits { bits, margins }
335}
336
337/// Step 6 — packs the 64 bits and hashes them into a salt.
338///
339/// Bits are packed most-significant-first inside each byte: bit `i` of the hash
340/// lands in bit `7 - (i % 8)` of byte `i / 8`. The order is arbitrary but it is
341/// part of the wire format, so both sides must agree on it.
342fn salt_from_bits(bits: &[bool; N_HASH_BITS]) -> PHashSalt {
343    let mut packed = [0u8; N_HASH_BITS / 8];
344    for (index, &bit) in bits.iter().enumerate() {
345        if bit {
346            packed[index / 8] |= 1 << (7 - (index % 8));
347        }
348    }
349
350    let mut hasher = Sha3_256::new();
351    hasher.update(packed);
352    hasher.update(PHASH_SALT_DOMAIN);
353    let digest: [u8; PHASH_SALT_LEN] = hasher.finalize().into();
354
355    PHashSalt::new(digest)
356}
357
358/// Every salt a container's perceptual hash could have produced.
359///
360/// One when the hash is fully determined; two when a single coefficient sits
361/// inside the margin, in which case the sender used one of the pair and nothing
362/// in the image says which.
363pub(crate) struct PHashHypotheses {
364    /// Salt of the bits exactly as they measure on this image.
365    ///
366    /// The one the sender used whenever the image is the unmodified cover, and
367    /// the overwhelmingly likely one even for a stego image: the uncertain
368    /// coefficient only moves if embedding happened to push it across the
369    /// median.
370    pub(crate) primary: PHashSalt,
371    /// Salt of the same bits with the uncertain one flipped, when there is one.
372    pub(crate) alternative: Option<PHashSalt>,
373}
374
375/// Enumerates the salts a container's hash can take, applying the same
376/// stability limit as [`compute_stable_phash`].
377///
378/// The extraction path needs this rather than a single salt. Its problem is
379/// circular: the payload is what tells the two hypotheses apart, and reading the
380/// payload needs the permutation seed, which is derived from the salt. It can
381/// only be broken by trying a hypothesis, extracting under it and letting
382/// [`recover_phash_salt`] judge the result — which requires being able to name
383/// the other hypothesis, and that is what this function provides.
384///
385/// # Errors
386///
387/// Returns [`PHashError::InsufficientStability`] when more than
388/// [`MAX_UNSTABLE_BITS`] coefficients sit within [`DELTA_MIN`] of the median.
389pub(crate) fn phash_salt_hypotheses(img: &ImageBuffer) -> Result<PHashHypotheses, PHashError> {
390    let hash = compute_hash_bits(img);
391    let unstable = hash.unstable_indices();
392
393    if unstable.len() > MAX_UNSTABLE_BITS {
394        return Err(PHashError::InsufficientStability {
395            unstable_bits: unstable.len(),
396            threshold: DELTA_MIN,
397        });
398    }
399
400    let alternative = unstable.first().map(|&index| {
401        let mut flipped = hash.bits;
402        // In bounds: the indices come from enumerating the margin array, which
403        // has exactly as many entries as the bit array. The fallback keeps the
404        // function total without an index panic.
405        if let Some(bit) = flipped.get_mut(index) {
406            *bit = !*bit;
407        }
408
409        salt_from_bits(&flipped)
410    });
411
412    Ok(PHashHypotheses {
413        primary: salt_from_bits(&hash.bits),
414        alternative,
415    })
416}
417
418/// Computes the perceptual hash salt of a container image, refusing images
419/// whose hash would not survive embedding.
420///
421/// Public so that the stability of a container can be asked about on its own,
422/// without running an embedding. The returned salt is opaque outside this
423/// crate — it has no public accessor — so exposing this function grants the
424/// verdict and not the value.
425///
426/// # Errors
427///
428/// Returns [`PHashError::InsufficientStability`] when more than
429/// `MAX_UNSTABLE_BITS` coefficients sit within `DELTA_MIN` of the median.
430/// Note that a single unstable bit is *accepted* here: the sender does not
431/// care which value it takes, because the receiver resolves the ambiguity
432/// during extraction.
433pub fn compute_stable_phash(img: &ImageBuffer) -> Result<PHashSalt, PHashError> {
434    Ok(phash_salt_hypotheses(img)?.primary)
435}
436
437/// Tests whether `ciphertext_prefix` was produced under the keys derived from a
438/// candidate salt.
439///
440/// # Why this is not a MAC check
441///
442/// Poly1305 authenticates the ciphertext as a whole and its tag lives at the
443/// very end of it, so a 64-byte prefix carries nothing that can be verified.
444/// Extracting the entire payload before the salt is known is not an option
445/// either: the STC decoder needs the permutation seed, which is derived from
446/// the very key we are trying to pin down.
447///
448/// What a prefix *does* allow is decrypting it with the raw XChaCha20
449/// keystream and looking at the plaintext. The payload is Zstandard-compressed
450/// before it is encrypted, so the correct key uncovers a Zstandard frame magic
451/// number and a wrong key uncovers uniformly random bytes — a discriminator
452/// that is wrong with probability `2^-32`.
453///
454/// That is sound because this function decides nothing security-relevant. It
455/// only picks which of two hypotheses to try first; the actual authentication
456/// still happens when the pipeline decrypts the full payload under the chosen
457/// salt, and a mistake here surfaces there as an ordinary authentication
458/// failure.
459fn prefix_matches_key(keys: &DerivedKeys, ciphertext_prefix: &[u8]) -> bool {
460    use chacha20::cipher::{KeyIvInit, StreamCipher, StreamCipherSeek};
461    use chacha20::XChaCha20;
462
463    let Some(head) = ciphertext_prefix.get(..ZSTD_FRAME_MAGIC.len()) else {
464        return false;
465    };
466
467    let mut cipher = XChaCha20::new(keys.enc_key().into(), keys.nonce().into());
468    if cipher.try_seek(AEAD_KEYSTREAM_OFFSET).is_err() {
469        return false;
470    }
471
472    let mut plaintext = Zeroizing::new(head.to_vec());
473
474    // The fallible form: as of `cipher` 0.5 the plain `apply_keystream_b2b`
475    // panics when the two buffers differ in length, and a panic is not
476    // something this crate is allowed to reach for. The lengths do match here —
477    // `plaintext` was built from `head` — so the branch is unreachable, and it
478    // stays as a branch rather than an assertion for exactly that reason.
479    if cipher
480        .try_apply_keystream_b2b(head, &mut plaintext)
481        .is_err()
482    {
483        return false;
484    }
485
486    plaintext.as_slice() == ZSTD_FRAME_MAGIC
487}
488
489/// Derives the keys for one hypothesis and checks them against the payload.
490///
491/// Returns `Ok(Some(salt))` when the hypothesis matches, `Ok(None)` when it
492/// does not, and `Err` only when key derivation itself failed — a condition
493/// that says nothing about which hypothesis was right and must not be confused
494/// with a rejection.
495fn try_hypothesis(
496    bits: &[bool; N_HASH_BITS],
497    password: &[u8],
498    kdf: &impl KeyDeriver,
499    ciphertext_prefix: &[u8],
500) -> Result<Option<PHashSalt>, PHashError> {
501    let salt = salt_from_bits(bits);
502
503    let master_key = kdf
504        .derive(password, &salt)
505        .map_err(|err| PHashError::KdfError(err.to_string()))?;
506    let keys =
507        expand_master_key(&master_key).map_err(|err| PHashError::KdfError(err.to_string()))?;
508    // Wiped here rather than at the end of the scope: nothing below needs it,
509    // and the two hypotheses run concurrently, so one of the two master keys is
510    // always live material that no longer has any use.
511    drop(master_key);
512
513    let matched = prefix_matches_key(&keys, ciphertext_prefix);
514    drop(keys);
515
516    if matched {
517        Ok(Some(salt))
518    } else {
519        Ok(None)
520    }
521}
522
523/// Recovers the salt of a stego image whose hash has one uncertain bit.
524///
525/// `ciphertext_prefix` is the first 64 bytes of the provisionally extracted
526/// payload. It is used to tell the two hypotheses apart; see
527/// [`prefix_matches_key`] for what "tell apart" means here and why it is not,
528/// and does not need to be, an authentication step.
529///
530/// # Errors
531///
532/// Returns [`PHashError::InsufficientStability`] if the stego image has more
533/// than [`MAX_UNSTABLE_BITS`] uncertain bits, [`PHashError::RecoveryFailed`] if
534/// neither hypothesis matches the payload, and [`PHashError::KdfError`] if the
535/// key derivation used to test a hypothesis failed.
536pub(crate) fn recover_phash_salt(
537    stego_img: &ImageBuffer,
538    password: &[u8],
539    kdf: &impl KeyDeriver,
540    ciphertext_prefix: &[u8],
541) -> Result<PHashSalt, PHashError> {
542    let hash = compute_hash_bits(stego_img);
543    let unstable = hash.unstable_indices();
544
545    let uncertain = match unstable.as_slice() {
546        // k = 0: every coefficient is far enough from the median that
547        // embedding cannot have moved it across. The hash the receiver just
548        // computed is the one the sender used, and no search is needed. This
549        // is exactly what `compute_stable_phash` would return, minus a second
550        // pass over the image.
551        [] => return Ok(salt_from_bits(&hash.bits)),
552        [index] => *index,
553        more => {
554            return Err(PHashError::InsufficientStability {
555                unstable_bits: more.len(),
556                threshold: DELTA_MIN,
557            })
558        }
559    };
560
561    let mut cleared = hash.bits;
562    cleared[uncertain] = false;
563    let mut set = hash.bits;
564    set[uncertain] = true;
565
566    // Argon2id at the production parameters costs 128 MiB and several hundred
567    // milliseconds. Running the two hypotheses on separate threads makes the
568    // recovery cost the same wall time as an ordinary derivation, at the price
569    // of holding two memory blocks at once.
570    let (cleared, set) = rayon::join(
571        || try_hypothesis(&cleared, password, kdf, ciphertext_prefix),
572        || try_hypothesis(&set, password, kdf, ciphertext_prefix),
573    );
574
575    // Both hypotheses can only match if the discriminator collided, which is a
576    // `2^-32` event; picking either one is then as good as picking the other,
577    // and the full-payload authentication downstream catches the mistake.
578    match (cleared?, set?) {
579        (Some(salt), _) | (None, Some(salt)) => Ok(salt),
580        (None, None) => Err(PHashError::RecoveryFailed),
581    }
582}
583
584#[cfg(test)]
585mod tests {
586    // The crate-wide bans on panicking helpers reach into `cfg(test)` code as
587    // well. A test that cannot panic cannot fail, so they are lifted here and
588    // only here.
589    #![allow(clippy::expect_used)]
590    #![allow(clippy::panic)]
591
592    use super::*;
593
594    use rand::rngs::StdRng;
595    use rand::{RngExt, SeedableRng};
596
597    use crate::crypto::aead::{compress_and_encrypt, XChaCha20Poly1305Cipher};
598    use crate::crypto::kdf::{Argon2Kdf, MasterKey};
599
600    /// Side length of the synthetic containers below.
601    ///
602    /// Exactly the thumbnail size, so the resize step is the identity and the
603    /// DCT reads the samples these tests wrote. Building the hash out of a
604    /// multi-megapixel image instead would test the resampler, not the filter.
605    const SIDE: u32 = PHASH_THUMBNAIL_SIZE as u32;
606
607    /// The password the recovery tests stretch.
608    const PASSWORD: &[u8] = b"a-container-passphrase";
609
610    /// A square of uncorrelated colour noise.
611    ///
612    /// White noise spreads the 64 AC coefficients over a range of some
613    /// thousands of units, which is three orders of magnitude above the `5.0`
614    /// stability margin — so the hash of such an image is reproducible for
615    /// almost any seed.
616    fn noise_image(seed: u64) -> ImageBuffer {
617        let mut rng = StdRng::seed_from_u64(seed);
618        let pixel_count = (SIDE * SIDE) as usize;
619        let pixels = (0..pixel_count * 3).map(|_| rng.random()).collect();
620
621        ImageBuffer::new(pixels, SIDE, SIDE, ColorSpace::Rgb8)
622    }
623
624    /// The first noise square at or after `seed` whose hash is reproducible.
625    ///
626    /// Almost every seed produces one, but not quite every seed: the two
627    /// central AC coefficients are what the verdict turns on, and they land
628    /// within `2 * DELTA_MIN` of each other often enough that pinning a
629    /// literal seed would make these tests brittle for no reason.
630    fn stable_noise_image(seed: u64) -> ImageBuffer {
631        match (seed..seed + 64)
632            .map(noise_image)
633            .find(|image| compute_hash_bits(image).unstable_indices().is_empty())
634        {
635            Some(image) => image,
636            None => panic!("no stable container in sixty-four candidates from {seed}"),
637        }
638    }
639
640    /// A container with no structure at all, and therefore no usable hash.
641    fn flat_image() -> ImageBuffer {
642        let pixel_count = (SIDE * SIDE) as usize;
643
644        ImageBuffer::new(vec![128u8; pixel_count * 3], SIDE, SIDE, ColorSpace::Rgb8)
645    }
646
647    /// Sixty-four bytes of ciphertext produced under the keys of `bits`.
648    ///
649    /// The head of a real payload: compressed with Zstandard and then
650    /// encrypted, which is exactly what [`prefix_matches_key`] is written to
651    /// recognise.
652    fn ciphertext_prefix_for(bits: &[bool; N_HASH_BITS]) -> Vec<u8> {
653        let salt = salt_from_bits(bits);
654        let master_key = Argon2Kdf::low_cost_for_tests()
655            .derive(PASSWORD, &salt)
656            .expect("a non-empty password must stretch");
657        let keys = expand_master_key(&master_key).expect("expansion must succeed");
658
659        let ciphertext = compress_and_encrypt(
660            &b"a payload long enough to fill a whole keystream block and then some".repeat(4),
661            keys.enc_key(),
662            keys.nonce(),
663            &XChaCha20Poly1305Cipher::new(),
664        )
665        .expect("encryption must succeed");
666
667        ciphertext.iter().copied().take(64).collect()
668    }
669
670    /// Every layout reduces to the same notion of brightness.
671    #[test]
672    fn luminance_reads_each_layout_on_the_same_scale() {
673        // Grayscale is passed through untouched.
674        assert_eq!(luminance(&[110], ColorSpace::Luma8), 110);
675
676        // BT.601 weights the green channel most and the blue least.
677        assert_eq!(luminance(&[255, 0, 0], ColorSpace::Rgb8), 76);
678        assert_eq!(luminance(&[0, 255, 0], ColorSpace::Rgb8), 149);
679        assert_eq!(luminance(&[0, 0, 255], ColorSpace::Rgb8), 29);
680
681        // The alpha channel is not part of the brightness of a pixel.
682        assert_eq!(
683            luminance(&[255, 0, 0, 17], ColorSpace::Rgba8),
684            luminance(&[255, 0, 0], ColorSpace::Rgb8)
685        );
686
687        // Sixteen-bit samples are read as little-endian pairs and normalised to
688        // the same 0..=255 range, so `DELTA_MIN` means one thing at both depths.
689        assert_eq!(
690            luminance(&[0xFF, 0xFF, 0, 0, 0, 0], ColorSpace::Rgb16),
691            luminance(&[255, 0, 0], ColorSpace::Rgb8)
692        );
693    }
694
695    /// A textured container hashes reproducibly and leaves nothing to guess.
696    #[test]
697    fn a_textured_container_has_a_single_hypothesis() {
698        let image = stable_noise_image(1);
699
700        let hypotheses = match phash_salt_hypotheses(&image) {
701            Ok(hypotheses) => hypotheses,
702            Err(error) => panic!("colour noise must hash stably: {error}"),
703        };
704
705        assert!(
706            hypotheses.alternative.is_none(),
707            "a fully determined hash has nothing to disambiguate"
708        );
709
710        let salt = compute_stable_phash(&image).expect("the same image must hash again");
711        assert_eq!(salt.as_bytes(), hypotheses.primary.as_bytes());
712    }
713
714    /// A container with no structure is refused, and the refusal counts the
715    /// bits it could not pin down.
716    #[test]
717    fn a_flat_container_is_refused_as_unstable() {
718        let error = phash_salt_hypotheses(&flat_image())
719            .map(|_| ())
720            .expect_err("a uniform image cannot hash reproducibly");
721
722        match error {
723            PHashError::InsufficientStability {
724                unstable_bits,
725                threshold,
726            } => {
727                assert!(unstable_bits > MAX_UNSTABLE_BITS);
728                assert_eq!(threshold, DELTA_MIN);
729            }
730            other => panic!("expected an instability verdict, got: {other:?}"),
731        }
732    }
733
734    /// Uncertain bits always come in pairs, so `k` is never exactly one.
735    ///
736    /// A structural property of the filter rather than an accident of the
737    /// images below. The median of an even-sized sample is the midpoint of its
738    /// two central values, so those two are equidistant from it and no other
739    /// coefficient can be closer: whenever the smallest margin falls under
740    /// [`DELTA_MIN`], two bits do at once.
741    ///
742    /// The consequence is worth stating where it can be checked: the branch of
743    /// [`recover_phash_salt`] that resolves a single uncertain bit describes a
744    /// case this filter cannot produce, and every container it accepts has a
745    /// hash that is fully determined.
746    #[test]
747    fn uncertain_bits_never_appear_alone() {
748        for seed in 0..64u64 {
749            let unstable = compute_hash_bits(&noise_image(seed))
750                .unstable_indices()
751                .len();
752
753            assert_ne!(unstable, 1, "seed {seed} produced a lone uncertain bit");
754        }
755
756        assert_ne!(compute_hash_bits(&flat_image()).unstable_indices().len(), 1);
757    }
758
759    /// The hash is a function of the image and of nothing else.
760    #[test]
761    fn the_salt_is_deterministic_and_image_dependent() {
762        let image = stable_noise_image(100);
763        let other_image = stable_noise_image(200);
764
765        let first = compute_stable_phash(&image).expect("noise must hash");
766        let again = compute_stable_phash(&image).expect("noise must hash");
767        let other = compute_stable_phash(&other_image).expect("noise must hash");
768
769        assert_eq!(first.as_bytes(), again.as_bytes());
770        assert_ne!(first.as_bytes(), other.as_bytes());
771    }
772
773    /// Flipping one hash bit changes the whole salt.
774    ///
775    /// The bits are packed and then hashed, so the salt is not a rearrangement
776    /// of them and a neighbouring hypothesis shares nothing with its partner.
777    #[test]
778    fn one_flipped_bit_changes_the_whole_salt() {
779        let mut bits = [false; N_HASH_BITS];
780        let base = salt_from_bits(&bits);
781
782        bits[17] = true;
783        let flipped = salt_from_bits(&bits);
784
785        assert_ne!(base.as_bytes(), flipped.as_bytes());
786    }
787
788    /// The median of an even sample is the midpoint of its two central values.
789    #[test]
790    fn the_median_is_the_midpoint_of_the_two_central_coefficients() {
791        let mut coefficients = [0.0f32; N_HASH_BITS];
792        for (index, value) in coefficients.iter_mut().enumerate() {
793            *value = index as f32;
794        }
795
796        assert_eq!(median(&coefficients), 31.5);
797    }
798
799    /// A correct key uncovers a Zstandard frame; a wrong one uncovers noise.
800    #[test]
801    fn the_zstd_magic_number_tells_the_keys_apart() {
802        let bits = [true; N_HASH_BITS];
803        let prefix = ciphertext_prefix_for(&bits);
804
805        let salt = salt_from_bits(&bits);
806        let master_key = Argon2Kdf::low_cost_for_tests()
807            .derive(PASSWORD, &salt)
808            .expect("a non-empty password must stretch");
809        let keys = expand_master_key(&master_key).expect("expansion must succeed");
810
811        assert!(prefix_matches_key(&keys, &prefix));
812
813        // Any other key decrypts the same prefix to bytes that are not a frame
814        // header, which is the whole discriminator.
815        let other = expand_master_key(&MasterKey::new([0x5Au8; 32])).expect("expansion");
816        assert!(!prefix_matches_key(&other, &prefix));
817
818        // A prefix too short to hold the magic number decides nothing.
819        assert!(!prefix_matches_key(&keys, &prefix[..3]));
820    }
821
822    /// One hypothesis is confirmed by the payload, the other is not.
823    #[test]
824    fn a_hypothesis_is_judged_by_the_payload_it_explains() {
825        let kdf = Argon2Kdf::low_cost_for_tests();
826        let bits = [true; N_HASH_BITS];
827        let prefix = ciphertext_prefix_for(&bits);
828
829        // `PHashSalt` implements neither `Debug` nor `PartialEq` — it is derived
830        // from the container and must not reach a log — so every arm below is
831        // spelled out rather than compared with `assert!(matches!(..))`.
832        match try_hypothesis(&bits, PASSWORD, &kdf, &prefix) {
833            Ok(Some(salt)) => assert_eq!(salt.as_bytes(), salt_from_bits(&bits).as_bytes()),
834            Ok(None) => panic!("the hypothesis the payload was made under must match"),
835            Err(error) => panic!("derivation must succeed: {error}"),
836        }
837
838        let mut wrong = bits;
839        wrong[0] = false;
840        match try_hypothesis(&wrong, PASSWORD, &kdf, &prefix) {
841            Ok(None) => {}
842            Ok(Some(_)) => panic!("a hypothesis that explains nothing must be rejected"),
843            Err(error) => panic!("derivation must succeed: {error}"),
844        }
845
846        // A derivation that fails is not a rejection: it says nothing about
847        // which hypothesis was right and must not be reported as one.
848        match try_hypothesis(&bits, &[], &kdf, &prefix) {
849            Err(PHashError::KdfError(_)) => {}
850            Err(error) => panic!("expected a derivation failure, got: {error}"),
851            Ok(_) => panic!("an empty password must fail derivation"),
852        }
853    }
854
855    /// A hash with no uncertain bit needs no search.
856    #[test]
857    fn recovery_short_circuits_on_a_certain_hash() {
858        let image = stable_noise_image(4);
859        let expected = compute_stable_phash(&image).expect("noise must hash");
860
861        let recovered = recover_phash_salt(&image, PASSWORD, &Argon2Kdf::low_cost_for_tests(), &[])
862            .expect("a fully determined hash needs no payload to be recovered");
863
864        assert_eq!(recovered.as_bytes(), expected.as_bytes());
865    }
866
867    /// An image whose hash is not reproducible is refused by the recovery path
868    /// on the same terms as by the sender's path.
869    #[test]
870    fn recovery_refuses_an_unstable_image() {
871        let error = recover_phash_salt(
872            &flat_image(),
873            PASSWORD,
874            &Argon2Kdf::low_cost_for_tests(),
875            &[],
876        )
877        .map(|_| ())
878        .expect_err("a uniform image has no hash to recover");
879
880        assert!(
881            matches!(error, PHashError::InsufficientStability { .. }),
882            "got: {error:?}"
883        );
884    }
885
886    /// Every failure of this layer explains itself.
887    #[test]
888    fn every_failure_explains_itself() {
889        let unstable = PHashError::InsufficientStability {
890            unstable_bits: 7,
891            threshold: DELTA_MIN,
892        }
893        .to_string();
894        assert!(unstable.contains('7') && unstable.contains("texture"));
895
896        assert!(!PHashError::RecoveryFailed.to_string().is_empty());
897        assert!(PHashError::KdfError("empty password".to_owned())
898            .to_string()
899            .contains("empty password"));
900    }
901}