1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
use bytes::{BufMut, Bytes, BytesMut};

use crate::util::{from_u16_bytes, from_u32_bytes};

/*
WasmRS operations list header

 0                   1                   2
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|       |   |       |                           |
+-------+---+-------+---------------------------+
|"\0wrs"| v | # ops |    operations...          |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

v (u16): Version
# Ops (u32): The number of operations to parse.

Operations

 0                   1
 0 1 2 3 4 5 6 7 8 9 0 ...
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
| | |       |   |                   |   |                     |
+-+-+-------+---+-------------------+---+---------------------+
|A|B| Index | N |  Namespace[0..N]  | O |  OpLen[N+2..N+2+O]  |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
A (u8): Operation type
B (u8): Direction (import/export)
Index (u32): The operation index to use when calling
N (u16): The length of the Namespace buffer
Namespace (u8[]): The Namespace String as UTF-8 bytes
O (u16): The length of the Operation buffer
Operation (u8[]): The Operation String as UTF-8 bytes
*/

static WASMRS_MAGIC: [u8; 4] = [0x00, 0x77, 0x72, 0x73];

#[derive(Debug, Copy, Clone)]
/// The types of RSocket operations supported by wasmRS.
pub enum OperationType {
  /// A request-response operation.
  RequestResponse,
  /// A fire-and-forget operation.
  RequestFnF,
  /// A request -> stream operation
  RequestStream,
  /// A stream -> stream operation
  RequestChannel,
}

impl From<u8> for OperationType {
  fn from(v: u8) -> Self {
    match v {
      1 => Self::RequestResponse,
      2 => Self::RequestFnF,
      3 => Self::RequestStream,
      4 => Self::RequestChannel,
      _ => unreachable!("Bad Operation Type {}", v),
    }
  }
}

impl From<OperationType> for u8 {
  fn from(op: OperationType) -> Self {
    match op {
      OperationType::RequestResponse => 1,
      OperationType::RequestFnF => 2,
      OperationType::RequestStream => 3,
      OperationType::RequestChannel => 4,
    }
  }
}

#[derive(Debug, Copy, Clone)]
pub enum Error {
  Magic,
  Version,
  Utf8String,
}

impl std::error::Error for Error {}
impl std::fmt::Display for Error {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.write_str(match self {
      Error::Magic => "Bad magic bytes",
      Error::Version => "Bad version",
      Error::Utf8String => "Could not convert bytes to UTF-8 String",
    })
  }
}
impl From<std::string::FromUtf8Error> for Error {
  fn from(_: std::string::FromUtf8Error) -> Self {
    Self::Utf8String
  }
}

#[derive(Debug, Clone)]
/// An operation record.
pub struct Operation {
  index: u32,
  kind: OperationType,
  namespace: String,
  operation: String,
}

#[derive(Debug, Default, Clone)]
/// A list of imports/exports for a wasmRS module.
#[must_use]
pub struct OperationList {
  imports: Vec<Operation>,
  exports: Vec<Operation>,
}

impl OperationList {
  #[must_use]
  /// Get the index for the imported namespace/operation.
  pub fn get_import(&self, namespace: &str, operation: &str) -> Option<u32> {
    Self::get_op(&self.imports, namespace, operation)
  }

  #[must_use]
  /// Get the index for the exported namespace/operation.
  pub fn get_export(&self, namespace: &str, operation: &str) -> Option<u32> {
    Self::get_op(&self.exports, namespace, operation)
  }

  #[must_use]
  /// Get a list of the exports by name.
  pub fn get_exports(&self) -> Vec<String> {
    self.exports.iter().map(|op| op.operation.clone()).collect()
  }

  fn get_op(list: &[Operation], namespace: &str, operation: &str) -> Option<u32> {
    list
      .iter()
      .find(|op| op.namespace == namespace && op.operation == operation)
      .map(|op| op.index)
  }

  /// Add an exported operation.
  pub fn add_export(
    &mut self,
    index: u32,
    kind: OperationType,
    namespace: impl AsRef<str>,
    operation: impl AsRef<str>,
  ) {
    Self::add_op(&mut self.exports, index, kind, namespace, operation);
  }

  /// Add an imported operation.
  pub fn add_import(
    &mut self,
    index: u32,
    kind: OperationType,
    namespace: impl AsRef<str>,
    operation: impl AsRef<str>,
  ) {
    Self::add_op(&mut self.imports, index, kind, namespace, operation);
  }

  fn add_op(
    list: &mut Vec<Operation>,
    index: u32,
    kind: OperationType,
    namespace: impl AsRef<str>,
    operation: impl AsRef<str>,
  ) {
    list.push(Operation {
      index,
      kind,
      namespace: namespace.as_ref().to_owned(),
      operation: operation.as_ref().to_owned(),
    });
  }

  #[must_use]
  /// Encode the operation list into a byte buffer.
  pub fn encode(&self) -> Bytes {
    let mut buff = BytesMut::new();
    let num_ops: u32 = (self.imports.len() + self.exports.len()) as u32;
    let version = 1u16;
    buff.put(WASMRS_MAGIC.as_slice());
    buff.put(version.to_be_bytes().as_slice());
    buff.put(num_ops.to_be_bytes().as_slice());
    for op in &self.exports {
      buff.put(Self::encode_op(op, 1));
    }
    for op in &self.imports {
      buff.put(Self::encode_op(op, 2));
    }
    buff.freeze()
  }

  fn encode_op(op: &Operation, dir: u8) -> Bytes {
    let mut buff = BytesMut::new();

    let kind: u8 = op.kind.into();
    buff.put([kind].as_slice());
    buff.put([dir].as_slice());
    buff.put(op.index.to_be_bytes().as_slice());
    buff.put((op.namespace.len() as u16).to_be_bytes().as_slice());
    buff.put(op.namespace.as_bytes());
    buff.put((op.operation.len() as u16).to_be_bytes().as_slice());
    buff.put(op.operation.as_bytes());
    buff.put(0_u16.to_be_bytes().as_slice());
    buff.freeze()
  }

  /// Decode bytes into an Operation List.
  pub fn decode(mut buf: Bytes) -> Result<Self, Error> {
    let magic = buf.split_to(4);
    if magic != WASMRS_MAGIC.as_slice() {
      return Err(Error::Magic);
    }
    let version = from_u16_bytes(&buf.split_to(2));
    match version {
      1 => Self::decode_v1(buf),
      _ => Err(Error::Version),
    }
  }

  fn decode_v1(mut buf: Bytes) -> Result<Self, Error> {
    let num_ops = from_u32_bytes(&buf.split_to(4));
    let mut imports = Vec::new();
    let mut exports = Vec::new();
    for _ in 0..num_ops {
      let kind = buf.split_to(1)[0];
      let kind: OperationType = kind.into();
      let dir = buf.split_to(1)[0];
      let index = from_u32_bytes(&buf.split_to(4));
      let ns_len = from_u16_bytes(&buf.split_to(2));
      let namespace = String::from_utf8(buf.split_to(ns_len as _).to_vec())?;
      let op_len = from_u16_bytes(&buf.split_to(2));
      let operation = String::from_utf8(buf.split_to(op_len as _).to_vec())?;
      let _reserved_len = from_u16_bytes(&buf.split_to(2));
      let op = Operation {
        index,
        kind,
        namespace,
        operation,
      };
      if dir == 1 {
        exports.push(op);
      } else {
        imports.push(op);
      }
    }
    Ok(Self { imports, exports })
  }
}