From ceebbdb200bc149476e70a4d1a202249154c23c3 Mon Sep 17 00:00:00 2001
From: EAimTY <ea.imty@gmail.com>
Date: Wed, 25 Jan 2023 18:37:04 +0900
Subject: [PATCH] adding recv methods for auth & dissoc & heartbeat

---
 tuic/src/prototype/authenticate.rs | 20 +++++++++++++++--
 tuic/src/prototype/dissociate.rs   | 20 +++++++++++++++--
 tuic/src/prototype/heartbeat.rs    | 13 +++++++++--
 tuic/src/prototype/mod.rs          | 35 +++++++++++++++++++++++++-----
 tuic/src/prototype/packet.rs       |  5 +----
 5 files changed, 77 insertions(+), 16 deletions(-)

diff --git a/tuic/src/prototype/authenticate.rs b/tuic/src/prototype/authenticate.rs
index 118ee20..ca54f39 100644
--- a/tuic/src/prototype/authenticate.rs
+++ b/tuic/src/prototype/authenticate.rs
@@ -10,8 +10,6 @@ pub struct Tx {
     header: Header,
 }
 
-pub struct Rx;
-
 impl Authenticate<side::Tx> {
     pub(super) fn new(token: [u8; 8]) -> Self {
         Self {
@@ -27,3 +25,21 @@ impl Authenticate<side::Tx> {
         &tx.header
     }
 }
+
+pub struct Rx {
+    token: [u8; 8],
+}
+
+impl Authenticate<side::Rx> {
+    pub(super) fn new(token: [u8; 8]) -> Self {
+        Self {
+            inner: Side::Rx(Rx { token }),
+            _marker: side::Rx,
+        }
+    }
+
+    pub fn token(&self) -> &[u8; 8] {
+        let Side::Rx(rx) = &self.inner else { unreachable!() };
+        &rx.token
+    }
+}
diff --git a/tuic/src/prototype/dissociate.rs b/tuic/src/prototype/dissociate.rs
index bf7f728..e6e4072 100644
--- a/tuic/src/prototype/dissociate.rs
+++ b/tuic/src/prototype/dissociate.rs
@@ -10,8 +10,6 @@ pub struct Tx {
     header: Header,
 }
 
-pub struct Rx;
-
 impl Dissociate<side::Tx> {
     pub(super) fn new(assoc_id: u16) -> Self {
         Self {
@@ -27,3 +25,21 @@ impl Dissociate<side::Tx> {
         &tx.header
     }
 }
+
+pub struct Rx {
+    assoc_id: u16,
+}
+
+impl Dissociate<side::Rx> {
+    pub(super) fn new(assoc_id: u16) -> Self {
+        Self {
+            inner: Side::Rx(Rx { assoc_id }),
+            _marker: side::Rx,
+        }
+    }
+
+    pub fn assoc_id(&self) -> &u16 {
+        let Side::Rx(rx) = &self.inner else { unreachable!() };
+        &rx.assoc_id
+    }
+}
diff --git a/tuic/src/prototype/heartbeat.rs b/tuic/src/prototype/heartbeat.rs
index b5b96eb..8fce8b0 100644
--- a/tuic/src/prototype/heartbeat.rs
+++ b/tuic/src/prototype/heartbeat.rs
@@ -10,8 +10,6 @@ pub struct Tx {
     header: Header,
 }
 
-pub struct Rx;
-
 impl Heartbeat<side::Tx> {
     pub(super) fn new() -> Self {
         Self {
@@ -27,3 +25,14 @@ impl Heartbeat<side::Tx> {
         &tx.header
     }
 }
+
+pub struct Rx;
+
+impl Heartbeat<side::Rx> {
+    pub(super) fn new() -> Self {
+        Self {
+            inner: Side::Rx(Rx),
+            _marker: side::Rx,
+        }
+    }
+}
diff --git a/tuic/src/prototype/mod.rs b/tuic/src/prototype/mod.rs
index b79cd7e..358f456 100644
--- a/tuic/src/prototype/mod.rs
+++ b/tuic/src/prototype/mod.rs
@@ -1,4 +1,7 @@
-use crate::protocol::{Address, Connect as ConnectHeader, Packet as PacketHeader};
+use crate::protocol::{
+    Address, Authenticate as AuthenticateHeader, Connect as ConnectHeader,
+    Dissociate as DissociateHeader, Heartbeat as HeartbeatHeader, Packet as PacketHeader,
+};
 use parking_lot::Mutex;
 use std::{
     collections::HashMap,
@@ -45,7 +48,12 @@ where
     }
 
     pub fn send_authenticate(&self, token: [u8; 8]) -> Authenticate<side::Tx> {
-        Authenticate::new(token)
+        Authenticate::<side::Tx>::new(token)
+    }
+
+    pub fn recv_authenticate(&self, header: AuthenticateHeader) -> Authenticate<side::Rx> {
+        let (token,) = header.into();
+        Authenticate::<side::Rx>::new(token)
     }
 
     pub fn send_connect(&self, addr: Address) -> Connect<side::Tx> {
@@ -82,11 +90,21 @@ where
     }
 
     pub fn send_dissociate(&self, assoc_id: u16) -> Dissociate<side::Tx> {
-        self.udp_sessions.lock().dissociate(assoc_id)
+        self.udp_sessions.lock().send_dissociate(assoc_id)
+    }
+
+    pub fn recv_dissociate(&self, header: DissociateHeader) -> Dissociate<side::Rx> {
+        let (assoc_id,) = header.into();
+        self.udp_sessions.lock().recv_dissociate(assoc_id)
     }
 
     pub fn send_heartbeat(&self) -> Heartbeat<side::Tx> {
-        Heartbeat::new()
+        Heartbeat::<side::Tx>::new()
+    }
+
+    pub fn recv_heartbeat(&self, header: HeartbeatHeader) -> Heartbeat<side::Rx> {
+        let () = header.into();
+        Heartbeat::<side::Rx>::new()
     }
 
     pub fn task_connect_count(&self) -> usize {
@@ -174,9 +192,14 @@ where
             .recv_packet(sessions, assoc_id, pkt_id, frag_total, frag_id, size, addr)
     }
 
-    fn dissociate(&mut self, assoc_id: u16) -> Dissociate<side::Tx> {
+    fn send_dissociate(&mut self, assoc_id: u16) -> Dissociate<side::Tx> {
         self.sessions.remove(&assoc_id);
-        Dissociate::new(assoc_id)
+        Dissociate::<side::Tx>::new(assoc_id)
+    }
+
+    fn recv_dissociate(&mut self, assoc_id: u16) -> Dissociate<side::Rx> {
+        self.sessions.remove(&assoc_id);
+        Dissociate::<side::Rx>::new(assoc_id)
     }
 
     fn insert<A>(
diff --git a/tuic/src/prototype/packet.rs b/tuic/src/prototype/packet.rs
index 5afd443..7b757a8 100644
--- a/tuic/src/prototype/packet.rs
+++ b/tuic/src/prototype/packet.rs
@@ -18,10 +18,7 @@ pub struct Tx {
     max_pkt_size: usize,
 }
 
-impl<B> Packet<side::Tx, B>
-where
-    B: AsRef<[u8]>,
-{
+impl<B> Packet<side::Tx, B> {
     pub(super) fn new(assoc_id: u16, pkt_id: u16, addr: Address, max_pkt_size: usize) -> Self {
         Self {
             inner: Side::Tx(Tx {