TLA  Line data    Source code

       1                 : //! Definitions for SASL messages.
       2                 : 
       3                 : use crate::parse::{split_at_const, split_cstr};
       4                 : use pq_proto::{BeAuthenticationSaslMessage, BeMessage};
       5                 : 
       6                 : /// SASL-specific payload of [`PasswordMessage`](pq_proto::FeMessage::PasswordMessage).
       7 UBC           0 : #[derive(Debug)]
       8                 : pub struct FirstMessage<'a> {
       9                 :     /// Authentication method, e.g. `"SCRAM-SHA-256"`.
      10                 :     pub method: &'a str,
      11                 :     /// Initial client message.
      12                 :     pub message: &'a str,
      13                 : }
      14                 : 
      15                 : impl<'a> FirstMessage<'a> {
      16                 :     // NB: FromStr doesn't work with lifetimes
      17 CBC          32 :     pub fn parse(bytes: &'a [u8]) -> Option<Self> {
      18              32 :         let (method_cstr, tail) = split_cstr(bytes)?;
      19              32 :         let method = method_cstr.to_str().ok()?;
      20                 : 
      21              32 :         let (len_bytes, bytes) = split_at_const(tail)?;
      22              32 :         let len = u32::from_be_bytes(*len_bytes) as usize;
      23              32 :         if len != bytes.len() {
      24 UBC           0 :             return None;
      25 CBC          32 :         }
      26                 : 
      27              32 :         let message = std::str::from_utf8(bytes).ok()?;
      28              32 :         Some(Self { method, message })
      29              32 :     }
      30                 : }
      31                 : 
      32                 : /// A single SASL message.
      33                 : /// This struct is deliberately decoupled from lower-level
      34                 : /// [`BeAuthenticationSaslMessage`].
      35 UBC           0 : #[derive(Debug)]
      36                 : pub(super) enum ServerMessage<T> {
      37                 :     /// We expect to see more steps.
      38                 :     Continue(T),
      39                 :     /// This is the final step.
      40                 :     Final(T),
      41                 : }
      42                 : 
      43                 : impl<'a> ServerMessage<&'a str> {
      44 CBC          58 :     pub(super) fn to_reply(&self) -> BeMessage<'a> {
      45              58 :         use BeAuthenticationSaslMessage::*;
      46              58 :         BeMessage::AuthenticationSasl(match self {
      47              31 :             ServerMessage::Continue(s) => Continue(s.as_bytes()),
      48              27 :             ServerMessage::Final(s) => Final(s.as_bytes()),
      49                 :         })
      50              58 :     }
      51                 : }
      52                 : 
      53                 : #[cfg(test)]
      54                 : mod tests {
      55                 :     use super::*;
      56                 : 
      57               1 :     #[test]
      58               1 :     fn parse_sasl_first_message() {
      59               1 :         let proto = "SCRAM-SHA-256";
      60               1 :         let sasl = "n,,n=,r=KHQ2Gjc7NptyB8aov5/TnUy4";
      61               1 :         let sasl_len = (sasl.len() as u32).to_be_bytes();
      62               1 :         let bytes = [proto.as_bytes(), &[0], sasl_len.as_ref(), sasl.as_bytes()].concat();
      63               1 : 
      64               1 :         let password = FirstMessage::parse(&bytes).unwrap();
      65               1 :         assert_eq!(password.method, proto);
      66               1 :         assert_eq!(password.message, sasl);
      67               1 :     }
      68                 : }