Move WireMessage to Packet, remove existing Message

crates/net/src/manager.rs

@@ -2,15 +2,15 @@ use std::collections::{HashMap, HashSet};
 use tokio::sync::mpsc;
 use tokio_util::task::JoinMap;
 
-use crate::{error::NetError, message::MsgId, packet::{ConnId, Message, Packet}, peer::PeerId, session::{ActiveSession, EphemeralState, PendingSession, SessionId}};
+use crate::{error::NetError, message::MsgId, packet::{ConnId, WireMessage, Packet}, peer::PeerId, session::{ActiveSession, EphemeralState, PendingSession, SessionId}};
 
 struct NetworkManager {
     sessions_by_peer: HashMap<PeerId, HashSet<SessionId>>,
     peer_by_session: HashMap<SessionId, PeerId>,
     active_sessions: HashMap<SessionId, ActiveSession>,
     pending_sessions: HashMap<SessionId, PendingSession>,
     ephemeral_states: HashMap<PeerId, EphemeralState>,
-    message_callbacks: HashMap<(SessionId, MsgId), mpsc::Sender<Message>>,
+    message_callbacks: HashMap<(SessionId, MsgId), mpsc::Sender<WireMessage>>,
     ttl_sessions: HashMap<u64, SessionId>,
 
     conns_by_peer: HashMap<PeerId, HashSet<ConnId>>,

crates/net/src/packet.rs

@@ -1,4 +1,4 @@
-use crypto::{ciphertext::Ciphertext, kem::{Capsule, PublicKey}};
+use crypto::{ciphertext::Ciphertext, kem::{Capsule, PublicKey}, id::Id};
 use serde::{Deserialize, Serialize};
 use serde_big_array::BigArray;
 
@@ -19,15 +19,17 @@ pub struct HandshakeAck {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct Message {
+pub struct WireMessage {
     pub seq: u64,
+    pub session_id: SessionId,
     pub ciphertext: Ciphertext,
-    pub conn_id: ConnId,
 }
 
+pub type SessionId = Id;
+
 #[derive(Serialize, Deserialize)]
 pub enum Packet {
     Handshake(Handshake),
     HandshakeAck(HandshakeAck),
-    Message(Message),
+    Message(WireMessage),
 }

crates/net/src/session.rs

@@ -1,17 +1,11 @@
 use bytes::Bytes;
 use crypto::{blake3, ciphertext::{Ciphertext, Nonce}, id::Id, kem::{SecretKey, SharedSecret}};
 
-use crate::{error::NetError, packet::Message};
+use crate::{error::NetError, packet::WireMessage};
 
 const WINDOW: usize = 32;
 const VERSION: u8 = 1;
 
-pub struct WireMessage { // TODO: Move it to the packet.rs
-    seq: u64,
-    session_id: SessionId,
-    ciphertext: Ciphertext,
-}
-
 pub type SessionId = Id;
 
 pub struct PendingSession(SharedSecret);
@@ -83,7 +77,7 @@ impl ActiveSession {
         Ok(WireMessage { seq: self.seq, ciphertext, session_id: self.session_id.clone() })
     }
 
-    pub fn recv(&mut self, msg: Message) -> Result<Bytes, NetError> {
+    pub fn recv(&mut self, msg: WireMessage) -> Result<Bytes, NetError> {
         if !self.check_seq(msg.seq) {
             return Err(NetError::InvalidSeq);
         }