Speed up PE checksum calculation

cmake/CMakeTest.cmake

@@ -43,6 +43,7 @@ if(Python3_FOUND)
             "${CMAKE_CURRENT_SOURCE_DIR}/tests/conf"
             "${CMAKE_CURRENT_SOURCE_DIR}/tests/client_http.py"
             "${CMAKE_CURRENT_SOURCE_DIR}/tests/make_certificates.py"
+            "${CMAKE_CURRENT_SOURCE_DIR}/tests/check_pe_checksum.py"
             "${CMAKE_CURRENT_SOURCE_DIR}/tests/start_server.py"
             "${CMAKE_CURRENT_SOURCE_DIR}/tests/exec.py"
             DESTINATION "${TEST_DIR}/")
@@ -248,6 +249,40 @@ if(Python3_FOUND AND NOT cryptography_error)
     endforeach(ext ${extensions_all})
 
 
+### PE checksum ###
+
+    # The CheckSum field of a signed PE file is compared with an independent
+    # computation, for overlays of odd and even sizes and of different contents.
+    set(checksum_cases "1_255" "3_0" "65537_255" "300001_127")
+    foreach(case ${checksum_cases})
+        string(REPLACE "_" ";" case_args ${case})
+        list(GET case_args 0 case_size)
+        list(GET case_args 1 case_fill)
+        add_test(NAME "checksum_${case}_create"
+            COMMAND ${Python3_EXECUTABLE} "${TEST_DIR}/check_pe_checksum.py"
+            "make" "${FILES}/unsigned.exe" "${FILES}/checksum_${case}.exe"
+            ${case_size} ${case_fill})
+        add_test(NAME "checksum_${case}_sign"
+            COMMAND ${Python3_EXECUTABLE} ${EXEC} ${OSSLSIGNCODE} "sign"
+            "-pkcs12" "${CERTS}/cert.p12"
+            "-readpass" "${CERTS}/password.txt"
+            "-in" "${FILES}/checksum_${case}.exe"
+            "-out" "${FILES}/checksum_${case}_signed.exe")
+        add_test(NAME "checksum_${case}_check"
+            COMMAND ${Python3_EXECUTABLE} "${TEST_DIR}/check_pe_checksum.py"
+            "check" "${FILES}/checksum_${case}_signed.exe")
+        set_tests_properties("checksum_${case}_create" PROPERTIES
+            DEPENDS "start_server")
+        set_tests_properties("checksum_${case}_sign" PROPERTIES
+            DEPENDS "checksum_${case}_create")
+        set_tests_properties("checksum_${case}_check" PROPERTIES
+            DEPENDS "checksum_${case}_sign")
+        list(APPEND ALL_TESTS "checksum_${case}_create" "checksum_${case}_sign" "checksum_${case}_check")
+        set(OUTPUT_FILES ${OUTPUT_FILES}
+            "${FILES}/checksum_${case}.exe" "${FILES}/checksum_${case}_signed.exe")
+    endforeach(case ${checksum_cases})
+
+
 ### Verify signature ###
 
     # Verify PE/MSI/CAB files signed in the catalog file

pe.c

@@ -83,6 +83,8 @@ FILE_FORMAT file_format_pe = {
 /* Prototypes */
 static PE_CTX *pe_ctx_get(char *indata, uint32_t filesize);
 static PKCS7 *pe_pkcs7_get_file(char *indata, PE_CTX *pe_ctx);
+static uint64_t pe_checksum_sum(const u_char *data, size_t nwords, uint32_t offset, uint32_t header_size);
+static uint32_t pe_checksum_fold(uint64_t sum, uint32_t offset);
 static uint32_t pe_calc_checksum(BIO *bio, uint32_t header_size);
 static uint32_t pe_calc_realchecksum(FILE_FORMAT_CTX *ctx);
 static int pe_modify_header(FILE_FORMAT_CTX *ctx, BIO *hash, BIO *outdata);
@@ -625,6 +627,49 @@ static PKCS7 *pe_pkcs7_get_file(char *indata, PE_CTX *pe_ctx)
     return NULL; /* FAILED */
 }
 
