Tee transaction filtering by viewing key

tee/Cargo.toml

@@ -12,11 +12,14 @@ categories = ["cryptography"]
 repository = "https://github.com/Cardinal-Cryptography/zkOS-monorepo"
 
 [workspace.dependencies]
+alloy-primitives = { version = "0.8.15" }
+base64 = "0.22.1"
 env_logger = "0.11.8"
 futures = "0.3.31"
 log = "0.4.27"
 serde = "1.0.219"
 serde_json = "1.0.140"
+shielder-circuits = { git = "https://github.com/Cardinal-Cryptography/zkOS-circuits", rev = "df31437" }
 shielder-rewards-common = { path = "crates/shielder-rewards-common" }
 thiserror = "2.0.12"
 tokio = "1.45.0"

tee/crates/shielder-rewards-common/src/protocol.rs

@@ -7,11 +7,14 @@ pub const VSOCK_PORT: u32 = 5000;
 #[derive(Serialize, Deserialize)]
 pub enum Request {
     Ping,
+    CalculateReward { viewing_key_base64: String },
 }
 
 #[derive(Serialize, Deserialize)]
 pub enum Response {
     Pong,
+    CalculateReward { reward: String },
+    Error { message: String },
 }
 
 pub type RewardServer = VsockServer<Request, Response>;

tee/crates/shielder-rewards-tee/Cargo.toml

@@ -9,8 +9,12 @@ categories = ["cryptography"]
 repository = "https://github.com/Cardinal-Cryptography/zkOS-monorepo"
 
 [dependencies]
+alloy-primitives = { workspace = true }
+base64 = { workspace = true }
 env_logger = { workspace = true }
 log = { workspace = true }
+shielder-circuits = { workspace = true }
 shielder-rewards-common = { workspace = true }
+thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] }
 tokio-vsock = { workspace = true }

tee/crates/shielder-rewards-tee/src/crypto.rs

@@ -0,0 +1,41 @@
+use std::borrow::Borrow;
+
+use alloy_primitives::U256;
+use shielder_circuits::{Fr, PrimeField};
+
+use crate::Error;
+
+/// Convert a U256 value to a field element.
+pub fn u256_to_field<Fr: From<[u64; 4]>>(value: impl Borrow<U256>) -> Fr {
+    Fr::from(*value.borrow().as_limbs())
+}
+
+/// Convert a field element to a U256 value.
+pub fn field_to_u256<F: PrimeField<Repr = [u8; BYTE_LENGTH]>, const BYTE_LENGTH: usize>(
+    value: impl Borrow<F>,
+) -> U256 {
+    U256::from_le_bytes(value.borrow().to_repr())
+}
+
+pub fn bytes32_to_field(bytes: &[u8; 32]) -> Result<Fr, Error> {
+    Fr::from_bytes(bytes)
+        .into_option()
+        .ok_or(Error::FieldConversion(
+            "Failed to convert to Fr".to_string(),
+        ))
+}
+
+pub fn blob_to_field(blob: &[u8]) -> Result<Fr, Error> {
+    if blob.len() != 32 {
+        return Err(Error::FieldConversion(format!(
+            "Expected 32 bytes, but got {} bytes",
+            blob.len()
+        )));
+    }
+
+    let bytes: [u8; 32] = blob
+        .try_into()
+        .map_err(|_| Error::FieldConversion("Failed to convert &[u8] to [u8; 32]".to_string()))?;
+
+    bytes32_to_field(&bytes)
+}

tee/crates/shielder-rewards-tee/src/lib.rs

