hashx: Rust hook for inspecting and modifying the random number stream

This patch has no effect on the C tor build.

Adds a function hashx_rng_callback() to the hashx API, defined only
when HASHX_RNG_CALLBACK is defined. This is then used in the Rust
wrapper to implement a similar rng_callback().

Included some minimal test cases. This code is intented for
use in cross-compatibility fuzzing tests which drive multiple
implementations of hashx with the same custom Rng stream.

Signed-off-by: Micah Elizabeth Scott <beth@torproject.org>

9 files changed, 156 insertions, 16 deletions

diff --git a/src/ext/equix/Cargo.toml b/src/ext/equix/Cargo.toml
index 45b24dcea1..bedb1ed565 100644
--- a/src/ext/equix/Cargo.toml
+++ b/src/ext/equix/Cargo.toml
@@ -9,7 +9,7 @@
 
 [package]
 name = "tor-c-equix"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 license = "LGPL-3.0-only"
 
diff --git a/src/ext/equix/build.rs b/src/ext/equix/build.rs
index f2825a50cf..b53f08e899 100644
--- a/src/ext/equix/build.rs
+++ b/src/ext/equix/build.rs
@@ -16,6 +16,8 @@ fn main() {
             "hashx/src/siphash_rng.c",
             "hashx/src/virtual_memory.c",
         ])
+        // Activate our patch for hashx_rng_callback
+        .define("HASHX_RNG_CALLBACK", "1")
         // Equi-X always uses HashX size 8 (64-bit output)
         .define("HASHX_SIZE", "8")
         // Avoid shared library API declarations, link statically
