Handle encryption/peer-id in a stateless manner

Fixes #30

TunnelKit/Sources/Core/Crypto.h

@@ -43,6 +43,12 @@ NS_ASSUME_NONNULL_BEGIN
 @protocol DataPathEncrypter;
 @protocol DataPathDecrypter;
 
+typedef struct {
+    uint32_t packetId;
+    const uint8_t *ad;
+    int adLength;
+} CryptoFlags;
+
 // WARNING: dest must be able to hold ciphertext
 @protocol Encrypter
 
@@ -50,7 +56,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (int)digestLength;
 
 - (NSInteger)encryptionCapacityWithLength:(NSInteger)length;
-- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(nullable const uint8_t *)extra error:(NSError **)error;
+- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags *_Nullable)flags error:(NSError **)error;
 
 - (id<DataPathEncrypter>)dataPathEncrypter;
 
@@ -63,8 +69,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (int)digestLength;
 
 - (NSInteger)encryptionCapacityWithLength:(NSInteger)length;
-- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(nullable const uint8_t *)extra error:(NSError **)error;
+- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags *_Nullable)flags error:(NSError **)error;
-- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length extra:(nullable const uint8_t *)extra error:(NSError **)error;
+- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length flags:(const CryptoFlags *_Nullable)flags error:(NSError **)error;
 
 - (id<DataPathDecrypter>)dataPathDecrypter;
 

TunnelKit/Sources/Core/CryptoAEAD.h

@@ -43,14 +43,14 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface CryptoAEAD : NSObject <Encrypter, Decrypter>
 
-@property (nonatomic, assign) int extraLength;
-
 - (instancetype)initWithCipherName:(NSString *)cipherName;
 
 @end
 
 @interface DataPathCryptoAEAD : NSObject <DataPathEncrypter, DataPathDecrypter>
 
+@property (nonatomic, assign) uint32_t peerId;
+
 - (instancetype)initWithCrypto:(CryptoAEAD *)crypto;
 
 @end

TunnelKit/Sources/Core/CryptoAEAD.m

@@ -50,7 +50,6 @@ const NSInteger CryptoAEADTagLength     = 16;
 @property (nonatomic, unsafe_unretained) const EVP_CIPHER *cipher;
 @property (nonatomic, assign) int cipherKeyLength;
 @property (nonatomic, assign) int cipherIVLength; // 12 (AD packetId + HMAC key)
-@property (nonatomic, assign) int extraPacketIdOffset;
 
 @property (nonatomic, unsafe_unretained) EVP_CIPHER_CTX *cipherCtxEnc;
 @property (nonatomic, unsafe_unretained) EVP_CIPHER_CTX *cipherCtxDec;
@@ -72,8 +71,6 @@ const NSInteger CryptoAEADTagLength     = 16;
         
         self.cipherKeyLength = EVP_CIPHER_key_length(self.cipher);
         self.cipherIVLength = EVP_CIPHER_iv_length(self.cipher);
-        self.extraLength = PacketIdLength;
-        self.extraPacketIdOffset = 0;
         
         self.cipherCtxEnc = EVP_CIPHER_CTX_new();
         self.cipherCtxDec = EVP_CIPHER_CTX_new();
@@ -118,19 +115,19 @@ const NSInteger CryptoAEADTagLength     = 16;
     [self prepareIV:self.cipherIVEnc withHMACKey:hmacKey];
 }
 
-- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
-    NSParameterAssert(extra);
+    NSParameterAssert(flags);
 
     int l1 = 0, l2 = 0;
     int x = 0;
     int code = 1;
 
