DTLS: crypto_ccm.rs implementation

dtls/src/crypto/crypto_ccm.rs

@@ -7,3 +7,154 @@
 
 // https://github.com/RustCrypto/AEADs
 // https://docs.rs/ccm/0.3.0/ccm/ Or https://crates.io/crates/aes-ccm?
+
+use util::Error;
+
+use rand::Rng;
+
+use std::io::Cursor;
+
+use super::*;
+use crate::content::*;
+use crate::errors::*;
+use crate::record_layer::record_layer_header::*;
+
+use aes::Aes256;
+use ccm::aead::{generic_array::GenericArray, AeadInPlace, NewAead};
+use ccm::{
+    consts::{U12, U16, U8},
+    Ccm,
+};
+
+const CRYPTO_CCM_8_TAG_LENGTH: usize = 8;
+const CRYPTO_CCM_TAG_LENGTH: usize = 16;
+const CRYPTO_CCM_NONCE_LENGTH: usize = 12;
+
+type AesCcm8 = Ccm<Aes256, U8, U12>;
+type AesCcm = Ccm<Aes256, U16, U12>;
+
+pub enum CryptoCcmTagLen {
+    CryptoCcm8TagLength,
+    CryptoCcmTagLength,
+}
+
+enum CryptoCcm {
+    CryptoCcm8(AesCcm8),
+    CryptoCcm(AesCcm),
+}
+
+// State needed to handle encrypted input/output
+pub struct CryptCcm {
+    local_ccm: CryptoCcm,
+    remote_ccm: CryptoCcm,
+    local_write_iv: Vec<u8>,
+    remote_write_iv: Vec<u8>,
+}
+
+impl CryptCcm {
+    pub fn new(
+        tag_len: CryptoCcmTagLen,
+        local_key: &[u8],
+        local_write_iv: &[u8],
+        remote_key: &[u8],
+        remote_write_iv: &[u8],
+    ) -> Self {
+        let key = GenericArray::from_slice(local_key);
+        let local_ccm = match tag_len {
+            CryptoCcmTagLen::CryptoCcmTagLength => CryptoCcm::CryptoCcm(AesCcm::new(key)),
+            CryptoCcmTagLen::CryptoCcm8TagLength => CryptoCcm::CryptoCcm8(AesCcm8::new(key)),
+        };
+
+        let key = GenericArray::from_slice(remote_key);
+        let remote_ccm = match tag_len {
+            CryptoCcmTagLen::CryptoCcmTagLength => CryptoCcm::CryptoCcm(AesCcm::new(key)),
+            CryptoCcmTagLen::CryptoCcm8TagLength => CryptoCcm::CryptoCcm8(AesCcm8::new(key)),
+        };
+
+        CryptCcm {
+            local_ccm,
+            local_write_iv: local_write_iv.to_vec(),
+            remote_ccm,
+            remote_write_iv: remote_write_iv.to_vec(),
+        }
+    }
+
+    pub fn encrypt(&self, pkt_rlh: &RecordLayerHeader, raw: &[u8]) -> Result<Vec<u8>, Error> {
+        let payload = &raw[RECORD_LAYER_HEADER_SIZE..];
+        let raw = &raw[..RECORD_LAYER_HEADER_SIZE];
+
+        let mut nonce = vec![0u8; CRYPTO_CCM_NONCE_LENGTH];
+        nonce[..4].copy_from_slice(&self.local_write_iv[..4]);
+        rand::thread_rng().fill(&mut nonce[4..]);
+        let nonce = GenericArray::from_slice(&nonce);
+
+        let additional_data = generate_aead_additional_data(pkt_rlh, payload.len());
+
+        let mut buffer: Vec<u8> = Vec::new();
+        buffer.extend_from_slice(payload);
+
+        match &self.local_ccm {
+            CryptoCcm::CryptoCcm(ccm) => {
+                ccm.encrypt_in_place(nonce, &additional_data, &mut buffer)?;
+            }
+            CryptoCcm::CryptoCcm8(ccm8) => {
+                ccm8.encrypt_in_place(nonce, &additional_data, &mut buffer)?;
+            }
+        }
+
+        let mut r = Vec::with_capacity(raw.len() + nonce.len() + buffer.len());
+
+        r.extend_from_slice(raw);
+        r.extend_from_slice(&nonce[4..]);
+        r.extend_from_slice(&buffer);
+
+        // Update recordLayer size to include explicit nonce
+        let r_len = (r.len() - RECORD_LAYER_HEADER_SIZE) as u16;
+        r[RECORD_LAYER_HEADER_SIZE - 2..RECORD_LAYER_HEADER_SIZE]
+            .copy_from_slice(&r_len.to_be_bytes());
+
+        Ok(r)
+    }
+
+    pub fn decrypt(&self, r: &[u8]) -> Result<Vec<u8>, Error> {
+        let mut reader = Cursor::new(r);
+        let h = RecordLayerHeader::unmarshal(&mut reader)?;
+        if h.content_type == ContentType::ChangeCipherSpec {
+            // Nothing to encrypt with ChangeCipherSpec
+            return Ok(r.to_vec());
+        }
+
+        if r.len() <= (RECORD_LAYER_HEADER_SIZE + 8) {
+            return Err(ERR_NOT_ENOUGH_ROOM_FOR_NONCE.clone());
+        }
+
+        let mut nonce = vec![];
+        nonce.extend_from_slice(&self.remote_write_iv[..4]);
+        nonce.extend_from_slice(&r[RECORD_LAYER_HEADER_SIZE..RECORD_LAYER_HEADER_SIZE + 8]);
+        let nonce = GenericArray::from_slice(&nonce);
+
+        let out = &r[RECORD_LAYER_HEADER_SIZE + 8..];
+
+        let mut buffer: Vec<u8> = Vec::new();
+        buffer.extend_from_slice(out);
+
+        match &self.remote_ccm {
+            CryptoCcm::CryptoCcm(ccm) => {
+                let additional_data =
+                    generate_aead_additional_data(&h, out.len() - CRYPTO_CCM_TAG_LENGTH);
+                ccm.decrypt_in_place(nonce, &additional_data, &mut buffer)?;
+            }
+            CryptoCcm::CryptoCcm8(ccm8) => {
+                let additional_data =
+                    generate_aead_additional_data(&h, out.len() - CRYPTO_CCM_8_TAG_LENGTH);
+                ccm8.decrypt_in_place(nonce, &additional_data, &mut buffer)?;
+            }
+        }
+
+        let mut d = Vec::with_capacity(RECORD_LAYER_HEADER_SIZE + buffer.len());
+        d.extend_from_slice(&r[..RECORD_LAYER_HEADER_SIZE]);
+        d.extend_from_slice(&buffer);
+
+        Ok(d)
+    }
+}

