From 81997ed93745b1c00c51b8e9fef3c7019fffeb4f Mon Sep 17 00:00:00 2001
From: Jedihy <huanyi@microsoft.com>
Date: Wed, 10 Feb 2021 11:21:18 -0800
Subject: [PATCH] Fix TotalLength calculation bug that leads to corrupt blocks
 (#34)

* fix comment length calc bug

* update readme
---
 README.md    |  2 ++
 src/main.c   |  4 ++--
 src/pcapng.h | 14 +++++++-------
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index a0e5627..7335ae9 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,8 @@ msbuild -t:rebuild -p:configuration=release -p:platform=x64
 
 # History
 
+1.4.1 - Fix a bug leading to writing corrupt packets.
+
 1.4.0 - Automatically infer original fragment length if captured fragments were truncated.
 
 1.3.0 - Add a comment to each packet containing the process id (PID).
diff --git a/src/main.c b/src/main.c
index cab55a7..7f7ee3d 100644
--- a/src/main.c
+++ b/src/main.c
@@ -452,7 +452,7 @@ void WINAPI EventCallback(PEVENT_RECORD ev)
             TimeStamp.LowPart,
             CommentLength > 0 ? (char*)&Comment : NULL,
             (unsigned short)CommentLength);
-        
+
         AuxFragBufOffset = 0;
         NumFramesConverted++;
     } else {
@@ -471,7 +471,7 @@ int __cdecl wmain(int argc, wchar_t** argv)
     if (argc == 2 &&
         (!wcscmp(argv[1], L"-v") ||
          !wcscmp(argv[1], L"--version"))) {
-        printf("etl2pcapng version 1.4.0\n");
+        printf("etl2pcapng version 1.4.1\n");
         return 0;
     }
 
diff --git a/src/pcapng.h b/src/pcapng.h
index 03bcff9..25fe970 100644
--- a/src/pcapng.h
+++ b/src/pcapng.h
@@ -155,11 +155,11 @@ inline int
 PcapNgWriteCommentOption(
     __in HANDLE File,
     __in PCHAR CommentBuffer,
-    __in unsigned short CommentLength
+    __in unsigned short CommentLength,
+    __in int CommentPadLength
     )
 {
     int Err = NO_ERROR;
-    int CommentPadLength = 4 - (CommentLength % 4 == 0 ? 4 : CommentLength % 4);
     struct PCAPNG_BLOCK_OPTION_COMMENT Comment;
     char Pad[4] = { 0 };
 
@@ -211,14 +211,13 @@ PcapNgWriteEnhancedPacket(
     struct PCAPNG_BLOCK_TAIL Tail;
     char Pad[4] = {0};
     BOOLEAN CommentProvided = (CommentLength > 0 && Comment != NULL);
-    int FragPadLength = (4 - ((sizeof(Body) + FragLength) & 3)) & 3; // pad to 4 bytes per the spec.
+    int CommentPadLength = (4 - (CommentLength & 3)) & 3; // pad to 4 bytes per the spec.
+    int FragPadLength = (4 - ((sizeof(Body) + FragLength) & 3)) & 3;
     int TotalLength =
         sizeof(Head) + sizeof(Body) + FragLength + FragPadLength +
         sizeof(EpbFlagsOption) + sizeof(EndOption) + sizeof(Tail) +
         (CommentProvided ?
-            sizeof(struct PCAPNG_BLOCK_OPTION_COMMENT) + sizeof(EndOption) + CommentLength +
-            (4 - (CommentLength % 4 == 0 ? 4 : CommentLength % 4)) //Comment Padding
-            : 0);
+            sizeof(struct PCAPNG_BLOCK_OPTION_COMMENT) + CommentLength + CommentPadLength : 0);
 
     Head.Type = PCAPNG_BLOCKTYPE_ENHANCED_PACKET;
     Head.Length = TotalLength;
@@ -264,7 +263,8 @@ PcapNgWriteEnhancedPacket(
         Err = PcapNgWriteCommentOption(
             File,
             Comment,
-            CommentLength);
+            CommentLength,
+            CommentPadLength);
         if (Err != NO_ERROR) {
             printf("WriteFile failed with %u\n", Err);
             goto Done;