-    assert(self.extraLength >= PacketIdLength);
+    assert(flags->adLength >= PacketIdLength);
-    memcpy(self.cipherIVEnc, extra + self.extraPacketIdOffset, PacketIdLength);
+    memcpy(self.cipherIVEnc, &flags->packetId, PacketIdLength);
 
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherInit(self.cipherCtxEnc, NULL, NULL, self.cipherIVEnc, -1);
-    TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxEnc, NULL, &x, extra, self.extraLength);
+    TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxEnc, NULL, &x, flags->ad, flags->adLength);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxEnc, dest + CryptoAEADTagLength, &l1, bytes, (int)length);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherFinal(self.cipherCtxEnc, dest + CryptoAEADTagLength + l1, &l2);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CIPHER_CTX_ctrl(self.cipherCtxEnc, EVP_CTRL_GCM_GET_TAG, CryptoAEADTagLength, dest);
@@ -164,20 +161,20 @@ const NSInteger CryptoAEADTagLength     = 16;
     [self prepareIV:self.cipherIVDec withHMACKey:hmacKey];
 }
 
-- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
-    NSParameterAssert(extra);
+    NSParameterAssert(flags);
 
     int l1 = 0, l2 = 0;
     int x = 0;
     int code = 1;
     
-    assert(self.extraLength >= PacketIdLength);
+    assert(flags->adLength >= PacketIdLength);
-    memcpy(self.cipherIVDec, extra + self.extraPacketIdOffset, PacketIdLength);
+    memcpy(self.cipherIVDec, &flags->packetId, PacketIdLength);
 
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherInit(self.cipherCtxDec, NULL, NULL, self.cipherIVDec, -1);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CIPHER_CTX_ctrl(self.cipherCtxDec, EVP_CTRL_GCM_SET_TAG, CryptoAEADTagLength, (uint8_t *)bytes);
-    TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxDec, NULL, &x, extra, self.extraLength);
+    TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxDec, NULL, &x, flags->ad, flags->adLength);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherUpdate(self.cipherCtxDec, dest, &l1, bytes + CryptoAEADTagLength, (int)length - CryptoAEADTagLength);
     TUNNEL_CRYPTO_TRACK_STATUS(code) EVP_CipherFinal(self.cipherCtxDec, dest + l1, &l2);
 
@@ -192,7 +189,7 @@ const NSInteger CryptoAEADTagLength     = 16;
     TUNNEL_CRYPTO_RETURN_STATUS(code)
 }
 
-- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
     [NSException raise:NSInvalidArgumentException format:@"Verification not supported"];
     return NO;
@@ -218,9 +215,6 @@ const NSInteger CryptoAEADTagLength     = 16;
 @interface DataPathCryptoAEAD ()
 
 @property (nonatomic, strong) CryptoAEAD *crypto;
-@property (nonatomic, assign) int headerLength;
-@property (nonatomic, copy) void (^setDataHeader)(uint8_t *, uint8_t);
-@property (nonatomic, copy) BOOL (^checkPeerId)(const uint8_t *);
 
 @end
 
@@ -239,28 +233,7 @@ const NSInteger CryptoAEADTagLength     = 16;
 
 - (void)setPeerId:(uint32_t)peerId
 {
-    peerId &= 0xffffff;
+    _peerId = peerId & 0xffffff;
-    
-    if (peerId == PacketPeerIdDisabled) {
-        self.headerLength = PacketOpcodeLength;
-        self.crypto.extraLength = PacketIdLength;
-        self.crypto.extraPacketIdOffset = 0;
-        self.setDataHeader = ^(uint8_t *to, uint8_t key) {
-            PacketHeaderSet(to, PacketCodeDataV1, key, nil);
-        };
-        self.checkPeerId = NULL;
-    }
-    else {
-        self.headerLength = PacketOpcodeLength + PacketPeerIdLength;
-        self.crypto.extraLength = self.headerLength + PacketIdLength;
-        self.crypto.extraPacketIdOffset = self.headerLength;
-        self.setDataHeader = ^(uint8_t *to, uint8_t key) {
-            PacketHeaderSetDataV2(to, key, peerId);
-        };
-        self.checkPeerId = ^BOOL(const uint8_t *ptr) {
-            return (PacketHeaderGetDataV2PeerId(ptr) == peerId);
-        };
-    }
 }
 
 - (NSInteger)encryptionCapacityWithLength:(NSInteger)length
