1 files changed, 33 insertions, 22 deletions

diff --git a/tunnat/include/tunnel.h b/tunnat/include/tunnel.h
index 1d83af0..fd2c314 100644
--- a/tunnat/include/tunnel.h
+++ b/tunnat/include/tunnel.h
@@ -1,4 +1,4 @@
-#pragma once
+#pragma once
 #include <array>
 
 extern "C"
@@ -13,14 +13,19 @@ extern "C"
 #include <marsio.h>
 #include <mrtunnat.h>
 }
-
-struct AddressInfo
+enum addr_type
+{
+    IPV4_ADDR,
+    IPV6_ADDR
+};
+struct ipv4_addr_t
 {
-    uint8_t src_mac_addr[6];
-    uint8_t dst_mac_addr[6];
-
     struct in_addr src_in_addr;
     struct in_addr dst_in_addr;
+};
+
+struct ipv6_addr_t
+{
     struct in6_addr src_in6_addr;
     struct in6_addr dst_in6_addr;
 
@@ -111,8 +116,7 @@ protected:
     struct udp_hdr udp_hdr_;
     struct g_vxlan_hdr vxlan_hdr_;
 
-    struct in_addr src_in_addr_;
-    struct in_addr dst_in_addr_;
+    struct ipv4_addr_t ipv4_addr_;
 
 public:
     int PacketParse(const char * pkt, unsigned int pkt_len);
@@ -121,14 +125,13 @@ public:
 
     int GetAddressInfo(struct AddressInfo & addr_info)
     {
-        memcpy(addr_info.src_mac_addr, &ether_hdr_.s_addr, sizeof(ether_hdr_.s_addr));
-        memcpy(addr_info.dst_mac_addr, &ether_hdr_.d_addr, sizeof(ether_hdr_.d_addr));
-
-        addr_info.src_in_addr = src_in_addr_;
-        addr_info.dst_in_addr = dst_in_addr_;
+        addr_info.type = IPV4_ADDR;
+        memcpy(&(addr_info.addr.ipv4_addr), &ipv4_addr_, sizeof(struct ipv4_addr_t)); 
         addr_info.s_port = udp_hdr_.src_port;
         addr_info.d_port = udp_hdr_.dst_port;
-        addr_info.vpn_id = (vxlan_hdr_.vlan_id_half_low & 0xf) | (vxlan_hdr_.vlan_id_half_high & 0xff) << 4;
+        addr_info.vpn_id = (vxlan_hdr_.vlan_id_half_low & 0xf) |
+            (vxlan_hdr_.vlan_id_half_high & 0xff) << 4;
+
         return 0;
     }
 
@@ -167,32 +170,40 @@ public:
 class TunInnerIPv4 : public Tunnel
 {
 protected:
-    struct in_addr src_in_addr_;
-    struct in_addr dst_in_addr_;
+    struct ipv4_addr_t ipv4_addr;
 public:
     int PacketParse(const char * pkt, unsigned int pkt_len);
 
     int GetAddressInfo(struct AddressInfo & addr_info)
     {
-        addr_info.src_in_addr = src_in_addr_;
-        addr_info.dst_in_addr = dst_in_addr_;
+        addr_info.type = IPV4_ADDR;
+        memcpy(&(addr_info.addr.ipv4_addr), &ipv4_addr, sizeof(struct ipv4_addr_t));
         return RT_SUCCESS;
     }
 
     void ReverseAddress()
     {
-        std::swap<in_addr>(src_in_addr_, dst_in_addr_);
+        std::swap<in_addr>(ipv4_addr.src_in_addr, ipv4_addr.dst_in_addr);
     }
 };
 
 class TunInnerIPv6 : public Tunnel
 {
 protected:
-    struct in6_addr src_in6_addr_;
-    struct in6_addr dst_in6_addr_;
+    struct ipv6_addr_t ipv6_addr;
 public:
     int PacketParse(const char * pkt, unsigned int pkt_len);
-    int GetAddressInfo(struct AddressInfo & addr_info);
+    int GetAddressInfo(struct AddressInfo & addr_info)
+    {
+        addr_info.type = IPV6_ADDR;
+        memcpy(&(addr_info.addr.ipv6_addr),&ipv6_addr,sizeof(struct ipv6_addr_t));
+        return RT_SUCCESS;
+    }
+    void ReverseAddress()
+    {
+        struct ipv6_addr_t ipv6_addr;
+        std::swap<in6_addr>(ipv6_addr.src_in6_addr, ipv6_addr.dst_in6_addr);
+    }
 };
 
 class TunInnerPPP : public Tunnel