1use bytes::{BufMut, Bytes, BytesMut};
2
3use crate::util::{from_u16_bytes, from_u32_bytes};
4
5static WASMRS_MAGIC: [u8; 4] = [0x00, 0x77, 0x72, 0x73];
38
39#[derive(Debug, Copy, Clone)]
40pub enum OperationType {
42  RequestResponse,
44  RequestFnF,
46  RequestStream,
48  RequestChannel,
50}
51
52impl From<u8> for OperationType {
53  fn from(v: u8) -> Self {
54    match v {
55      1 => Self::RequestResponse,
56      2 => Self::RequestFnF,
57      3 => Self::RequestStream,
58      4 => Self::RequestChannel,
59      _ => unreachable!("Bad Operation Type {}", v),
60    }
61  }
62}
63
64impl From<OperationType> for u8 {
65  fn from(op: OperationType) -> Self {
66    match op {
67      OperationType::RequestResponse => 1,
68      OperationType::RequestFnF => 2,
69      OperationType::RequestStream => 3,
70      OperationType::RequestChannel => 4,
71    }
72  }
73}
74
75#[derive(Debug, Copy, Clone)]
76pub enum Error {
77  Magic,
78  Version,
79  Utf8String,
80}
81
82impl std::error::Error for Error {}
83impl std::fmt::Display for Error {
84  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85    f.write_str(match self {
86      Error::Magic => "Bad magic bytes",
87      Error::Version => "Bad version",
88      Error::Utf8String => "Could not convert bytes to UTF-8 String",
89    })
90  }
91}
92impl From<std::string::FromUtf8Error> for Error {
93  fn from(_: std::string::FromUtf8Error) -> Self {
94    Self::Utf8String
95  }
96}
97
98#[derive(Debug, Clone)]
99pub struct Operation {
101  index: u32,
102  kind: OperationType,
103  namespace: String,
104  operation: String,
105}
106
107#[derive(Debug, Default, Clone)]
108#[must_use]
110pub struct OperationList {
111  imports: Vec<Operation>,
112  exports: Vec<Operation>,
113}
114
115impl OperationList {
116  #[must_use]
117  pub fn get_import(&self, namespace: &str, operation: &str) -> Option<u32> {
119    Self::get_op(&self.imports, namespace, operation)
120  }
121
122  #[must_use]
123  pub fn get_export(&self, namespace: &str, operation: &str) -> Option<u32> {
125    Self::get_op(&self.exports, namespace, operation)
126  }
127
128  #[must_use]
129  pub fn get_exports(&self) -> Vec<String> {
131    self.exports.iter().map(|op| op.operation.clone()).collect()
132  }
133
134  fn get_op(list: &[Operation], namespace: &str, operation: &str) -> Option<u32> {
135    list
136      .iter()
137      .find(|op| op.namespace == namespace && op.operation == operation)
138      .map(|op| op.index)
139  }
140
141  pub fn add_export(
143    &mut self,
144    index: u32,
145    kind: OperationType,
146    namespace: impl AsRef<str>,
147    operation: impl AsRef<str>,
148  ) {
149    Self::add_op(&mut self.exports, index, kind, namespace, operation);
150  }
151
152  pub fn add_import(
154    &mut self,
155    index: u32,
156    kind: OperationType,
157    namespace: impl AsRef<str>,
158    operation: impl AsRef<str>,
159  ) {
160    Self::add_op(&mut self.imports, index, kind, namespace, operation);
161  }
162
163  fn add_op(
164    list: &mut Vec<Operation>,
165    index: u32,
166    kind: OperationType,
167    namespace: impl AsRef<str>,
168    operation: impl AsRef<str>,
169  ) {
170    list.push(Operation {
171      index,
172      kind,
173      namespace: namespace.as_ref().to_owned(),
174      operation: operation.as_ref().to_owned(),
175    });
176  }
177
178  #[must_use]
179  pub fn encode(&self) -> Bytes {
181    let mut buff = BytesMut::new();
182    let num_ops: u32 = (self.imports.len() + self.exports.len()) as u32;
183    let version = 1u16;
184    buff.put(WASMRS_MAGIC.as_slice());
185    buff.put(version.to_be_bytes().as_slice());
186    buff.put(num_ops.to_be_bytes().as_slice());
187    for op in &self.exports {
188      buff.put(Self::encode_op(op, 1));
189    }
190    for op in &self.imports {
191      buff.put(Self::encode_op(op, 2));
192    }
193    buff.freeze()
194  }
195
196  fn encode_op(op: &Operation, dir: u8) -> Bytes {
197    let mut buff = BytesMut::new();
198
199    let kind: u8 = op.kind.into();
200    buff.put([kind].as_slice());
201    buff.put([dir].as_slice());
202    buff.put(op.index.to_be_bytes().as_slice());
203    buff.put((op.namespace.len() as u16).to_be_bytes().as_slice());
204    buff.put(op.namespace.as_bytes());
205    buff.put((op.operation.len() as u16).to_be_bytes().as_slice());
206    buff.put(op.operation.as_bytes());
207    buff.put(0_u16.to_be_bytes().as_slice());
208    buff.freeze()
209  }
210
211  pub fn decode(mut buf: Bytes) -> Result<Self, Error> {
213    let magic = buf.split_to(4);
214    if magic != WASMRS_MAGIC.as_slice() {
215      return Err(Error::Magic);
216    }
217    let version = from_u16_bytes(&buf.split_to(2));
218    match version {
219      1 => Self::decode_v1(buf),
220      _ => Err(Error::Version),
221    }
222  }
223
224  fn decode_v1(mut buf: Bytes) -> Result<Self, Error> {
225    let num_ops = from_u32_bytes(&buf.split_to(4));
226    let mut imports = Vec::new();
227    let mut exports = Vec::new();
228    for _ in 0..num_ops {
229      let kind = buf.split_to(1)[0];
230      let kind: OperationType = kind.into();
231      let dir = buf.split_to(1)[0];
232      let index = from_u32_bytes(&buf.split_to(4));
233      let ns_len = from_u16_bytes(&buf.split_to(2));
234      let namespace = String::from_utf8(buf.split_to(ns_len as _).to_vec())?;
235      let op_len = from_u16_bytes(&buf.split_to(2));
236      let operation = String::from_utf8(buf.split_to(op_len as _).to_vec())?;
237      let _reserved_len = from_u16_bytes(&buf.split_to(2));
238      let op = Operation {
239        index,
240        kind,
241        namespace,
242        operation,
243      };
244      if dir == 1 {
245        exports.push(op);
246      } else {
247        imports.push(op);
248      }
249    }
250    Ok(Self { imports, exports })
251  }
252}