@@ -285,24 +258,32 @@ const NSInteger CryptoAEADTagLength     = 16;
 
 - (NSData *)encryptedDataPacketWithKey:(uint8_t)key packetId:(uint32_t)packetId packetBytes:(const uint8_t *)packetBytes packetLength:(NSInteger)packetLength error:(NSError *__autoreleasing *)error
 {
-    const int capacity = self.headerLength + PacketIdLength + (int)[self.crypto encryptionCapacityWithLength:packetLength];
+    DATA_PATH_ENCRYPT_INIT(self.peerId)
+    
+    const int capacity = headerLength + PacketIdLength + (int)[self.crypto encryptionCapacityWithLength:packetLength];
     NSMutableData *encryptedPacket = [[NSMutableData alloc] initWithLength:capacity];
     uint8_t *ptr = encryptedPacket.mutableBytes;
     NSInteger encryptedPacketLength = INT_MAX;
 
-    self.setDataHeader(ptr, key);
+    CryptoFlags flags;
-    *(uint32_t *)(ptr + self.headerLength) = htonl(packetId);
+    flags.packetId = htonl(packetId);
-
+    if (hasPeerId) {
-    const uint8_t *extra = ptr; // AD = header + peer id + packet id
+        PacketHeaderSetDataV2(ptr, key, self.peerId);
-    if (!self.checkPeerId) {
+        flags.ad = ptr;
-        extra += self.headerLength; // AD = packet id only
+        flags.adLength = headerLength + PacketIdLength;
     }
+    else {
+        PacketHeaderSet(ptr, PacketCodeDataV1, key, nil);
+        flags.ad = ptr + headerLength;
+        flags.adLength = PacketIdLength;
+    }
+    *(uint32_t *)(ptr + headerLength) = flags.packetId;
 
     const BOOL success = [self.crypto encryptBytes:packetBytes
                                             length:packetLength
-                                              dest:(ptr + self.headerLength + PacketIdLength) // skip header and packet id
+                                              dest:(ptr + headerLength + PacketIdLength) // skip header and packet id
                                         destLength:&encryptedPacketLength
-                                             extra:extra
+                                             flags:&flags
                                              error:error];
     
     NSAssert(encryptedPacketLength <= capacity, @"Did not allocate enough bytes for payload");
@@ -311,7 +292,7 @@ const NSInteger CryptoAEADTagLength     = 16;
         return nil;
     }
     
-    encryptedPacket.length = self.headerLength + PacketIdLength + encryptedPacketLength;
+    encryptedPacket.length = headerLength + PacketIdLength + encryptedPacketLength;
     return encryptedPacket;
 }
 
