1 files changed, 93 insertions, 0 deletions
diff --git a/toolkit/components/telemetry/dap/ffi/src/prg.rs b/toolkit/components/telemetry/dap/ffi/src/prg.rs
new file mode 100644
index 0000000000..a7ebeb11cb
--- /dev/null
+++ b/toolkit/components/telemetry/dap/ffi/src/prg.rs
@@ -0,0 +1,93 @@
+/* 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 https://mozilla.org/MPL/2.0/. */
+
+use std::ffi::c_void;
+
+use prio::vdaf::prg::{Prg, SeedStream};
+
+extern "C" {
+    pub fn dapStartCmac(aSeed: *mut u8) -> *mut c_void;
+    pub fn dapUpdateCmac(aContext: *mut c_void, aData: *const u8, aDataLen: u32);
+    pub fn dapFinalizeCmac(aContext: *mut c_void, aMacOutput: *mut u8);
+    pub fn dapReleaseCmac(aContext: *mut c_void);
+
+    pub fn dapStartAesCtr(aKey: *const u8) -> *mut c_void;
+    pub fn dapCtrFillBuffer(aContext: *mut c_void, aBuffer: *mut u8, aBufferSize: i32);
+    pub fn dapReleaseCtrCtx(aContext: *mut c_void);
+}
+
+#[derive(Clone, Debug)]
+pub struct PrgAes128Alt {
+    nss_context: *mut c_void,
+}
+
+impl Prg<16> for PrgAes128Alt {
+    type SeedStream = SeedStreamAes128Alt;
+
+    fn init(seed_bytes: &[u8; 16]) -> Self {
+        let mut my_seed_bytes = *seed_bytes;
+        let ctx = unsafe { dapStartCmac(my_seed_bytes.as_mut_ptr()) };
+        assert!(!ctx.is_null());
+
+        Self { nss_context: ctx }
+    }
+
+    fn update(&mut self, data: &[u8]) {
+        unsafe {
+            dapUpdateCmac(
+                self.nss_context,
+                data.as_ptr(),
+                u32::try_from(data.len()).unwrap(),
+            );
+        }
+    }
+
+    fn into_seed_stream(self) -> Self::SeedStream {
+        // finish the MAC and create a new random data stream using the result as key and 0 as IV for AES-CTR
+        let mut key = [0u8; 16];
+        unsafe {
+            dapFinalizeCmac(self.nss_context, key.as_mut_ptr());
+        }
+
+        SeedStreamAes128Alt::new(&mut key, &[0; 16])
+    }
+}
+
+impl Drop for PrgAes128Alt {
+    fn drop(&mut self) {
+        unsafe {
+            dapReleaseCmac(self.nss_context);
+        }
+    }
+}
+
+pub struct SeedStreamAes128Alt {
+    nss_context: *mut c_void,
+}
+
+impl SeedStreamAes128Alt {
+    pub(crate) fn new(key: &mut [u8; 16], iv: &[u8; 16]) -> Self {
+        debug_assert_eq!(iv, &[0; 16]);
+        let ctx = unsafe { dapStartAesCtr(key.as_ptr()) };
+        Self { nss_context: ctx }
+    }
+}
+
+impl SeedStream for SeedStreamAes128Alt {
+    fn fill(&mut self, buf: &mut [u8]) {
+        unsafe {
+            dapCtrFillBuffer(
+                self.nss_context,
+                buf.as_mut_ptr(),
+                i32::try_from(buf.len()).unwrap(),
+            );
+        }
+    }
+}
+
+impl Drop for SeedStreamAes128Alt {
+    fn drop(&mut self) {
+        unsafe { dapReleaseCtrCtx(self.nss_context) };
+    }
+}