github/Sunscreen
1
1
Fork 0
mirror of https://github.com/Sunscreen-tech/Sunscreen.git synced 2026-04-19 03:00:06 -04:00
Code Issues Packages Projects Releases Wiki Activity

Remove unneccessary unsafe code. Add more validation and document some stuff

Browse Source
This commit is contained in:
Rick Weber
2022-02-18 15:45:46 -08:00
parent c46f583bc8
commit e84f8da88b
15 changed files with 121 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/simple_multiply/src/main.rs
View File

@@ -1,7 +1,7 @@
use sunscreen_compiler::{
fhe_program,
types::{bfv::Signed, Cipher},
Compiler, PlainModulusConstraint, Runtime
Compiler, PlainModulusConstraint, Runtime,
};
/**

5
sunscreen_compiler/src/error.rs
View File

@@ -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<seal::Error> for Error {

2
sunscreen_compiler/src/params.rs
View File

@@ -143,6 +143,8 @@ where
let ir = fhe_program_fn.build(&params)?.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

27
sunscreen_compiler/src/types/intern/fhe_program_node.rs
View File

@@ -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<T: NumCiphertexts> {
@@ -48,19 +52,21 @@ impl<T: NumCiphertexts> FheProgramNode<T> {
/**
* 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<T: NumCiphertexts> FheProgramNode<T> {
// 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,

9
sunscreen_compiler/src/types/ops/add.rs
View File

@@ -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 {
/**

15
sunscreen_compiler/src/types/ops/div.rs
View File

@@ -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 {
/**

9
sunscreen_compiler/src/types/ops/mul.rs
View File

@@ -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 {
/**

3
sunscreen_compiler/src/types/ops/neg.rs
View File

@@ -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 {
/**

3
sunscreen_compiler/src/types/ops/rotate.rs
View File

@@ -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

3
sunscreen_compiler/src/types/ops/sub.rs
View File

@@ -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 {
/**

16
sunscreen_compiler/tests/edge_cases.rs
View File

@@ -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());
}
}

4
sunscreen_compiler/tests/typename_tests.rs
View File

@@ -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)]

2
sunscreen_fhe_program/src/lib.rs
View File

@@ -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 {

58
sunscreen_runtime/src/run.rs
View File

@@ -365,16 +365,11 @@ where
let ir = ir.as_ref();
// Initialize the number of incomplete dependencies.
let mut deps: Vec<AtomicUsize> = 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::<Vec<AtomicUsize>>();
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::<Vec<usize>>();
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::<Vec<usize>>();
let returned_result = AtomicCell::new(Ok(()));
@@ -402,8 +401,9 @@ where
deps: &[AtomicUsize],
returned_result: &AtomicCell<Result<(), FheProgramRunFailure>>,
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,
);
});
}
});

5
sunscreen_runtime/src/runtime.rs
View File

@@ -206,6 +206,11 @@ impl Runtime {
where
I: Into<FheProgramInput>,
{
// 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.