1 files changed, 230 insertions, 0 deletions
diff --git a/toolkit/components/telemetry/dap/ffi-gtest/test.rs b/toolkit/components/telemetry/dap/ffi-gtest/test.rs
new file mode 100644
index 0000000000..8db79f5f05
--- /dev/null
+++ b/toolkit/components/telemetry/dap/ffi-gtest/test.rs
@@ -0,0 +1,230 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use serde::Deserialize;
+use std::ffi::c_void;
+use std::fs::File;
+use std::io::Cursor;
+
+use thin_vec::ThinVec;
+
+use dap_ffi::prg::PrgAes128Alt;
+use dap_ffi::types::Report;
+
+use prio::codec::{Decode, Encode};
+use prio::vdaf::prg::{Prg, SeedStream};
+
+#[no_mangle]
+pub extern "C" fn dap_test_encoding() {
+    let r = Report::new_dummy();
+    let mut encoded = Vec::<u8>::new();
+    Report::encode(&r, &mut encoded);
+    let decoded = Report::decode(&mut Cursor::new(&encoded)).expect("Report decoding failed!");
+    if r != decoded {
+        println!("Report:");
+        println!("{:?}", r);
+        println!("Encoded Report:");
+        println!("{:?}", encoded);
+        println!("Decoded Report:");
+        println!("{:?}", decoded);
+        panic!("Report changed after encoding & decoding.");
+    }
+}
+
+#[derive(Deserialize, Debug)]
+struct PrgTestCase {
+    seed: [u8; 16],
+    info_string: Vec<u8>,
+    buffer1_out: Vec<u8>,
+    buffer2_out: Vec<u8>,
+}
+
+#[no_mangle]
+pub extern "C" fn dap_test_prg() {
+    let file = File::open("PrgAes128_tests.json").unwrap();
+    let testcases: Vec<PrgTestCase> = serde_json::from_reader(file).unwrap();
+    for testcase in testcases {
+        let mut p = PrgAes128Alt::init(&testcase.seed);
+        p.update(&testcase.info_string);
+        let mut s = p.into_seed_stream();
+        let mut b1 = vec![0u8; testcase.buffer1_out.len()];
+        s.fill(&mut b1);
+        assert_eq!(b1, testcase.buffer1_out);
+        let mut b2 = vec![0u8; testcase.buffer2_out.len()];
+        s.fill(&mut b2);
+        assert_eq!(b2, testcase.buffer2_out);
+    }
+}
+
+extern "C" {
+    pub fn dapHpkeEncrypt(
+        aContext: *mut c_void,
+        aAad: *mut u8,
+        aAadLength: u32,
+        aPlaintext: *mut u8,
+        aPlaintextLength: u32,
+        aOutputShare: &mut ThinVec<u8>,
+    ) -> bool;
+    pub fn dapSetupHpkeContextForTesting(
+        aKey: *const u8,
+        aKeyLength: u32,
+        aInfo: *mut u8,
+        aInfoLength: u32,
+        aPkEm: *const u8,
+        aPkEmLength: u32,
+        aSkEm: *const u8,
+        aSkEmLength: u32,
+        aOutputEncapsulatedKey: &mut ThinVec<u8>,
+    ) -> *mut c_void;
+    pub fn dapDestroyHpkeContext(aContext: *mut c_void);
+}
+
+struct HpkeContext(*mut c_void);
+
+impl Drop for HpkeContext {
+    fn drop(&mut self) {
+        unsafe {
+            dapDestroyHpkeContext(self.0);
+        }
+    }
+}
+
+type Testsuites = Vec<CiphersuiteTest>;
+
+#[derive(Debug, Deserialize)]
+pub struct HexString(#[serde(with = "hex")] Vec<u8>);
+impl AsRef<[u8]> for HexString {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}
+#[derive(Debug, Deserialize)]
+struct CiphersuiteTest {
+    mode: i64,
+    kem_id: i64,
+    kdf_id: i64,
+    aead_id: i64,
+    info: HexString,
+    #[serde(rename = "ikmR")]
+    ikm_r: HexString,
+    #[serde(rename = "ikmE")]
+    ikm_e: HexString,
+    #[serde(rename = "skRm")]
+    sk_r_m: HexString,
+    #[serde(rename = "skEm")]
+    sk_e_m: HexString,
+    #[serde(rename = "pkRm")]
+    pk_r_m: HexString,
+    #[serde(rename = "pkEm")]
+    pk_e_m: HexString,
+    enc: HexString,
+    shared_secret: HexString,
+    key_schedule_context: HexString,
+    secret: HexString,
+    key: HexString,
+    base_nonce: HexString,
+    exporter_secret: HexString,
+    encryptions: Vec<Encryption>,
+    exports: Vec<Export>,
+    psk: Option<HexString>,
+    psk_id: Option<HexString>,
+    ikm_s: Option<HexString>,
+    sk_sm: Option<HexString>,
+    pk_sm: Option<HexString>,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct Encryption {
+    pub aad: HexString,
+    pub ciphertext: HexString,
+    pub nonce: HexString,
+    pub plaintext: HexString,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct Export {
+    pub exporter_context: HexString,
+    #[serde(rename = "L")]
+    pub length: i64,
+    pub exported_value: HexString,
+}
+
+#[no_mangle]
+pub extern "C" fn dap_test_hpke_encrypt() {
+    let file = File::open("hpke-vectors.json").unwrap();
+    let tests: Testsuites = serde_json::from_reader(file).unwrap();
+
+    let mut have_tested = false;
+
+    for (test_idx, test) in tests.into_iter().enumerate() {
+        // Mode must be "Base"
+        if test.mode != 0
+         // KEM must be DHKEM(X25519, HKDF-SHA256)
+         || test.kem_id != 32
+         // KDF must be HKDF-SHA256
+         || test.kdf_id != 1
+         // AEAD must be AES-128-GCM
+         || test.aead_id != 1
+        {
+            continue;
+        }
+
+        have_tested = true;
+
+        let mut pk_r_serialized = test.pk_r_m.0;
+        let mut info = test.info.0;
+        let mut pk_e_serialized = test.pk_e_m.0;
+        let mut sk_e_serialized = test.sk_e_m.0;
+
+        let mut encapsulated_key = ThinVec::<u8>::new();
+
+        let ctx = HpkeContext(unsafe {
+            dapSetupHpkeContextForTesting(
+                pk_r_serialized.as_mut_ptr(),
+                pk_r_serialized.len().try_into().unwrap(),
+                info.as_mut_ptr(),
+                info.len().try_into().unwrap(),
+                pk_e_serialized.as_mut_ptr(),
+                pk_e_serialized.len().try_into().unwrap(),
+                sk_e_serialized.as_mut_ptr(),
+                sk_e_serialized.len().try_into().unwrap(),
+                &mut encapsulated_key,
+            )
+        });
+        if ctx.0.is_null() {
+            panic!("Failed to set up HPKE context.");
+        }
+        if encapsulated_key != test.enc.0 {
+            panic!("Encapsulated key is wrong!");
+        }
+
+        for (encryption_idx, encryption) in test.encryptions.into_iter().enumerate() {
+            let mut encrypted_share = ThinVec::<u8>::new();
+
+            let mut aad = encryption.aad.0.clone();
+            let mut pt = encryption.plaintext.0.clone();
+            unsafe {
+                dapHpkeEncrypt(
+                    ctx.0,
+                    aad.as_mut_ptr(),
+                    aad.len().try_into().unwrap(),
+                    pt.as_mut_ptr(),
+                    pt.len().try_into().unwrap(),
+                    &mut encrypted_share,
+                );
+            }
+
+            if encrypted_share != encryption.ciphertext.0 {
+                println!("Test: {}, Encryption: {}", test_idx, encryption_idx);
+                println!("Expected:");
+                println!("{:?}", encryption.ciphertext.0);
+                println!("Actual:");
+                println!("{:?}", encrypted_share);
+                panic!("Encryption outputs did not match!");
+            }
+        }
+    }
+
+    assert!(have_tested);
+}