@@ -320,31 +301,40 @@ 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
+    DATA_PATH_DECRYPT_INIT(packet.bytes)
-    if (!self.checkPeerId) {
+    if (packet.length < headerLength) {
-        extra += self.headerLength; // AD = packet id only
+        return NO;
+    }
+
+    CryptoFlags flags;
+    flags.packetId = *(const uint32_t *)(packet.bytes + headerLength);
+    if (hasPeerId) {
+        if (peerId != self.peerId) {
+            if (error) {
+                *error = TunnelKitErrorWithCode(TunnelKitErrorCodeDataPathPeerIdMismatch);
+            }
+            return NO;
+        }
+        flags.ad = packet.bytes;
+        flags.adLength = headerLength + PacketIdLength;
+    }
+    else {
+        flags.ad = packet.bytes + headerLength;
+        flags.adLength = PacketIdLength;
     }
 
     // skip header + packet id
-    const BOOL success = [self.crypto decryptBytes:(packet.bytes + self.headerLength + PacketIdLength)
+    const BOOL success = [self.crypto decryptBytes:(packet.bytes + headerLength + PacketIdLength)
-                                            length:(int)(packet.length - (self.headerLength + PacketIdLength))
+                                            length:(int)(packet.length - (headerLength + PacketIdLength))
                                               dest:packetBytes
                                         destLength:packetLength
-                                             extra:extra
+                                             flags:&flags
                                              error:error];
     if (!success) {
         return NO;
     }
-    if (self.checkPeerId && !self.checkPeerId(packet.bytes)) {
+    *packetId = ntohl(flags.packetId);
-        if (error) {
-            *error = TunnelKitErrorWithCode(TunnelKitErrorCodeDataPathPeerIdMismatch);
-        }
-        return NO;
-    }
-    *packetId = ntohl(*(const uint32_t *)(packet.bytes + self.headerLength));
     return YES;
 }
 

TunnelKit/Sources/Core/CryptoCBC.h

@@ -49,6 +49,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface DataPathCryptoCBC : NSObject <DataPathEncrypter, DataPathDecrypter>
 
+@property (nonatomic, assign) uint32_t peerId;
+
 - (instancetype)initWithCrypto:(CryptoCBC *)crypto;
 
 @end

TunnelKit/Sources/Core/CryptoCBC.m

@@ -141,7 +141,7 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
     HMAC_Init_ex(self.hmacCtxEnc, hmacKey.bytes, self.hmacKeyLength, self.digest, NULL);
 }
 
-- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)encryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
     uint8_t *outIV = dest + self.digestLength;
     uint8_t *outEncrypted = dest + self.digestLength + self.cipherIVLength;
@@ -200,7 +200,7 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
     HMAC_Init_ex(self.hmacCtxDec, hmacKey.bytes, self.hmacKeyLength, self.digest, NULL);
 }
 
-- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)decryptBytes:(const uint8_t *)bytes length:(NSInteger)length dest:(uint8_t *)dest destLength:(NSInteger *)destLength flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
     NSAssert(self.cipher, @"No cipher provided");
 
@@ -229,7 +229,7 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
     TUNNEL_CRYPTO_RETURN_STATUS(code)
 }
 
-- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length extra:(const uint8_t *)extra error:(NSError *__autoreleasing *)error
+- (BOOL)verifyBytes:(const uint8_t *)bytes length:(NSInteger)length flags:(const CryptoFlags * _Nullable)flags error:(NSError * _Nullable __autoreleasing * _Nullable)error
 {
     int l1 = 0;
     int code = 1;
@@ -260,9 +260,6 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
 @interface DataPathCryptoCBC ()
 
 @property (nonatomic, strong) CryptoCBC *crypto;
-@property (nonatomic, assign) int headerLength;
-@property (nonatomic, copy) void (^setDataHeader)(uint8_t *, uint8_t);
-@property (nonatomic, copy) BOOL (^checkPeerId)(const uint8_t *);
 
 @end
 
@@ -281,24 +278,7 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
 
 - (void)setPeerId:(uint32_t)peerId
 {
-    peerId &= 0xffffff;
+    _peerId = peerId & 0xffffff;
-
-    if (peerId == PacketPeerIdDisabled) {
-        self.headerLength = PacketOpcodeLength;
-        self.setDataHeader = ^(uint8_t *to, uint8_t key) {
-            PacketHeaderSet(to, PacketCodeDataV1, key, nil);
-        };
-        self.checkPeerId = NULL;
-    }
-    else {
-        self.headerLength = PacketOpcodeLength + PacketPeerIdLength;
-        self.setDataHeader = ^(uint8_t *to, uint8_t key) {
-            PacketHeaderSetDataV2(to, key, peerId);
-        };
-        self.checkPeerId = ^BOOL(const uint8_t *ptr) {
-            return (PacketHeaderGetDataV2PeerId(ptr) == peerId);
-        };
-    }
 }
 
 - (NSInteger)encryptionCapacityWithLength:(NSInteger)length
@@ -326,15 +306,17 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
 
 - (NSData *)encryptedDataPacketWithKey:(uint8_t)key packetId:(uint32_t)packetId packetBytes:(const uint8_t *)packetBytes packetLength:(NSInteger)packetLength error:(NSError *__autoreleasing *)error
 {
-    const int capacity = self.headerLength + (int)[self.crypto encryptionCapacityWithLength:packetLength];
+    DATA_PATH_ENCRYPT_INIT(self.peerId)
+
+    const int capacity = headerLength + (int)[self.crypto encryptionCapacityWithLength:packetLength];
     NSMutableData *encryptedPacket = [[NSMutableData alloc] initWithLength:capacity];
     uint8_t *ptr = encryptedPacket.mutableBytes;
     NSInteger encryptedPacketLength = INT_MAX;
     const BOOL success = [self.crypto encryptBytes:packetBytes
                                             length:packetLength
-                                              dest:(ptr + self.headerLength) // skip header byte
+                                              dest:(ptr + headerLength) // skip header bytes
                                         destLength:&encryptedPacketLength
-                                             extra:NULL
+                                             flags:NULL
                                              error:error];
     
     NSAssert(encryptedPacketLength <= capacity, @"Did not allocate enough bytes for payload");
@@ -343,8 +325,13 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
         return nil;
     }
 
-    self.setDataHeader(ptr, key);
+    if (hasPeerId) {
-    encryptedPacket.length = self.headerLength + encryptedPacketLength;
+        PacketHeaderSetDataV2(ptr, key, self.peerId);
+    }
+    else {
+        PacketHeaderSet(ptr, PacketCodeDataV1, key, nil);
+    }
+    encryptedPacket.length = headerLength + encryptedPacketLength;
     return encryptedPacket;
 }
 
