From 5535617b74fa8c3376a7a6bc919f457c23a373ac Mon Sep 17 00:00:00 2001 From: ef3d0c3e Date: Sat, 17 Aug 2024 20:29:41 +0200 Subject: [PATCH] Refactor in progress --- src/block.rs | 228 +++++++++++++++++++++++++++++++++++++++++++++----- src/embed.rs | 58 +++++++------ src/header.rs | 70 +++++++++++++++- src/main.rs | 8 +- 4 files changed, 315 insertions(+), 49 deletions(-) diff --git a/src/block.rs b/src/block.rs index 977bb44..3618523 100644 --- a/src/block.rs +++ b/src/block.rs @@ -1,8 +1,10 @@ use std::fmt::Formatter; use std::str::FromStr; -use rand::Rng; +use bitvec::slice::BitSlice; +use bitvec::vec::BitVec; use rand::prelude::SliceRandom; +use rand::Rng; use crate::embed::EmbedAlgorithm; @@ -42,7 +44,7 @@ impl BlockMode { Self { len: 1usize << len, - crc_len: (crc * len) as usize + crc_len: (crc * len) as usize, } } } @@ -76,50 +78,238 @@ impl FromStr for BlockMode { #[derive(Debug)] pub struct BlockPlacement<'a> { + algorithm: &'a EmbedAlgorithm, data: &'a mut [u8], block_size: usize, blocks: Vec, } -impl<'a> BlockPlacement<'a> -{ +impl<'a> BlockPlacement<'a> { // Attempts to create a new block placement // // # Errors // // Will fail if the data is too small to hold all the blocks - pub fn new(data: &'a mut [u8], block_size: usize, algorithm: &EmbedAlgorithm, embed_size: usize, rng: &mut R) -> Result - where R: Rng + ?Sized + pub fn new( + data: &'a mut [u8], + algorithm: &'a EmbedAlgorithm, + block_size: usize, + embed_size: usize, + rng: &mut R, + ) -> Result + where + R: Rng + ?Sized, { // Total size of the embed (crc included) let embedded_size = algorithm.embedded_size(embed_size); + // Maximum number of blocks + let max_blocks = data.len() / block_size; + // Number of blocks let blocks_num = (embedded_size as f64 / block_size as f64).ceil() as usize; - - // Safe distance for spacing the blocks equally - let safe_spacing = data.len() / blocks_num; - if safe_spacing*blocks_num < embedded_size { - return Err(format!("Failed to determine a safe spacing size: {safe_spacing} < {}", embedded_size as f64 / blocks_num as f64)) + if blocks_num > max_blocks { + return Err(format!( + "Too many blocks required: {blocks_num}, maximum: {max_blocks}" + )); } // Blocks in the resulting data - let mut blocks = Vec::with_capacity(blocks_num); - for i in 0 .. blocks_num { - // Choose a random position within [0, safe_spacing[ for the block - let pos = rng.gen_range(i*safe_spacing..(i+1)*safe_spacing); - - blocks.push(pos); - } + let mut blocks = (0..max_blocks).collect::>(); // Shuffle the block order blocks.shuffle(rng); + // Only keep the first blocks_num blocks + blocks.resize(blocks_num, 0); + Ok(Self { + algorithm, data, block_size, - blocks + blocks, }) } + + pub fn write_embed(&mut self) +} + +// Iterator over blocks in the resulting image +pub struct BlockPlacementIterator<'a> { + algorithm: &'a EmbedAlgorithm, + data: &'a [u8], + block_size: usize, + + // Block index + index: usize, + // Position of the blocks + blocks: Vec, +} + +impl<'a> BlockPlacementIterator<'a> { + pub fn new( + algorithm: &'a EmbedAlgorithm, + data: &'a [u8], + block_size: usize, + rng: &mut R, + ) -> Self { + // Maximum number of blocks + let max_blocks = data.len() / block_size; + + // Blocks + let mut blocks = (0..max_blocks).collect::>(); + + // Shuffle the block order + blocks.shuffle(rng); + + Self { + algorithm, + data, + block_size, + index: 0, + blocks, + } + } +} + +impl<'a> Iterator for BlockPlacementIterator<'a> { + type Item = Block<'a>; + + fn next(&mut self) -> Option { + if self.index == self.blocks.len() { + return None; + } + + let pos = self.blocks[self.index] * self.block_size; + let slice = &self.data[pos..pos + self.block_size]; + self.index += 1; + + Some(Block(self.algorithm, slice)) + } +} + +// Block of data in the resulting image +#[derive(Debug)] +pub struct Block<'a>(&'a EmbedAlgorithm, &'a [u8]); + +// Iterator to read embedded data inside a block +pub struct BlockIterator<'a> { + // Block of the iterator + block: &'a Block<'a>, + + // Byte position in [`data`] + index: usize, + + // Remainder, i.e. bits that have been read and will be part of the `next` byte + remainder: BitVec, +} + +impl<'a> BlockIterator<'a> { + pub fn new(block: &'a Block, previous: Option) -> Self { + if let Some(previous) = previous { + Self { + block, + index: 0, + remainder: previous.remainder, + } + } else { + Self { + block, + index: 0, + remainder: BitVec::with_capacity(7), + } + } + } +} + +impl<'a> Iterator for BlockIterator<'a> { + type Item = u8; + + fn next(&mut self) -> Option { + let mut byte = 0u8; + + // Read remainder + let mut bit_idx = 0; + for bit in &self.remainder { + byte |= (*bit as u8) << bit_idx; + bit_idx += 1; + } + self.remainder.clear(); + + match self.block.0 { + EmbedAlgorithm::Lo(bits) => { + while bit_idx < 8 { + // End of data + if self.index == self.block.1.len() { + return None; + } + + // Read next byte + let next = self.block.1[self.index]; + self.index += 1; + + for i in 0..*bits { + if bit_idx < 8 { + // Prepend bit to result + byte |= ((next >> i) & 0b1) << bit_idx; + } else { + // Append bit to remainder + self.remainder.push((next >> i) & 0b1 == 0b1) + } + + bit_idx += 1; + } + } + + Some(byte) + } + } + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + use super::*; + + #[test] + fn block_iterator() { + let algorithm = EmbedAlgorithm::Lo(3); + let data = vec![ + 0b10111000, 0b11111001, 0b01101010, 0b00111011, 0b11011100, 0b11100110, 0b01100111, + 0b01100000, + ]; + + let block = Block(&algorithm, &data); + let mut it = BlockIterator::new(&block, None); + + assert_eq!(it.next(), Some(0b10_001_000)); + assert_eq!(it.next(), Some(0b0_100_011_0)); + assert_eq!(it.next(), Some(0b000_111_11)); + } + + #[test] + fn blockplacement_iterator() { + let algorithm = EmbedAlgorithm::Lo(4); + let mut data = vec![ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, + 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, + 0x1c, 0x1d, 0x1e, 0x1f, + ]; + + let mut rand = ChaCha8Rng::from_seed([1u8; 32]); + let mut it = BlockPlacementIterator::new::<_>(&algorithm, &data, 4, &mut rand); + + let mut rand = ChaCha8Rng::from_seed([1u8; 32]); + let mut positions = (0..8).collect::>(); + positions.shuffle(&mut rand); + + for i in 0..8 { + let block = it.next().unwrap(); + assert_eq!(block.1[0] / 4, positions[i]); + } + } } diff --git a/src/embed.rs b/src/embed.rs index fa3afd9..350ba0b 100644 --- a/src/embed.rs +++ b/src/embed.rs @@ -1,12 +1,13 @@ use std::fmt::Formatter; use std::str::FromStr; -use bitvec::prelude::*; +use bitvec::slice::BitSlice; use bitvec::vec::BitVec; use crate::block::BlockMode; use crate::image::ImageInfo; +#[derive(Debug)] pub enum EmbedAlgorithm { Lo(u8), } @@ -15,9 +16,7 @@ impl EmbedAlgorithm { /// Get the size of the data (in bytes) once embedded by the algorithm pub fn embedded_size(&self, size: usize) -> usize { match self { - EmbedAlgorithm::Lo(bits) => { - ((size * 8) as f64 / *bits as f64).ceil() as usize - } + EmbedAlgorithm::Lo(bits) => ((size * 8) as f64 / *bits as f64).ceil() as usize, } } @@ -26,35 +25,43 @@ impl EmbedAlgorithm { match self { EmbedAlgorithm::Lo(bits) => { - (((blockmode.len - blockmode.crc_len)*blocks_num) as f64 * (*bits as f64) / 8f64).floor() as usize + (((blockmode.len - blockmode.crc_len) * blocks_num) as f64 * (*bits as f64) / 8f64) + .floor() as usize } } } - pub fn next_block(&self, original_data: &mut [u8], mut data_pos: usize, embed_data: &BitVec, mut embed_offset: usize, blockmode: &BlockMode) -> (usize, usize) { + pub fn next_block( + &self, + original_data: &mut [u8], + mut data_pos: usize, + embed_data: &BitVec, + mut embed_offset: usize, + blockmode: &BlockMode, + ) -> (usize, usize) { match self { EmbedAlgorithm::Lo(bits) => { - let mask = (1<, bits: u8) -> u8 - { - let mut result : u8 = 0; - for i in 0..bits - { + fn bits_to_byte(slice: &BitSlice, bits: u8) -> u8 { + let mut result: u8 = 0; + for i in 0..bits { result |= (slice[i as usize] as u8) << i; } result } let start = embed_offset; - while embed_offset-start < (blockmode.len-blockmode.crc_len)*8 - { + while embed_offset - start < (blockmode.len - blockmode.crc_len) * 8 { let hi = std::cmp::min(*bits as usize, embed_data.len() - embed_offset); - let embed = bits_to_byte(embed_data.get(embed_offset..embed_offset+hi).unwrap(), hi as u8); + let embed = bits_to_byte( + embed_data.get(embed_offset..embed_offset + hi).unwrap(), + hi as u8, + ); original_data[data_pos] &= !mask; original_data[data_pos] |= embed; - + data_pos += 1; embed_offset += hi; } @@ -62,24 +69,27 @@ impl EmbedAlgorithm { // TODO: WRITE CRC } } - + (data_pos, embed_offset) } - pub fn read_block(&self, encoded_data: &[u8], mut data_pos: usize, incoming: &mut BitVec, blockmode: &BlockMode) -> usize { + pub fn read_block( + &self, + encoded_data: &[u8], + mut data_pos: usize, + incoming: &mut BitVec, + blockmode: &BlockMode, + ) -> usize { match self { EmbedAlgorithm::Lo(bits) => { - fn push(vec: &mut BitVec, bits: u8, b: u8) - { - for i in 0..bits - { + fn push(vec: &mut BitVec, bits: u8, b: u8) { + for i in 0..bits { vec.push((b >> i) & 0b1 == 0b1) } } let start = incoming.len(); - while incoming.len()-start < (blockmode.len-blockmode.crc_len)*8 - { + while incoming.len() - start < (blockmode.len - blockmode.crc_len) * 8 { push(incoming, *bits, encoded_data[data_pos]); data_pos += 1; } diff --git a/src/header.rs b/src/header.rs index 1189f5e..59eb53f 100644 --- a/src/header.rs +++ b/src/header.rs @@ -1,13 +1,49 @@ -use bitvec::slice::BitSlice; +use bitvec::{slice::BitSlice, vec::BitVec}; +use crc::Crc; use crate::block::BlockMode; -pub struct Header { - pub blockmode: BlockMode, - pub comment: Option, +#[repr(u16)] +#[allow(non_camel_case_types)] +#[derive(Debug, Clone, Copy)] +pub enum Version { + VERSION_1, } +pub trait Encode { + // Encode the data + fn encode(&self, vec: &mut Vec); +} + +//pub trait Decode { +// fn decode(incoming: &mut EmbedIterator) -> (usize, Self); +//} + +#[derive(Debug)] +pub struct Header { + version: Version, + block_size: usize, + data_len: u32, + data_crc: u32, + comment: Option, +} + + impl Header { + pub fn new(version: Version, block_size: usize, data: &[u8], comment: Option) -> Self { + assert_eq!((data.len() as u32) as usize, data.len()); + assert_eq!(1 << usize::trailing_zeros(block_size), block_size); + assert!(comment.as_ref().map_or(0, |c| c.len()) < u16::MAX as usize); + + Self { + version, + block_size, + data_len: data.len() as u32, + data_crc: Crc::::new(&crc::CRC_32_CKSUM).checksum(data), + comment, + } + } + /* pub fn to_data(&self, version: u16, embed_len: u32) -> Vec { let mut header = vec![]; @@ -55,4 +91,30 @@ impl Header { (version, blockmode, len, comment_len) } + */ +} + +impl Encode for Header { + fn encode(&self, vec: &mut Vec) { + // Version + vec.extend_from_slice((self.version as u16).to_le_bytes().as_slice()); + + // Block size + vec.push((usize::trailing_zeros(self.block_size) as u8).to_le()); + + // Data Len + vec.extend_from_slice(self.data_len.to_le_bytes().as_slice()); + + // Data CRC + vec.extend_from_slice(self.data_crc.to_le_bytes().as_slice()); + + // Comment length + let comment_length = self.comment.as_ref().map_or(0u16, |c| c.len() as u16); + vec.extend_from_slice(comment_length.to_le_bytes().as_slice()); + + // Comment + if let Some(comment) = &self.comment { + vec.extend_from_slice(comment.as_bytes()); + } + } } diff --git a/src/main.rs b/src/main.rs index 56666ac..22252b1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -43,6 +43,7 @@ Public domain\n" ); } +/* impl ImageInfo for png::OutputInfo { fn width(&self) -> u32 { self.width } @@ -202,9 +203,9 @@ fn encode(image: String, matches: Matches) -> Result, String> { .map_err(|err| format!("Failed to read embed file `{embed_file}`: {err}"))?; let mut rand = ChaCha8Rng::from_seed(seed); - //let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?; + let placement = BlockPlacement::new(data.as_mut_slice(), blockmode.len, &algorithm, embed_data.len(), &mut rand)?; - //return Ok(vec![]); + return Ok(vec![]); // Get header let header = Header { @@ -295,6 +296,7 @@ fn encode(image: String, matches: Matches) -> Result, String> { Ok(vec![]) } +*/ fn main() -> ExitCode { let args: Vec = env::args().collect(); @@ -331,6 +333,7 @@ fn main() -> ExitCode { return ExitCode::FAILURE; } + /* let input = matches.free[0].clone(); if matches.opt_present("z") { @@ -350,6 +353,7 @@ fn main() -> ExitCode { } } } + */ ExitCode::SUCCESS }