From 925bc61ddefa7e6f30795d77a551eed6df30ccd1 Mon Sep 17 00:00:00 2001 From: lonerapier Date: Tue, 24 Sep 2024 02:36:57 +0530 Subject: [PATCH] add tests --- src/codegen/http.rs | 267 ++++++++++++++++++++++++-------------- src/codegen/integrated.rs | 38 ++---- src/codegen/json.rs | 71 ++++++++-- src/witness.rs | 53 ++++---- 4 files changed, 264 insertions(+), 165 deletions(-) diff --git a/src/codegen/http.rs b/src/codegen/http.rs index 27fd5ec..e1f885b 100644 --- a/src/codegen/http.rs +++ b/src/codegen/http.rs @@ -1,7 +1,4 @@ -use crate::{ - circuit_config::CircomkitCircuitConfig, witness::read_input_file_as_bytes, ExtractorArgs, - FileType, -}; +use crate::{circuit_config::CircomkitCircuitConfig, ExtractorArgs, FileType}; use regex::Regex; use serde::{Deserialize, Serialize}; @@ -98,6 +95,106 @@ impl HttpData { inputs } + + pub fn parse_input( + &self, + input: Vec, + ) -> Result<(HttpData, Vec), Box> { + let input_string = String::from_utf8(input)?; + + let parts: Vec<&str> = input_string.split("\r\n\r\n").collect(); + assert!(parts.len() <= 2); + + let mut body = vec![]; + if parts.len() == 2 { + body = parts[1].as_bytes().to_vec(); + } + + let headers: Vec<&str> = parts[0].split("\r\n").collect(); + let start_line: Vec<&str> = headers[0].split(" ").collect(); + assert_eq!(start_line.len(), 3); + + let (_, headers) = headers.split_at(1); + let mut headers_map = HashMap::::new(); + let re = Regex::new(r":\s+").unwrap(); + for &header in headers { + println!("header: {:?}", header); + let key_value: Vec<&str> = re.split(header).collect(); + assert_eq!(key_value.len(), 2); + println!("key: {:?}", key_value); + headers_map.insert(key_value[0].to_string(), key_value[1].to_string()); + } + + let http_data = match self { + HttpData::Request(_) => HttpData::Request(Request { + method: start_line[0].to_string(), + target: start_line[1].to_string(), + version: start_line[2].to_string(), + headers: headers_map, + }), + HttpData::Response(_) => HttpData::Response(Response { + version: start_line[0].to_string(), + status: start_line[1].to_string(), + message: start_line[2].to_string(), + headers: headers_map, + }), + }; + + Ok((http_data, body)) + } + + pub fn populate_params( + &self, + input: Vec, + ) -> Result, Box> { + let (_, http_body) = self.parse_input(input.clone())?; + + let mut params = vec![input.len()]; + + match self { + HttpData::Request(request) => { + params.push(request.method.len()); + params.push(request.target.len()); + params.push(request.version.len()); + for (key, value) in request.headers.iter() { + params.push(key.len()); + params.push(value.len()); + } + } + HttpData::Response(response) => { + params.push(http_body.len()); + params.push(response.version.len()); + params.push(response.status.len()); + params.push(response.message.len()); + for (key, value) in response.headers.iter() { + params.push(key.len()); + params.push(value.len()); + } + } + } + + Ok(params) + } + + fn build_circuit_config( + &self, + input_file: &Path, + codegen_filename: &str, + ) -> Result> { + println!("input_ifle: {:?}", input_file); + let input = FileType::Http.read_input(input_file)?; + + let circuit_template_name = match self { + HttpData::Request(_) => String::from("LockHTTPRequest"), + HttpData::Response(_) => String::from("LockHTTPResponse"), + }; + + Ok(CircomkitCircuitConfig { + file: format!("main/{}", codegen_filename), + template: circuit_template_name, + params: self.populate_params(input)?, + }) + } } impl std::fmt::Debug for HttpData { @@ -168,8 +265,6 @@ where Ok(map) } -const PRAGMA: &str = "pragma circom 2.1.9;\n\n"; - fn build_http_circuit( config: &CircomkitCircuitConfig, data: &HttpData, @@ -184,7 +279,7 @@ fn build_http_circuit( circuit_buffer += "\n*/\n"; // Version and includes - circuit_buffer += PRAGMA; + circuit_buffer += "pragma circom 2.1.9;\n\n"; circuit_buffer += "include \"../http/interpreter.circom\";\n"; circuit_buffer += "include \"../http/parser/machine.circom\";\n"; circuit_buffer += "include \"../utils/bytes.circom\";\n"; @@ -529,96 +624,6 @@ fn build_http_circuit( Ok(()) } -pub fn parse_http_file( - locfile: &HttpData, - input: Vec, -) -> Result<(HttpData, Vec), Box> { - let input_string = String::from_utf8(input)?; - - let parts: Vec<&str> = input_string.split("\r\n\r\n").collect(); - assert!(parts.len() <= 2); - - let mut body = vec![]; - if parts.len() == 2 { - body = parts[1].as_bytes().to_vec(); - } - - let headers: Vec<&str> = parts[0].split("\r\n").collect(); - let start_line: Vec<&str> = headers[0].split(" ").collect(); - assert_eq!(start_line.len(), 3); - - let (_, headers) = headers.split_at(1); - let mut headers_map = HashMap::::new(); - let re = Regex::new(r":\s(.+)").unwrap(); - for &header in headers { - let key_value: Vec<&str> = re.split(header).collect(); - assert_eq!(key_value.len(), 2); - headers_map.insert(key_value[0].to_string(), key_value[1].to_string()); - } - - let http_data = match locfile { - HttpData::Request(_) => HttpData::Request(Request { - method: start_line[0].to_string(), - target: start_line[1].to_string(), - version: start_line[2].to_string(), - headers: headers_map, - }), - HttpData::Response(_) => HttpData::Response(Response { - version: start_line[0].to_string(), - status: start_line[1].to_string(), - message: start_line[2].to_string(), - headers: headers_map, - }), - }; - - Ok((http_data, body)) -} - -fn build_circuit_config( - input_file: &Path, - lockfile: &HttpData, - codegen_filename: &str, -) -> Result> { - let input = read_input_file_as_bytes(&FileType::Http, input_file)?; - - let (_, http_body) = parse_http_file(lockfile, input.clone())?; - - let circuit_template_name = match lockfile { - HttpData::Request(_) => String::from("LockHTTPRequest"), - HttpData::Response(_) => String::from("LockHTTPResponse"), - }; - - let mut params = vec![input.len()]; - - match lockfile { - HttpData::Request(request) => { - params.push(request.method.len()); - params.push(request.target.len()); - params.push(request.version.len()); - for (key, value) in request.headers.iter() { - params.push(key.len()); - params.push(value.len()); - } - } - HttpData::Response(response) => { - params.push(http_body.len()); - params.push(response.version.len()); - params.push(response.status.len()); - params.push(response.message.len()); - for (key, value) in response.headers.iter() { - params.push(key.len()); - params.push(value.len()); - } - } - } - - Ok(CircomkitCircuitConfig { - file: format!("main/{}", codegen_filename), - template: circuit_template_name, - params, - }) -} - pub fn http_circuit_from_args( args: &ExtractorArgs, ) -> Result> { @@ -642,9 +647,75 @@ pub fn http_circuit_from_lockfile( codegen_filename: &str, debug: bool, ) -> Result> { - let config = build_circuit_config(input_file, http_data, codegen_filename)?; + let config = http_data.build_circuit_config(input_file, codegen_filename)?; build_http_circuit(&config, http_data, codegen_filename, debug)?; Ok(config) } + +#[cfg(test)] +mod test { + use std::path::PathBuf; + + use super::*; + + #[test] + fn params() { + let lockfile: HttpData = serde_json::from_slice(include_bytes!( + "../../examples/http/lockfile/spotify.lock.json" + )) + .unwrap(); + + let params = lockfile.params(); + + assert_eq!(params.len(), 7); + assert_eq!(params[0], "DATA_BYTES"); + assert_eq!(params[1], "maxContentLength"); + } + + #[test] + fn inputs() { + let lockfile: HttpData = serde_json::from_slice(include_bytes!( + "../../examples/http/lockfile/spotify.lock.json" + )) + .unwrap(); + + let inputs = lockfile.inputs(); + + assert_eq!(inputs.len(), 6); + assert_eq!(inputs[1], "version"); + assert_eq!(inputs[2], "status"); + assert_eq!(inputs[3], "message"); + } + + #[test] + fn populate_params() { + let lockfile: HttpData = serde_json::from_slice(include_bytes!( + "../../examples/http/lockfile/request.lock.json" + )) + .unwrap(); + + let input = include_bytes!("../../examples/http/get_request.http"); + + let params = lockfile.populate_params(input.to_vec()).unwrap(); + + assert_eq!(params.len(), 8); + assert_eq!(params, [input.len(), 3, 4, 8, 4, 9, 6, 16]); + } + + #[test] + fn parse_input() { + let lockfile: HttpData = serde_json::from_slice(include_bytes!( + "../../examples/http/lockfile/request.lock.json" + )) + .unwrap(); + + let input = include_bytes!("../../examples/http/get_request.http"); + + let (http, body) = lockfile.parse_input(input.to_vec()).unwrap(); + + assert_eq!(body.len(), 0); + assert_eq!(http.headers()["Accept"], "application/json"); + } +} diff --git a/src/codegen/integrated.rs b/src/codegen/integrated.rs index 89630c1..0d2c354 100644 --- a/src/codegen/integrated.rs +++ b/src/codegen/integrated.rs @@ -1,11 +1,10 @@ use crate::{ circuit_config::CircomkitCircuitConfig, codegen::{ - http::{parse_http_file, HttpData}, + http::HttpData, json::{Key, Lockfile as JsonLockfile}, }, - witness::read_input_file_as_bytes, - ExtractorArgs, + ExtractorArgs, FileType, }; use serde::{Deserialize, Serialize}; use std::path::Path; @@ -158,33 +157,12 @@ fn build_circuit_config( json_lockfile: &JsonLockfile, output_filename: &str, ) -> Result> { - let input = read_input_file_as_bytes(&crate::FileType::Http, &args.input_file)?; + let input = FileType::Http.read_input(&args.input_file)?; - let (_, http_body) = parse_http_file(http_data, input.clone())?; + let (_, http_body) = http_data.parse_input(input.clone())?; - let mut params = vec![input.len()]; - - match http_data { - HttpData::Request(request) => { - params.push(request.method.len()); - params.push(request.target.len()); - params.push(request.version.len()); - for (key, value) in request.headers.iter() { - params.push(key.len()); - params.push(value.len()); - } - } - HttpData::Response(response) => { - params.push(http_body.len()); - params.push(response.version.len()); - params.push(response.status.len()); - params.push(response.message.len()); - for (key, value) in response.headers.iter() { - params.push(key.len()); - params.push(value.len()); - } - } - } + // populate http params + let mut params = http_data.populate_params(input)?; // add json params and remove first param: `DATA_BYTES` let mut json_params = json_lockfile.populate_params(&http_body)?; @@ -217,8 +195,8 @@ pub fn integrated_circuit(args: &ExtractorArgs) -> Result<(), Box HashMap> { + pub fn keys_as_bytes(&self) -> HashMap> { let mut keys = HashMap::>::new(); for (i, key) in self.keys.iter().enumerate() { if let Key::String(key) = key { @@ -80,7 +81,7 @@ impl Lockfile { &self, input: &[u8], output_filename: &str, - ) -> Result> { + ) -> Result> { let circuit_template_name = match self.value_type { ValueType::String => String::from("ExtractStringValue"), ValueType::Number => String::from("ExtractNumValue"), @@ -95,7 +96,7 @@ impl Lockfile { /// Builds circuit arguments /// `[DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1, ..., maxValueLen]` - pub fn populate_params(&self, input: &[u8]) -> Result, Box> { + pub fn populate_params(&self, input: &[u8]) -> Result, Box> { let mut params = vec![input.len(), json_max_stack_height(input)]; for (i, key) in self.keys.iter().enumerate() { @@ -112,7 +113,7 @@ impl Lockfile { Ok(params) } - pub fn get_value(&self, input: &[u8]) -> Result> { + pub fn get_value(&self, input: &[u8]) -> Result> { let mut current_value: Value = serde_json::from_slice(input)?; for key in self.keys.iter() { match key { @@ -290,7 +291,7 @@ fn build_json_circuit( data: &Lockfile, output_filename: &str, debug: bool, -) -> Result<(), Box> { +) -> Result<(), Box> { let mut circuit_buffer = String::new(); // Dump out the contents of the lockfile used into the circuit @@ -641,12 +642,12 @@ fn build_json_circuit( /// - writes file pub fn json_circuit_from_args( args: &ExtractorArgs, -) -> Result> { - let lockfile: Lockfile = serde_json::from_slice(&std::fs::read(&args.lockfile)?)?; +) -> Result> { + let lockfile: Lockfile = serde_json::from_slice(&fs::read(&args.lockfile)?)?; let circuit_filename = format!("json_{}", args.circuit_name); - let input = std::fs::read(&args.input_file)?; + let input = fs::read(&args.input_file)?; let config = json_circuit_from_lockfile(&input, &lockfile, &circuit_filename, args.debug)?; config.write(&args.circuit_name)?; @@ -659,7 +660,7 @@ pub fn json_circuit_from_lockfile( lockfile: &Lockfile, output_filename: &str, debug: bool, -) -> Result> { +) -> Result> { let config = lockfile.build_circuit_config(input, output_filename)?; build_json_circuit(&config, lockfile, output_filename, debug)?; @@ -668,7 +669,7 @@ pub fn json_circuit_from_lockfile( #[cfg(test)] mod test { - use super::Lockfile; + use super::*; #[test] fn params() { @@ -696,4 +697,54 @@ mod test { assert_eq!(inputs.len(), 2); assert_eq!(inputs[0], "data"); } + + #[test] + fn populate_params() { + let input = include_bytes!("../../examples/json/test/spotify.json"); + let lockfile: Lockfile = + serde_json::from_slice(include_bytes!("../../examples/json/lockfile/spotify.json")) + .unwrap(); + + let params = lockfile.populate_params(input).unwrap(); + + assert_eq!(params.len(), lockfile.params().len()); + assert_eq!(params[0], input.len()); + } + + #[test] + fn build_circuit_config() { + let input = include_bytes!("../../examples/json/test/spotify.json"); + let lockfile: Lockfile = + serde_json::from_slice(include_bytes!("../../examples/json/lockfile/spotify.json")) + .unwrap(); + + let config = lockfile + .build_circuit_config(input, "output_filename") + .unwrap(); + + assert_eq!(config.template, "ExtractStringValue"); + assert_eq!(config.file, "main/output_filename"); + } + + #[test] + fn json_value() { + let input = include_bytes!("../../examples/json/test/spotify.json"); + let lockfile: Lockfile = + serde_json::from_slice(include_bytes!("../../examples/json/lockfile/spotify.json")) + .unwrap(); + + let value = lockfile.get_value(input).unwrap(); + + assert_eq!(value, "Taylor Swift"); + } + + #[test] + fn max_stack_height() { + let input = include_bytes!("../../examples/json/test/two_keys.json"); + + assert_eq!(json_max_stack_height(input), 1); + + let input = include_bytes!("../../examples/json/test/spotify.json"); + assert_eq!(json_max_stack_height(input), 5); + } } diff --git a/src/witness.rs b/src/witness.rs index 79c2bdb..20acaf9 100644 --- a/src/witness.rs +++ b/src/witness.rs @@ -55,31 +55,30 @@ fn print_boxed_output(lines: Vec) { println!("{}", bottom_border); } -pub fn read_input_file_as_bytes( - file_type: &FileType, - file_path: &Path, -) -> Result, Box> { - match file_type { - FileType::Json => Ok(std::fs::read(file_path)?), - FileType::Http | FileType::Extended => { - let mut data = std::fs::read(file_path)?; - let mut i = 0; - // convert LF to CRLF - while i < data.len() { - if data[i] == 10 && (i == 0 || data[i - 1] != 13) { - data.insert(i, 13); - i += 2; - } else { - i += 1; +impl FileType { + pub fn read_input(&self, input: &Path) -> Result, Box> { + match self { + FileType::Json => Ok(std::fs::read(input)?), + FileType::Http | FileType::Extended => { + let mut data = std::fs::read(input)?; + let mut i = 0; + // convert LF to CRLF + while i < data.len() { + if data[i] == 10 && (i == 0 || data[i - 1] != 13) { + data.insert(i, 13); + i += 2; + } else { + i += 1; + } } + Ok(data) } - Ok(data) } } } pub fn parser_witness(args: ParserWitnessArgs) -> Result<(), Box> { - let data = read_input_file_as_bytes(&args.subcommand, &args.input_file)?; + let data = args.subcommand.read_input(&args.input_file)?; let witness = ParserWitness { data: data.clone() }; @@ -114,7 +113,7 @@ pub fn parser_witness(args: ParserWitnessArgs) -> Result<(), Box Result<(), Box> { // read input and lockfile - let input_data = read_input_file_as_bytes(&args.subcommand, &args.input_file)?; + let input_data = args.subcommand.read_input(&args.input_file)?; let lockfile_data = std::fs::read(&args.lockfile)?; let lockfile: Lockfile = serde_json::from_slice(&lockfile_data)?; @@ -122,7 +121,7 @@ fn json_extractor_witness(args: ExtractorWitnessArgs) -> Result<(), Box Result<(), Box Result<(), Box> { // read input and lockfile - let input_data = read_input_file_as_bytes(&args.subcommand, &args.input_file)?; + let data = args.subcommand.read_input(&args.input_file)?; let lockfile_data = std::fs::read(&args.lockfile)?; let http_data: HttpData = serde_json::from_slice(&lockfile_data)?; // create witness data let witness = HttpExtractorWitness { - data: input_data.clone(), + data: data.clone(), http_data, }; @@ -180,7 +179,7 @@ fn http_extractor_witness(args: ExtractorWitnessArgs) -> Result<(), Box Result<(), Box> { // read input and lockfile - let input_data = read_input_file_as_bytes(&args.subcommand, &args.input_file)?; + let data = args.subcommand.read_input(&args.input_file)?; let lockfile_data = std::fs::read(&args.lockfile)?; let lockfile: ExtendedLockfile = serde_json::from_slice(&lockfile_data)?; @@ -200,10 +199,10 @@ fn extended_extractor_witness( // create witness data let witness = ExtendedWitness { http_witness: HttpExtractorWitness { - data: input_data.clone(), + data: data.clone(), http_data: lockfile.http, }, - keys: lockfile.json.as_bytes(), + keys: lockfile.json.keys_as_bytes(), }; // create witness dir @@ -221,7 +220,7 @@ fn extended_extractor_witness( // Prepare lines to print let mut lines = Vec::new(); - lines.push(format!("Data length: {}", input_data.len())); + lines.push(format!("Data length: {}", data.len())); // Print the output inside a nicely formatted box print_boxed_output(lines);