1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use crate::io::BufMut;
use crate::postgres::connection::PgConnection;
use crate::postgres::protocol::authentication::Authentication::SaslContinue;
use crate::postgres::protocol::Encode;
use crate::postgres::protocol::Message;
use crate::Result;
use byteorder::NetworkEndian;
use hmac::{Hmac, Mac};
use sha2::{Digest, Sha256};

pub struct SaslInitialResponse<'a>(pub &'a str);

impl<'a> Encode for SaslInitialResponse<'a> {
    fn encode(&self, buf: &mut Vec<u8>) {
        let len = self.0.as_bytes().len() as u32;
        buf.push(b'p');
        buf.put_u32::<NetworkEndian>(4u32 + len + 14u32 + 4u32);
        buf.put_str_nul("SCRAM-SHA-256");
        buf.put_u32::<NetworkEndian>(len);
        buf.extend_from_slice(self.0.as_bytes());
    }
}

pub struct SaslResponse<'a>(pub &'a str);

impl<'a> Encode for SaslResponse<'a> {
    fn encode(&self, buf: &mut Vec<u8>) {
        buf.push(b'p');
        buf.put_u32::<NetworkEndian>(4u32 + self.0.as_bytes().len() as u32);
        buf.extend_from_slice(self.0.as_bytes());
    }
}

// Hi(str, salt, i):
pub fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32]> {
    let mut mac = Hmac::<Sha256>::new_varkey(s.as_bytes())
        .map_err(|_| protocol_err!("HMAC can take key of any size"))?;

    mac.input(&salt);
    mac.input(&1u32.to_be_bytes());

    let mut u = mac.result().code();
    let mut hi = u;

    for _ in 1..iter_count {
        let mut mac = Hmac::<Sha256>::new_varkey(s.as_bytes())
            .map_err(|_| protocol_err!(" HMAC can take key of any size"))?;
        mac.input(u.as_slice());
        u = mac.result().code();
        hi = hi.iter().zip(u.iter()).map(|(&a, &b)| a ^ b).collect();
    }

    Ok(hi.into())
}