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  
 19  use darkfi_serial::{deserialize_partial, VarInt};
 20  
 21  use super::{
 22      compiler::MAGIC_BYTES,
 23      constants::{MAX_K, MAX_NS_LEN, MIN_BIN_SIZE},
 24      types::HeapType,
 25      LitType, Opcode, VarType,
 26  };
 27  use 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)]
 32  pub 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
 42  fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
 43      haystack.windows(needle.len()).position(|window| window == needle)
 44  }
 45  
 46  impl 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)]
264  mod 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  }