Move compression framing to stateless blocks

Prepare payload blocks in DataPath.

TunnelKit/Sources/Core/CryptoAEAD.h

@@ -51,8 +51,6 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface DataPathCryptoAEAD : NSObject <DataPathEncrypter, DataPathDecrypter>
 
-@property (nonatomic, assign) CompressionFraming compressionFraming;
-
 - (instancetype)initWithCrypto:(nonnull CryptoAEAD *)crypto;
 
 @end

TunnelKit/Sources/Core/CryptoAEAD.m

@@ -290,32 +290,17 @@ const NSInteger CryptoAEADTagLength     = 16;
 
 #pragma mark DataPathEncrypter
 
-- (void)assembleDataPacketWithPacketId:(uint32_t)packetId payload:(NSData *)payload into:(uint8_t *)dest length:(NSInteger *)length
+- (void)assembleDataPacketWithBlock:(DataPathAssembleBlock)block packetId:(uint32_t)packetId payload:(NSData *)payload into:(uint8_t *)dest length:(NSInteger *)length
 {
-    uint8_t *ptr = dest;
+    *length = payload.length;
-    *length = (int)(ptr - dest + payload.length);
+    if (!block) {
-
+        memcpy(dest, payload.bytes, payload.length);
-    switch (self.compressionFraming) {
+        return;
-        case CompressionFramingDisabled:
-            memcpy(ptr, payload.bytes, payload.length);
-            break;
-            
-        case CompressionFramingCompress:
-            memcpy(ptr, payload.bytes, payload.length);
-            ptr[payload.length] = *ptr;
-            *ptr = CompressionFramingNoCompressSwap;
-            *length += sizeof(uint8_t);
-            break;
-            
-        case CompressionFramingCompLZO:
-            memcpy(ptr + sizeof(uint8_t), payload.bytes, payload.length);
-            *ptr = CompressionFramingNoCompress;
-            *length += sizeof(uint8_t);
-            break;
-            
-        default:
-            break;
     }
+
+    NSInteger packetLengthOffset;
+    block(dest, &packetLengthOffset, payload);
+    *length += packetLengthOffset;
 }
 
 - (NSData *)encryptedDataPacketWithKey:(uint8_t)key packetId:(uint32_t)packetId payload:(const uint8_t *)payload payloadLength:(NSInteger)payloadLength error:(NSError *__autoreleasing *)error
@@ -379,27 +364,19 @@ const NSInteger CryptoAEADTagLength     = 16;
     return YES;
 }
 
