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}