1 files changed, 157 insertions, 0 deletions
diff --git a/aero-sasl/src/encode.rs b/aero-sasl/src/encode.rs
new file mode 100644
index 0000000..625d035
--- /dev/null
+++ b/aero-sasl/src/encode.rs
@@ -0,0 +1,157 @@
+use anyhow::Result;
+use base64::Engine;
+use tokio_util::bytes::{BufMut, BytesMut};
+
+use super::types::*;
+
+pub trait Encode {
+    fn encode(&self, out: &mut BytesMut) -> Result<()>;
+}
+
+fn tab_enc(out: &mut BytesMut) {
+    out.put(&[0x09][..])
+}
+
+fn lf_enc(out: &mut BytesMut) {
+    out.put(&[0x0A][..])
+}
+
+impl Encode for Mechanism {
+    fn encode(&self, out: &mut BytesMut) -> Result<()> {
+        match self {
+            Self::Plain => out.put(&b"PLAIN"[..]),
+            Self::Login => out.put(&b"LOGIN"[..]),
+        }
+        Ok(())
+    }
+}
+
+impl Encode for MechanismParameters {
+    fn encode(&self, out: &mut BytesMut) -> Result<()> {
+        match self {
+            Self::Anonymous => out.put(&b"anonymous"[..]),
+            Self::PlainText => out.put(&b"plaintext"[..]),
+            Self::Dictionary => out.put(&b"dictionary"[..]),
+            Self::Active => out.put(&b"active"[..]),
+            Self::ForwardSecrecy => out.put(&b"forward-secrecy"[..]),
+            Self::MutualAuth => out.put(&b"mutual-auth"[..]),
+            Self::Private => out.put(&b"private"[..]),
+        }
+        Ok(())
+    }
+}
+
+impl Encode for FailCode {
+    fn encode(&self, out: &mut BytesMut) -> Result<()> {
+        match self {
+            Self::TempFail => out.put(&b"temp_fail"[..]),
+            Self::AuthzFail => out.put(&b"authz_fail"[..]),
+            Self::UserDisabled => out.put(&b"user_disabled"[..]),
+            Self::PassExpired => out.put(&b"pass_expired"[..]),
+        };
+        Ok(())
+    }
+}
+
+impl Encode for ServerCommand {
+    fn encode(&self, out: &mut BytesMut) -> Result<()> {
+        match self {
+            Self::Version(Version { major, minor }) => {
+                out.put(&b"VERSION"[..]);
+                tab_enc(out);
+                out.put(major.to_string().as_bytes());
+                tab_enc(out);
+                out.put(minor.to_string().as_bytes());
+                lf_enc(out);
+            }
+            Self::Spid(pid) => {
+                out.put(&b"SPID"[..]);
+                tab_enc(out);
+                out.put(pid.to_string().as_bytes());
+                lf_enc(out);
+            }
+            Self::Cuid(pid) => {
+                out.put(&b"CUID"[..]);
+                tab_enc(out);
+                out.put(pid.to_string().as_bytes());
+                lf_enc(out);
+            }
+            Self::Cookie(cval) => {
+                out.put(&b"COOKIE"[..]);
+                tab_enc(out);
+                out.put(hex::encode(cval).as_bytes());
+                lf_enc(out);
+            }
+            Self::Mech { kind, parameters } => {
+                out.put(&b"MECH"[..]);
+                tab_enc(out);
+                kind.encode(out)?;
+                for p in parameters.iter() {
+                    tab_enc(out);
+                    p.encode(out)?;
+                }
+                lf_enc(out);
+            }
+            Self::Done => {
+                out.put(&b"DONE"[..]);
+                lf_enc(out);
+            }
+            Self::Cont { id, data } => {
+                out.put(&b"CONT"[..]);
+                tab_enc(out);
+                out.put(id.to_string().as_bytes());
+                tab_enc(out);
+                if let Some(rdata) = data {
+                    let b64 = base64::engine::general_purpose::STANDARD.encode(rdata);
+                    out.put(b64.as_bytes());
+                }
+                lf_enc(out);
+            }
+            Self::Ok {
+                id,
+                user_id,
+                extra_parameters,
+            } => {
+                out.put(&b"OK"[..]);
+                tab_enc(out);
+                out.put(id.to_string().as_bytes());
+                if let Some(user) = user_id {
+                    tab_enc(out);
+                    out.put(&b"user="[..]);
+                    out.put(user.as_bytes());
+                }
+                for p in extra_parameters.iter() {
+                    tab_enc(out);
+                    out.put(&p[..]);
+                }
+                lf_enc(out);
+            }
+            Self::Fail {
+                id,
+                user_id,
+                code,
+                extra_parameters,
+            } => {
+                out.put(&b"FAIL"[..]);
+                tab_enc(out);
+                out.put(id.to_string().as_bytes());
+                if let Some(user) = user_id {
+                    tab_enc(out);
+                    out.put(&b"user="[..]);
+                    out.put(user.as_bytes());
+                }
+                if let Some(code_val) = code {
+                    tab_enc(out);
+                    out.put(&b"code="[..]);
+                    code_val.encode(out)?;
+                }
+                for p in extra_parameters.iter() {
+                    tab_enc(out);
+                    out.put(&p[..]);
+                }
+                lf_enc(out);
+            }
+        }
+        Ok(())
+    }
+}