From 70b50a7a2e75b1df0349e37d098aa13dd153ae75 Mon Sep 17 00:00:00 2001
From: Davide De Rosa <keeshux@gmail.com>
Date: Fri, 19 Oct 2018 11:31:23 +0200
Subject: [PATCH] Parse data opcode when decrypting

Assume it could be DATA_V1/V2 regardless of peer-id.
---
 TunnelKit/Sources/Core/CryptoAEAD.m   |  4 ++++
 TunnelKit/Sources/Core/CryptoCBC.m    |  4 ++++
 TunnelKit/Sources/Core/PacketMacros.h | 10 ++++++++++
 3 files changed, 18 insertions(+)

diff --git a/TunnelKit/Sources/Core/CryptoAEAD.m b/TunnelKit/Sources/Core/CryptoAEAD.m
index 85185a8..6b43146 100644
--- a/TunnelKit/Sources/Core/CryptoAEAD.m
+++ b/TunnelKit/Sources/Core/CryptoAEAD.m
@@ -319,6 +319,10 @@ const NSInteger CryptoAEADTagLength     = 16;
 
 - (BOOL)decryptDataPacket:(NSData *)packet into:(uint8_t *)packetBytes length:(NSInteger *)packetLength packetId:(uint32_t *)packetId error:(NSError *__autoreleasing *)error
 {
+    NSAssert(packet.length > 0, @"Decrypting an empty packet, how did it get this far?");
+    PacketCode code;
+    PacketOpcodeGet(packet.bytes, &code, NULL);
+
     const uint8_t *extra = packet.bytes; // AD = header + peer id + packet id
     if (!self.checkPeerId) {
         extra += self.headerLength; // AD = packet id only
diff --git a/TunnelKit/Sources/Core/CryptoCBC.m b/TunnelKit/Sources/Core/CryptoCBC.m
index 5b69e10..ae27263 100644
--- a/TunnelKit/Sources/Core/CryptoCBC.m
+++ b/TunnelKit/Sources/Core/CryptoCBC.m
@@ -352,6 +352,10 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
 
 - (BOOL)decryptDataPacket:(NSData *)packet into:(uint8_t *)packetBytes length:(NSInteger *)packetLength packetId:(uint32_t *)packetId error:(NSError *__autoreleasing *)error
 {
+    NSAssert(packet.length > 0, @"Decrypting an empty packet, how did it get this far?");
+    PacketCode code;
+    PacketOpcodeGet(packet.bytes, &code, NULL);
+    
     // skip header = (code, key)
     const BOOL success = [self.crypto decryptBytes:(packet.bytes + self.headerLength)
                                             length:(int)(packet.length - self.headerLength)
diff --git a/TunnelKit/Sources/Core/PacketMacros.h b/TunnelKit/Sources/Core/PacketMacros.h
index 019a170..f4f540e 100644
--- a/TunnelKit/Sources/Core/PacketMacros.h
+++ b/TunnelKit/Sources/Core/PacketMacros.h
@@ -62,6 +62,16 @@ typedef NS_ENUM(uint8_t, PacketCode) {
 
 extern const uint8_t DataPacketPingData[16];
 
+static inline void PacketOpcodeGet(const uint8_t *from, PacketCode *_Nullable code, uint8_t *_Nullable key)
+{
+    if (code) {
+        *code = (PacketCode)(*from >> 3);
+    }
+    if (key) {
+        *key = *from & 0b111;
+    }
+}
+
 // Ruby: header
 static inline int PacketHeaderSet(uint8_t *to, PacketCode code, uint8_t key, const uint8_t *_Nullable sessionId)
 {