@@ -353,24 +340,29 @@ 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);
+    DATA_PATH_DECRYPT_INIT(packet.bytes)
+    if (packet.length < headerLength) {
+        return NO;
+    }
 
     // skip header = (code, key)
-    const BOOL success = [self.crypto decryptBytes:(packet.bytes + self.headerLength)
+    const BOOL success = [self.crypto decryptBytes:(packet.bytes + headerLength)
-                                            length:(int)(packet.length - self.headerLength)
+                                            length:(int)(packet.length - headerLength)
                                               dest:packetBytes
                                         destLength:packetLength
-                                             extra:NULL
+                                             flags:NULL
                                              error:error];
     if (!success) {
         return NO;
     }
-    if (self.checkPeerId && !self.checkPeerId(packet.bytes)) {
+    if (hasPeerId) {
-        if (error) {
+        if (peerId != self.peerId) {
-            *error = TunnelKitErrorWithCode(TunnelKitErrorCodeDataPathPeerIdMismatch);
+            if (error) {
+                *error = TunnelKitErrorWithCode(TunnelKitErrorCodeDataPathPeerIdMismatch);
+            }
+            return NO;
         }
-        return NO;
     }
     *packetId = ntohl(*(uint32_t *)packetBytes);
     return YES;

TunnelKit/Sources/Core/DataPath.m

@@ -86,6 +86,8 @@
     NSParameterAssert(decrypter);
     NSParameterAssert(maxPackets > 0);
     
+    peerId &= 0xffffff;
+    
     if ((self = [super init])) {
         self.encrypter = encrypter;
         self.decrypter = decrypter;

TunnelKit/Sources/Core/DataPathCrypto.h

@@ -38,11 +38,30 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+#define DATA_PATH_ENCRYPT_INIT(peerId) \
+    const BOOL hasPeerId = (peerId != PacketPeerIdDisabled); \
+    int headerLength = PacketOpcodeLength; \
+    if (hasPeerId) { \
+        headerLength += PacketPeerIdLength; \
+    }
+
+#define DATA_PATH_DECRYPT_INIT(ptr) \
+    PacketCode code; \
+    PacketOpcodeGet(ptr, &code, NULL); \
+    uint32_t peerId = PacketPeerIdDisabled; \
+    const BOOL hasPeerId = (code == PacketCodeDataV2); \
+    int headerLength = PacketOpcodeLength; \
+    if (hasPeerId) { \
+        headerLength += PacketPeerIdLength; \
+        peerId = PacketHeaderGetDataV2PeerId(packet.bytes); \
+    }
+
 typedef void (^DataPathAssembleBlock)(uint8_t *packetDest, NSInteger *packetLengthOffset, NSData *payload);
 typedef void (^DataPathParseBlock)(uint8_t *payload, NSInteger *payloadOffset, NSInteger *headerLength, const uint8_t *packet, NSInteger packetLength);
 
 @protocol DataPathChannel
 
+- (uint32_t)peerId;
 - (void)setPeerId:(uint32_t)peerId;
 - (NSInteger)encryptionCapacityWithLength:(NSInteger)length;
 

TunnelKitTests/DataPathEncryptionTests.swift

@@ -99,15 +99,17 @@ class DataPathEncryptionTests: XCTestCase {
         if let peerId = peerId {
             enc.setPeerId(peerId)
             dec.setPeerId(peerId)
-//            XCTAssertEqual(enc.peerId(), peerId & 0xffffff)
+            XCTAssertEqual(enc.peerId(), peerId & 0xffffff)
-//            XCTAssertEqual(dec.peerId(), peerId & 0xffffff)
+            XCTAssertEqual(dec.peerId(), peerId & 0xffffff)
         }
 
         let expectedPayload = Data(hex: "00112233445566778899")
         let key: UInt8 = 4
 
         let encrypted = try! path.encryptPackets([expectedPayload], key: key)
+        print(encrypted.map { $0.toHex() })
         let decrypted = try! path.decryptPackets(encrypted, keepAlive: nil)
+        print(decrypted.map { $0.toHex() })
         let payload = decrypted.first!
 
         XCTAssertEqual(payload, expectedPayload)
@@ -117,8 +119,8 @@ class DataPathEncryptionTests: XCTestCase {
         if let peerId = peerId {
             enc.setPeerId(peerId)
             dec.setPeerId(peerId)
-//            XCTAssertEqual(enc.peerId(), peerId & 0xffffff)
+            XCTAssertEqual(enc.peerId(), peerId & 0xffffff)
-//            XCTAssertEqual(dec.peerId(), peerId & 0xffffff)
+            XCTAssertEqual(dec.peerId(), peerId & 0xffffff)
         }
 
         let expectedPayload = Data(hex: "00112233445566778899")

TunnelKitTests/EncryptionPerformanceTests.swift

@@ -72,7 +72,7 @@ class EncryptionPerformanceTests: XCTestCase {
         let suite = TestUtils.generateDataSuite(1000, 100000)
         measure {
             for data in suite {
-                let _ = try! self.cbcEncrypter.encryptData(data, extra: nil)
+                let _ = try! self.cbcEncrypter.encryptData(data, flags: nil)
             }
         }
     }
@@ -80,10 +80,11 @@ class EncryptionPerformanceTests: XCTestCase {
     // 0.684s
     func testGCMEncryption() {
         let suite = TestUtils.generateDataSuite(1000, 100000)
-        let extra: [UInt8] = [0x11, 0x22, 0x33, 0x44]
+        let ad: [UInt8] = [0x11, 0x22, 0x33, 0x44]
+        var flags = CryptoFlags(packetId: 0, ad: ad, adLength: 4)
         measure {
             for data in suite {
-                let _ = try! self.gcmEncrypter.encryptData(data, extra: extra)
+                let _ = try! self.gcmEncrypter.encryptData(data, flags: &flags)
             }
         }
     }

TunnelKitTests/EncryptionTests.swift

@@ -63,8 +63,8 @@ class EncryptionTests: XCTestCase {
         let (client, server) = clientServer("aes-128-cbc", "sha256")
 
         let plain = Data(hex: "00112233445566778899")
-        let encrypted = try! client.encrypter().encryptData(plain, extra: nil)
+        let encrypted = try! client.encrypter().encryptData(plain, flags: nil)
-        let decrypted = try! server.decrypter().decryptData(encrypted, extra: nil)
+        let decrypted = try! server.decrypter().decryptData(encrypted, flags: nil)
         XCTAssertEqual(plain, decrypted)
     }
 
@@ -72,18 +72,19 @@ class EncryptionTests: XCTestCase {
         let (client, server) = clientServer(nil, "sha256")
 
         let plain = Data(hex: "00112233445566778899")
-        let encrypted = try! client.encrypter().encryptData(plain, extra: nil)
+        let encrypted = try! client.encrypter().encryptData(plain, flags: nil)
-        XCTAssertNoThrow(try server.decrypter().verifyData(encrypted, extra: nil))
+        XCTAssertNoThrow(try server.decrypter().verifyData(encrypted, flags: nil))
     }
     
     func testGCM() {
         let (client, server) = clientServer("aes-256-gcm", nil)
         
 //        let packetId: UInt32 = 0x56341200
-        let extra: [UInt8] = [0x00, 0x12, 0x34, 0x56]
+        let ad: [UInt8] = [0x00, 0x12, 0x34, 0x56]
+        var flags = CryptoFlags(packetId: 0, ad: ad, adLength: 4)
         let plain = Data(hex: "00112233445566778899")
-        let encrypted = try! client.encrypter().encryptData(plain, extra: extra)
+        let encrypted = try! client.encrypter().encryptData(plain, flags: &flags)
-        let decrypted = try! server.decrypter().decryptData(encrypted, extra: extra)
+        let decrypted = try! server.decrypter().decryptData(encrypted, flags: &flags)
         XCTAssertEqual(plain, decrypted)
     }
     

TunnelKitTests/TestUtils.swift

@@ -57,12 +57,12 @@ class TestUtils {
 }
 
 extension Encrypter {
-    func encryptData(_ data: Data, extra: [UInt8]?) throws -> Data {
+    func encryptData(_ data: Data, flags: UnsafePointer<CryptoFlags>?) throws -> Data {
         let srcLength = data.count
         var dest: [UInt8] = Array(repeating: 0, count: srcLength + 256)
         var destLength = 0
         try data.withUnsafeBytes {
-            try encryptBytes($0, length: srcLength, dest: &dest, destLength: &destLength, extra: extra)
+            try encryptBytes($0, length: srcLength, dest: &dest, destLength: &destLength, flags: flags)
         }
         dest.removeSubrange(destLength..<dest.count)
         return Data(dest)
@@ -70,21 +70,21 @@ extension Encrypter {
 }
 
 extension Decrypter {
-    func decryptData(_ data: Data, extra: [UInt8]?) throws -> Data {
+    func decryptData(_ data: Data, flags: UnsafePointer<CryptoFlags>?) throws -> Data {
         let srcLength = data.count
         var dest: [UInt8] = Array(repeating: 0, count: srcLength + 256)
         var destLength = 0
         try data.withUnsafeBytes {
-            try decryptBytes($0, length: srcLength, dest: &dest, destLength: &destLength, extra: extra)
+            try decryptBytes($0, length: srcLength, dest: &dest, destLength: &destLength, flags: flags)
         }
         dest.removeSubrange(destLength..<dest.count)
         return Data(dest)
     }
     
-    func verifyData(_ data: Data, extra: [UInt8]?) throws {
+    func verifyData(_ data: Data, flags: UnsafePointer<CryptoFlags>?) throws {
         let srcLength = data.count
         try data.withUnsafeBytes {
-            try verifyBytes($0, length: srcLength, extra: extra)
+            try verifyBytes($0, length: srcLength, flags: flags)
         }
     }
 }