b64_rs/src/encode.rs
2025-01-04 19:04:56 -05:00

119 lines
3.5 KiB
Rust

use {
crate::{error::Error, B64Alphabet},
std::{
fmt::Write,
io::{ErrorKind, Read},
},
};
pub struct Encoder<R: Read, W: Write> {
reader: R,
writer: W,
alphabet: B64Alphabet,
wrap: Option<usize>,
}
impl<R: Read, W: Write> Encoder<R, W> {
/// Creates a new encoder with the given reader and writer. If alphabet is
/// `None` the encoder will use the default rfc4648 alphabet.
pub fn new(reader: R, writer: W, alphabet: Option<B64Alphabet>, wrap: Option<usize>) -> Self {
Self {
reader,
writer,
alphabet: alphabet.unwrap_or_default(),
wrap,
}
}
pub fn encode(&mut self) -> Result<(), Error> {
let mut total: usize = 0;
loop {
let mut ibuf = [0; 3];
let mut obuf = [self.alphabet.pad(); 4];
let mut n_bytes = 0;
let mut num: u64 = 0;
loop {
n_bytes += match self.reader.read(&mut ibuf) {
Ok(n) => n,
Err(e) if e.kind() == ErrorKind::Interrupted => continue,
Err(e) => return Err(e.into()),
};
break;
}
if n_bytes == 0 {
break;
}
for (idx, n) in ibuf.iter().enumerate() {
num <<= 8;
if idx < n_bytes {
num |= *n as u64;
}
}
let mut outlen = n_bytes * 8 / 6;
if n_bytes * 8 % 6 > 0 {
outlen += 1;
}
for idx in (0..4).rev() {
if outlen == 4 || idx < outlen {
let b = num & 0b111111;
obuf[idx] = self.alphabet.items[b as usize];
}
num >>= 6;
}
if let Some(wrap) = self.wrap {
for idx in 0..=3 {
write!(self.writer, "{}", obuf[idx])?;
total += 1;
if total % wrap == 0 {
writeln!(self.writer)?;
}
}
} else {
write!(self.writer, "{}{}{}{}", obuf[0], obuf[1], obuf[2], obuf[3])?;
}
if outlen < 4 {
break;
}
}
Ok(())
}
}
impl<R: Read> Encoder<R, String> {
pub fn output(self) -> String {
self.writer
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs::{self, File};
#[test]
fn encode() {
let mut encoder = Encoder::new("Hello, World".as_bytes(), String::new(), None, None);
encoder.encode().unwrap();
assert_eq!(encoder.output(), "SGVsbG8sIFdvcmxk");
encoder = Encoder {
reader: "Hello, World!".as_bytes(),
writer: String::new(),
alphabet: B64Alphabet::default(),
wrap: None,
};
encoder.encode().unwrap();
assert_eq!(encoder.output(), "SGVsbG8sIFdvcmxkIQ==");
encoder = Encoder::new("Hello, World!\n".as_bytes(), String::new(), None, None);
encoder.encode().unwrap();
assert_eq!(encoder.output(), "SGVsbG8sIFdvcmxkIQo=");
}
#[test]
fn wrapping() {
let infile = File::open("src/testdata/lorem.txt").unwrap();
let outfile = fs::read_to_string("src/testdata/lorem_b64.txt").unwrap();
let mut encoder = Encoder::new(infile, String::new(), None, Some(76));
encoder.encode().unwrap();
assert_eq!(encoder.output(), outfile);
}
}