darkfi/zkas/
decoder.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2025 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use darkfi_serial::{deserialize_partial, VarInt};
20
21use super::{
22    compiler::MAGIC_BYTES,
23    constants::{MAX_K, MAX_NS_LEN, MIN_BIN_SIZE},
24    types::HeapType,
25    LitType, Opcode, VarType,
26};
27use crate::{Error::ZkasDecoderError as ZkasErr, Result};
28
29/// A ZkBinary decoded from compiled zkas code.
30/// This is used by the zkvm.
31#[derive(Clone, Debug)]
32pub struct ZkBinary {
33    pub namespace: String,
34    pub k: u32,
35    pub constants: Vec<(VarType, String)>,
36    pub literals: Vec<(LitType, String)>,
37    pub witnesses: Vec<VarType>,
38    pub opcodes: Vec<(Opcode, Vec<(HeapType, usize)>)>,
39}
40
41// https://stackoverflow.com/questions/35901547/how-can-i-find-a-subsequence-in-a-u8-slice
42fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
43    haystack.windows(needle.len()).position(|window| window == needle)
44}
45
46impl ZkBinary {
47    pub fn decode(bytes: &[u8]) -> Result<Self> {
48        // Ensure that bytes is a certain minimum length. Otherwise the code
49        // below will panic due to an index out of bounds error.
50        if bytes.len() < MIN_BIN_SIZE {
51            return Err(ZkasErr("Not enough bytes".to_string()))
52        }
53        let magic_bytes = &bytes[0..4];
54        if magic_bytes != MAGIC_BYTES {
55            return Err(ZkasErr("Magic bytes are incorrect".to_string()))
56        }
57
58        let _binary_version = &bytes[4];
59
60        // Deserialize the k param
61        let (k, _): (u32, _) = deserialize_partial(&bytes[5..9])?;
62
63        // For now, we'll limit k.
64        if k > MAX_K {
65            return Err(ZkasErr("k param is too high, max allowed is 16".to_string()))
66        }
67
68        // After the binary version and k, we're supposed to have the witness namespace
69        let (namespace, _): (String, _) = deserialize_partial(&bytes[9..])?;
70
71        // Enforce a limit on the namespace string length
72        if namespace.len() > MAX_NS_LEN {
73            return Err(ZkasErr("Namespace too long".to_string()))
74        }
75
76        let constants_offset = match find_subslice(bytes, b".constant") {
77            Some(v) => v,
78            None => return Err(ZkasErr("Could not find .constant section".to_string())),
79        };
80
81        let literals_offset = match find_subslice(bytes, b".literal") {
82            Some(v) => v,
83            None => return Err(ZkasErr("Could not find .literal section".to_string())),
84        };
85
86        let witness_offset = match find_subslice(bytes, b".witness") {
87            Some(v) => v,
88            None => return Err(ZkasErr("Could not find .witness section".to_string())),
89        };
90
91        let circuit_offset = match find_subslice(bytes, b".circuit") {
92            Some(v) => v,
93            None => return Err(ZkasErr("Could not find .circuit section".to_string())),
94        };
95
96        let debug_offset = match find_subslice(bytes, b".debug") {
97            Some(v) => v,
98            None => bytes.len(),
99        };
100
101        if constants_offset > literals_offset {
102            return Err(ZkasErr(".literal section appeared before .constant".to_string()))
103        }
104
105        if literals_offset > witness_offset {
106            return Err(ZkasErr(".witness section appeared before .literal".to_string()))
107        }
108
109        if witness_offset > circuit_offset {
110            return Err(ZkasErr(".circuit section appeared before .witness".to_string()))
111        }
112
113        if circuit_offset > debug_offset {
114            return Err(ZkasErr(".debug section appeared before .circuit or EOF".to_string()))
115        }
116
117        let constants_section = &bytes[constants_offset + b".constant".len()..literals_offset];
118        let literals_section = &bytes[literals_offset + b".literal".len()..witness_offset];
119        let witness_section = &bytes[witness_offset + b".witness".len()..circuit_offset];
120        let circuit_section = &bytes[circuit_offset + b".circuit".len()..debug_offset];
121
122        let constants = ZkBinary::parse_constants(constants_section)?;
123        let literals = ZkBinary::parse_literals(literals_section)?;
124        let witnesses = ZkBinary::parse_witness(witness_section)?;
125        let opcodes = ZkBinary::parse_circuit(circuit_section)?;
126
127        // TODO: Debug info
128
129        Ok(Self { namespace, k, constants, literals, witnesses, opcodes })
130    }
131
132    fn parse_constants(bytes: &[u8]) -> Result<Vec<(VarType, String)>> {
133        let mut constants = vec![];
134
135        let mut iter_offset = 0;
136        while iter_offset < bytes.len() {
137            let c_type = match VarType::from_repr(bytes[iter_offset]) {
138                Some(v) => v,
139                None => {
140                    return Err(ZkasErr(format!(
141                        "Could not decode constant VarType from {}",
142                        bytes[iter_offset],
143                    )))
144                }
145            };
146            iter_offset += 1;
147            let (name, offset) = deserialize_partial::<String>(&bytes[iter_offset..])?;
148            iter_offset += offset;
149
150            constants.push((c_type, name));
151        }
152
153        Ok(constants)
154    }
155
156    fn parse_literals(bytes: &[u8]) -> Result<Vec<(LitType, String)>> {
157        let mut literals = vec![];
158
159        let mut iter_offset = 0;
160        while iter_offset < bytes.len() {
161            let l_type = match LitType::from_repr(bytes[iter_offset]) {
162                Some(v) => v,
163                None => {
164                    return Err(ZkasErr(format!(
165                        "Could not decode literal LitType from {}",
166                        bytes[iter_offset],
167                    )))
168                }
169            };
170            iter_offset += 1;
171            let (name, offset) = deserialize_partial::<String>(&bytes[iter_offset..])?;
172            iter_offset += offset;
173
174            literals.push((l_type, name));
175        }
176
177        Ok(literals)
178    }
179
180    fn parse_witness(bytes: &[u8]) -> Result<Vec<VarType>> {
181        let mut witnesses = vec![];
182
183        let mut iter_offset = 0;
184        while iter_offset < bytes.len() {
185            let w_type = match VarType::from_repr(bytes[iter_offset]) {
186                Some(v) => v,
187                None => {
188                    return Err(ZkasErr(format!(
189                        "Could not decode witness VarType from {}",
190                        bytes[iter_offset],
191                    )))
192                }
193            };
194
195            iter_offset += 1;
196
197            witnesses.push(w_type);
198        }
199
200        Ok(witnesses)
201    }
202
203    #[allow(clippy::type_complexity)]
204    fn parse_circuit(bytes: &[u8]) -> Result<Vec<(Opcode, Vec<(HeapType, usize)>)>> {
205        let mut opcodes = vec![];
206
207        let mut iter_offset = 0;
208        while iter_offset < bytes.len() {
209            let opcode = match Opcode::from_repr(bytes[iter_offset]) {
210                Some(v) => v,
211                None => {
212                    return Err(ZkasErr(format!(
213                        "Could not decode Opcode from {}",
214                        bytes[iter_offset]
215                    )))
216                }
217            };
218            iter_offset += 1;
219
220            // TODO: Check that the types and arg number are correct
221
222            let (arg_num, offset) = deserialize_partial::<VarInt>(&bytes[iter_offset..])?;
223            iter_offset += offset;
224
225            let mut args = vec![];
226            for _ in 0..arg_num.0 {
227                // Check bounds each time bytes[iter_offset] is accessed to prevent panics.
228                if iter_offset >= bytes.len() {
229                    return Err(ZkasErr(format!(
230                        "Bad offset for circuit: offset {} is >= circuit length {}",
231                        iter_offset,
232                        bytes.len()
233                    )))
234                }
235                let heap_type = bytes[iter_offset];
236                iter_offset += 1;
237
238                if iter_offset >= bytes.len() {
239                    return Err(ZkasErr(format!(
240                        "Bad offset for circuit: offset {} is >= circuit length {}",
241                        iter_offset,
242                        bytes.len()
243                    )))
244                }
245                let (heap_index, offset) = deserialize_partial::<VarInt>(&bytes[iter_offset..])?;
246                iter_offset += offset;
247                let heap_type = match HeapType::from_repr(heap_type) {
248                    Some(v) => v,
249                    None => {
250                        return Err(ZkasErr(format!("Could not decode HeapType from {}", heap_type)))
251                    }
252                };
253                args.push((heap_type, heap_index.0 as usize));
254            }
255
256            opcodes.push((opcode, args));
257        }
258
259        Ok(opcodes)
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use crate::zkas::ZkBinary;
266
267    #[test]
268    fn panic_regression_001() {
269        // Out-of-memory panic from string deserialization.
270        // Read `doc/src/zkas/bincode.md` to understand the input.
271        let data = vec![11u8, 1, 177, 53, 1, 0, 0, 0, 0, 255, 0, 204, 200, 72, 72, 72, 72, 1];
272        let _dec = ZkBinary::decode(&data);
273    }
274
275    #[test]
276    fn panic_regression_002() {
277        // Index out of bounds panic in parse_circuit().
278        // Read `doc/src/zkas/bincode.md` to understand the input.
279        let data = vec![
280            11u8, 1, 177, 53, 2, 13, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 83, 105,
281            109, 112, 108, 101, 46, 99, 111, 110, 115, 116, 97, 110, 116, 3, 18, 86, 65, 76, 85,
282            69, 95, 67, 79, 77, 77, 73, 84, 95, 86, 65, 76, 85, 69, 2, 19, 86, 65, 76, 85, 69, 95,
283            67, 79, 77, 77, 73, 84, 95, 82, 65, 77, 68, 79, 77, 46, 108, 105, 116, 101, 114, 97,
284            108, 46, 119, 105, 116, 110, 101, 115, 115, 16, 18, 46, 99, 105, 114, 99, 117, 105,
285            116, 4, 2, 0, 2, 0, 0, 2, 2, 0, 3, 0, 1, 8, 2, 0, 4, 0, 5, 8, 1, 0, 6, 9, 1, 0, 6, 240,
286            1, 0, 7, 240, 41, 0, 0, 0, 1, 0, 8,
287        ];
288        let _dec = ZkBinary::decode(&data);
289    }
290}