+/*
+ * Sum of the 16-bit little-endian words of a data block, excluding the two
+ * words of the CheckSum field (header_size + 88 and header_size + 90).
+ * [in] data: data block
+ * [in] nwords: number of 16-bit words
+ * [in] offset: file offset of the first word
+ * [in] header_size: PE header size
+ * [returns] sum of the words
+ */
+static uint64_t pe_checksum_sum(const u_char *data, size_t nwords, uint32_t offset, uint32_t header_size)
+{
+    uint64_t sum = 0;
+    uint32_t pos[2];
+    size_t i;
+
+    for (i = 0; i < nwords; i++)
+        sum += (uint32_t)data[2 * i] | ((uint32_t)data[2 * i + 1] << 8);
+    pos[0] = header_size + 88;
+    pos[1] = header_size + 90;
+    for (i = 0; i < 2; i++) {
+        /* words start at even offsets of the file */
+        if (pos[i] >= offset && (uint64_t)(pos[i] - offset) / 2 < nwords && (pos[i] - offset) % 2 == 0) {
+            const u_char *w = data + (pos[i] - offset);
+            sum -= (uint32_t)w[0] | ((uint32_t)w[1] << 8);
+        }
+    }
+    return sum;
+}
+
+/*
+ * Fold the sum of words into a 16-bit value (one's complement sum)
+ * and add the file length.
+ * [in] sum: sum of 16-bit words
+ * [in] offset: number of bytes summed
+ * [returns] checksum
+ */
+static uint32_t pe_checksum_fold(uint64_t sum, uint32_t offset)
+{
+    uint32_t checkSum = sum == 0 ? 0 : (uint32_t)(1 + (sum - 1) % 0xFFFF);
+
+    return checkSum + offset;
+}
+
 /*
  * Calculate checksum.
  * A signed PE file is padded (with 0's) to 8 byte boundary,
@@ -635,28 +680,19 @@ static PKCS7 *pe_pkcs7_get_file(char *indata, PE_CTX *pe_ctx)
  */
 static uint32_t pe_calc_checksum(BIO *outdata, uint32_t header_size)
 {
-    uint32_t checkSum = 0, offset = 0;
+    uint64_t sum = 0;
+    uint32_t offset = 0;
     int nread;
-    unsigned short *buf = OPENSSL_malloc(SIZE_64K);
+    u_char *buf = OPENSSL_malloc(SIZE_64K);
 
     /* recalculate the checksum */
     (void)BIO_seek(outdata, 0);
     while ((nread = BIO_read(outdata, buf, SIZE_64K)) > 0) {
-        unsigned short val;
-        int i;
-        for (i = 0; i < nread / 2; i++) {
-            val = LE_UINT16(buf[i]);
-            if (offset == header_size + 88 || offset == header_size + 90)
-                val = 0;
-            checkSum += val;
-            checkSum = LOWORD(LOWORD(checkSum) + HIWORD(checkSum));
-            offset += 2;
-        }
+        sum += pe_checksum_sum(buf, (size_t)nread / 2, offset, header_size);
+        offset += (uint32_t)(nread / 2) * 2;
     }
     OPENSSL_free(buf);
-    checkSum = LOWORD(LOWORD(checkSum) + HIWORD(checkSum));
-    checkSum += offset;
-    return checkSum;
+    return pe_checksum_fold(sum, offset);
 }
 
 /*
@@ -666,40 +702,11 @@ static uint32_t pe_calc_checksum(BIO *outdata, uint32_t header_size)
  */
 static uint32_t pe_calc_realchecksum(FILE_FORMAT_CTX *ctx)
 {
-    uint32_t n = 0, checkSum = 0, offset = 0;
-    BIO *bio = BIO_new(BIO_s_mem());
-    unsigned short *buf = OPENSSL_malloc(SIZE_64K);
-
-    /* calculate the checkSum */
-    while (n < ctx->pe_ctx->fileend) {
-        size_t i, written, nread;
-        size_t left = ctx->pe_ctx->fileend - n;
-        unsigned short val;
-        if (left > SIZE_64K)
-            left = SIZE_64K;
-        if (!BIO_write_ex(bio, ctx->options->indata + n, left, &written))
-            goto err; /* FAILED */
-        (void)BIO_seek(bio, 0);
-        n += (uint32_t)written;
-        if (!BIO_read_ex(bio, buf, written, &nread))
-            goto err; /* FAILED */
-        for (i = 0; i < nread / 2; i++) {
-            val = LE_UINT16(buf[i]);
-            if (offset == ctx->pe_ctx->header_size + 88
-                || offset == ctx->pe_ctx->header_size + 90) {
-                val = 0;
-            }
-            checkSum += val;
-            checkSum = LOWORD(LOWORD(checkSum) + HIWORD(checkSum));
-            offset += 2;
-        }
-    }
-    checkSum = LOWORD(LOWORD(checkSum) + HIWORD(checkSum));
-    checkSum += offset;
-err:
-    OPENSSL_free(buf);
-    BIO_free(bio);
-    return checkSum;
+    size_t nwords = ctx->pe_ctx->fileend / 2;
+    uint64_t sum = pe_checksum_sum((const u_char *)ctx->options->indata, nwords, 0,
+        ctx->pe_ctx->header_size);
+
+    return pe_checksum_fold(sum, (uint32_t)(nwords * 2));
 }
 
 /*

tests/check_pe_checksum.py

@@ -0,0 +1,70 @@
+#!/usr/bin/python3
+"""Check the CheckSum field of a PE file against an independent computation."""
+
+import struct
+import sys
+
+
+def pe_checksum(data: bytes) -> int:
+    """Compute the PE checksum as documented for the CheckSum field.
+
+    Add the 16-bit little-endian words of the file, without the CheckSum field,
+    with an end-around carry, then add the length of the file."""
+    pe_offset = struct.unpack_from("<I", data, 60)[0]
+    field = pe_offset + 24 + 64  # OptionalHeader.CheckSum
+    total = 0
+    for pos in range(0, len(data) - 1, 2):
+        if pos in (field, field + 2):
+            continue
+        total += data[pos] | (data[pos + 1] << 8)
+        total = (total & 0xFFFF) + (total >> 16)
+    total = (total & 0xFFFF) + (total >> 16)
+    return total + len(data)
+
+
+def check(path: str) -> int:
+    """Return 0 when the stored checksum is correct."""
+    with open(path, "rb") as file:
+        data = file.read()
+    pe_offset = struct.unpack_from("<I", data, 60)[0]
+    stored = struct.unpack_from("<I", data, pe_offset + 24 + 64)[0]
+    expected = pe_checksum(data)
+    if stored != expected:
+        print("{}: checksum 0x{:08X}, expected 0x{:08X}".format(path, stored, expected),
+              file=sys.stderr)
+        return 1
+    return 0
+
+
+def make(src: str, dst: str, size: int, fill: int) -> None:
+    """Append size bytes with the value fill to a PE file, as an overlay."""
+    with open(src, "rb") as file:
+        data = file.read()
+    with open(dst, "wb") as file:
+        file.write(data + bytes([fill]) * size)
+
+
+def main() -> None:
+    """Parse the command line."""
+    if len(sys.argv) == 3 and sys.argv[1] == "check":
+        sys.exit(check(sys.argv[2]))
+    if len(sys.argv) == 6 and sys.argv[1] == "make":
+        make(sys.argv[2], sys.argv[3], int(sys.argv[4]), int(sys.argv[5]))
+        return
+    print("Usage:\n\t{0} check FILE\n\t{0} make IN OUT SIZE FILL".format(sys.argv[0]),
+          file=sys.stderr)
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
+
+
+# pylint: disable=pointless-string-statement
+"""Local Variables:
+    c-basic-offset: 4
+    tab-width: 4
+    indent-tabs-mode: no
+End:
+    vim: set ts=4 expandtab:
+"""