1use crate::{
6 connection::{Connection, ConnectionDetails, SendQueueStatus, WriteStatus},
7 handshake::Handshake,
8 node_table::{NodeEndpoint, NodeEntry, NodeId},
9 parse_msg_id_leb128_2_bytes_at_most,
10 service::{NetworkServiceInner, ProtocolVersion},
11 DisconnectReason, Error, ProtocolId, ProtocolInfo, SessionMetadata,
12 UpdateNodeOperation, PROTOCOL_ID_SIZE,
13};
14use bytes::Bytes;
15use cfx_util_macros::bail;
16use diem_crypto::{bls::BLS_PUBLIC_KEY_LENGTH, ValidCryptoMaterial};
17use diem_types::validator_config::{ConsensusPublicKey, ConsensusVRFPublicKey};
18use io::{IoContext, StreamToken};
19use log::{debug, trace};
20use mio::{net::TcpStream, Registry, Token};
21use priority_send_queue::SendQueuePriority;
22use rlp::{Rlp, RlpStream};
23use serde::Deserialize;
24use serde_derive::Serialize;
25use std::{
26 collections::HashSet,
27 convert::TryFrom,
28 fmt,
29 net::SocketAddr,
30 str,
31 time::{Duration, Instant},
32};
33
34pub struct Session {
45 pub metadata: SessionMetadata,
47 address: SocketAddr,
49 state: State,
51 sent_hello: Instant,
53 had_hello: Option<Instant>,
55 expired: Option<Instant>,
57
58 last_read: Instant,
60 last_write: (Instant, WriteStatus),
61 pos_public_key: Option<(ConsensusPublicKey, ConsensusVRFPublicKey)>,
62}
63
64enum State {
66 Handshake(MovableWrapper<Handshake>),
70 Session(Connection),
72}
73
74pub enum SessionData {
76 None,
78 Ready {
80 pos_public_key: Option<(ConsensusPublicKey, ConsensusVRFPublicKey)>,
81 },
82 Message { data: Vec<u8>, protocol: ProtocolId },
85 Continue,
87}
88
89pub struct SessionDataWithDisconnectInfo {
90 pub session_data: SessionData,
91 pub token_to_disconnect: Option<(StreamToken, String)>,
92}
93
94const PACKET_HELLO: u8 = 0x80;
96const MAX_PEER_PROTOCOLS_IN_HELLO: usize = 64;
99const PACKET_DISCONNECT: u8 = 0x01;
101pub const PACKET_USER: u8 = 0x10;
103pub const PACKET_HEADER_VERSION: u8 = 0;
106const HEADER_VERSION_WITH_EXTENSION: u8 = 0;
108
109impl Session {
110 #[allow(clippy::too_many_arguments)]
113 pub fn new<Message: Send + Sync + Clone + 'static>(
114 io: &IoContext<Message>, socket: TcpStream, address: SocketAddr,
115 id: Option<&NodeId>, peer_header_version: u8, token: StreamToken,
116 host: &NetworkServiceInner,
117 pos_public_key: Option<(ConsensusPublicKey, ConsensusVRFPublicKey)>,
118 ) -> Result<Session, Error> {
119 let originated = id.is_some();
120
121 let mut handshake = Handshake::new(token, id, socket);
122 handshake.start(io, &host.metadata)?;
123
124 Ok(Session {
125 metadata: SessionMetadata {
126 id: id.cloned(),
127 peer_protocols: Vec::new(),
128 originated,
129 peer_header_version,
130 },
131 address,
132 state: State::Handshake(MovableWrapper::new(handshake)),
133 sent_hello: Instant::now(),
134 had_hello: None,
135 expired: None,
136 last_read: Instant::now(),
137 last_write: (Instant::now(), WriteStatus::Complete),
138 pos_public_key,
139 })
140 }
141
142 pub fn have_capability(&self, protocol: ProtocolId) -> bool {
143 self.metadata
144 .peer_protocols
145 .iter()
146 .any(|c| c.protocol == protocol)
147 }
148
149 pub fn id(&self) -> Option<&NodeId> { self.metadata.id.as_ref() }
151
152 pub fn originated(&self) -> bool { self.metadata.originated }
153
154 pub fn is_ready(&self) -> bool { self.had_hello.is_some() }
155
156 pub fn expired(&self) -> bool { self.expired.is_some() }
157
158 pub fn set_expired(&mut self) { self.expired = Some(Instant::now()); }
159
160 pub fn done(&self) -> bool {
161 self.expired() && !self.connection().is_sending()
162 }
163
164 fn connection(&self) -> &Connection {
165 match self.state {
166 State::Handshake(ref h) => &h.get().connection,
167 State::Session(ref c) => c,
168 }
169 }
170
171 fn connection_mut(&mut self) -> &mut Connection {
172 match self.state {
173 State::Handshake(ref mut h) => &mut h.get_mut().connection,
174 State::Session(ref mut c) => c,
175 }
176 }
177
178 pub fn token(&self) -> StreamToken { self.connection().token() }
179
180 pub fn address(&self) -> SocketAddr { self.address }
181
182 pub fn register_socket(
185 &mut self, reg: Token, poll_registry: &Registry,
186 ) -> Result<(), Error> {
187 if !self.expired() {
188 self.connection_mut().register_socket(reg, poll_registry)?;
189 }
190
191 Ok(())
192 }
193
194 pub fn update_socket(
196 &mut self, reg: Token, poll_registry: &Registry,
197 ) -> Result<(), Error> {
198 self.connection_mut().update_socket(reg, poll_registry)?;
199 Ok(())
200 }
201
202 pub fn deregister_socket(
204 &mut self, poll_registry: &Registry,
205 ) -> Result<(), Error> {
206 self.connection_mut().deregister_socket(poll_registry)?;
207 Ok(())
208 }
209
210 fn complete_handshake<Message>(
215 &mut self, io: &IoContext<Message>, host: &NetworkServiceInner,
216 ) -> Result<(), Error>
217 where Message: Send + Sync + Clone {
218 let wrapper = match self.state {
219 State::Handshake(ref mut h) => h,
220 State::Session(_) => panic!("Unexpected session state"),
221 };
222
223 if self.metadata.id.is_none() {
225 let id = wrapper.get().id;
226
227 if host.node_db.write().evaluate_blacklisted(&id) {
229 return Err(self.send_disconnect(DisconnectReason::Blacklisted));
230 }
231
232 self.metadata.id = Some(id);
233 }
234
235 self.state = State::Session(wrapper.take().connection);
237 self.write_hello(io, host)?;
238
239 Ok(())
240 }
241
242 pub fn readable<Message: Send + Sync + Clone>(
244 &mut self, io: &IoContext<Message>, host: &NetworkServiceInner,
245 ) -> Result<SessionDataWithDisconnectInfo, Error> {
246 self.last_read = Instant::now();
248
249 if self.expired() {
250 debug!("cannot read data due to expired, session = {:?}", self);
251 return Ok(SessionDataWithDisconnectInfo {
252 session_data: SessionData::None,
253 token_to_disconnect: None,
254 });
255 }
256
257 match self.state {
258 State::Handshake(ref mut h) => {
259 let h = h.get_mut();
260
261 if !h.readable(io, &host.metadata)? {
262 return Ok(SessionDataWithDisconnectInfo {
263 session_data: SessionData::None,
264 token_to_disconnect: None,
265 });
266 }
267
268 if h.done() {
269 self.complete_handshake(io, host)?;
270 io.update_registration(self.token()).unwrap_or_else(|e| {
271 debug!("Token registration error: {:?}", e)
272 });
273 }
274
275 Ok(SessionDataWithDisconnectInfo {
276 session_data: SessionData::Continue,
277 token_to_disconnect: None,
278 })
279 }
280 State::Session(ref mut c) => match c.readable()? {
281 Some(data) => Ok(self.read_packet(data, host)?),
282 None => Ok(SessionDataWithDisconnectInfo {
283 session_data: SessionData::None,
284 token_to_disconnect: None,
285 }),
286 },
287 }
288 }
289
290 fn read_packet(
292 &mut self, data: Bytes, host: &NetworkServiceInner,
293 ) -> Result<SessionDataWithDisconnectInfo, Error> {
294 let packet = SessionPacket::parse(data)?;
295
296 if packet.id != PACKET_HELLO
299 && packet.id != PACKET_DISCONNECT
300 && self.had_hello.is_none()
301 {
302 return Err(Error::BadProtocol);
303 }
304
305 match packet.id {
306 PACKET_HELLO => {
307 debug!("Read HELLO in session {:?}", self);
308 if self.is_ready() {
312 debug!("Duplicate HELLO on ready session {:?}", self);
313 return Err(Error::BadProtocol);
314 }
315 self.metadata.peer_header_version = packet.header_version;
316
317 let rlp = Rlp::new(&packet.data);
320 let (peer_protocols, node_entry, pos_public_key) =
321 self.validate_hello(&rlp, host)?;
322
323 let token_to_disconnect = self.update_ingress_node_id(host)?;
324
325 if token_to_disconnect.as_ref().map(|(t, _)| *t)
328 == Some(self.token())
329 {
330 return Ok(SessionDataWithDisconnectInfo {
331 session_data: SessionData::None,
332 token_to_disconnect,
333 });
334 }
335
336 self.metadata.peer_protocols = peer_protocols;
337 host.node_db
338 .write()
339 .insert_with_token(node_entry, self.token());
340 self.had_hello = Some(Instant::now());
341 Ok(SessionDataWithDisconnectInfo {
342 session_data: SessionData::Ready { pos_public_key },
343 token_to_disconnect,
344 })
345 }
346 PACKET_DISCONNECT => {
347 let rlp = Rlp::new(&packet.data);
348 let reason: DisconnectReason = rlp.as_val()?;
349 debug!(
350 "read packet DISCONNECT, reason = {}, session = {:?}",
351 reason, self
352 );
353 Err(Error::Disconnect(reason))
354 }
355 PACKET_USER => Ok(SessionDataWithDisconnectInfo {
356 session_data: SessionData::Message {
357 data: packet.data.to_vec(),
358 protocol: packet
359 .protocol
360 .expect("protocol should available for USER packet"),
361 },
362 token_to_disconnect: None,
363 }),
364 _ => {
365 debug!(
366 "read packet UNKNOWN, packet_id = {:?}, session = {:?}",
367 packet.id, self
368 );
369 Err(Error::BadProtocol)
370 }
371 }
372 }
373
374 fn update_ingress_node_id(
378 &mut self, host: &NetworkServiceInner,
379 ) -> Result<Option<(usize, String)>, Error> {
380 if self.metadata.originated {
382 return Ok(None);
383 }
384
385 let token = self.token();
386 let node_id = self
387 .metadata
388 .id
389 .expect("should have node id after handshake");
390
391 let result = host
392 .sessions
393 .update_ingress_node_id(token, &node_id)
394 .map_err(|reason| {
395 debug!(
396 "failed to update node id of ingress session, reason = {:?}, session = {:?}",
397 reason, self
398 );
399 self.send_disconnect(DisconnectReason::UpdateNodeIdFailed)
400 })?;
401
402 match result {
403 crate::session_manager::UpdateIngressResult::Inserted => Ok(None),
404 crate::session_manager::UpdateIngressResult::Replaced(old) => {
405 Ok(Some((
406 old,
407 String::from("Remove old session from the same node"),
408 )))
409 }
410 crate::session_manager::UpdateIngressResult::DropNew => {
411 debug!(
414 "lost simultaneous-dial tie-break, dropping new ingress session = {:?}",
415 self
416 );
417 let _ = self.send_disconnect(DisconnectReason::Custom(
418 "simultaneous dial: drop new connection".into(),
419 ));
420 Ok(Some((
421 token,
422 String::from("simultaneous dial: drop new connection"),
423 )))
424 }
425 }
426 }
427
428 fn validate_hello(
432 &mut self, rlp: &Rlp, host: &NetworkServiceInner,
433 ) -> Result<
434 (
435 Vec<ProtocolInfo>,
436 NodeEntry,
437 Option<(ConsensusPublicKey, ConsensusVRFPublicKey)>,
438 ),
439 Error,
440 > {
441 let remote_network_id: u64 = rlp.val_at(0)?;
442 if remote_network_id != host.metadata.network_id {
443 debug!(
444 "failed to read hello, network id mismatch, self = {}, remote = {}",
445 host.metadata.network_id, remote_network_id);
446 return Err(self.send_disconnect(DisconnectReason::Custom(
447 "network id mismatch".into(),
448 )));
449 }
450
451 let peer_caps_count = rlp
455 .at(1)?
456 .iter()
457 .take(MAX_PEER_PROTOCOLS_IN_HELLO + 1)
458 .count();
459 if peer_caps_count > MAX_PEER_PROTOCOLS_IN_HELLO {
460 debug!(
461 "Too many protocols in hello: {}, remote = {}",
462 peer_caps_count, remote_network_id
463 );
464 return Err(self.send_disconnect(DisconnectReason::Custom(
465 "Invalid protocol list: too many protocols.".into(),
466 )));
467 }
468
469 let mut peer_caps: Vec<ProtocolInfo> = rlp.list_at(1)?;
470 let mut seen_protocols = HashSet::with_capacity(peer_caps.len());
471 for cap in &peer_caps {
472 if !seen_protocols.insert(cap.protocol) {
473 debug!(
474 "Invalid protocol list from hello. Duplication: {:?}, \
475 remote = {}",
476 cap.protocol, remote_network_id
477 );
478 return Err(self.send_disconnect(DisconnectReason::Custom(
479 "Invalid protocol list: duplication.".into(),
480 )));
481 }
482 }
483
484 peer_caps.retain(|c| {
485 host.metadata
486 .minimum_peer_protocol_version
487 .read()
488 .iter()
489 .any(|hc| hc.protocol == c.protocol && hc.version <= c.version)
490 });
491
492 if peer_caps.is_empty() {
493 debug!("No common capabilities with remote peer, peer_node_id = {:?}, session = {:?}", self.metadata.id, self);
494 return Err(self.send_disconnect(DisconnectReason::UselessPeer));
495 }
496
497 let mut hello_from = NodeEndpoint::from_rlp(&rlp.at(2)?)?;
498 hello_from.address.set_ip(self.address.ip());
502
503 let ping_to = NodeEndpoint {
504 address: hello_from.address,
505 udp_port: hello_from.udp_port,
506 };
507
508 let entry = NodeEntry {
509 id: self
510 .metadata
511 .id
512 .expect("should have node ID after handshake"),
513 endpoint: ping_to,
514 };
515 if !entry.endpoint.is_valid() {
516 debug!("Got invalid endpoint {:?}, session = {:?}", entry, self);
517 return Err(
518 self.send_disconnect(DisconnectReason::WrongEndpointInfo)
519 );
520 } else if !(entry.endpoint.is_allowed(host.get_ip_filter())
521 && entry.id != *host.metadata.id())
522 {
523 debug!(
524 "Address not allowed, endpoint = {:?}, session = {:?}",
525 entry, self
526 );
527 return Err(self.send_disconnect(DisconnectReason::IpLimited));
528 }
529 debug!("Received valid endpoint {:?}, session = {:?}", entry, self);
530
531 let pos_public_key = match rlp.item_count()? {
532 3 => None,
533 4 => {
534 let pos_public_key_bytes: Vec<u8> = rlp.val_at(3)?;
536 trace!("pos_public_key_bytes: {:?}", pos_public_key_bytes);
537 if pos_public_key_bytes.len() < BLS_PUBLIC_KEY_LENGTH {
538 bail!("pos public key bytes is too short!");
539 }
540 let bls_pub_key = ConsensusPublicKey::try_from(
541 &pos_public_key_bytes[..BLS_PUBLIC_KEY_LENGTH],
542 )
543 .map_err(|e| Error::Decoder(format!("{:?}", e)))?;
544 let vrf_pub_key = ConsensusVRFPublicKey::try_from(
545 &pos_public_key_bytes[BLS_PUBLIC_KEY_LENGTH..],
546 )
547 .map_err(|e| Error::Decoder(format!("{:?}", e)))?;
548
549 Some((bls_pub_key, vrf_pub_key))
550 }
551 length => {
552 return Err(Error::Decoder(format!(
553 "Hello has incorrect rlp length: {:?}",
554 length
555 )))
556 }
557 };
558
559 Ok((peer_caps, entry, pos_public_key))
560 }
561
562 fn prepare_packet(
566 &self, protocol: Option<ProtocolId>, packet_id: u8, data: Vec<u8>,
567 ) -> Result<Vec<u8>, Error> {
568 if protocol.is_some() && self.had_hello.is_none() {
569 debug!(
570 "Sending to unconfirmed session {}, protocol: {:?}, packet: {}",
571 self.token(),
572 protocol
573 .as_ref()
574 .map(|p| str::from_utf8(&p[..]).unwrap_or("???")),
575 packet_id
576 );
577 bail!(Error::Expired);
578 }
579
580 if self.expired() {
581 return Err(Error::Expired);
582 }
583
584 Ok(SessionPacket::assemble(
585 packet_id,
586 self.metadata.peer_header_version,
587 protocol,
588 data,
589 ))
590 }
591
592 #[inline]
593 pub fn check_message_protocol_version(
594 &self, protocol: Option<ProtocolId>,
595 min_protocol_version: ProtocolVersion, mut msg: &[u8],
596 ) -> Result<(), Error> {
597 if let Some(protocol) = protocol {
600 for peer_protocol in &self.metadata.peer_protocols {
601 if protocol.eq(&peer_protocol.protocol) {
602 if min_protocol_version <= peer_protocol.version {
603 break;
604 } else {
605 bail!(Error::SendUnsupportedMessage {
606 protocol,
607 msg_id: parse_msg_id_leb128_2_bytes_at_most(
608 &mut msg
609 )
610 .map_err(|_| Error::Msg(
611 "msg_id parse failed when checking protocol version".into()
612 ))?,
613 peer_protocol_version: Some(peer_protocol.version),
614 min_supported_version: None,
615 });
616 }
617 }
618 }
619 }
620
621 Ok(())
622 }
623
624 pub fn send_packet<Message: Send + Sync + Clone>(
626 &mut self, io: &IoContext<Message>, protocol: Option<ProtocolId>,
627 min_proto_version: ProtocolVersion, packet_id: u8, data: Vec<u8>,
628 priority: SendQueuePriority,
629 ) -> Result<SendQueueStatus, Error> {
630 self.check_message_protocol_version(
631 protocol,
632 min_proto_version,
633 &data,
634 )?;
635 let packet = self.prepare_packet(protocol, packet_id, data)?;
636 self.connection_mut().send(io, packet, priority)
637 }
638
639 pub fn send_packet_immediately(
641 &mut self, protocol: Option<ProtocolId>,
642 min_proto_version: ProtocolVersion, packet_id: u8, data: Vec<u8>,
643 ) -> Result<usize, Error> {
644 self.check_message_protocol_version(
645 protocol,
646 min_proto_version,
647 &data,
648 )?;
649 let packet = self.prepare_packet(protocol, packet_id, data)?;
650 self.connection_mut().write_raw_data(packet)
651 }
652
653 pub fn send_disconnect(&mut self, reason: DisconnectReason) -> Error {
655 let packet = rlp::encode(&reason).to_vec();
656 let _ = self.send_packet_immediately(
657 None,
658 ProtocolVersion::default(),
659 PACKET_DISCONNECT,
660 packet,
661 );
662 Error::Disconnect(reason)
663 }
664
665 fn write_hello<Message: Send + Sync + Clone>(
667 &mut self, io: &IoContext<Message>, host: &NetworkServiceInner,
668 ) -> Result<(), Error> {
669 debug!("Sending Hello, session = {:?}", self);
670 let mut rlp = RlpStream::new_list(4);
671 rlp.append(&host.metadata.network_id);
672 rlp.append_list(&host.metadata.protocols.read());
673 host.metadata.public_endpoint.to_rlp_list(&mut rlp);
674 let mut key_bytes =
675 self.pos_public_key.as_ref().unwrap().0.to_bytes().to_vec();
676 key_bytes.append(
677 &mut self.pos_public_key.as_ref().unwrap().1.to_bytes().to_vec(),
678 );
679 rlp.append(&key_bytes);
680 self.send_packet(
681 io,
682 None,
683 ProtocolVersion::default(),
684 PACKET_HELLO,
685 rlp.out().to_vec(),
686 SendQueuePriority::High,
687 )
688 .map(|_| ())
689 }
690
691 pub fn writable<Message: Send + Sync + Clone>(
693 &mut self, io: &IoContext<Message>,
694 ) -> Result<(), Error> {
695 let status = self.connection_mut().writable(io)?;
696 self.last_write = (Instant::now(), status);
697 Ok(())
698 }
699
700 pub fn details(&self) -> SessionDetails {
703 SessionDetails {
704 originated: self.metadata.originated,
705 node_id: self.metadata.id,
706 address: self.address,
707 connection: self.connection().details(),
708 status: if let Some(time) = self.expired {
709 format!("expired ({:?})", time.elapsed())
710 } else if let Some(time) = self.had_hello {
711 format!("communicating ({:?})", time.elapsed())
712 } else {
713 format!("handshaking ({:?})", self.sent_hello.elapsed())
714 },
715 last_read: format!("{:?}", self.last_read.elapsed()),
716 last_write: format!("{:?}", self.last_write.0.elapsed()),
717 last_write_status: format!("{:?}", self.last_write.1),
718 }
719 }
720
721 pub fn check_timeout(&self) -> (bool, Option<UpdateNodeOperation>) {
730 if let Some(time) = self.expired {
731 if time.elapsed() > Duration::from_secs(5) {
733 return (true, None);
734 }
735 } else if self.had_hello.is_none() {
736 if self.sent_hello.elapsed() > Duration::from_secs(300) {
738 return (true, Some(UpdateNodeOperation::Failure));
739 }
740 }
741
742 (false, None)
743 }
744}
745
746impl fmt::Debug for Session {
747 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
748 write!(f, "Session {{ token: {}, id: {:?}, originated: {}, address: {:?}, had_hello: {}, expired: {} }}",
749 self.token(), self.id(), self.metadata.originated, self.address, self.had_hello.is_some(), self.expired())
750 }
751}
752
753#[derive(Serialize, Deserialize, Clone)]
755#[serde(rename_all = "camelCase")]
756pub struct SessionDetails {
757 pub originated: bool,
758 pub node_id: Option<NodeId>,
759 pub address: SocketAddr,
760 pub connection: ConnectionDetails,
761 pub status: String,
762 pub last_read: String,
763 pub last_write: String,
764 pub last_write_status: String,
765}
766
767struct MovableWrapper<T> {
770 item: Option<T>,
771}
772
773impl<T> MovableWrapper<T> {
774 fn new(item: T) -> Self { MovableWrapper { item: Some(item) } }
775
776 fn get(&self) -> &T {
777 match self.item {
778 Some(ref item) => item,
779 None => panic!("cannot get moved item"),
780 }
781 }
782
783 fn get_mut(&mut self) -> &mut T {
784 match self.item {
785 Some(ref mut item) => item,
786 None => panic!("cannot get_mut moved item"),
787 }
788 }
789
790 fn take(&mut self) -> T {
791 if self.item.is_none() {
792 panic!("cannot take moved item")
793 }
794
795 self.item.take().expect("should have value")
796 }
797}
798
799#[derive(Eq, PartialEq)]
820struct SessionPacket {
821 pub id: u8,
822 pub protocol: Option<ProtocolId>,
823 pub data: Bytes,
824 pub header_version: u8,
825 pub extensions: Vec<Vec<u8>>,
826}
827
828impl SessionPacket {
829 fn assemble(
831 id: u8, header_version: u8, protocol: Option<ProtocolId>,
832 mut data: Vec<u8>,
833 ) -> Vec<u8> {
834 let mut protocol_flag = 0;
835 if let Some(protocol) = protocol {
836 data.extend_from_slice(&protocol);
837 protocol_flag = 1;
838 }
839
840 let header_byte = (header_version << 1) + protocol_flag;
841 data.push(header_byte);
842 data.push(id);
843
844 data
845 }
846
847 fn parse(mut data: Bytes) -> Result<Self, Error> {
848 if data.is_empty() {
850 debug!("failed to parse session packet, packet id missed");
851 return Err(Error::BadProtocol);
852 }
853
854 let packet_id = data.split_off(data.len() - 1)[0];
855
856 if data.is_empty() {
858 debug!("failed to parse session packet, protocol flag missed");
859 return Err(Error::BadProtocol);
860 }
861
862 let header_byte = data.split_off(data.len() - 1)[0];
863 let protocol_flag = header_byte & 1;
864 let header_version = (header_byte & 0x0f) >> 1;
865 if header_version > HEADER_VERSION_WITH_EXTENSION {
866 debug!("unsupported header_version {}", header_version);
867 return Err(Error::BadProtocol);
868 }
869 let has_extension = (header_byte & 0x10) >> 4;
870
871 if protocol_flag == 0 {
873 if packet_id == PACKET_USER {
874 debug!("failed to parse session packet, no protocol for user packet");
875 return Err(Error::BadProtocol);
876 }
877
878 let (data, extensions) =
879 Self::parse_extensions(data, has_extension != 0)?;
880
881 return Ok(SessionPacket {
882 id: packet_id,
883 header_version,
884 protocol: None,
885 data,
886 extensions,
887 });
888 }
889
890 if packet_id != PACKET_USER {
891 debug!("failed to parse session packet, invalid packet id");
892 return Err(Error::BadProtocol);
893 }
894
895 if data.len() < PROTOCOL_ID_SIZE {
897 debug!("failed to parse session packet, protocol missed");
898 return Err(Error::BadProtocol);
899 }
900
901 let protocol_bytes = data.split_off(data.len() - PROTOCOL_ID_SIZE);
902 let mut protocol = ProtocolId::default();
903 protocol.copy_from_slice(&protocol_bytes);
904
905 let (data, extensions) =
907 Self::parse_extensions(data, has_extension != 0)?;
908
909 Ok(SessionPacket {
910 id: packet_id,
911 protocol: Some(protocol),
912 header_version,
913 data,
914 extensions,
915 })
916 }
917
918 fn parse_extensions(
919 mut data: Bytes, mut has_extension: bool,
920 ) -> Result<(Bytes, Vec<Vec<u8>>), Error> {
921 let mut extensions = Vec::new();
922 while has_extension {
923 if data.is_empty() {
924 debug!(
925 "failed to parse session packet, extension data exhausted"
926 );
927 bail!(Error::BadProtocol);
928 }
929 let extension_byte = data.split_off(data.len() - 1)[0];
930 let extension_len = (extension_byte >> 1) as usize;
931 has_extension = (extension_byte & 1) != 0;
932 if data.len() < extension_len {
933 debug!("failed to parse session packet, not enough bytes for extension.");
934 bail!(Error::BadProtocol);
935 }
936 extensions
937 .push(data.split_off(data.len() - extension_len).to_vec());
938 }
939
940 Ok((data, extensions))
941 }
942}
943
944impl fmt::Debug for SessionPacket {
945 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
946 write!(
947 f,
948 "SessionPacket {{ id: {}, protocol: {:?}, date_len: {} }}",
949 self.id,
950 self.protocol,
951 self.data.len()
952 )
953 }
954}
955
956#[cfg(test)]
957mod tests {
958 use super::*;
959
960 #[test]
961 fn test_packet_assemble() {
962 let packet =
963 SessionPacket::assemble(5, PACKET_HEADER_VERSION, None, vec![1, 3]);
964 assert_eq!(packet, vec![1, 3, 0, 5]);
965
966 let packet = SessionPacket::assemble(
967 6,
968 PACKET_HEADER_VERSION,
969 Some([8; 3]),
970 vec![2, 4],
971 );
972 assert_eq!(packet, vec![2, 4, 8, 8, 8, 1, 6]);
973 }
974
975 #[test]
976 fn test_packet_parse() {
977 assert!(SessionPacket::parse(vec![].into()).is_err());
979
980 assert!(SessionPacket::parse(vec![1].into()).is_err());
982
983 assert!(SessionPacket::parse(vec![2, 1].into()).is_err());
985
986 assert!(SessionPacket::parse(vec![0, PACKET_USER].into()).is_err());
988
989 let packet = SessionPacket::parse(vec![1, 2, 0, 20].into()).unwrap();
991 assert_eq!(
992 packet,
993 SessionPacket {
994 id: 20,
995 header_version: 0,
996 protocol: None,
997 data: vec![1, 2].into(),
998 extensions: vec![],
999 }
1000 );
1001
1002 assert!(SessionPacket::parse(vec![6, 6, 6, 1, 7].into()).is_err());
1004
1005 assert!(
1007 SessionPacket::parse(vec![6, 6, 1, PACKET_USER].into()).is_err()
1008 );
1009
1010 let packet =
1012 SessionPacket::parse(vec![1, 9, 3, 3, 3, 1, PACKET_USER].into())
1013 .unwrap();
1014 assert_eq!(
1015 packet,
1016 SessionPacket {
1017 id: PACKET_USER,
1018 header_version: 0,
1019 protocol: Some([3; 3]),
1020 data: vec![1, 9].into(),
1021 extensions: vec![],
1022 }
1023 );
1024 }
1025}