init

1 files changed, 372 insertions, 0 deletions

diff --git a/src/key.rs b/src/key.rs
new file mode 100644
index 0000000..19bc127
--- /dev/null
+++ b/src/key.rs
@@ -0,0 +1,372 @@
+use base64::Engine;
+use netlink_packet_wireguard::constants::WG_KEY_LEN;
+use rand::{rngs::OsRng, RngCore};
+
+// Code from: https://git.zx2c4.com/wireguard-tools/tree/contrib/embeddable-wg-library/wireguard.c
+
+type Fe = [i64; 16];
+
+#[derive(Debug)]
+pub struct KeyDecodeError;
+
+impl std::fmt::Display for KeyDecodeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Key decode error")
+    }
+}
+
+impl std::error::Error for KeyDecodeError {}
+
+#[derive(Clone, Default, Copy, PartialEq, Eq, Hash)]
+pub struct Key([u8; WG_KEY_LEN]);
+
+impl Key {
+    pub fn into_array(self) -> [u8; WG_KEY_LEN] {
+        self.0
+    }
+
+    pub fn as_array(&self) -> &[u8; WG_KEY_LEN] {
+        &self.0
+    }
+
+    pub fn as_slice(&self) -> &[u8] {
+        &self.0
+    }
+
+    pub fn encode(&self) -> String {
+        base64::engine::general_purpose::STANDARD.encode(&self.0)
+    }
+
+    pub fn decode(encoded: &str) -> Result<Self, KeyDecodeError> {
+        let decoded = base64::engine::general_purpose::STANDARD
+            .decode(encoded)
+            .map_err(|_| KeyDecodeError)?;
+        if decoded.len() != WG_KEY_LEN {
+            return Err(KeyDecodeError);
+        }
+        let mut key = [0u8; WG_KEY_LEN];
+        key.copy_from_slice(&decoded);
+        Ok(Key(key))
+    }
+
+    pub fn generate_pub_priv() -> (Self, Self) {
+        let private_key = Self::generate_private();
+        let public_key = Self::generate_public(&private_key);
+        (public_key, private_key)
+    }
+
+    pub fn generate_public(private: &Key) -> Self {
+        let mut r: i32 = Default::default();
+        let mut public_key: [u8; WG_KEY_LEN] = Default::default();
+        let mut z: [u8; WG_KEY_LEN] = private.0;
+        let mut a = fe_new_one(1);
+        let mut b = fe_new_one(9);
+        let mut c = fe_new_one(0);
+        let mut d = fe_new_one(1);
+        let mut e = fe_new_default();
+        let mut f = fe_new_default();
+
+        clamp_key(&mut z);
+
+        for i in (0..=254i32).rev() {
+            r = ((z[(i >> 3) as usize] >> (i & 7)) & 1) as i32;
+            cswap(&mut a, &mut b, r);
+            cswap(&mut c, &mut d, r);
+            add(&mut e, &a, &c);
+            {
+                let a_clone = a;
+                subtract(&mut a, &a_clone, &c);
+            }
+            add(&mut c, &b, &d);
+            {
+                let b_clone = b;
+                subtract(&mut b, &b_clone, &d);
+            }
+            multmod(&mut d, &e, &e);
+            multmod(&mut f, &a, &a);
+            {
+                let a_clone = a;
+                multmod(&mut a, &c, &a_clone);
+            }
+            multmod(&mut c, &b, &e);
+            add(&mut e, &a, &c);
+            {
+                let a_clone = a;
+                subtract(&mut a, &a_clone, &c);
+            }
+            multmod(&mut b, &a, &a);
+            subtract(&mut c, &d, &f);
+            //multmod(&mut a, &c, (const fe){ 0xdb41, 1 });
+            multmod(&mut a, &c, &fe_new_two(0xdb41, 1));
+            {
+                let a_clone = a;
+                add(&mut a, &a_clone, &d);
+            }
+            {
+                let c_clone = c;
+                multmod(&mut c, &c_clone, &a);
+            }
+            multmod(&mut a, &d, &f);
+            multmod(&mut d, &b, &fe_new_one(9));
+            multmod(&mut b, &e, &e);
+            cswap(&mut a, &mut b, r);
+            cswap(&mut c, &mut d, r);
+        }
+        {
+            let c_clone = c;
+            invert(&mut c, &c_clone);
+        }
+        {
+            let a_clone = a;
+            multmod(&mut a, &a_clone, &c);
+        }
+        pack(&mut public_key, &a);
+
+        memzero_explicit(&mut r);
+        memzero_explicit(&mut z);
+        memzero_explicit(&mut a);
+        memzero_explicit(&mut b);
+        memzero_explicit(&mut c);
+        memzero_explicit(&mut d);
+        memzero_explicit(&mut e);
+        memzero_explicit(&mut f);
+
+        Self(public_key)
+    }
+
+    pub fn generate_private() -> Self {
+        let mut preshared = Self::generate_preshared();
+        clamp_key(&mut preshared.0);
+        preshared
+    }
+
+    pub fn generate_preshared() -> Self {
+        let mut key = [0u8; WG_KEY_LEN];
+        OsRng.fill_bytes(&mut key);
+        Self(key)
+    }
+}
+
+impl From<&[u8; WG_KEY_LEN]> for Key {
+    fn from(k: &[u8; WG_KEY_LEN]) -> Self {
+        Self(*k)
+    }
+}
+
+impl From<[u8; WG_KEY_LEN]> for Key {
+    fn from(k: [u8; WG_KEY_LEN]) -> Self {
+        Self(k)
+    }
+}
+
+impl std::str::FromStr for Key {
+    type Err = KeyDecodeError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Key::decode(s)
+    }
+}
+
+impl std::fmt::Debug for Key {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut buf = [0; WG_KEY_LEN * 2];
+        let len = base64::engine::general_purpose::STANDARD
+            .encode_slice(&self.0, &mut buf)
+            .expect("base64 should encode");
+        let b64 = std::str::from_utf8(&buf[..len]).expect("base64 should be valid utf-8");
+        f.debug_tuple("Key").field(&b64).finish()
+    }
+}
+
+impl std::fmt::Display for Key {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut buf = [0; WG_KEY_LEN * 2];
+        let len = base64::engine::general_purpose::STANDARD
+            .encode_slice(&self.0, &mut buf)
+            .expect("base64 should encode");
+        let b64 = std::str::from_utf8(&buf[..len]).expect("base64 should be valid utf-8");
+        f.write_str(b64)
+    }
+}
+
+impl serde::Serialize for Key {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serializer.serialize_str(&self.encode())
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for Key {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let s = String::deserialize(deserializer)?;
+        Key::decode(&s).map_err(serde::de::Error::custom)
+    }
+}
+
+fn fe_new_default() -> Fe {
+    Default::default()
+}
+
+fn fe_new_one(x: i64) -> Fe {
+    let mut fe = fe_new_default();
+    fe[0] = x;
+    fe
+}
+
+fn fe_new_two(x: i64, y: i64) -> Fe {
+    let mut fe = fe_new_default();
+    fe[0] = x;
+    fe[1] = y;
+    fe
+}
+
+fn clamp_key(key: &mut [u8]) {
+    key[31] = (key[31] & 127) | 64;
+    key[0] &= 248;
+}
+
+fn carry(o: &mut Fe) {
+    for i in 0..16 {
+        let x = if i == 15 { 38 } else { 1 };
+        o[(i + 1) % 16] += x * (o[i] >> 16);
+        o[i] &= 0xffff;
+    }
+}
+
+fn cswap(p: &mut Fe, q: &mut Fe, mut b: i32) {
+    let mut t: i64 = 0;
+    let mut c: i64 = !i64::from(b).wrapping_sub(1);
+
+    for i in 0..16 {
+        t = c & (p[i] ^ q[i]);
+        p[i] ^= t;
+        q[i] ^= t;
+    }
+
+    memzero_explicit(&mut t);
+    memzero_explicit(&mut c);
+    memzero_explicit(&mut b);
+}
+
+fn pack(o: &mut [u8; WG_KEY_LEN], n: &Fe) {
+    let mut b: i32 = 0;
+    let mut t: Fe = fe_new_default();
+    let mut m: Fe = fe_new_default();
+
+    t.copy_from_slice(n);
+    carry(&mut t);
+    carry(&mut t);
+    carry(&mut t);
+    for _ in 0..2 {
+        m[0] = t[0] - 0xffed;
+        for i in 1..15 {
+            m[i] = t[i] - 0xffff - ((m[i - 1] >> 16) & 1);
+            m[i - 1] &= 0xffff;
+        }
+        m[15] = t[15] - 0x7fff - ((m[14] >> 16) & 1);
+        b = ((m[15] >> 16) & 1) as i32;
+        m[14] &= 0xffff;
+        cswap(&mut t, &mut m, 1 - b);
+    }
+    for i in 0..16 {
+        o[2 * i] = (t[i] & 0xff) as u8;
+        o[2 * i + 1] = (t[i] >> 8) as u8;
+    }
+
+    memzero_explicit(&mut m);
+    memzero_explicit(&mut t);
+    memzero_explicit(&mut b);
+}
+
+fn add(o: &mut Fe, a: &Fe, b: &Fe) {
+    for i in 0..16 {
+        o[i] = a[i] + b[i];
+    }
+}
+
+fn subtract(o: &mut Fe, a: &Fe, b: &Fe) {
+    for i in 0..16 {
+        o[i] = a[i] - b[i];
+    }
+}
+
+fn multmod(o: &mut Fe, a: &Fe, b: &Fe) {
+    let mut t: [i64; 31] = [0; 31];
+
+    for i in 0..16 {
+        for j in 0..16 {
+            t[i + j] += a[i] * b[j];
+        }
+    }
+    for i in 0..15 {
+        t[i] += 38 * t[i + 16];
+    }
+    o.copy_from_slice(&t[..16]);
+    carry(o);
+    carry(o);
+
+    memzero_explicit(&mut t);
+}
+
+fn invert(o: &mut Fe, i: &Fe) {
+    let mut c: Fe = fe_new_default();
+
+    c.copy_from_slice(i);
+    for a in (0..=253).rev() {
+        {
+            let c_clone = c;
+            multmod(&mut c, &c_clone, &c_clone);
+        }
+        if a != 2 && a != 4 {
+            {
+                let c_clone = c;
+                multmod(&mut c, &c_clone, i);
+            }
+        }
+    }
+    o.copy_from_slice(&c);
+
+    memzero_explicit(&mut c);
+}
+
+fn memzero_explicit<T>(v: &mut T) {
+    unsafe {
+        let zeroed = std::mem::zeroed();
+        std::ptr::write_volatile(v as *mut _, zeroed);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Key;
+
+    #[test]
+    fn decode_encode_key() {
+        let key = "6F5rOtYE5A2KcXTKf9jdzWa9Y/kuV5gPS3LcKlxmOnY=";
+        let key = super::Key::decode(key).unwrap();
+        let key = key.encode();
+        assert_eq!(key, "6F5rOtYE5A2KcXTKf9jdzWa9Y/kuV5gPS3LcKlxmOnY=");
+    }
+
+    #[test]
+    fn generate_public_key() {
+        assert_eq!(
+            Key::decode("3D5lgnI9ztvnuyWDm7dlBDgm6xr0+WVWPoo6HIfzHRU=").unwrap(),
+            Key::generate_public(
+                &Key::decode("+Op7voRskU0Zm2fHFR/5tVE+PJtnwn6cbnme71jXt0E=").unwrap()
+            )
+        );
+
+        assert_eq!(
+            Key::decode("//eq/raPUE4+sOlTlozx76XEE+W8L0bUqNfyg9IpX0Q=").unwrap(),
+            Key::generate_public(
+                &Key::decode("8OD8QPWH/a0D5LmbWVnb7bwFq4Fghy/QUEFkIhyL/EI=").unwrap()
+            )
+        );
+    }
+}