pozk_utils/
utils.rs

1use anyhow::{anyhow, Result};
2use once_cell::sync::OnceCell;
3use std::path::PathBuf;
4use tokio::fs;
5
6static BASE_PATH: OnceCell<PathBuf> = OnceCell::new();
7static API_SERVER: OnceCell<String> = OnceCell::new();
8
9pub fn init_path_and_server(path: &str, server: &str) {
10    BASE_PATH
11        .set(PathBuf::from(path))
12        .expect("Unable set BASE_PATH");
13    API_SERVER
14        .set(server.to_owned())
15        .expect("Unable set API_SERVER");
16}
17
18pub async fn write_task_input(tid: &str, inputs: Vec<u8>, publics: Vec<u8>) -> Result<()> {
19    let mut path = BASE_PATH.get().expect("Missing BASE PATH").clone();
20    path.push(tid);
21
22    let mut bytes = (inputs.len() as u32).to_be_bytes().to_vec();
23    bytes.extend(inputs);
24    bytes.extend(publics);
25
26    fs::write(path, bytes).await?;
27    Ok(())
28}
29
30pub async fn read_task_input(tid: &str) -> Result<Vec<u8>> {
31    let mut path = BASE_PATH.get().expect("Missing BASE PATH").clone();
32    path.push(tid);
33
34    let bytes = fs::read(path).await?;
35    Ok(bytes)
36}
37
38pub async fn parse_task_input(data: Vec<u8>) -> Result<(Vec<u8>, Vec<u8>)> {
39    let mut inputs_len_bytes = [0u8; 4];
40    inputs_len_bytes.copy_from_slice(&data[0..4]);
41    let inputs_len = u32::from_be_bytes(inputs_len_bytes) as usize;
42    if data.len() < inputs_len + 4 {
43        return Err(anyhow!("Invalid proof length"));
44    }
45
46    let raw_data = &data[4..];
47    let inputs = raw_data[..inputs_len].to_vec();
48    let publics = raw_data[inputs_len..].to_vec();
49
50    Ok((inputs, publics))
51}
52
53pub async fn remove_task_input(tid: &str) -> Result<()> {
54    let mut path = BASE_PATH.get().expect("Missing BASE PATH").clone();
55    path.push(tid);
56
57    fs::remove_file(path).await?;
58    Ok(())
59}
60
61pub async fn write_task_proof(tid: &str, proof: Vec<u8>) -> Result<()> {
62    let mut path = BASE_PATH.get().expect("Missing BASE PATH").clone();
63    path.push(format!("proof-{}", tid));
64
65    fs::write(path, proof).await?;
66    Ok(())
67}
68
69pub async fn read_task_proof(tid: &str) -> Result<Vec<u8>> {
70    let mut path = BASE_PATH.get().expect("Missing BASE PATH").clone();
71    path.push(format!("proof-{}", tid));
72
73    let bytes = fs::read(&path).await?;
74    fs::remove_file(path).await?;
75
76    Ok(bytes)
77}
78
79pub fn get_task_api(tid: &str) -> String {
80    let server = API_SERVER.get().expect("Missing API SERVER");
81    format!("{}/inner/tasks/{}", server, tid)
82}
83
84pub fn convert_task_to_connect_api(url: &str) -> (String, &str) {
85    let mut v: Vec<&str> = url.split('/').collect();
86    let id = v.pop().unwrap_or("");
87    let _ = v.pop(); // pop tasks
88    v.push("connect"); // push connect
89    (v.join("/"), id)
90}
91
92pub async fn download_input() -> Result<Vec<u8>> {
93    let uri = std::env::var("INPUT").unwrap();
94    download_input_with_uri(&uri).await
95}
96
97pub async fn download_input_with_uri(uri: &str) -> Result<Vec<u8>> {
98    let body = reqwest::get(uri).await?.bytes().await?;
99    Ok(body.to_vec())
100}
101
102pub async fn upload_proof(output: Vec<u8>, proof: Vec<u8>) -> Result<()> {
103    let uri = std::env::var("INPUT").unwrap();
104    upload_proof_with_uri(&uri, output, proof).await
105}
106
107pub async fn upload_proof_with_uri(uri: &str, output: Vec<u8>, proof: Vec<u8>) -> Result<()> {
108    let mut bytes = vec![];
109    bytes.extend((output.len() as u32).to_be_bytes());
110    bytes.extend(output);
111    bytes.extend(proof);
112
113    let client = reqwest::Client::new();
114    let _ = client.post(uri).body(bytes).send().await?;
115
116    Ok(())
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[tokio::test]
124    async fn test_download() {
125        let res = download_input_with_uri("http://localhost:9098/inner/tasks/1")
126            .await
127            .unwrap();
128        println!("{} {:?}", res.len(), res);
129    }
130
131    #[tokio::test]
132    async fn test_upload() {
133        upload_proof_with_uri(
134            "http://localhost:9098/inner/tasks/1",
135            vec![1, 2, 3],
136            vec![1, 2, 3, 4],
137        )
138        .await
139        .unwrap();
140    }
141}
pozk_utils/utils.rs

pozk_utils/
utils.rs
