@@ -6,6 +6,8 @@ use crate::delta::Delta;
66use crate :: digest:: Digest ;
77use 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
7375impl 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
112116impl 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