-- (const uint8_t *)parsePayloadWithDataPacket:(uint8_t *)packet packetLength:(NSInteger)packetLength length:(NSInteger *)length
+- (const uint8_t *)parsePayloadWithBlock:(DataPathParseBlock)block dataPacket:(uint8_t *)packet packetLength:(NSInteger)packetLength length:(NSInteger *)length
 {
-    uint8_t *ptr = packet;
+    uint8_t *payload = packet;
-    *length = packetLength - (int)(ptr - packet);
+    *length = packetLength - (int)(payload - packet);
-    if (self.compressionFraming != CompressionFramingDisabled) {
+    if (!block) {
-        switch (*ptr) {
+        return payload;
-            case CompressionFramingNoCompress:
-                ptr += sizeof(uint8_t);
-                break;
-                
-            case CompressionFramingNoCompressSwap:
-                *ptr = packet[packetLength - 1];
-                break;
-                
-            default:
-                NSAssert(NO, @"Compression not supported (found %X)", *ptr);
-                break;
-        }
-        *length -= sizeof(uint8_t);
     }
-    return ptr;
+    
+    NSInteger payloadOffset;
+    NSInteger payloadHeaderLength;
+    block(payload, &payloadOffset, &payloadHeaderLength, packet, packetLength);
+    *length -= payloadHeaderLength;
+    return payload + payloadOffset;
 }
 
 @end

TunnelKit/Sources/Core/CryptoCBC.h

@@ -50,8 +50,6 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface DataPathCryptoCBC : NSObject <DataPathEncrypter, DataPathDecrypter>
 
-@property (nonatomic, assign) CompressionFraming compressionFraming;
-
 - (instancetype)initWithCrypto:(nonnull CryptoCBC *)crypto;
 
 @end

TunnelKit/Sources/Core/CryptoCBC.m

@@ -288,34 +288,20 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
 
 #pragma mark DataPathEncrypter
 
-- (void)assembleDataPacketWithPacketId:(uint32_t)packetId payload:(NSData *)payload into:(uint8_t *)dest length:(NSInteger *)length
+- (void)assembleDataPacketWithBlock:(DataPathAssembleBlock)block packetId:(uint32_t)packetId payload:(NSData *)payload into:(uint8_t *)dest length:(NSInteger *)length
 {
     uint8_t *ptr = dest;
     *(uint32_t *)ptr = htonl(packetId);
     ptr += sizeof(uint32_t);
     *length = (int)(ptr - dest + payload.length);
-
+    if (!block) {
-    switch (self.compressionFraming) {
+        memcpy(ptr, payload.bytes, payload.length);
-        case CompressionFramingDisabled:
+        return;
-            memcpy(ptr, payload.bytes, payload.length);
-            break;
-            
-        case CompressionFramingCompress:
-            memcpy(ptr, payload.bytes, payload.length);
-            ptr[payload.length] = *ptr;
-            *ptr = CompressionFramingNoCompressSwap;
-            *length += sizeof(uint8_t);
-            break;
-        
-        case CompressionFramingCompLZO:
-            memcpy(ptr + sizeof(uint8_t), payload.bytes, payload.length);
-            *ptr = CompressionFramingNoCompress;
-            *length += sizeof(uint8_t);
-            break;
-            
-        default:
-            break;
     }
+
+    NSInteger packetLengthOffset;
+    block(ptr, &packetLengthOffset, payload);
+    *length += packetLengthOffset;
 }
 
 - (NSData *)encryptedDataPacketWithKey:(uint8_t)key packetId:(uint32_t)packetId payload:(const uint8_t *)payload payloadLength:(NSInteger)payloadLength error:(NSError *__autoreleasing *)error
@@ -366,28 +352,20 @@ const NSInteger CryptoCBCMaxHMACLength = 100;
     return YES;
 }
 
-- (const uint8_t *)parsePayloadWithDataPacket:(uint8_t *)packet packetLength:(NSInteger)packetLength length:(NSInteger *)length
+- (const uint8_t *)parsePayloadWithBlock:(DataPathParseBlock)block dataPacket:(uint8_t *)packet packetLength:(NSInteger)packetLength length:(NSInteger *)length
 {
-    uint8_t *ptr = packet;
+    uint8_t *payload = packet;
-    ptr += sizeof(uint32_t); // packet id
+    payload += sizeof(uint32_t); // packet id
-    *length = packetLength - (int)(ptr - packet);
+    *length = packetLength - (int)(payload - packet);
-    if (self.compressionFraming != CompressionFramingDisabled) {
+    if (!block) {
-        switch (*ptr) {
+        return payload;
-            case CompressionFramingNoCompress:
-                ptr += sizeof(uint8_t);
-                break;
-
-            case CompressionFramingNoCompressSwap:
-                *ptr = packet[packetLength - 1];
-                break;
-                
-            default:
-                NSAssert(NO, @"Compression not supported (found %X)", *ptr);
-                break;
-        }
-        *length -= sizeof(uint8_t);
     }
-    return ptr;
+
+    NSInteger payloadOffset;
+    NSInteger payloadHeaderLength;
+    block(payload, &payloadOffset, &payloadHeaderLength, packet, packetLength);
+    *length -= payloadHeaderLength;
+    return payload + payloadOffset;
 }
 
 @end

TunnelKit/Sources/Core/DataPath.m

@@ -64,6 +64,9 @@
 @property (nonatomic, assign) int decBufferCapacity;
 @property (nonatomic, strong) ReplayProtector *inReplay;
 
+@property (nonatomic, copy) DataPathAssembleBlock assemblePayloadBlock;
+@property (nonatomic, copy) DataPathParseBlock parsePayloadBlock;
+
 @end
 
 @implementation DataPath
@@ -105,6 +108,8 @@
         if (usesReplayProtection) {
             self.inReplay = [[ReplayProtector alloc] init];
         }
+
+        self.compressionFraming = CompressionFramingDisabled;
     }
     return self;
 }
@@ -162,11 +167,47 @@
 
 - (void)setCompressionFraming:(CompressionFraming)compressionFraming
 {
-    NSAssert(self.encrypter, @"Setting compressionFraming to nil encrypter");
+    switch (compressionFraming) {
-    NSAssert(self.decrypter, @"Setting compressionFraming to nil decrypter");
+        case CompressionFramingDisabled: {
-    
+            self.assemblePayloadBlock = ^(uint8_t * _Nonnull packetDest, NSInteger * _Nonnull packetLengthOffset, NSData * _Nonnull payload) {
-    [self.encrypter setCompressionFraming:compressionFraming];
+                memcpy(packetDest, payload.bytes, payload.length);
-    [self.decrypter setCompressionFraming:compressionFraming];
+                *packetLengthOffset = 0;
+            };
+            self.parsePayloadBlock = ^(uint8_t * _Nonnull payload, NSInteger *payloadOffset, NSInteger * _Nonnull headerLength, const uint8_t * _Nonnull packet, NSInteger packetLength) {
+                *payloadOffset = 0;
+                *headerLength = 0;
+            };
+            break;
+        }
+        case CompressionFramingCompress: {
+            self.assemblePayloadBlock = ^(uint8_t * _Nonnull packetDest, NSInteger * _Nonnull packetLengthOffset, NSData * _Nonnull payload) {
+                memcpy(packetDest, payload.bytes, payload.length);
+                packetDest[payload.length] = packetDest[0];
+                packetDest[0] = CompressionFramingNoCompressSwap;
+                *packetLengthOffset = 1;
+            };
+            self.parsePayloadBlock = ^(uint8_t * _Nonnull payload, NSInteger *payloadOffset, NSInteger * _Nonnull headerLength, const uint8_t * _Nonnull packet, NSInteger packetLength) {
+                NSCAssert(payload[0] == CompressionFramingNoCompressSwap, @"Expected NO_COMPRESS_SWAP (found %X != %X)", payload[0], CompressionFramingNoCompressSwap);
+                payload[0] = packet[packetLength - 1];
+                *payloadOffset = 0;
+                *headerLength = 1;
+            };
+            break;
+        }
+        case CompressionFramingCompLZO: {
+            self.assemblePayloadBlock = ^(uint8_t * _Nonnull packetDest, NSInteger * _Nonnull packetLengthOffset, NSData * _Nonnull payload) {
+                memcpy(packetDest + 1, payload.bytes, payload.length);
+                packetDest[0] = CompressionFramingNoCompress;
+                *packetLengthOffset = 1;
+            };
+            self.parsePayloadBlock = ^(uint8_t * _Nonnull payload, NSInteger *payloadOffset, NSInteger * _Nonnull headerLength, const uint8_t * _Nonnull packet, NSInteger packetLength) {
+                NSCAssert(payload[0] == CompressionFramingNoCompress, @"Expected NO_COMPRESS (found %X != %X)", payload[0], CompressionFramingNoCompress);
+                *payloadOffset = 1;
+                *headerLength = 1;
+            };
+            break;
+        }
+    }
 }
 
 #pragma mark DataPath
