From e84f8da88b60127cdc752f28572300e94ed2fa54 Mon Sep 17 00:00:00 2001 From: Rick Weber Date: Fri, 18 Feb 2022 15:45:46 -0800 Subject: [PATCH] Remove unneccessary unsafe code. Add more validation and document some stuff --- examples/simple_multiply/src/main.rs | 2 +- sunscreen_compiler/src/error.rs | 5 ++ sunscreen_compiler/src/params.rs | 2 + .../src/types/intern/fhe_program_node.rs | 27 +++++---- sunscreen_compiler/src/types/ops/add.rs | 9 +++ sunscreen_compiler/src/types/ops/div.rs | 15 +++++ sunscreen_compiler/src/types/ops/mul.rs | 9 +++ sunscreen_compiler/src/types/ops/neg.rs | 3 + sunscreen_compiler/src/types/ops/rotate.rs | 3 + sunscreen_compiler/src/types/ops/sub.rs | 3 + sunscreen_compiler/tests/edge_cases.rs | 16 +++-- sunscreen_compiler/tests/typename_tests.rs | 4 +- sunscreen_fhe_program/src/lib.rs | 2 +- sunscreen_runtime/src/run.rs | 58 ++++++++++++------- sunscreen_runtime/src/runtime.rs | 5 ++ 15 files changed, 121 insertions(+), 42 deletions(-) diff --git a/examples/simple_multiply/src/main.rs b/examples/simple_multiply/src/main.rs index 69bbba9a0..606915eab 100644 --- a/examples/simple_multiply/src/main.rs +++ b/examples/simple_multiply/src/main.rs @@ -1,7 +1,7 @@ use sunscreen_compiler::{ fhe_program, types::{bfv::Signed, Cipher}, - Compiler, PlainModulusConstraint, Runtime + Compiler, PlainModulusConstraint, Runtime, }; /** diff --git a/sunscreen_compiler/src/error.rs b/sunscreen_compiler/src/error.rs index 3336558c9..72ee20186 100644 --- a/sunscreen_compiler/src/error.rs +++ b/sunscreen_compiler/src/error.rs @@ -28,6 +28,11 @@ pub enum Error { * An Error occurred in the Sunscreen runtime. */ RuntimeError(crate::RuntimeError), + + /** + * The compiled Sunscreen FHE program is malformed. + */ + FheProgramError(sunscreen_fhe_program::Error), } impl From for Error { diff --git a/sunscreen_compiler/src/params.rs b/sunscreen_compiler/src/params.rs index 1914ff5df..55033150a 100644 --- a/sunscreen_compiler/src/params.rs +++ b/sunscreen_compiler/src/params.rs @@ -143,6 +143,8 @@ where let ir = fhe_program_fn.build(¶ms)?.compile(); + ir.validate().map_err(|e| Error::FheProgramError(e))?; + // From a noise standpoint, it doesn't matter what is in the plaintext or if the output // is meaningful or not. Just run a bunch of 1 values through the fhe_program and measure the // noise. We choose 1, as it avoids transparent ciphertexts when diff --git a/sunscreen_compiler/src/types/intern/fhe_program_node.rs b/sunscreen_compiler/src/types/intern/fhe_program_node.rs index 4115a4120..640224761 100644 --- a/sunscreen_compiler/src/types/intern/fhe_program_node.rs +++ b/sunscreen_compiler/src/types/intern/fhe_program_node.rs @@ -28,10 +28,14 @@ use std::ops::{Add, Div, Mul, Neg, Shl, Shr, Sub}; * but rather until we clear the arena. We clean the arena in the FHE program macro after * FHE program construction and thus after all FheProgramNodes have gone out of scope. * + * You should never explicitly construct these outside of e.g. + * FHE type GraphCipherPlainAdd traits, which run during graph + * construction. + * * # Undefined behavior * These types must be constructed while a [`crate::CURRENT_CTX`] refers to a valid * [`crate::Context`]. Furthermore, no [`FheProgramNode`] should outlive the said context. - * Violating any of these condicitions may result in memory corruption or + * Violating any of these conditions may result in memory corruption or * use-after-free. */ pub struct FheProgramNode { @@ -48,19 +52,21 @@ impl FheProgramNode { /** * Creates a new FHE program node with the given node index. * - * These are an implementation detail needed while constructing the FHE program graph - * and should not be constructed at any other time. Thus, you should never - * directly create a [`FheProgramNode`]. + * These are an implementation detail needed while constructing the + * FHE program graph + * and should not be constructed at any other time. You should never + * need to directly create a [`FheProgramNode`]. * * # Remarks - * This type internally captures a slice rather than directly storing its own Vec. We do this - * so the type can impl Copy and composing FHE programs is natural without the user needing to call - * clone() all the time. + * This type internally captures a slice rather than directly + * storing its own Vec. We do this so the type can impl Copy and + * composing FHE programs is natural without the user needing to + * call clone() all the time. * * # Undefined behavior - * This type references memory in a backing [`crate::Context`] and without carefully ensuring FheProgramNodes - * never outlive the backing context, use-after-free can occur. - * + * This type references memory in a bump allocator. Failing to + * ensure FheProgramNodes never outlive the backing context, will + * result in use-after-free. */ pub fn new(ids: &[NodeIndex]) -> Self { INDEX_ARENA.with(|allocator| { @@ -72,6 +78,7 @@ impl FheProgramNode { // The memory in the bump allocator is valid until we call reset, which // we do after creating the FHE program. At this time, no FheProgramNodes should // remain. + // We invoke the dark transmutation ritual to turn a finite lifetime into a 'static. Self { ids: unsafe { std::mem::transmute(ids_dest) }, _phantom: std::marker::PhantomData, diff --git a/sunscreen_compiler/src/types/ops/add.rs b/sunscreen_compiler/src/types/ops/add.rs index 95efffc2b..ed2c1c35c 100644 --- a/sunscreen_compiler/src/types/ops/add.rs +++ b/sunscreen_compiler/src/types/ops/add.rs @@ -6,6 +6,9 @@ use crate::types::{ /** * Called when an Fhe Program encounters a + operation on two encrypted * types. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherAdd { /** @@ -30,6 +33,9 @@ pub trait GraphCipherAdd { /** * Called when an Fhe Program encounters a + operation on one encrypted * and one unencrypted type. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherPlainAdd { /** @@ -54,6 +60,9 @@ pub trait GraphCipherPlainAdd { /** * Called when an Fhe Program encounters a + operation on one encrypted * and a literal. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherConstAdd { /** diff --git a/sunscreen_compiler/src/types/ops/div.rs b/sunscreen_compiler/src/types/ops/div.rs index 45d05c6a9..0a0c0fa8c 100644 --- a/sunscreen_compiler/src/types/ops/div.rs +++ b/sunscreen_compiler/src/types/ops/div.rs @@ -5,6 +5,9 @@ use crate::types::{ /** * Called when an Fhe Program encounters a / operation on two encrypted types. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherDiv { /** @@ -29,6 +32,9 @@ pub trait GraphCipherDiv { /** * Called when an Fhe Program encounters a / operation with a * ciphertext numerator and plaintext denominator. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherPlainDiv { /** @@ -53,6 +59,9 @@ pub trait GraphCipherPlainDiv { /** * Called when an Fhe Program encounters a / operation with a * plaintext numerator and ciphertext denominator. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphPlainCipherDiv { /** @@ -76,6 +85,9 @@ pub trait GraphPlainCipherDiv { /** * Called when an Fhe Program encounters a / operation on an encrypted numerator and literal denominator. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherConstDiv { /** @@ -100,6 +112,9 @@ pub trait GraphCipherConstDiv { /** * Called when an Fhe Program encounters a / operation on a * literal numerator and encrypted denominator. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphConstCipherDiv { /** diff --git a/sunscreen_compiler/src/types/ops/mul.rs b/sunscreen_compiler/src/types/ops/mul.rs index 21c20dc84..c6c26a0cd 100644 --- a/sunscreen_compiler/src/types/ops/mul.rs +++ b/sunscreen_compiler/src/types/ops/mul.rs @@ -5,6 +5,9 @@ use crate::types::{ /** * Called when an Fhe Program encounters a * operation on two encrypted types. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherMul { /** @@ -29,6 +32,9 @@ pub trait GraphCipherMul { /** * Called when an Fhe Program encounters a * operation on an encrypted * and plaintext data type. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherPlainMul { /** @@ -53,6 +59,9 @@ pub trait GraphCipherPlainMul { /** * Called when an Fhe Program encounters a + operation on one encrypted * and a literal. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherConstMul { /** diff --git a/sunscreen_compiler/src/types/ops/neg.rs b/sunscreen_compiler/src/types/ops/neg.rs index 2f45d342b..09a59b016 100644 --- a/sunscreen_compiler/src/types/ops/neg.rs +++ b/sunscreen_compiler/src/types/ops/neg.rs @@ -5,6 +5,9 @@ use crate::types::{ /** * Called when the user performs unary negation (-) on a ciphertext. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherNeg { /** diff --git a/sunscreen_compiler/src/types/ops/rotate.rs b/sunscreen_compiler/src/types/ops/rotate.rs index 5fbccc6e0..3086ff94b 100644 --- a/sunscreen_compiler/src/types/ops/rotate.rs +++ b/sunscreen_compiler/src/types/ops/rotate.rs @@ -2,6 +2,9 @@ use crate::types::{intern::FheProgramNode, Cipher, FheType}; /** * Swaps the rows of the given ciphertext. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherSwapRows where diff --git a/sunscreen_compiler/src/types/ops/sub.rs b/sunscreen_compiler/src/types/ops/sub.rs index 4f3f65f69..d567fe3b1 100644 --- a/sunscreen_compiler/src/types/ops/sub.rs +++ b/sunscreen_compiler/src/types/ops/sub.rs @@ -74,6 +74,9 @@ pub trait GraphPlainCipherSub { /** * Called when an Fhe Program encounters a - operation on two encrypted types. + * + * This trait is an implementation detail of FHE program compilation; + * you should not directly call methods on this trait. */ pub trait GraphCipherConstSub { /** diff --git a/sunscreen_compiler/tests/edge_cases.rs b/sunscreen_compiler/tests/edge_cases.rs index f8858dd84..2480169fb 100644 --- a/sunscreen_compiler/tests/edge_cases.rs +++ b/sunscreen_compiler/tests/edge_cases.rs @@ -1,6 +1,6 @@ use sunscreen_compiler::{ + types::{bfv::Signed, Cipher}, *, - types::{bfv::{Signed}, Cipher} }; #[test] @@ -22,7 +22,9 @@ fn unused_cipher_parameter_1() { let a = runtime.encrypt(Signed::from(15), &public_key).unwrap(); let b = runtime.encrypt(Signed::from(5), &public_key).unwrap(); - let result = runtime.run(&program, vec![a.clone(), a, b], &public_key).unwrap(); + let result = runtime + .run(&program, vec![a.clone(), a, b], &public_key) + .unwrap(); let c: Signed = runtime.decrypt(&result[0], &private_key).unwrap(); @@ -48,7 +50,9 @@ fn unused_cipher_parameter_2() { let a = runtime.encrypt(Signed::from(15), &public_key).unwrap(); let b = runtime.encrypt(Signed::from(5), &public_key).unwrap(); - let result = runtime.run(&program, vec![a.clone(), a, b], &public_key).unwrap(); + let result = runtime + .run(&program, vec![a.clone(), a, b], &public_key) + .unwrap(); let c: Signed = runtime.decrypt(&result[0], &private_key).unwrap(); @@ -74,7 +78,9 @@ fn unused_cipher_parameter_3() { let a = runtime.encrypt(Signed::from(15), &public_key).unwrap(); let b = runtime.encrypt(Signed::from(5), &public_key).unwrap(); - let result = runtime.run(&program, vec![a, b.clone(), b], &public_key).unwrap(); + let result = runtime + .run(&program, vec![a, b.clone(), b], &public_key) + .unwrap(); let c: Signed = runtime.decrypt(&result[0], &private_key).unwrap(); @@ -163,4 +169,4 @@ fn unused_plain_parameter_3() { let c: Signed = runtime.decrypt(&result[0], &private_key).unwrap(); assert_eq!(c, 20.into()); -} \ No newline at end of file +} diff --git a/sunscreen_compiler/tests/typename_tests.rs b/sunscreen_compiler/tests/typename_tests.rs index 0daab9d9c..bf4ac3da8 100644 --- a/sunscreen_compiler/tests/typename_tests.rs +++ b/sunscreen_compiler/tests/typename_tests.rs @@ -1,11 +1,9 @@ use sunscreen_compiler::{ crate_version, - TypeName as DeriveTypeName, types::{Type, TypeName, TypeNameInstance, Version}, + TypeName as DeriveTypeName, }; - - #[test] fn derive_typename_example() { #[derive(DeriveTypeName)] diff --git a/sunscreen_fhe_program/src/lib.rs b/sunscreen_fhe_program/src/lib.rs index fc80a7a56..b2f86e8f2 100644 --- a/sunscreen_fhe_program/src/lib.rs +++ b/sunscreen_fhe_program/src/lib.rs @@ -598,7 +598,7 @@ impl FheProgram { let is_input = match n.operation { Operation::InputPlaintext(_) => true, Operation::InputCiphertext(_) => true, - _ => false + _ => false, }; if closure_set.contains(&revmap[id.index()]) || is_input { diff --git a/sunscreen_runtime/src/run.rs b/sunscreen_runtime/src/run.rs index 67a41c3a4..4986d7a9a 100644 --- a/sunscreen_runtime/src/run.rs +++ b/sunscreen_runtime/src/run.rs @@ -365,16 +365,11 @@ where let ir = ir.as_ref(); // Initialize the number of incomplete dependencies. - let mut deps: Vec = Vec::with_capacity(ir.graph.node_count()); - - for n in ir.graph.node_indices() { - unsafe { - *deps.get_unchecked_mut(n.index()) = - AtomicUsize::new(ir.graph.neighbors_directed(n, Direction::Incoming).count()); - } - } - - unsafe { deps.set_len(ir.graph.node_count()) }; + let deps = ir + .graph + .node_indices() + .map(|n| AtomicUsize::new(ir.graph.neighbors_directed(n, Direction::Incoming).count())) + .collect::>(); let items_remaining = AtomicUsize::new(ir.graph.node_count()); @@ -383,14 +378,18 @@ where // iteration causes a race condition between the filter_map closer // evaluating and the deps counts being decremented, potentially // resulting in nodes being run more than once. - let initial_ready = deps.iter().enumerate().filter_map(|(id, count)| { - if count.load(Ordering::Relaxed) == 0 { - log::trace!("parallel_traverse: Initial node {}", id); - Some(id) - } else { - None - } - }).collect::>(); + let initial_ready = deps + .iter() + .enumerate() + .filter_map(|(id, count)| { + if count.load(Ordering::Relaxed) == 0 { + log::trace!("parallel_traverse: Initial node {}", id); + Some(id) + } else { + None + } + }) + .collect::>(); let returned_result = AtomicCell::new(Ok(())); @@ -402,8 +401,9 @@ where deps: &[AtomicUsize], returned_result: &AtomicCell>, items_remaining: &AtomicUsize, - callback: &F - ) where F: Fn(NodeIndex) -> Result<(), FheProgramRunFailure> + Sync + Send + callback: &F, + ) where + F: Fn(NodeIndex) -> Result<(), FheProgramRunFailure> + Sync + Send, { log::trace!("parallel_traverse: Running node {}", node_id.index()); @@ -427,7 +427,14 @@ where if old_val == 1 { s.spawn(move |_| { log::trace!("Node {} ready", e.index()); - run_internal(e, ir, deps, returned_result, items_remaining, callback); + run_internal( + e, + ir, + deps, + returned_result, + items_remaining, + callback, + ); }); } } @@ -440,7 +447,14 @@ where let callback = &callback; s.spawn(move |_| { - run_internal(NodeIndex::from(node_id as u32), ir, deps,returned_result, items_remaining, callback); + run_internal( + NodeIndex::from(node_id as u32), + ir, + deps, + returned_result, + items_remaining, + callback, + ); }); } }); diff --git a/sunscreen_runtime/src/runtime.rs b/sunscreen_runtime/src/runtime.rs index c0412a0ef..d186feabb 100644 --- a/sunscreen_runtime/src/runtime.rs +++ b/sunscreen_runtime/src/runtime.rs @@ -206,6 +206,11 @@ impl Runtime { where I: Into, { + // We're going to call run_program_unchecked, which + // can result in undefined behavior, non-termination, + // or panics on malformed programs. Since this method is safe, + // it must guard against calling run_program_unchecked with + // inputs that result in undefined behavior. fhe_program.fhe_program_fn.validate()?; // Aside from FHE program correctness, check that the required keys are given.