Skip to main content

pg/protocol/
mod.rs

1pub mod backend;
2pub mod frontend;
3pub mod scram;
4
5pub use backend::{BackendDecoder, BackendMessage, FieldDescription};
6pub use frontend::FrontendMessage;
7pub(crate) use scram::ScramClient;
8
9pub(crate) fn base64_encode(data: &[u8]) -> String {
10    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
11    let mut result = Vec::new();
12    for chunk in data.chunks(3) {
13        let b0 = chunk[0] as u32;
14        let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
15        let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
16        let triple = (b0 << 16) | (b1 << 8) | b2;
17
18        result.push(CHARS[((triple >> 18) & 0x3F) as usize]);
19        result.push(CHARS[((triple >> 12) & 0x3F) as usize]);
20
21        if chunk.len() > 1 {
22            result.push(CHARS[((triple >> 6) & 0x3F) as usize]);
23        } else {
24            result.push(b'=');
25        }
26        if chunk.len() > 2 {
27            result.push(CHARS[(triple & 0x3F) as usize]);
28        } else {
29            result.push(b'=');
30        }
31    }
32    String::from_utf8(result).unwrap()
33}
34
35pub(crate) fn base64_decode(input: &str) -> std::result::Result<Vec<u8>, String> {
36    let input = input.trim_end_matches('=');
37    let mut result = Vec::new();
38    let mut buffer = 0u32;
39    let mut bits_collected = 0;
40
41    for &byte in input.as_bytes() {
42        let value = match byte {
43            b'A'..=b'Z' => byte - b'A',
44            b'a'..=b'z' => byte - b'a' + 26,
45            b'0'..=b'9' => byte - b'0' + 52,
46            b'+' => 62,
47            b'/' => 63,
48            _ => return Err(format!("invalid base64 character: {}", byte as char)),
49        } as u32;
50
51        buffer = (buffer << 6) | value;
52        bits_collected += 6;
53
54        if bits_collected >= 8 {
55            bits_collected -= 8;
56            result.push((buffer >> bits_collected) as u8);
57            buffer &= (1 << bits_collected) - 1;
58        }
59    }
60
61    Ok(result)
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_base64_roundtrip() {
70        let data = b"hello world";
71        let encoded = base64_encode(data);
72        let decoded = base64_decode(&encoded).unwrap();
73        assert_eq!(data, &decoded[..]);
74    }
75
76    #[test]
77    fn test_base64_scram() {
78        let data = b"Client Key";
79        let encoded = base64_encode(data);
80        let decoded = base64_decode(&encoded).unwrap();
81        assert_eq!(data, &decoded[..]);
82    }
83}