@@ -192,10 +233,11 @@
         
         uint8_t *payload = self.encBufferAligned;
         NSInteger payloadLength;
-        [self.encrypter assembleDataPacketWithPacketId:self.outPacketId
+        [self.encrypter assembleDataPacketWithBlock:self.assemblePayloadBlock
-                                               payload:raw
+                                           packetId:self.outPacketId
-                                                  into:payload
+                                            payload:raw
-                                                length:&payloadLength];
+                                               into:payload
+                                             length:&payloadLength];
         MSSFix(payload, payloadLength);
         
         NSData *encryptedPacket = [self.encrypter encryptedDataPacketWithKey:key
@@ -246,9 +288,10 @@
         }
         
         NSInteger payloadLength;
-        const uint8_t *payload = [self.decrypter parsePayloadWithDataPacket:packet
+        const uint8_t *payload = [self.decrypter parsePayloadWithBlock:self.parsePayloadBlock
-                                                               packetLength:packetLength
+                                                            dataPacket:packet
-                                                                     length:&payloadLength];
+                                                          packetLength:packetLength
+                                                                length:&payloadLength];
         
         if ((payloadLength == sizeof(DataPacketPingData)) && !memcmp(payload, DataPacketPingData, payloadLength)) {
             if (keepAlive) {

TunnelKit/Sources/Core/DataPathEncryption.h

@@ -36,18 +36,19 @@
 
 #import <Foundation/Foundation.h>
 
+typedef void (^DataPathAssembleBlock)(uint8_t *_Nonnull packetDest, NSInteger *_Nonnull packetLengthOffset, NSData *_Nonnull payload);
+typedef void (^DataPathParseBlock)(uint8_t *_Nonnull payload, NSInteger *_Nonnull payloadOffset, NSInteger *_Nonnull headerLength, const uint8_t *_Nonnull packet, NSInteger packetLength);
+
 @protocol DataPathChannel
 
 - (int)overheadLength;
 - (void)setPeerId:(uint32_t)peerId;
-- (CompressionFraming)compressionFraming;
-- (void)setCompressionFraming:(CompressionFraming)compressionFraming;
 
 @end
 
 @protocol DataPathEncrypter <DataPathChannel>
 
-- (void)assembleDataPacketWithPacketId:(uint32_t)packetId payload:(NSData *)payload into:(nonnull uint8_t *)dest length:(nonnull NSInteger *)length;
+- (void)assembleDataPacketWithBlock:(DataPathAssembleBlock)block packetId:(uint32_t)packetId payload:(NSData *)payload into:(nonnull uint8_t *)dest length:(nonnull NSInteger *)length;
 - (NSData *)encryptedDataPacketWithKey:(uint8_t)key packetId:(uint32_t)packetId payload:(const uint8_t *)payload payloadLength:(NSInteger)payloadLength error:(NSError **)error;
 
 @end
@@ -55,6 +56,6 @@
 @protocol DataPathDecrypter <DataPathChannel>
 
 - (BOOL)decryptDataPacket:(nonnull NSData *)packet into:(nonnull uint8_t *)dest length:(nonnull NSInteger *)length packetId:(nonnull uint32_t *)packetId error:(NSError **)error;
-- (nonnull const uint8_t *)parsePayloadWithDataPacket:(nonnull uint8_t *)packet packetLength:(NSInteger)packetLength length:(nonnull NSInteger *)length;
+- (nonnull const uint8_t *)parsePayloadWithBlock:(DataPathParseBlock)block dataPacket:(nonnull uint8_t *)packet packetLength:(NSInteger)packetLength length:(nonnull NSInteger *)length;
 
 @end