@@ -31,6 +33,7 @@ fn main() {
         .header_contents(
             "wrapper.h",
             r#"
+                #define HASHX_RNG_CALLBACK 1
                 #define HASHX_SIZE 8
                 #define HASHX_SHARED 1
                 #define EQUIX_SHARED 1
diff --git a/src/ext/equix/hashx/include/hashx.h b/src/ext/equix/hashx/include/hashx.h
index 2910515d9a..3f6d059b92 100644
--- a/src/ext/equix/hashx/include/hashx.h
+++ b/src/ext/equix/hashx/include/hashx.h
@@ -169,6 +169,25 @@ HASHX_API hashx_result hashx_exec(const hashx_ctx* ctx,
 */
 HASHX_API void hashx_free(hashx_ctx* ctx);
 
+#ifdef HASHX_RNG_CALLBACK
+/*
+ * Set a callback for inspecting or modifying the HashX random number stream.
+ *
+ * The callback and its user pointer are associated with the provided context
+ * even if it's re-used for another hash program. A callback value of NULL
+ * disables the callback.
+ *
+ * @param ctx is pointer to a HashX instance.
+ * @param callback is invoked after each new 64-bit pseudorandom value
+ *        is generated in a buffer. The callback may record it and/or replace
+ *        it. A NULL pointer here disables the callback.
+ * @param user_data is an opaque parameter given to the callback
+ */
+HASHX_API void hashx_rng_callback(hashx_ctx* ctx,
+                                  void (*callback)(uint64_t*, void*),
+                                  void* user_data);
+#endif
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/ext/equix/hashx/src/context.c b/src/ext/equix/hashx/src/context.c
index 03a9de57fd..da1b997e16 100644
--- a/src/ext/equix/hashx/src/context.c
+++ b/src/ext/equix/hashx/src/context.c
@@ -55,3 +55,13 @@ void hashx_free(hashx_ctx* ctx) {
 		free(ctx);
 	}
 }
+
+#ifdef HASHX_RNG_CALLBACK
+void hashx_rng_callback(hashx_ctx* ctx,
+                        void (*callback)(uint64_t*, void*),
+                        void* callback_user_data)
+{
+	ctx->program.rng_callback = callback;
+	ctx->program.rng_callback_user_data = callback_user_data;
+}
+#endif
diff --git a/src/ext/equix/hashx/src/program.c b/src/ext/equix/hashx/src/program.c
index b44bdb855a..1017d4070a 100644
--- a/src/ext/equix/hashx/src/program.c
+++ b/src/ext/equix/hashx/src/program.c
@@ -554,6 +554,10 @@ bool hashx_program_generate(const siphash_state* key, hashx_program* program) {
 		.ports = {{ 0 }}
 	};
 	hashx_siphash_rng_init(&ctx.gen, key);
+#ifdef HASHX_RNG_CALLBACK
+	ctx.gen.callback = program->rng_callback;
+	ctx.gen.callback_user_data = program->rng_callback_user_data;
+#endif
 	for (int i = 0; i < 8; ++i) {
 		ctx.registers[i].last_op = -1;
 		ctx.registers[i].latency = 0;
diff --git a/src/ext/equix/hashx/src/program.h b/src/ext/equix/hashx/src/program.h
index 096cc4ee0a..78dbb8b6e3 100644
--- a/src/ext/equix/hashx/src/program.h
+++ b/src/ext/equix/hashx/src/program.h
@@ -29,6 +29,10 @@ typedef struct hashx_program {
 	int branch_count;
 	int branches[16];
 #endif
+#ifdef HASHX_RNG_CALLBACK
+	void (*rng_callback)(uint64_t *buffer, void *user_data);
+	void *rng_callback_user_data;
+#endif
 } hashx_program;
 
 #ifdef __cplusplus
diff --git a/src/ext/equix/hashx/src/siphash_rng.c b/src/ext/equix/hashx/src/siphash_rng.c
index 89ed8fc845..c0f457be76 100644
--- a/src/ext/equix/hashx/src/siphash_rng.c
+++ b/src/ext/equix/hashx/src/siphash_rng.c
@@ -15,6 +15,11 @@ uint8_t hashx_siphash_rng_u8(siphash_rng* gen) {
 		gen->buffer8 = hashx_siphash13_ctr(gen->counter, &gen->keys);
 		gen->counter++;
 		gen->count8 = sizeof(gen->buffer8);
+#ifdef HASHX_RNG_CALLBACK
+		if (gen->callback) {
+			gen->callback(&gen->buffer8, gen->callback_user_data);
+		}
+#endif
 	}
 	gen->count8--;
 	return gen->buffer8 >> (gen->count8 * 8);
@@ -25,6 +30,11 @@ uint32_t hashx_siphash_rng_u32(siphash_rng* gen) {
 		gen->buffer32 = hashx_siphash13_ctr(gen->counter, &gen->keys);
 		gen->counter++;
 		gen->count32 = sizeof(gen->buffer32) / sizeof(uint32_t);
+#ifdef HASHX_RNG_CALLBACK
+		if (gen->callback) {
+			gen->callback(&gen->buffer32, gen->callback_user_data);
+		}
+#endif
 	}
 	gen->count32--;
 	return (uint32_t)(gen->buffer32 >> (gen->count32 * 32));
diff --git a/src/ext/equix/hashx/src/siphash_rng.h b/src/ext/equix/hashx/src/siphash_rng.h
index 638b177e06..7b402fdc6e 100644
--- a/src/ext/equix/hashx/src/siphash_rng.h
+++ b/src/ext/equix/hashx/src/siphash_rng.h
@@ -13,6 +13,10 @@ typedef struct siphash_rng {
 	uint64_t counter;
 	uint64_t buffer8, buffer32;
 	unsigned count8, count32;
+#ifdef HASHX_RNG_CALLBACK
+	void (*callback)(uint64_t *buffer, void *user_data);
+	void *callback_user_data;
+#endif
 } siphash_rng;
 
 #ifdef __cplusplus
diff --git a/src/ext/equix/src/lib.rs b/src/ext/equix/src/lib.rs
index 0db1fc1bb3..8eb163075a 100644
--- a/src/ext/equix/src/lib.rs
+++ b/src/ext/equix/src/lib.rs
@@ -12,6 +12,10 @@
 //! See `LICENSE` for licensing information.
 //!
 
+use core::ffi::c_void;
+use core::mem;
+use core::ptr::null_mut;
+
 pub mod ffi {
     //! Low-level access to the C API
 
@@ -34,8 +38,14 @@ pub const HASHX_SIZE: usize = ffi::HASHX_SIZE as usize;
 /// Output value obtained by executing a HashX hash function
 pub type HashXOutput = [u8; HASHX_SIZE];
 
+/// Type for callback functions that inspect or replace the pseudorandom stream
+pub type RngCallback = Box<dyn FnMut(u64) -> u64>;
+
 /// Safe wrapper around a HashX context
-pub struct HashX(*mut ffi::hashx_ctx);
+pub struct HashX {
+    ctx: *mut ffi::hashx_ctx,
+    rng_callback: Option<RngCallback>,
+}
 
 impl HashX {
     /// Allocate a new HashX context
@@ -44,7 +54,10 @@ impl HashX {
         if ctx.is_null() {
             panic!("out of memory in hashx_alloc");
         }
-        Self(ctx)
+        Self {
+            ctx,
+            rng_callback: None,
+        }
     }
 
     /// Create a new hash function within this context, using the given seed
@@ -53,14 +66,15 @@ impl HashX {
     /// error occurs while the interpreter is disabled.
     #[inline(always)]
     pub fn make(&mut self, seed: &[u8]) -> HashXResult {
-        unsafe { ffi::hashx_make(self.0, seed.as_ptr() as *const std::ffi::c_void, seed.len()) }
+        unsafe { ffi::hashx_make(self.ctx, seed.as_ptr() as *const c_void, seed.len()) }
     }
 
     /// Check which implementation was selected by `make`
     #[inline(always)]
     pub fn query_type(&mut self) -> Result<HashXType, HashXResult> {
         let mut buffer = HashXType::HASHX_TYPE_INTERPRETED; // Arbitrary default
-        let result = unsafe { ffi::hashx_query_type(self.0, &mut buffer as *mut ffi::hashx_type) };
+        let result =
+            unsafe { ffi::hashx_query_type(self.ctx, &mut buffer as *mut ffi::hashx_type) };
         match result {
             HashXResult::HASHX_OK => Ok(buffer),
             e => Err(e),
@@ -71,23 +85,45 @@ impl HashX {
     #[inline(always)]
     pub fn exec(&mut self, input: u64) -> Result<HashXOutput, HashXResult> {
         let mut buffer: HashXOutput = Default::default();
-        let result = unsafe {
-            ffi::hashx_exec(
-                self.0,
-                input,
-                &mut buffer as *mut u8 as *mut std::ffi::c_void,
-            )
-        };
+        let result =
+            unsafe { ffi::hashx_exec(self.ctx, input, &mut buffer as *mut u8 as *mut c_void) };
         match result {
             HashXResult::HASHX_OK => Ok(buffer),
             e => Err(e),
         }
     }
+
+    /// Set a callback function that may inspect and/or modify the internal
+    /// pseudorandom number stream used by this context.
+    ///
+    /// The function will be owned by this context, and it replaces any
+    /// previous function that may have been set. Returns the previous callback
+    /// if any.
+    pub fn rng_callback(&mut self, callback: Option<RngCallback>) -> Option<RngCallback> {
+        // Keep ownership of our Rust value in the context wrapper, to match
+        // the lifetime of the mutable pointer that the C API saves.
+        let result = mem::replace(&mut self.rng_callback, callback);
+        match &mut self.rng_callback {
+            None => unsafe { ffi::hashx_rng_callback(self.ctx, None, null_mut()) },
+            Some(callback) => unsafe {
+                ffi::hashx_rng_callback(
+                    self.ctx,
+                    Some(wrapper),
+                    callback as *mut RngCallback as *mut c_void,
+                );
+            },
+        }
+        unsafe extern "C" fn wrapper(buffer: *mut u64, callback: *mut c_void) {
+            let callback: &mut RngCallback = unsafe { mem::transmute(callback) };
+            buffer.write(callback(buffer.read()));
+        }
+        result
+    }
 }
 
 impl Drop for HashX {
     fn drop(&mut self) {
-        let ctx = std::mem::replace(&mut self.0, std::ptr::null_mut());
+        let ctx = mem::replace(&mut self.ctx, null_mut());
         unsafe {
             ffi::hashx_free(ctx);
         }
@@ -146,7 +182,7 @@ impl EquiX {
         unsafe {
             ffi::equix_verify(
                 self.0,
-                challenge.as_ptr() as *const std::ffi::c_void,
+                challenge.as_ptr() as *const c_void,
                 challenge.len(),
                 solution as *const ffi::equix_solution,
             )
@@ -159,7 +195,7 @@ impl EquiX {
         unsafe {
             ffi::equix_solve(
                 self.0,
-                challenge.as_ptr() as *const std::ffi::c_void,
+                challenge.as_ptr() as *const c_void,
                 challenge.len(),
                 buffer as *mut ffi::equix_solutions_buffer,
             )
@@ -169,7 +205,7 @@ impl EquiX {
 
 impl Drop for EquiX {
     fn drop(&mut self) {
-        let ctx = std::mem::replace(&mut self.0, std::ptr::null_mut());
+        let ctx = mem::replace(&mut self.0, null_mut());
         unsafe {
             ffi::equix_free(ctx);
         }
@@ -180,6 +216,8 @@ impl Drop for EquiX {
 mod tests {
     use crate::*;
     use hex_literal::hex;
+    use std::cell::RefCell;
+    use std::sync::Arc;
 
     #[test]
     fn equix_context() {
@@ -290,4 +328,52 @@ mod tests {
         assert_eq!(ctx.exec(123456), Ok(hex!("ab3d155bf4bbb0aa")));
         assert_eq!(ctx.exec(987654321123456789), Ok(hex!("8dfef0497c323274")));
     }
+
+    #[test]
+    fn rng_callback_read() {
+        // Use a Rng callback to read the sequence of pseudorandom numbers
+        // without changing them, and spot check the list we get back.
+        let mut ctx = HashX::new(HashXType::HASHX_TRY_COMPILE);
+        let seq = Arc::new(RefCell::new(Vec::new()));
+        {
+            let seq = seq.clone();
+            ctx.rng_callback(Some(Box::new(move |value| {
+                seq.borrow_mut().push(value);
+                value
+            })));
+        }
+        assert_eq!(seq.borrow().len(), 0);
+        assert_eq!(ctx.make(b"abc"), HashXResult::HASHX_OK);
+        assert_eq!(ctx.exec(12345).unwrap(), hex!("c0bc95da7cc30f37"));
+        assert_eq!(seq.borrow().len(), 563);
+        assert_eq!(
+            seq.borrow()[..4],
+            [
+                0xf695edd02205449d,
+                0x51c1ac51cd19a7d1,
+                0xadf4cb303b9814cf,
+                0x79793a52d965083d
+            ]
+        );
+    }
+
+    #[test]
+    fn rng_callback_replace() {
+        // Use a Rng callback to replace the random number stream.
+        // We have to choose the replacement somewhat carefully since
+        // many stationary replacement values will cause infinite loops.
+        let mut ctx = HashX::new(HashXType::HASHX_TYPE_INTERPRETED);
+        let counter = Arc::new(RefCell::new(0u32));
+        {
+            let counter = counter.clone();
+            ctx.rng_callback(Some(Box::new(move |_value| {
+                *counter.borrow_mut() += 1;
+                0x0807060504030201
+            })));
+        }
+        assert_eq!(*counter.borrow(), 0);
+        assert_eq!(ctx.make(b"abc"), HashXResult::HASHX_OK);
+        assert_eq!(ctx.exec(12345).unwrap(), hex!("825a9b6dd5d074af"));
+        assert_eq!(*counter.borrow(), 575);
+    }
 }