@@ -0,0 +1,41 @@
+use alloy_primitives::U256;
+use thiserror::Error;
+
+pub mod crypto;
+pub mod rewards;
+
+#[derive(Clone, Debug)]
+pub struct ShielderTransaction {
+    pub value: U256,
+    pub mac_salt: U256,
+    pub mac_commitment: U256,
+}
+
+#[derive(Clone, Debug)]
+
+pub struct AppState {
+    pub txs: Vec<ShielderTransaction>,
+}
+
+#[derive(Debug, Error)]
+#[error(transparent)]
+#[non_exhaustive]
+pub enum Error {
+    #[error("Error reading AR private key file")]
+    ARKeyRead(#[from] std::io::Error),
+
+    #[error("Error converting from a little-endian byte representation to grumpkin::Fr")]
+    NotAGrumpkinBaseFieldElement,
+
+    #[error("Event is missing some data")]
+    MissingData,
+
+    #[error("Field conversion")]
+    FieldConversion(String),
+
+    #[error("Error while deserializing public key")]
+    DeserializePubKey,
+
+    #[error("Public key does not satisfy y^2 = x^3 - 17")]
+    PubkeyNotOnCurve,
+}

tee/crates/shielder-rewards-tee/src/main.rs

@@ -1,39 +1,62 @@
 use log::info;
 use shielder_rewards_common::protocol::{Request, Response, RewardServer, VSOCK_PORT};
+use shielder_rewards_tee::{rewards::calculate_reward, AppState};
 use tokio::spawn;
 use tokio_vsock::{VsockAddr, VsockListener, VsockStream, VMADDR_CID_ANY};
 
 #[tokio::main]
 async fn main() {
     env_logger::init();
 
-    if let Err(e) = run_server().await {
+    // TODO: Load app state from a file
+
+    let app_state = AppState { txs: vec![] };
+
+    if let Err(e) = run_server(&app_state).await {
         eprintln!("VSOCK Server error: {}", e);
     }
 }
 
-async fn run_server() -> Result<(), Box<dyn std::error::Error>> {
+async fn run_server(app_state: &AppState) -> Result<(), Box<dyn std::error::Error>> {
     let listener = VsockListener::bind(VsockAddr::new(VMADDR_CID_ANY, VSOCK_PORT))?;
 
     loop {
         let (stream, _) = listener.accept().await?;
-        spawn(handle_client(stream));
+        spawn(handle_client(stream, app_state.clone()));
     }
 }
 
-async fn handle_client(stream: VsockStream) {
-    let result = do_handle_client(stream).await;
+async fn handle_client(stream: VsockStream, app_state: AppState) {
+    let result = do_handle_client(stream, &app_state).await;
     info!("Client disconnected: {:?}", result);
 }
 
-async fn do_handle_client(stream: VsockStream) -> Result<(), Box<dyn std::error::Error>> {
+async fn do_handle_client(
+    stream: VsockStream,
+    app_state: &AppState,
+) -> Result<(), Box<dyn std::error::Error>> {
     let mut server: RewardServer = stream.into();
 
     loop {
         server
             .handle_request(async |command| match command {
                 Request::Ping => Response::Pong,
+                Request::CalculateReward { viewing_key_base64 } => {
+                    handle_error(calculate_reward(viewing_key_base64, app_state).await)
+                }
             })
             .await?;
     }
 }
+
+fn handle_error(result: Result<Response, Box<dyn std::error::Error>>) -> Response {
+    match result {
+        Ok(response) => response,
+        Err(e) => {
+            eprintln!("Error handling request: {}", e);
+            Response::Error {
+                message: "Internal server error".to_string(),
+            }
+        }
+    }
+}

tee/crates/shielder-rewards-tee/src/rewards.rs

@@ -0,0 +1,41 @@
+use alloy_primitives::U256;
+use base64::prelude::*;
+use shielder_circuits::{poseidon::off_circuit::hash, Fr};
+use shielder_rewards_common::protocol::Response;
+
+use crate::{
+    crypto::{blob_to_field, u256_to_field},
+    AppState,
+};
+
+pub fn viewing_key_filter(viewing_key: Fr, mac_salt: U256, mac_commitment: U256) -> bool {
+    let expected_commitment = hash(&[u256_to_field(mac_salt), viewing_key]);
+    u256_to_field::<Fr>(mac_commitment).eq(&expected_commitment)
+}
+
+pub async fn calculate_reward(
+    viewing_key_base64: String,
+    app_state: &AppState,
+) -> Result<Response, Box<dyn std::error::Error>> {
+    let viewing_key_bytes = BASE64_STANDARD.decode(&viewing_key_base64)?;
+
+    let viewing_key = blob_to_field(&viewing_key_bytes)?;
+
+    let filtered_txs = app_state
+        .txs
+        .iter()
+        .filter(|tx| viewing_key_filter(viewing_key, tx.mac_salt, tx.mac_commitment))
+        .collect::<Vec<_>>();
+
+    // Here you would calculate the reward based on the filtered transactions.
+    let sum = filtered_txs
+        .iter()
+        .map(|tx| tx.value)
+        .fold(U256::ZERO, |acc, value| acc + value);
+
+    // For now, we just return a sum.
+
+    Ok(Response::CalculateReward {
+        reward: sum.to_string(),
+    })
+}