dumbgoy/soulseek-rs-modded
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: seperate into lib and cli packages

Browse Source
This commit is contained in:
Michel de Graaf
2025-11-11 21:55:53 +01:00
parent 00046f2d59
commit 7bf2748e1f
51 changed files with 334 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
soulseek-rs-lib/src/utils/logger.rs
+149
View File
@@ -0,0 +1,149 @@
use std::env;
use std::sync::Once;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum LogLevel {
Error = 0,
Warn = 1,
Info = 2,
Debug = 3,
Trace = 4,
}
static INIT: Once = Once::new();
static mut LOG_LEVEL: LogLevel = LogLevel::Warn;
pub fn init() {
INIT.call_once(|| {
let level = env::var("LOG_LEVEL")
.or_else(|_| env::var("RUST_LOG"))
.unwrap_or_else(|_| "WARN".to_string())
.to_uppercase();
unsafe {
LOG_LEVEL = match level.as_str() {
"ERROR" => LogLevel::Error,
"WARN" => LogLevel::Warn,
"INFO" => LogLevel::Info,
"DEBUG" => LogLevel::Debug,
"TRACE" => LogLevel::Trace,
"VERBOSE" => LogLevel::Debug, // Map VERBOSE to DEBUG
_ => LogLevel::Warn, // Default to WARN
};
}
});
}
pub fn log(level: LogLevel, message: &str) {
unsafe {
if level <= LOG_LEVEL {
let level_str = match level {
LogLevel::Error => "\x1b[31mERROR\x1b[0m", // Red
LogLevel::Warn => "\x1b[33mWARN\x1b[0m", // Yellow
LogLevel::Info => "\x1b[32mINFO\x1b[0m", // Green
LogLevel::Debug => "\x1b[34mDEBUG\x1b[0m", // Blue
LogLevel::Trace => "\x1b[35mTRACE\x1b[0m", // Magenta
};
let now = std::time::SystemTime::now();
let datetime = now.duration_since(std::time::UNIX_EPOCH).unwrap();
let secs = datetime.as_secs();
let subsec_millis = datetime.subsec_millis();
// Format as YYYY-MM-DD HH:MM:SS.mmm
let days_since_epoch = secs / 86400;
let days_since_1970 = days_since_epoch as i32;
// Calculate year (approximately)
let mut year = 1970;
let mut remaining_days = days_since_1970;
while remaining_days >= 365 {
let is_leap =
(year % 4 == 0 && year % 100 != 0) || (year % 400 == 0);
let days_in_year = if is_leap { 366 } else { 365 };
if remaining_days >= days_in_year {
remaining_days -= days_in_year;
year += 1;
} else {
break;
}
}
// Calculate month and day (simplified)
let month_days = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
let is_leap =
(year % 4 == 0 && year % 100 != 0) || (year % 400 == 0);
let mut month = 1;
let mut day = remaining_days + 1;
for &days_in_month in &month_days {
let actual_days = if month == 2 && is_leap {
29
} else {
days_in_month
};
if day > actual_days {
day -= actual_days;
month += 1;
} else {
break;
}
}
// Calculate time of day
let seconds_in_day = secs % 86400;
let hours = seconds_in_day / 3600;
let minutes = (seconds_in_day % 3600) / 60;
let seconds = seconds_in_day % 60;
eprintln!(
"[{:04}-{:02}-{:02} {:02}:{:02}:{:02}.{:03}] [{}] {}",
year,
month,
day,
hours,
minutes,
seconds,
subsec_millis,
level_str,
message
);
}
}
}
#[macro_export]
macro_rules! error {
($($arg:tt)*) => {
$crate::utils::logger::log($crate::utils::logger::LogLevel::Error, &format!($($arg)*))
};
}
#[macro_export]
macro_rules! warn {
($($arg:tt)*) => {
$crate::utils::logger::log($crate::utils::logger::LogLevel::Warn, &format!($($arg)*))
};
}
#[macro_export]
macro_rules! info {
($($arg:tt)*) => {
$crate::utils::logger::log($crate::utils::logger::LogLevel::Info, &format!($($arg)*))
};
}
#[macro_export]
macro_rules! debug {
($($arg:tt)*) => {
$crate::utils::logger::log($crate::utils::logger::LogLevel::Debug, &format!($($arg)*))
};
}
#[macro_export]
macro_rules! trace {
($($arg:tt)*) => {
$crate::utils::logger::log($crate::utils::logger::LogLevel::Trace, &format!($($arg)*))
};
}
soulseek-rs-lib/src/utils/md5.rs
+474
View File
@@ -0,0 +1,474 @@
use std::convert::TryInto;
// The MD5 Message-Digest Algorithm, written in Rust.
// Implemented by michael dimovich (@mdimovich)
//
// The purpose of this code is to act as a learning resource
// for the md5 message digest algorithm. I wanted to provide
// a straightforward walk-through of the md5 algorithm using
// Rust.
//
// This code was implemented using the original md5 message digest
// RFC, written by R. Rivest in April 1992.
//
// For more information about the md5 algorithm, or the original C
// source code implementation, please reference RFC 1321.
//
// (source: https://datatracker.ietf.org/doc/html/rfc1321)
/**
* This function is specified in the rfc as a way to generate the
* table that we use in rounds 1-4 as part of the hashing operation.
* The function is specified as:
* K[i] = floor(2^32 * abs(sin(i+1)))
*/
fn table_construction_function(i: u32) -> u32 {
let x: f64 = i as f64;
let sin_eval = x.sin().abs();
// note: 4294967296 == 2^32
(4294967296.0 * sin_eval) as u32
}
/**
* 4 auxiliary functions that take 3 32-bit
* integers as input and return 1 32 bit integer as output.
*
* these bitwise operations are the crux of the md5 hashing
* algorithm, and they are performed a number of times to end
* up with our final result.
*/
fn f(x: u32, y: u32, z: u32) -> u32 {
x & y | !x & z
}
fn g(x: u32, y: u32, z: u32) -> u32 {
x & z | y & !z
}
fn h(x: u32, y: u32, z: u32) -> u32 {
x ^ y ^ z
}
fn i(x: u32, y: u32, z: u32) -> u32 {
y ^ (x | !z)
}
// utility function to convert a vector of type T with size N into an array of type T with size N.
fn vec_to_array<T, const N: usize>(v: Vec<T>) -> [T; N] {
v.try_into().unwrap_or_else(|_v: Vec<T>| {
panic!("error converting vector to array - sizes don't match")
})
}
/**
* Round 1 function is specified as 16 operations using the following:
* Let [abcd, k s i] denote the operation -
* a = b + ((a + f(b,c,d) + x[k] + table[i]) <<< s);
* where x[k] is a 32 bit chunk of our original message,
* table[i] is an entry generated by our sine function, s is an unsigned
* integer, and a,b,c,d are our 32 bit registers we perform operations on.
*/
fn round_one_operations(
mut a: u32,
mut b: u32,
mut c: u32,
mut d: u32,
table: &[u32],
x: &[u32],
) -> [u32; 4] {
let _ = x;
let _ = table;
macro_rules! round1 {
( $a:ident, $b:ident, $c:ident, $d:ident, $k:expr, $s:expr, $i: expr ) => {
$a = $b.wrapping_add(
($a.wrapping_add(f($b, $c, $d))
.wrapping_add(x[$k])
.wrapping_add(table[$i]))
.rotate_left($s),
)
};
}
round1!(a, b, c, d, 0, 7, 1);
round1!(d, a, b, c, 1, 12, 2);
round1!(c, d, a, b, 2, 17, 3);
round1!(b, c, d, a, 3, 22, 4);
round1!(a, b, c, d, 4, 7, 5);
round1!(d, a, b, c, 5, 12, 6);
round1!(c, d, a, b, 6, 17, 7);
round1!(b, c, d, a, 7, 22, 8);
round1!(a, b, c, d, 8, 7, 9);
round1!(d, a, b, c, 9, 12, 10);
round1!(c, d, a, b, 10, 17, 11);
round1!(b, c, d, a, 11, 22, 12);
round1!(a, b, c, d, 12, 7, 13);
round1!(d, a, b, c, 13, 12, 14);
round1!(c, d, a, b, 14, 17, 15);
round1!(b, c, d, a, 15, 22, 16);
[a, b, c, d]
}
/**
* Round 2 is equivalent to round 1, except we swap out our f function for g as such:
* a = b + ((a + g(b,c,d) + x[k] + table[i]) <<< s);
*/
fn round_two_operations(
mut a: u32,
mut b: u32,
mut c: u32,
mut d: u32,
table: &[u32],
x: &[u32],
) -> [u32; 4] {
macro_rules! round2 {
( $a:ident, $b:ident, $c:ident, $d:ident, $k:expr, $s:expr, $i:expr) => {
$a = $b.wrapping_add(
($a.wrapping_add(g($b, $c, $d))
.wrapping_add(x[$k])
.wrapping_add(table[$i]))
.rotate_left($s),
)
};
}
round2!(a, b, c, d, 1, 5, 17);
round2!(d, a, b, c, 6, 9, 18);
round2!(c, d, a, b, 11, 14, 19);
round2!(b, c, d, a, 0, 20, 20);
round2!(a, b, c, d, 5, 5, 21);
round2!(d, a, b, c, 10, 9, 22);
round2!(c, d, a, b, 15, 14, 23);
round2!(b, c, d, a, 4, 20, 24);
round2!(a, b, c, d, 9, 5, 25);
round2!(d, a, b, c, 14, 9, 26);
round2!(c, d, a, b, 3, 14, 27);
round2!(b, c, d, a, 8, 20, 28);
round2!(a, b, c, d, 13, 5, 29);
round2!(d, a, b, c, 2, 9, 30);
round2!(c, d, a, b, 7, 14, 31);
round2!(b, c, d, a, 12, 20, 32);
[a, b, c, d]
}
/**
* Round 3 is equivalent to round 1, except we swap out our f function for h as such:
* a = b + ((a + h(b,c,d) + x[k] + table[i]) <<< s);
*/
fn round_three_operations(
mut a: u32,
mut b: u32,
mut c: u32,
mut d: u32,
table: &[u32],
x: &[u32],
) -> [u32; 4] {
macro_rules! round3 {
( $a:ident, $b:ident, $c:ident, $d:ident, $k:expr, $s:expr, $i:expr ) => {
$a = $b.wrapping_add(
($a.wrapping_add(h($b, $c, $d))
.wrapping_add(x[$k])
.wrapping_add(table[$i]))
.rotate_left($s),
)
};
}
round3!(a, b, c, d, 5, 4, 33);
round3!(d, a, b, c, 8, 11, 34);
round3!(c, d, a, b, 11, 16, 35);
round3!(b, c, d, a, 14, 23, 36);
round3!(a, b, c, d, 1, 4, 37);
round3!(d, a, b, c, 4, 11, 38);
round3!(c, d, a, b, 7, 16, 39);
round3!(b, c, d, a, 10, 23, 40);
round3!(a, b, c, d, 13, 4, 41);
round3!(d, a, b, c, 0, 11, 42);
round3!(c, d, a, b, 3, 16, 43);
round3!(b, c, d, a, 6, 23, 44);
round3!(a, b, c, d, 9, 4, 45);
round3!(d, a, b, c, 12, 11, 46);
round3!(c, d, a, b, 15, 16, 47);
round3!(b, c, d, a, 2, 23, 48);
[a, b, c, d]
}
/**
* Round 4 is equivalent to round 1, except we swap out our f function for h as such:
* a = b + ((a + i(b,c,d) + x[k] + table[i]) <<< s);
*/
fn round_four_operations(
mut a: u32,
mut b: u32,
mut c: u32,
mut d: u32,
table: &[u32],
x: &[u32],
) -> [u32; 4] {
macro_rules! round4 {
( $a:ident, $b:ident, $c:ident, $d:ident, $k:expr, $s:expr, $i:expr ) => {
$a = $b.wrapping_add(
($a.wrapping_add(i($b, $c, $d))
.wrapping_add(x[$k])
.wrapping_add(table[$i]))
.rotate_left($s),
)
};
}
round4!(a, b, c, d, 0, 6, 49);
round4!(d, a, b, c, 7, 10, 50);
round4!(c, d, a, b, 14, 15, 51);
round4!(b, c, d, a, 5, 21, 52);
round4!(a, b, c, d, 12, 6, 53);
round4!(d, a, b, c, 3, 10, 54);
round4!(c, d, a, b, 10, 15, 55);
round4!(b, c, d, a, 1, 21, 56);
round4!(a, b, c, d, 8, 6, 57);
round4!(d, a, b, c, 15, 10, 58);
round4!(c, d, a, b, 6, 15, 59);
round4!(b, c, d, a, 13, 21, 60);
round4!(a, b, c, d, 4, 6, 61);
round4!(d, a, b, c, 11, 10, 62);
round4!(c, d, a, b, 2, 15, 63);
round4!(b, c, d, a, 9, 21, 64);
[a, b, c, d]
}
/**
* utility function to iterate over our slice of u8 ints
* and convert into a vector of unsigned 32 bit ints
*/
fn convert_u8_chunk_to_u32(chunk: &mut [u8]) -> Vec<u32> {
let mut x: Vec<u32> = Vec::new();
let mut count = 0;
let mut temporary_vec: Vec<u8> = Vec::new();
// iterate over our block and take
// our 8 bit ints and convert them to
// 32 bit integers
for item in chunk {
temporary_vec.push(*item);
count += 1;
if count == 4 {
let temp_arr: [u8; 4] = vec_to_array(temporary_vec.clone());
let value = u32::from_ne_bytes(temp_arr);
x.push(value);
count = 0;
temporary_vec.clear();
}
}
x
}
fn compute_md5_digest(mut v: Vec<u8>) -> String {
// as described in the rfc,
// 4 32-bit words initialized as fixed constants.
let mut word_a = 0x67452301u32;
let mut word_b = 0xefcdab89u32;
let mut word_c = 0x98badcfeu32;
let mut word_d = 0x10325476u32;
// construct the 64 element constant table.
let table = construct_value_table();
// let M[0 .. N-1] = words of resulting message, where N is multiple of 16
for chunk in v.chunks_exact_mut(64) {
let x = convert_u8_chunk_to_u32(chunk);
// set all values of a,b,c,d to aa,bb,cc, and dd respectively.
// this is to save the initial values.
let word_aa = word_a;
let word_bb = word_b;
let word_cc = word_c;
let word_dd = word_d;
// execute round 1
let result =
round_one_operations(word_a, word_b, word_c, word_d, &table, &x);
word_a = result[0];
word_b = result[1];
word_c = result[2];
word_d = result[3];
// execute round 2
let result =
round_two_operations(word_a, word_b, word_c, word_d, &table, &x);
word_a = result[0];
word_b = result[1];
word_c = result[2];
word_d = result[3];
// execute round 3
let result =
round_three_operations(word_a, word_b, word_c, word_d, &table, &x);
word_a = result[0];
word_b = result[1];
word_c = result[2];
word_d = result[3];
// execute round 4
let result =
round_four_operations(word_a, word_b, word_c, word_d, &table, &x);
word_a = result[0];
word_b = result[1];
word_c = result[2];
word_d = result[3];
// at end of loop iteration, add original word
// to the current word value
word_a = word_a.wrapping_add(word_aa);
word_b = word_b.wrapping_add(word_bb);
word_c = word_c.wrapping_add(word_cc);
word_d = word_d.wrapping_add(word_dd);
}
// format and return the final result, which
// is a 128-bit digest string.
let message_digest = format!(
"{:08x}{:08x}{:08x}{:08x}",
word_a.swap_bytes(),
word_b.swap_bytes(),
word_c.swap_bytes(),
word_d.swap_bytes()
);
message_digest
}
/*
* BIT PADDING STEP
* 1. Append a 1 bit to the end of original message.
* 2. Until the bit length of the data is 64 bits
* short of being a multiple of 512, append 0's
* to the data.
* 3. With the last 64 bits, append the length in 64 bits
* (in lower-order bits first).
*/
fn bit_padding(input: &str) -> Vec<u8> {
let mut input_vector: Vec<u8> = convert_str_to_vec(input);
let bit_length: u64 = (input.len() as u64) * 8u64; // todo - add support for > 2^64 bit size
// 128_u8 is the equivalent of padding 1 as an unsigned 8-bit integer
// with lower-order bits first
input_vector.push(128_u8);
//check if bit length % 512 is 448 (64 less than 512)
while (input_vector.len() * 8) % 512 != 448 {
input_vector.push(0_u8); // push in another 8-bit 0 padded value until the correct
// result is reached;
}
let length_bits_as_u8_array = split_u64_to_u8_array(bit_length);
input_vector.extend(length_bits_as_u8_array);
input_vector
}
fn split_u64_to_u8_array(s: u64) -> [u8; 8] {
[
s as u8,
(s >> 8) as u8,
(s >> 16) as u8,
(s >> 24) as u8,
(s >> 32) as u8,
(s >> 40) as u8,
(s >> 48) as u8,
(s >> 56) as u8,
]
}
fn construct_value_table() -> Vec<u32> {
let mut t: Vec<u32> = Vec::new();
t.push(0x00000000);
for i in 1..=64 {
t.push(table_construction_function(i));
}
t
}
// this should only work with utf-8 encoding and not full unicode support
// due to multi-byte unicode chars
fn convert_str_to_vec(input: &str) -> Vec<u8> {
let mut byte_vec: Vec<u8> = Vec::new();
byte_vec.extend(input.as_bytes());
byte_vec
}
/**
* A Basic Overview of this MD5 Implementation:
* 1. Take in a command line string and convert into
* a vector of unsigned 8 bit ints (Vec<u8>)
* 2. Pad the message with a single 1, followed by n 0's,
* such that the message length in bits % 512 is equal to 448
* 3. Pad the last 64 bits with a little-endian representation
* of the original message length.
* 4. Convert our vector of unsigned 8-bit integers into
* blocks of 32 bit unsigned ints (u32).
* 5. Using the 4 auxiliary functions described in rfc 1321,
* begin a series of 64 bitwise operations on each 512 bit
* block of our message.
* 6. Once we are done, combine our 4 32 bit registers used to
* during our bitwise operations to retrieve our message digest.
* (This is a simplification, but describes the core functionality)
*/
pub fn md5(input: &str) -> String {
let input_vec = bit_padding(input);
compute_md5_digest(input_vec)
}
#[test]
fn a_correct_hash() {
assert_eq!("0cc175b9c0f1b6a831c399e269772661", md5("a"));
}
#[test]
fn empty_string_correct_hash() {
assert_eq!("d41d8cd98f00b204e9800998ecf8427e", md5(""));
}
#[test]
fn abc_correct_hash() {
assert_eq!("900150983cd24fb0d6963f7d28e17f72", md5("abc"));
}
#[test]
fn message_digest_correct_hash() {
assert_eq!("f96b697d7cb7938d525a2f31aaf161d0", md5("message digest"));
}
#[test]
fn lowercase_alphabet_correct_hash() {
assert_eq!(
"c3fcd3d76192e4007dfb496cca67e13b",
md5("abcdefghijklmnopqrstuvwxyz")
)
}
#[test]
fn all_alphanumeric_correct_hash() {
let alpha_num_str =
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
assert_eq!("d174ab98d277d9f5a5611c2c9f419d9f", md5(alpha_num_str));
}
#[test]
fn repeated_numbers_correct_hash() {
let repeated_number_str =
"12345678901234567890123456789012345678901234567890123456789012345678901234567890";
assert_eq!("57edf4a22be3c955ac49da2e2107b67a", md5(repeated_number_str));
}
soulseek-rs-lib/src/utils/mod.rs
+8
View File
@@ -0,0 +1,8 @@
#[macro_use]
pub mod logger;
pub mod md5;
pub mod thread_pool;
pub mod zlib;
// Re-export commonly used items
pub use md5::md5;
soulseek-rs-lib/src/utils/thread_pool.rs
+132
View File
@@ -0,0 +1,132 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
type Job = Box<dyn FnOnce() + Send + 'static>;
enum Message {
NewJob(Job),
Terminate,
}
pub struct ThreadPool {
workers: Vec<Worker>,
sender: Option<mpsc::Sender<Message>>,
active_threads: Arc<AtomicUsize>,
}
impl ThreadPool {
pub fn new(size: usize) -> ThreadPool {
assert!(size > 0);
let (sender, receiver) = mpsc::channel();
let receiver = Arc::new(Mutex::new(receiver));
let active_threads = Arc::new(AtomicUsize::new(0));
let mut workers = Vec::with_capacity(size);
for id in 0..size {
workers.push(Worker::new(
id,
Arc::clone(&receiver),
Arc::clone(&active_threads),
size,
));
}
ThreadPool {
workers,
sender: Some(sender),
active_threads,
}
}
pub fn execute<F>(&self, f: F)
where
F: FnOnce() + Send + 'static,
{
let job = Box::new(f);
if let Some(ref sender) = self.sender {
sender.send(Message::NewJob(job)).unwrap();
}
}
pub fn thread_count(&self) -> usize {
self.workers.len()
}
pub fn active_thread_count(&self) -> usize {
self.active_threads.load(Ordering::SeqCst)
}
}
impl Drop for ThreadPool {
fn drop(&mut self) {
// Send termination message to all workers
if let Some(sender) = self.sender.take() {
for _ in &self.workers {
sender.send(Message::Terminate).unwrap();
}
}
// Wait for all workers to finish
for worker in &mut self.workers {
if let Some(thread) = worker.thread.take() {
thread.join().unwrap();
}
}
}
}
struct Worker {
_id: usize,
thread: Option<thread::JoinHandle<()>>,
}
impl Worker {
fn new(
id: usize,
receiver: Arc<Mutex<mpsc::Receiver<Message>>>,
_active_threads: Arc<AtomicUsize>,
_total_threads: usize,
) -> Worker {
let thread = thread::spawn(move || loop {
let message = receiver.lock().unwrap().recv();
match message {
Ok(Message::NewJob(job)) => {
// let active =
// active_threads.fetch_add(1, Ordering::SeqCst) + 1;
// trace!(
// "Thread {} started job (active: {}/{})",
// id,
// active,
// total_threads
// );
job();
// let active =
// active_threads.fetch_sub(1, Ordering::SeqCst) - 1;
// trace!(
// "Thread {} finished job (active: {}/{})",
// id,
// active,
// total_threads
// );
}
Ok(Message::Terminate) => {
// Received termination signal
break;
}
Err(_) => {
// Channel disconnected
break;
}
}
});
Worker {
_id: id,
thread: Some(thread),
}
}
}
soulseek-rs-lib/src/utils/zlib.rs
+465
View File
@@ -0,0 +1,465 @@
//https://www.rfc-editor.org/rfc/rfc1950
struct BitReader {
mem: Vec<u8>,
pos: usize,
b: u8,
numbits: i32,
}
impl BitReader {
fn new(mem: Vec<u8>) -> Self {
Self {
mem,
pos: 0,
b: 0,
numbits: 0,
}
}
fn read_byte(&mut self) -> std::result::Result<u8, String> {
self.numbits = 0; // discard unread bits
if self.pos >= self.mem.len() {
return Err("End of data".to_string());
}
let b = self.mem[self.pos];
self.pos += 1;
Ok(b)
}
fn read_bit(&mut self) -> std::result::Result<u8, String> {
if self.numbits <= 0 {
self.b = self.read_byte()?;
self.numbits = 8;
}
self.numbits -= 1;
// shift bit out of byte
let bit = self.b & 1;
self.b >>= 1;
Ok(bit)
}
fn read_bits(&mut self, n: usize) -> std::result::Result<u32, String> {
let mut o = 0u32;
for i in 0..n {
o |= (self.read_bit()? as u32) << i;
}
Ok(o)
}
fn read_bytes(&mut self, n: usize) -> std::result::Result<u32, String> {
// read bytes as an integer in little-endian
let mut o = 0u32;
for i in 0..n {
o |= (self.read_byte()? as u32) << (8 * i);
}
Ok(o)
}
}
use crate::error::{Result, SoulseekRs};
pub fn deflate(input: &[u8]) -> Result<Vec<u8>> {
let mut r = BitReader::new(input.to_vec());
let cmf = r.read_byte()?;
let cm = cmf & 15; // Compression method
if cm != 8 {
// only CM=8 is supported
return Err(SoulseekRs::CompressionError("invalid CM".to_string()));
}
let cinfo = (cmf >> 4) & 15; // Compression info
if cinfo > 7 {
return Err(SoulseekRs::CompressionError("invalid CINFO".to_string()));
}
let flg = r.read_byte()?;
if !((cmf as u32) * 256 + (flg as u32)).is_multiple_of(31) {
return Err(SoulseekRs::CompressionError(
"CMF+FLG checksum failed".to_string(),
));
}
let fdict = (flg >> 5) & 1; // preset dictionary?
if fdict != 0 {
return Err(SoulseekRs::CompressionError(
"preset dictionary not supported".to_string(),
));
}
let out = inflate(&mut r).map_err(SoulseekRs::CompressionError)?; // decompress DEFLATE data
let _adler32 = r.read_bytes(4)?; // Adler-32 checksum (for this exercise, we ignore it)
Ok(out)
}
fn inflate(r: &mut BitReader) -> std::result::Result<Vec<u8>, String> {
let mut bfinal = 0;
let mut out = Vec::new();
while bfinal == 0 {
bfinal = r.read_bit()?;
let btype = r.read_bits(2)?;
match btype {
0 => inflate_block_no_compression(r, &mut out)?,
1 => inflate_block_fixed(r, &mut out)?,
2 => inflate_block_dynamic(r, &mut out)?,
_ => return Err("invalid BTYPE".to_string()),
}
}
Ok(out)
}
fn inflate_block_no_compression(
r: &mut BitReader,
o: &mut Vec<u8>,
) -> std::result::Result<(), String> {
let len = r.read_bytes(2)?;
let _nlen = r.read_bytes(2)?;
for _ in 0..len {
o.push(r.read_byte()?);
}
Ok(())
}
#[derive(Clone)]
struct Node {
symbol: Option<u32>,
left: Option<Box<Node>>,
right: Option<Box<Node>>,
}
impl Node {
fn new() -> Self {
Self {
symbol: None,
left: None,
right: None,
}
}
}
struct HuffmanTree {
root: Node,
}
impl HuffmanTree {
fn new() -> Self {
Self { root: Node::new() }
}
fn insert(&mut self, codeword: u32, n: usize, symbol: u32) {
// Insert an entry into the tree mapping `codeword` of len `n` to `symbol`
let mut node = &mut self.root;
for i in (0..n).rev() {
let b = (codeword >> i) & 1;
if b != 0 {
if node.right.is_none() {
node.right = Some(Box::new(Node::new()));
}
node = node.right.as_mut().unwrap();
} else {
if node.left.is_none() {
node.left = Some(Box::new(Node::new()));
}
node = node.left.as_mut().unwrap();
}
}
node.symbol = Some(symbol);
}
}
fn decode_symbol(
r: &mut BitReader,
t: &HuffmanTree,
) -> std::result::Result<u32, String> {
let mut node = &t.root;
while node.left.is_some() || node.right.is_some() {
let b = r.read_bit()?;
node = if b != 0 {
node.right.as_ref().unwrap()
} else {
node.left.as_ref().unwrap()
};
}
node.symbol.ok_or("No symbol found".to_string())
}
const LENGTH_EXTRA_BITS: [usize; 29] = [
0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5,
5, 5, 5, 0,
];
const LENGTH_BASE: [u32; 29] = [
3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59,
67, 83, 99, 115, 131, 163, 195, 227, 258,
];
const DISTANCE_EXTRA_BITS: [usize; 30] = [
0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10,
11, 11, 12, 12, 13, 13,
];
const DISTANCE_BASE: [u32; 30] = [
1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513,
769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
];
fn inflate_block_data(
r: &mut BitReader,
literal_length_tree: &HuffmanTree,
distance_tree: &HuffmanTree,
out: &mut Vec<u8>,
) -> std::result::Result<(), String> {
loop {
let sym = decode_symbol(r, literal_length_tree)?;
if sym <= 255 {
// Literal byte
out.push(sym as u8);
} else if sym == 256 {
// End of block
return Ok(());
} else {
// <length, backward distance> pair
let sym_idx = (sym - 257) as usize;
if sym_idx >= LENGTH_EXTRA_BITS.len() {
return Err("Invalid length symbol".to_string());
}
let length =
r.read_bits(LENGTH_EXTRA_BITS[sym_idx])? + LENGTH_BASE[sym_idx];
let dist_sym = decode_symbol(r, distance_tree)?;
if dist_sym as usize >= DISTANCE_EXTRA_BITS.len() {
return Err("Invalid distance symbol".to_string());
}
let dist = r.read_bits(DISTANCE_EXTRA_BITS[dist_sym as usize])?
+ DISTANCE_BASE[dist_sym as usize];
if dist as usize > out.len() {
return Err("Distance too large".to_string());
}
for _ in 0..length {
let idx = out.len() - dist as usize;
let byte = out[idx];
out.push(byte);
}
}
}
}
fn bl_list_to_tree(bl: &[usize], alphabet: &[u32]) -> HuffmanTree {
let max_bits = *bl.iter().max().unwrap_or(&0);
let mut bl_count = vec![0; max_bits + 1];
for &bitlen in bl {
if bitlen != 0 {
bl_count[bitlen] += 1;
}
}
let mut next_code = vec![0; max_bits + 1];
for bits in 2..=max_bits {
next_code[bits] = (next_code[bits - 1] + bl_count[bits - 1]) << 1;
}
let mut t = HuffmanTree::new();
for (i, &bitlen) in bl.iter().enumerate() {
if bitlen != 0 && i < alphabet.len() {
t.insert(next_code[bitlen], bitlen, alphabet[i]);
next_code[bitlen] += 1;
}
}
t
}
const CODE_LENGTH_CODES_ORDER: [usize; 19] = [
16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
];
fn decode_trees(
r: &mut BitReader,
) -> std::result::Result<(HuffmanTree, HuffmanTree), String> {
// The number of literal/length codes
let hlit = r.read_bits(5)? + 257;
// The number of distance codes
let hdist = r.read_bits(5)? + 1;
// The number of code length codes
let hclen = r.read_bits(4)? + 4;
// Read code lengths for the code length alphabet
let mut code_length_tree_bl = vec![0; 19];
for i in 0..hclen as usize {
code_length_tree_bl[CODE_LENGTH_CODES_ORDER[i]] =
r.read_bits(3)? as usize;
}
// Construct code length tree
let code_length_alphabet: Vec<u32> = (0..19).collect();
let code_length_tree =
bl_list_to_tree(&code_length_tree_bl, &code_length_alphabet);
// Read literal/length + distance code length list
let mut bl = Vec::new();
while bl.len() < (hlit + hdist) as usize {
let sym = decode_symbol(r, &code_length_tree)?;
if sym <= 15 {
// literal value
bl.push(sym as usize);
} else if sym == 16 {
// copy the previous code length 3..6 times.
// the next 2 bits indicate repeat length ( 0 = 3, ..., 3 = 6 )
if bl.is_empty() {
return Err("No previous code length".to_string());
}
let prev_code_length = bl[bl.len() - 1];
let repeat_length = r.read_bits(2)? + 3;
for _ in 0..repeat_length {
bl.push(prev_code_length);
}
} else if sym == 17 {
// repeat code length 0 for 3..10 times. (3 bits of length)
let repeat_length = r.read_bits(3)? + 3;
bl.resize(bl.len() + repeat_length as usize, 0);
} else if sym == 18 {
// repeat code length 0 for 11..138 times. (7 bits of length)
let repeat_length = r.read_bits(7)? + 11;
bl.resize(bl.len() + repeat_length as usize, 0);
} else {
return Err("Invalid symbol".to_string());
}
}
// Construct trees
let literal_length_alphabet: Vec<u32> = (0..286).collect();
let literal_length_tree =
bl_list_to_tree(&bl[..hlit as usize], &literal_length_alphabet);
let distance_alphabet: Vec<u32> = (0..30).collect();
let distance_tree =
bl_list_to_tree(&bl[hlit as usize..], &distance_alphabet);
Ok((literal_length_tree, distance_tree))
}
fn inflate_block_dynamic(
r: &mut BitReader,
o: &mut Vec<u8>,
) -> std::result::Result<(), String> {
let (literal_length_tree, distance_tree) = decode_trees(r)?;
inflate_block_data(r, &literal_length_tree, &distance_tree, o)
}
fn inflate_block_fixed(
r: &mut BitReader,
o: &mut Vec<u8>,
) -> std::result::Result<(), String> {
let mut bl = Vec::new();
bl.extend(vec![8; 144]); // 0-143: 8 bits
bl.extend(vec![9; 112]); // 144-255: 9 bits
bl.extend(vec![7; 24]); // 256-279: 7 bits
bl.extend(vec![8; 8]); // 280-287: 8 bits
let literal_length_alphabet: Vec<u32> = (0..286).collect();
let literal_length_tree = bl_list_to_tree(&bl, &literal_length_alphabet);
let bl_dist = vec![5; 30];
let distance_alphabet: Vec<u32> = (0..30).collect();
let distance_tree = bl_list_to_tree(&bl_dist, &distance_alphabet);
inflate_block_data(r, &literal_length_tree, &distance_tree, o)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bitreader_read_bits() {
let data = vec![0b11010010, 0b10110101];
let mut reader = BitReader::new(data);
assert_eq!(reader.read_bits(3).unwrap(), 0b010); // First 3 bits: 010
assert_eq!(reader.read_bits(5).unwrap(), 0b11010); // Next 5 bits: 11010
}
#[test]
fn test_bitreader_read_bytes() {
let data = vec![0x12, 0x34, 0x56, 0x78];
let mut reader = BitReader::new(data);
assert_eq!(reader.read_bytes(2).unwrap(), 0x3412); // Little-endian: 0x3412
assert_eq!(reader.read_bytes(2).unwrap(), 0x7856); // Little-endian: 0x7856
}
#[test]
fn test_extract_header_success() {
let _data = [120, 156]; // Valid zlib header
let result = deflate(&[120, 156, 3, 0, 0, 0, 0, 1]); // Minimal valid zlib stream
assert!(result.is_ok());
}
#[test]
fn test_extract_header_fail_to_short() {
let data = vec![120]; // Too short
let mut reader = BitReader::new(data);
let result = reader.read_byte();
assert!(result.is_ok());
let result = reader.read_byte();
assert!(result.is_err());
}
#[test]
fn test_deflate() {
// Test data from the original test - this should work with our new implementation
let data = vec![
120, 156, 203, 72, 205, 201, 201, 87, 8, 207, 47, 202, 73, 1, 0,
24, 11, 4, 93,
];
let result = deflate(&data);
assert!(result.is_ok());
let decompressed = result.unwrap();
assert_eq!(decompressed, b"hello World");
}
#[test]
fn test_deflate2() {
let data = vec![
120, 156, 99, 103, 96, 96, 72, 201, 79, 201, 76, 79, 204, 203, 213,
158, 98, 194, 4, 228, 50, 250, 3, 9, 7, 135, 162, 156, 148, 194,
188, 152, 228, 252, 220, 130, 156, 212, 146, 212, 24, 231, 196,
188, 228, 204, 252, 188, 212, 226, 152, 144, 162, 210, 226, 226,
212, 28, 93, 67, 75, 115, 75, 93, 119, 160, 144, 130, 91, 126, 145,
66, 72, 70, 170, 66, 120, 106, 106, 118, 106, 94, 138, 174, 161,
89, 82, 102, 137, 174, 137, 137, 142, 161, 119, 70, 149, 94, 90,
78, 98, 114, 203, 175, 243, 32, 163, 193, 128, 25, 100, 7, 16, 23,
0, 9, 22, 32, 237, 178, 134, 129, 129, 21, 72, 11, 128, 196, 243,
176, 217, 29, 156, 153, 151, 158, 147, 90, 12, 54, 95, 193, 216,
84, 193, 200, 192, 200, 36, 198, 45, 181, 168, 40, 53, 57, 91, 193,
37, 177, 60, 79, 71, 193, 55, 177, 44, 181, 40, 19, 200, 13, 78,
76, 42, 74, 85, 80, 83, 240, 75, 45, 7, 10, 38, 103, 100, 2, 221,
167, 139, 238, 66, 5, 13, 144, 17, 154, 96, 167, 173, 228, 215, 98,
68, 119, 218, 74, 6, 76, 167, 49, 60, 153, 202, 200, 160, 199, 128,
0, 0, 161, 99, 76, 142,
];
let expect = vec![
7, 0, 0, 0, 100, 111, 100, 105, 103, 97, 110, 109, 43, 148, 52, 2,
0, 0, 0, 1, 79, 0, 0, 0, 64, 64, 114, 108, 100, 113, 110, 92, 99,
111, 109, 112, 108, 101, 116, 101, 92, 67, 97, 110, 99, 105, 111,
110, 101, 115, 92, 84, 114, 117, 115, 115, 101, 108, 45, 49, 57,
55, 57, 45, 71, 111, 110, 101, 32, 70, 111, 114, 32, 84, 104, 101,
32, 87, 101, 101, 107, 101, 110, 100, 45, 49, 54, 98, 105, 116, 45,
52, 52, 44, 49, 75, 104, 122, 46, 102, 108, 97, 99, 132, 250, 207,
2, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 112, 1, 0, 0, 4,
0, 0, 0, 68, 172, 0, 0, 5, 0, 0, 0, 16, 0, 0, 0, 1, 110, 0, 0, 0,
64, 64, 114, 108, 100, 113, 110, 92, 99, 111, 109, 112, 108, 101,
116, 101, 92, 83, 105, 110, 103, 108, 101, 115, 32, 87, 101, 101,
107, 32, 51, 53, 32, 50, 48, 50, 52, 92, 70, 101, 114, 114, 101,
99, 107, 32, 68, 97, 119, 110, 44, 32, 77, 97, 118, 101, 114, 105,
99, 107, 32, 83, 97, 98, 114, 101, 32, 38, 32, 78, 101, 119, 32,
77, 97, 99, 104, 105, 110, 101, 32, 45, 32, 70, 111, 114, 32, 84,
104, 101, 32, 87, 101, 101, 107, 101, 110, 100, 32, 40, 50, 48, 50,
52, 41, 46, 102, 108, 97, 99, 169, 15, 42, 1, 0, 0, 0, 0, 0, 0, 0,
0, 3, 0, 0, 0, 1, 0, 0, 0, 169, 0, 0, 0, 4, 0, 0, 0, 68, 172, 0, 0,
5, 0, 0, 0, 16, 0, 0, 0, 0, 228, 149, 1, 0, 46, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0,
]
.to_vec();
let result = deflate(&data);
assert!(result.is_ok());
let decompressed = result.unwrap();
assert_eq!(decompressed, expect);
}
}