Skip to content

Commit f783620

Browse files
authored
gAdded a magic number to the message serializing. (#141)
The point here is to make sure we ignore message coming from a different program, or coming from a version of chitchat that was not properly versioned. In that case, the enum discriminant could coincide with our version number.
1 parent d9049ab commit f783620

File tree

1 file changed

+32
-23
lines changed

1 file changed

+32
-23
lines changed

chitchat/src/message.rs

Lines changed: 32 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ use crate::delta::Delta;
66
use crate::digest::Digest;
77
use crate::serialize::{Deserializable, Serializable};
88

9+
const MAGIC_NUMBER: u16 = 45_139;
10+
911
/// Chitchat message.
1012
///
1113
/// Each variant represents a step of the gossip "handshake"
@@ -72,6 +74,7 @@ impl MessageType {
7274

7375
impl Serializable for ChitchatMessage {
7476
fn serialize(&self, buf: &mut Vec<u8>) {
77+
buf.extend(MAGIC_NUMBER.to_le_bytes());
7578
ProtocolVersion::V0.to_code().serialize(buf);
7679

7780
match self {
@@ -96,34 +99,39 @@ impl Serializable for ChitchatMessage {
9699
}
97100

98101
fn serialized_len(&self) -> usize {
99-
1 + match self {
100-
ChitchatMessage::Syn { cluster_id, digest } => {
101-
1 + cluster_id.serialized_len() + digest.serialized_len()
102+
2 + 1
103+
+ match self {
104+
ChitchatMessage::Syn { cluster_id, digest } => {
105+
1 + cluster_id.serialized_len() + digest.serialized_len()
106+
}
107+
ChitchatMessage::SynAck { digest, delta } => {
108+
1 + digest.serialized_len() + delta.serialized_len()
109+
}
110+
ChitchatMessage::Ack { delta } => 1 + delta.serialized_len(),
111+
ChitchatMessage::BadCluster => 1,
102112
}
103-
ChitchatMessage::SynAck { digest, delta } => {
104-
1 + digest.serialized_len() + delta.serialized_len()
105-
}
106-
ChitchatMessage::Ack { delta } => 1 + delta.serialized_len(),
107-
ChitchatMessage::BadCluster => 1,
108-
}
109113
}
110114
}
111115

112116
impl Deserializable for ChitchatMessage {
113117
fn deserialize(buf: &mut &[u8]) -> anyhow::Result<Self> {
114-
let protocol_version = buf
115-
.first()
116-
.copied()
117-
.and_then(ProtocolVersion::from_code)
118-
.context("invalid protocol version")?;
118+
if buf.len() < 3 {
119+
bail!("buffer too small to store the magic number and the protocol version");
120+
}
121+
let magic_number = u16::from_le_bytes(buf[0..2].try_into().unwrap());
122+
if magic_number != MAGIC_NUMBER {
123+
bail!("invalid chitchat magic number");
124+
}
125+
let protocol_version =
126+
ProtocolVersion::from_code(buf[2]).context("invalid protocol version")?;
119127

120128
if protocol_version != ProtocolVersion::V0 {
121129
bail!(
122130
"unsupported protocol version `{}`",
123131
protocol_version.to_code()
124132
)
125133
}
126-
buf.consume(1);
134+
buf.consume(3);
127135

128136
let message_type = buf
129137
.first()
@@ -164,7 +172,7 @@ mod tests {
164172
cluster_id: "cluster-a".to_string(),
165173
digest: Digest::default(),
166174
};
167-
test_serdeser_aux(&syn, 15);
175+
test_serdeser_aux(&syn, 17);
168176
}
169177
{
170178
let mut digest = Digest::default();
@@ -175,7 +183,7 @@ mod tests {
175183
cluster_id: "cluster-a".to_string(),
176184
digest,
177185
};
178-
test_serdeser_aux(&syn, 66);
186+
test_serdeser_aux(&syn, 68);
179187
}
180188
}
181189

@@ -186,8 +194,9 @@ mod tests {
186194
digest: Digest::default(),
187195
delta: Delta::default(),
188196
};
189-
// 1 (protocol version) + 1 (message tag) + 2 (digest len) + 1 (delta end op)
190-
test_serdeser_aux(&syn_ack, 5);
197+
// 2 (magic number) + 1 (protocol version) + 1 (message tag) + 2 (digest len) + 1 (delta
198+
// end op)
199+
test_serdeser_aux(&syn_ack, 7);
191200
}
192201
{
193202
// 2 bytes.
@@ -212,7 +221,7 @@ mod tests {
212221
let syn_ack = ChitchatMessage::SynAck { digest, delta };
213222
// 1 byte (protocol version) + 1 byte (message tag) + 53 bytes (digest) + 60 bytes
214223
// (delta).
215-
test_serdeser_aux(&syn_ack, 1 + 1 + 53 + 60);
224+
test_serdeser_aux(&syn_ack, 2 + 1 + 1 + 53 + 60);
216225
}
217226
}
218227

@@ -221,7 +230,7 @@ mod tests {
221230
{
222231
let delta = Delta::default();
223232
let ack = ChitchatMessage::Ack { delta };
224-
test_serdeser_aux(&ack, 3);
233+
test_serdeser_aux(&ack, 5);
225234
}
226235
{
227236
// 4 bytes.
@@ -233,12 +242,12 @@ mod tests {
233242
delta.add_kv(&node, "key", "value", 0, true);
234243
delta.set_serialized_len(60);
235244
let ack = ChitchatMessage::Ack { delta };
236-
test_serdeser_aux(&ack, 1 + 1 + 60);
245+
test_serdeser_aux(&ack, 2 + 1 + 1 + 60);
237246
}
238247
}
239248

240249
#[test]
241250
fn test_bad_cluster() {
242-
test_serdeser_aux(&ChitchatMessage::BadCluster, 2);
251+
test_serdeser_aux(&ChitchatMessage::BadCluster, 4);
243252
}
244253
}

0 commit comments

Comments
 (0)