PE Checksum Algorithm的較簡實作

這篇BLOG是我很早以前寫的，因為現在搬移到CNBLOGS了，經過整理後重新發出來。

工作之前的幾年一直都在搞計算機安全/病毒相關的東西（純學習，不作惡），其中PE檔案格式是必須知識。有些PE檔案，比如驅動，系統會在加載時對checksum進行校驗，確定驅動檔案的完整性。關于PE檔案如何校驗，網上有很多資料可以學習，這裡有一篇文章《An Analysis of the Windows PE Checksum Algorithm》是對WINDOWS API  CheckSumMappedFile進行逆向分析的。文章的結尾提到WINDOWS的這個校驗和算法和IP協定的校驗和算法類似，IP的校驗和算法實作是RFC1071,如果對其他的校驗和算法感興趣，可以閱讀WIKI的《Error Detection and correction》。

但是CheckSumMappedFile的實作略顯複雜，不夠直覺，後來讀WRK的代碼時，發現了WINDOWS核心在加載驅動時實作的校驗和算法要簡潔直覺很多，分享給大家：

uint32_t calc_checksum(uint32_t checksum, void *data, int length) {
    if (length && data != nullptr) {
        uint32_t sum = 0;
        do {
            sum = *(uint16_t *)data + checksum;
            checksum = (uint16_t)sum + (sum >> 16);
            data = (char *)data + 2;
        } while (--length);
    }

    return checksum + (checksum >> 16);
}

uint32_t generate_pe_checksum(void *file_base, uint32_t file_size) {
    uint32_t file_checksum = 0;
    PIMAGE_NT_HEADERS nt_headers = ImageNtHeader(file_base);
    if (nt_headers) {
        uint32_t header_size = (uintptr_t)nt_headers - (uintptr_t)file_base +
            ((uintptr_t)&nt_headers->OptionalHeader.CheckSum -
            (uintptr_t)nt_headers);
        uint32_t remain_size = (file_size - header_size - 4) >> 1;
        void *remain = &nt_headers->OptionalHeader.Subsystem;
        uint32_t header_checksum = calc_checksum(0, file_base, header_size >> 1);
        file_checksum = calc_checksum(header_checksum, remain, remain_size);
        if (file_size & 1){
            file_checksum += (uint16_t)*((char *)file_base + file_size - 1);
        }
    }

    return (file_size + file_checksum);
}