dtls/src/crypto/crypto_test.rs

@@ -1,5 +1,9 @@
+use super::crypto_ccm::*;
 use super::*;
 
+use crate::content::ContentType;
+use crate::record_layer::record_layer_header::{ProtocolVersion, RECORD_LAYER_HEADER_SIZE};
+
 use std::io::Cursor;
 
 use util::Error;
@@ -98,3 +102,53 @@ fn test_generate_key_signature() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[test]
+fn test_ccm_encryption_and_decryption() -> Result<(), Error> {
+    let key = vec![
+        0x18, 0x78, 0xac, 0xc2, 0x2a, 0xd8, 0xbd, 0xd8, 0xc6, 0x01, 0xa6, 0x17, 0x12, 0x6f, 0x63,
+        0x54, 0x18, 0x78, 0xac, 0xc2, 0x2a, 0xd8, 0xbd, 0xd8, 0xc6, 0x01, 0xa6, 0x17, 0x12, 0x6f,
+        0x63, 0x54,
+    ];
+    let iv = vec![0x0e, 0xb2, 0x09, 0x06];
+
+    let ccm = CryptCcm::new(CryptoCcmTagLen::CryptoCcmTagLength, &key, &iv, &key, &iv);
+
+    let rlh = RecordLayerHeader {
+        content_type: ContentType::ApplicationData,
+        protocol_version: ProtocolVersion {
+            major: 0xfe,
+            minor: 0xff,
+        },
+        epoch: 0,
+        sequence_number: 18,
+        content_len: 3,
+    };
+
+    let raw = vec![
+        0x17, 0xfe, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x12, 0x00, 0x03, 0xff, 0xaa,
+        0xbb,
+    ];
+
+    let cipher_text = ccm.encrypt(&rlh, &raw)?;
+
+    assert_eq!(
+        [0, 27],
+        &cipher_text[RECORD_LAYER_HEADER_SIZE - 2..RECORD_LAYER_HEADER_SIZE],
+        "RecordLayer size updating failed \nexp: {:?} \nactual {:?} ",
+        [0, 27],
+        &cipher_text[RECORD_LAYER_HEADER_SIZE - 2..RECORD_LAYER_HEADER_SIZE]
+    );
+
+    let plain_text = ccm.decrypt(&cipher_text)?;
+
+    assert_eq!(
+        raw[RECORD_LAYER_HEADER_SIZE..],
+        plain_text[RECORD_LAYER_HEADER_SIZE..],
+        "Decryption failed \nexp: {:?} \nactual {:?} ",
+        &raw[RECORD_LAYER_HEADER_SIZE..],
+        &plain_text[RECORD_LAYER_HEADER_SIZE..]
+    );
+
+    Ok(())
+}