1use alloc::vec::Vec;
7use serde::{Deserialize, Serialize};
8
9use crate::{
10 PROTOCOL_VERSION,
11 error::{Error, Result},
12 params::Params,
13};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16#[repr(u8)]
17pub enum MsgType {
18 KgMsg1 = 1,
19 KgMsg2 = 2,
20 KgMsg3 = 3,
21 SignMsg1 = 4,
22 SignMsg2 = 5,
23 SignMsg3 = 6,
24 SetupData = 7,
25}
26
27impl MsgType {
28 fn from_u8(v: u8) -> Result<Self> {
29 Ok(match v {
30 1 => Self::KgMsg1,
31 2 => Self::KgMsg2,
32 3 => Self::KgMsg3,
33 4 => Self::SignMsg1,
34 5 => Self::SignMsg2,
35 6 => Self::SignMsg3,
36 7 => Self::SetupData,
37 _ => return Err(Error::InvalidEncoding),
38 })
39 }
40
41 pub fn max_len(self, params: &Params) -> usize {
43 let n = params.n_rsa as usize / 8;
44 let c = params.n_client as usize / 8;
45 let ell = params.ell as usize;
46 let overhead = 4 + 64;
47 overhead
48 + match self {
49 Self::KgMsg1 => 2 * 32 + 8,
50 Self::KgMsg2 => 33 + 4 * c + 4 * (c * 2) + 64 * ell + 4096,
51 Self::KgMsg3 => 33 + 2 * n + 2 * c + 4 * n + 4096,
52 Self::SignMsg1 => 2 * 33 + 8,
53 Self::SignMsg2 => 2 * 33 + 2 * n + 8 * n + 4096,
54 Self::SignMsg3 => 2 * 32 + 1 + 8,
55 Self::SetupData => (40 * n + 1_000 * ell) + (1 << 12),
56 }
57 }
58}
59
60pub fn wrap(params: &Params, ty: MsgType, body: &[u8]) -> Vec<u8> {
62 let mut out = Vec::with_capacity(4 + body.len());
63 out.extend_from_slice(&PROTOCOL_VERSION.to_be_bytes());
64 out.push(params.id as u8);
65 out.push(ty as u8);
66 out.extend_from_slice(body);
67 out
68}
69
70pub fn unwrap<'a>(params: &Params, expected: MsgType, bytes: &'a [u8]) -> Result<&'a [u8]> {
72 if bytes.len() < 4 || bytes.len() > expected.max_len(params) {
73 return Err(Error::InvalidEncoding);
74 }
75 let version = u16::from_be_bytes([bytes[0], bytes[1]]);
76 if version != PROTOCOL_VERSION || bytes[2] != params.id as u8 {
77 return Err(Error::InvalidEncoding);
78 }
79 if MsgType::from_u8(bytes[3])? != expected {
80 return Err(Error::InvalidEncoding);
81 }
82 Ok(&bytes[4..])
83}
84
85fn to_body<T: Serialize>(v: &T) -> Result<Vec<u8>> {
86 postcard::to_allocvec(v).map_err(|_| Error::InvalidEncoding)
87}
88
89fn from_body<T: for<'de> Deserialize<'de>>(body: &[u8]) -> Result<T> {
90 let (v, rest) = postcard::take_from_bytes(body).map_err(|_| Error::InvalidEncoding)?;
92 if !rest.is_empty() {
93 return Err(Error::InvalidEncoding);
94 }
95 Ok(v)
96}
97
98#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
100pub struct KgMsg1 {
101 pub kg_nonce: Vec<u8>,
102 pub b: Vec<u8>,
103}
104
105#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
107pub struct KgMsg2 {
108 pub x1: Vec<u8>,
109 pub m_hat: Vec<u8>,
110 pub v: Vec<u8>,
111 pub u1: Vec<u8>,
112 pub u2: Vec<u8>,
113 pub psi_df_star: Vec<u8>,
114 pub psi_dlog: Vec<u8>,
115}
116
117#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
119pub struct KgMsg3 {
120 pub x2: Vec<u8>,
121 pub e: Vec<u8>,
122 pub psi_dlenc: Vec<u8>,
123}
124
125#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
127pub struct SignMsg1 {
128 pub r2: Vec<u8>,
129 pub y: Vec<u8>,
130}
131
132#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
134pub struct SignMsg2 {
135 pub r1: Vec<u8>,
136 pub r: Vec<u8>,
137 pub s_ct: Vec<u8>,
138 pub psi_sig: Vec<u8>,
139}
140
141#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
143pub struct SignMsg3 {
144 pub r: Vec<u8>,
145 pub s: Vec<u8>,
146 pub recid: u8,
147}
148
149macro_rules! codec {
150 ($t:ty, $ty:expr) => {
151 impl $t {
152 pub fn encode(&self, params: &Params) -> Result<Vec<u8>> {
153 Ok(wrap(params, $ty, &to_body(self)?))
154 }
155 pub fn decode(bytes: &[u8], params: &Params) -> Result<Self> {
156 from_body(unwrap(params, $ty, bytes)?)
157 }
158 }
159 };
160}
161codec!(KgMsg1, MsgType::KgMsg1);
162codec!(KgMsg2, MsgType::KgMsg2);
163codec!(KgMsg3, MsgType::KgMsg3);
164codec!(SignMsg1, MsgType::SignMsg1);
165codec!(SignMsg2, MsgType::SignMsg2);
166codec!(SignMsg3, MsgType::SignMsg3);