Protocols now always set the next layer protocol flag.

3 files changed, 12 insertions, 21 deletions

diff --git a/src/dot1q.cpp b/src/dot1q.cpp
index d0e045b..528030d 100644
--- a/src/dot1q.cpp
+++ b/src/dot1q.cpp
@@ -103,7 +103,7 @@ void Dot1Q::write_serialization(uint8_t *buffer, uint32_t total_sz, const PDU *)
     #ifdef TINS_DEBUG
     assert(total_sz >= sizeof(_header) + trailer);
     #endif
-    if ((payload_type() == 0) && inner_pdu()) {
+    if (inner_pdu()) {
         Constants::Ethernet::e flag = Internals::pdu_flag_to_ether_type(
             inner_pdu()->pdu_type()
         );
diff --git a/src/snap.cpp b/src/snap.cpp
index bcb66f5..662a54b 100644
--- a/src/snap.cpp
+++ b/src/snap.cpp
@@ -97,22 +97,13 @@ void Tins::SNAP::write_serialization(uint8_t *buffer, uint32_t total_sz, const P
     #ifdef TINS_DEBUG
     assert(total_sz >= sizeof(_snap));
     #endif
-    if (!_snap.eth_type && inner_pdu()) {
-        uint16_t type = Tins::Constants::Ethernet::IP;
-        switch (inner_pdu()->pdu_type()) {
-            case PDU::IP:
-                type = Tins::Constants::Ethernet::IP;
-                break;
-            case PDU::ARP:
-                type = Tins::Constants::Ethernet::ARP;
-                break;
-            case PDU::EAPOL:
-                type = Tins::Constants::Ethernet::EAPOL;
-                break;
-            default:
-                type = 0;
-        }
-        _snap.eth_type = Endian::host_to_be(type);
+    if (inner_pdu()) {
+        Constants::Ethernet::e flag = Internals::pdu_flag_to_ether_type(
+            inner_pdu()->pdu_type()
+        );
+        _snap.eth_type = Endian::host_to_be(
+            static_cast<uint16_t>(flag)
+        );
     }
     std::memcpy(buffer, &_snap, sizeof(_snap));
 }
diff --git a/tests/src/allocators.cpp b/tests/src/allocators.cpp
index 0292fde..465b517 100644
--- a/tests/src/allocators.cpp
+++ b/tests/src/allocators.cpp
@@ -85,13 +85,13 @@ TEST_F(AllocatorsTest, LinkLayerPDUs) {
         EXPECT_EQ(pkt.serialize(), link_layer_data);
     }
     {
-        SLL pkt(&link_layer_data[0], link_layer_data.size());
-        EXPECT_TRUE(pkt.find_pdu<DummyPDU<3> >());
+        Dot1Q pkt(&link_layer_data[0], link_layer_data.size());
+        EXPECT_TRUE(pkt.find_pdu<DummyPDU<2> >());
         EXPECT_EQ(pkt.serialize(), link_layer_data);
     }
     {
-        Dot1Q pkt(&link_layer_data[0], link_layer_data.size());
-        EXPECT_TRUE(pkt.find_pdu<DummyPDU<2> >());
+        SLL pkt(&link_layer_data[0], link_layer_data.size());
+        EXPECT_TRUE(pkt.find_pdu<DummyPDU<3> >());
         EXPECT_EQ(pkt.serialize(), link_layer_data);
     }
 }