From 51aa23b63571cad1ce981cd4ec1426ad88af0abc Mon Sep 17 00:00:00 2001
From: "sijie.sun" <sijie.sun@smartx.com>
Date: Sat, 11 May 2024 22:59:00 +0800
Subject: [PATCH] add ttl for packet

---
 easytier/src/peers/peer_manager.rs | 10 ++++++++--
 easytier/src/tunnel/packet_def.rs  |  6 +++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/easytier/src/peers/peer_manager.rs b/easytier/src/peers/peer_manager.rs
index b6fec0f..a2d7e34 100644
--- a/easytier/src/peers/peer_manager.rs
+++ b/easytier/src/peers/peer_manager.rs
@@ -324,14 +324,20 @@ impl PeerManager {
         self.tasks.lock().await.spawn(async move {
             log::trace!("start_peer_recv");
             while let Some(mut ret) = recv.next().await {
-                let Some(hdr) = ret.peer_manager_header() else {
+                let Some(hdr) = ret.mut_peer_manager_header() else {
                     tracing::warn!(?ret, "invalid packet, skip");
                     continue;
                 };
-                tracing::trace!(?hdr, ?ret, "peer recv a packet...");
+                tracing::trace!(?hdr, "peer recv a packet...");
                 let from_peer_id = hdr.from_peer_id.get();
                 let to_peer_id = hdr.to_peer_id.get();
                 if to_peer_id != my_peer_id {
+                    if hdr.ttl <= 1 {
+                        tracing::warn!(?hdr, "ttl is 0, drop packet");
+                        continue;
+                    }
+
+                    hdr.ttl -= 1;
                     tracing::trace!(?to_peer_id, ?my_peer_id, "need forward");
                     let ret = peers.send_msg(ret, to_peer_id).await;
                     if ret.is_err() {
diff --git a/easytier/src/tunnel/packet_def.rs b/easytier/src/tunnel/packet_def.rs
index 036b40d..dd0a106 100644
--- a/easytier/src/tunnel/packet_def.rs
+++ b/easytier/src/tunnel/packet_def.rs
@@ -7,6 +7,8 @@ use zerocopy::FromZeroes;
 
 type DefaultEndian = LittleEndian;
 
+pub const DEFAULT_TTL: u8 = 8;
+
 // TCP TunnelHeader
 #[repr(C, packed)]
 #[derive(AsBytes, FromBytes, FromZeroes, Clone, Debug, Default)]
@@ -69,7 +71,8 @@ pub struct PeerManagerHeader {
     pub to_peer_id: U32<DefaultEndian>,
     pub packet_type: u8,
     pub flags: u8,
-    reserved: U16<DefaultEndian>,
+    pub ttl: u8,
+    reserved: u8,
     pub len: U32<DefaultEndian>,
 }
 pub const PEER_MANAGER_HEADER_SIZE: usize = std::mem::size_of::<PeerManagerHeader>();
@@ -362,6 +365,7 @@ impl ZCPacket {
         hdr.to_peer_id.set(to_peer_id);
         hdr.packet_type = packet_type;
         hdr.flags = 0;
+        hdr.ttl = DEFAULT_TTL;
         hdr.len.set(payload_len as u32);
     }