imagehlp: Implement ImageGetDigestStream.

Mostly based on patch from Juan Lang.

dlls/imagehlp/integrity.c

@@ -4,6 +4,8 @@
  *	Copyright 1998	Patrik Stridvall
  *	Copyright 2003	Mike McCormack
  *	Copyright 2009  Owen Rudge for CodeWeavers
+ *	Copyright 2010  Juan Lang
+ *	Copyright 2010  Andrey Turkin
  *
  * This library is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
@@ -630,18 +632,256 @@ BOOL WINAPI ImageGetCertificateHeader(
     return TRUE;
 }
 
+/* Finds the section named section in the array of IMAGE_SECTION_HEADERs hdr.  If
+ * found, returns the offset to the section.  Otherwise returns 0.  If the section
+ * is found, optionally returns the size of the section (in size) and the base
+ * address of the section (in base.)
+ */
+static DWORD IMAGEHLP_GetSectionOffset( IMAGE_SECTION_HEADER *hdr,
+    DWORD num_sections, LPCSTR section, PDWORD size, PDWORD base )
+{
+    DWORD i, offset = 0;
+
+    for( i = 0; !offset && i < num_sections; i++, hdr++ )
+    {
+        if( !memcmp( hdr->Name, section, strlen(section) ) )
+        {
+            offset = hdr->PointerToRawData;
+            if( size )
+                *size = hdr->SizeOfRawData;
+            if( base )
+                *base = hdr->VirtualAddress;
+        }
+    }
+    return offset;
+}
+
+/* Calls DigestFunction e bytes at offset offset from the file mapped at map.
+ * Returns the return value of DigestFunction, or FALSE if the data is not available.
+ */
+static BOOL IMAGEHLP_ReportSectionFromOffset( DWORD offset, DWORD size,
+    BYTE *map, DWORD fileSize, DIGEST_FUNCTION DigestFunction, DIGEST_HANDLE DigestHandle )
+{
+    if( offset + size > fileSize )
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
+        return FALSE;
+    }
+    return DigestFunction( DigestHandle, map + offset, size );
+}
+
+/* Finds the section named section among the IMAGE_SECTION_HEADERs in
+ * section_headers and calls DigestFunction for this section.  Returns
+ * the return value from DigestFunction, or FALSE if the data could not be read.
+ */
+static BOOL IMAGEHLP_ReportSection( IMAGE_SECTION_HEADER *section_headers,
+    DWORD num_sections, LPCSTR section, BYTE *map, DWORD fileSize,
+    DIGEST_FUNCTION DigestFunction, DIGEST_HANDLE DigestHandle )
+{
+    DWORD offset, size = 0;
+
+    offset = IMAGEHLP_GetSectionOffset( section_headers, num_sections, section,
+        &size, NULL );
+    if( !offset )
+        return FALSE;
+    return IMAGEHLP_ReportSectionFromOffset( offset, size, map, fileSize,
+            DigestFunction, DigestHandle );
+}
+
+/* Calls DigestFunction for all sections with the IMAGE_SCN_CNT_CODE flag set.
+ * Returns the return value from * DigestFunction, or FALSE if a section could not be read.
+ */
+static BOOL IMAGEHLP_ReportCodeSections( IMAGE_SECTION_HEADER *hdr, DWORD num_sections,
+    BYTE *map, DWORD fileSize, DIGEST_FUNCTION DigestFunction, DIGEST_HANDLE DigestHandle )
+{
+    DWORD i;
+    BOOL ret = TRUE;
+
+    for( i = 0; ret && i < num_sections; i++, hdr++ )
+    {
+        if( hdr->Characteristics & IMAGE_SCN_CNT_CODE )
+            ret = IMAGEHLP_ReportSectionFromOffset( hdr->PointerToRawData,
+                hdr->SizeOfRawData, map, fileSize, DigestFunction, DigestHandle );
+    }
+    return ret;
+}
+
+/* Reports the import section from the file FileHandle.  If
+ * CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO is set in DigestLevel, reports the entire
+ * import section.
+ * FIXME: if it's not set, the function currently fails.
+ */
+static BOOL IMAGEHLP_ReportImportSection( IMAGE_SECTION_HEADER *hdr,
+    DWORD num_sections, BYTE *map, DWORD fileSize, DWORD DigestLevel,
+    DIGEST_FUNCTION DigestFunction, DIGEST_HANDLE DigestHandle )
+{
+    BOOL ret = FALSE;
+    DWORD offset, size, base;
+    PBYTE import;
+
+    /* Get import data */
+    offset = IMAGEHLP_GetSectionOffset( hdr, num_sections, ".idata", &size,
+        &base );
+    if( !offset )
+        return FALSE;
+
+    /* If CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO is set, the entire
+     * section is reported.  Otherwise, the debug info section is
+     * decoded and reported piecemeal.  See tests.  However, I haven't been
+     * able to figure out how the native implementation decides which values
+     * to report.  Either it's buggy or my understanding is flawed.
+     */
+    if( DigestLevel & CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO )
+        ret = IMAGEHLP_ReportSectionFromOffset( offset, size, map, fileSize,
+                DigestFunction, DigestHandle );
+    else
+    {
+        FIXME("not supported except for CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO\n");
+        SetLastError(ERROR_INVALID_PARAMETER);
+        ret = FALSE;
+    }
+    HeapFree( GetProcessHeap(), 0, import );
+    return ret;
+}
+
 /***********************************************************************
  *		ImageGetDigestStream (IMAGEHLP.@)
+ *
+ * Gets a stream of bytes from a PE file overwhich a hash might be computed to
+ * verify that the image has not changed.  Useful for creating a certificate to
+ * be added to the file with ImageAddCertificate.
+ *
+ * PARAMS
+ *  FileHandle     [In] File for which to return a stream.
+ *  DigestLevel    [In] Flags to control which portions of the file to return.
+ *                      0 is allowed, as is any combination of:
+ *                       CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO: reports the entire
+ *                        import section rather than selected portions of it.
+ *                       CERT_PE_IMAGE_DIGEST_DEBUG_INFO: reports the debug section.
+ *                       CERT_PE_IMAGE_DIGEST_RESOURCES: reports the resources
+                          section.
+ *  DigestFunction [In] Callback function.
+ *  DigestHandle   [In] Handle passed as first parameter to DigestFunction.
+ *
+ * RETURNS
+ *  TRUE if successful.
+ *  FALSE if unsuccessful.  GetLastError returns more about the error.
+ *
+ * NOTES
+ *  Only supports 32-bit PE files, not tested with any other format.
+ *  Reports data in the following order:
+ *  1. The file headers are reported first
+ *  2. Any code sections are reported next.
+ *  3. The data (".data" and ".rdata") sections are reported next.
+ *  4. The import section is reported next.
+ *  5. If CERT_PE_IMAGE_DIGEST_DEBUG_INFO is set in DigestLevel, the debug section is
+ *     reported next.
+ *  6. If CERT_PE_IMAGE_DIGEST_RESOURCES is set in DigestLevel, the resources section
+ *     is reported next.
+ *
+ * BUGS
+ *  CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO must be specified, returns an error if not.
  */
 BOOL WINAPI ImageGetDigestStream(
   HANDLE FileHandle, DWORD DigestLevel,
   DIGEST_FUNCTION DigestFunction, DIGEST_HANDLE DigestHandle)
 {
-  FIXME("(%p, %d, %p, %p): stub\n",
+    DWORD error = 0;
-    FileHandle, DigestLevel, DigestFunction, DigestHandle
+    BOOL ret = FALSE;
-  );
+    DWORD offset, size, num_sections, fileSize;
-  SetLastError(ERROR_CALL_NOT_IMPLEMENTED);
+    HANDLE hMap = INVALID_HANDLE_VALUE;
-  return FALSE;
+    BYTE *map = NULL;
+    IMAGE_DOS_HEADER *dos_hdr;
+    IMAGE_NT_HEADERS *nt_hdr;
+    IMAGE_SECTION_HEADER *section_headers;
+
+    TRACE("(%p, %d, %p, %p)\n", FileHandle, DigestLevel, DigestFunction,
+        DigestHandle);
+
+    /* Get the file size */
+    if( !FileHandle )
+        goto invalid_parameter;
+    fileSize = GetFileSize( FileHandle, NULL );
+    if(fileSize == INVALID_FILE_SIZE )
+        goto invalid_parameter;
+
+    /* map file */
+    hMap = CreateFileMappingW( FileHandle, NULL, PAGE_READONLY, 0, 0, NULL );
+    if( hMap == INVALID_HANDLE_VALUE )
+        goto invalid_parameter;
+    map = MapViewOfFile( hMap, FILE_MAP_COPY, 0, 0, 0 );
+    if( !map )
+        goto invalid_parameter;
+
+    /* Read the file header */
+    if( fileSize < sizeof(IMAGE_DOS_HEADER) )
+        goto invalid_parameter;
+    dos_hdr = (IMAGE_DOS_HEADER *)map;
+
+    if( dos_hdr->e_magic != IMAGE_DOS_SIGNATURE )
+        goto invalid_parameter;
+    offset = dos_hdr->e_lfanew;
+    if( !offset || offset > fileSize )
+        goto invalid_parameter;
+    ret = DigestFunction( DigestHandle, map, offset );
+    if( !ret )
+        goto end;
+
+    /* Read the NT header */
+    if( offset + sizeof(IMAGE_NT_HEADERS) > fileSize )
+        goto invalid_parameter;
+    nt_hdr = (IMAGE_NT_HEADERS *)(map + offset);
+    if( nt_hdr->Signature != IMAGE_NT_SIGNATURE )
+        goto invalid_parameter;
+    /* It's clear why the checksum is cleared, but why only these size headers?
+     */
+    nt_hdr->OptionalHeader.SizeOfInitializedData = 0;
+    nt_hdr->OptionalHeader.SizeOfImage = 0;
+    nt_hdr->OptionalHeader.CheckSum = 0;
+    size = sizeof(nt_hdr->Signature) + sizeof(nt_hdr->FileHeader) +
+        nt_hdr->FileHeader.SizeOfOptionalHeader;
+    ret = DigestFunction( DigestHandle, map + offset, size );
+    if( !ret )
+        goto end;
+
+    /* Read the section headers */
+    offset += size;
+    num_sections = nt_hdr->FileHeader.NumberOfSections;
+    size = num_sections * sizeof(IMAGE_SECTION_HEADER);
+    if( offset + size > fileSize )
+        goto invalid_parameter;
+    ret = DigestFunction( DigestHandle, map + offset, size );
+    if( !ret )
+        goto end;
+
+    section_headers = (IMAGE_SECTION_HEADER *)(map + offset);
+    IMAGEHLP_ReportCodeSections( section_headers, num_sections,
+        map, fileSize, DigestFunction, DigestHandle );
+    IMAGEHLP_ReportSection( section_headers, num_sections, ".data",
+        map, fileSize, DigestFunction, DigestHandle );
+    IMAGEHLP_ReportSection( section_headers, num_sections, ".rdata",
+        map, fileSize, DigestFunction, DigestHandle );
+    IMAGEHLP_ReportImportSection( section_headers, num_sections,
+        map, fileSize, DigestLevel, DigestFunction, DigestHandle );
+    if( DigestLevel & CERT_PE_IMAGE_DIGEST_DEBUG_INFO )
+        IMAGEHLP_ReportSection( section_headers, num_sections, ".debug",
+            map, fileSize, DigestFunction, DigestHandle );
+    if( DigestLevel & CERT_PE_IMAGE_DIGEST_RESOURCES )
+        IMAGEHLP_ReportSection( section_headers, num_sections, ".rsrc",
+            map, fileSize, DigestFunction, DigestHandle );
+
+end:
+    if( map )
+        UnmapViewOfFile( map );
+    if( hMap != INVALID_HANDLE_VALUE )
+        CloseHandle( hMap );
+    if( error )
+        SetLastError(error);
+    return ret;
+
+invalid_parameter:
+    error = ERROR_INVALID_PARAMETER;
+    goto end;
 }
 
 /***********************************************************************

dlls/imagehlp/tests/image.c

@@ -146,6 +146,7 @@ struct expected_update_accum
 {
     DWORD cUpdates;
     const struct expected_blob *updates;
+    BOOL  todo;
 };
 
 static BOOL WINAPI accumulating_stream_output(DIGEST_HANDLE handle, BYTE *pb,
@@ -180,7 +181,11 @@ static void check_updates(LPCSTR header, const struct expected_update_accum *exp
 {
     DWORD i;
 
-    ok(expected->cUpdates == got->cUpdates, "%s: expected %d updates, got %d\n",
+    if (expected->todo)
+        todo_wine ok(expected->cUpdates == got->cUpdates, "%s: expected %d updates, got %d\n",
+            header, expected->cUpdates, got->cUpdates);
+    else
+        ok(expected->cUpdates == got->cUpdates, "%s: expected %d updates, got %d\n",
             header, expected->cUpdates, got->cUpdates);
     for (i = 0; i < min(expected->cUpdates, got->cUpdates); i++)
     {
@@ -217,7 +222,7 @@ static const struct expected_blob b1[] = {
     {FILE_TOTAL-FILE_IDATA-FIELD_OFFSET(struct Imports, ibn),
         &bin.idata_section.ibn}
 };
-static const struct expected_update_accum a1 = { sizeof(b1) / sizeof(b1[0]), b1 };
+static const struct expected_update_accum a1 = { sizeof(b1) / sizeof(b1[0]), b1, TRUE };
 
 static const struct expected_blob b2[] = {
     {FILE_PE_START,  &bin},
@@ -227,7 +232,7 @@ static const struct expected_blob b2[] = {
     {FILE_IDATA-FILE_TEXT, &bin.text_section},
     {FILE_TOTAL-FILE_IDATA, &bin.idata_section}
 };
-static const struct expected_update_accum a2 = { sizeof(b2) / sizeof(b2[0]), b2 };
+static const struct expected_update_accum a2 = { sizeof(b2) / sizeof(b2[0]), b2, FALSE };
 
 /* Creates a test file and returns a handle to it.  The file's path is returned
  * in temp_file, which must be at least MAX_PATH characters in length.
@@ -278,7 +283,6 @@ static void test_get_digest_stream(void)
 
     SetLastError(0xdeadbeef);
     ret = pImageGetDigestStream(NULL, 0, NULL, NULL);
-    todo_wine
     ok(!ret && GetLastError() == ERROR_INVALID_PARAMETER,
      "expected ERROR_INVALID_PARAMETER, got %d\n", GetLastError());
     file = create_temp_file(temp_file);
@@ -289,18 +293,15 @@ static void test_get_digest_stream(void)
     }
     SetLastError(0xdeadbeef);
     ret = pImageGetDigestStream(file, 0, NULL, NULL);
-    todo_wine
     ok(!ret && GetLastError() == ERROR_INVALID_PARAMETER,
      "expected ERROR_INVALID_PARAMETER, got %d\n", GetLastError());
     SetLastError(0xdeadbeef);
     ret = pImageGetDigestStream(NULL, 0, accumulating_stream_output, &accum);
-    todo_wine
     ok(!ret && GetLastError() == ERROR_INVALID_PARAMETER,
      "expected ERROR_INVALID_PARAMETER, got %d\n", GetLastError());
     /* Even with "valid" parameters, it fails with an empty file */
     SetLastError(0xdeadbeef);
     ret = pImageGetDigestStream(file, 0, accumulating_stream_output, &accum);
-    todo_wine
     ok(!ret && GetLastError() == ERROR_INVALID_PARAMETER,
      "expected ERROR_INVALID_PARAMETER, got %d\n", GetLastError());
     /* Finally, with a valid executable in the file, it succeeds.  Note that
@@ -316,16 +317,12 @@ static void test_get_digest_stream(void)
     bin.nt_headers.OptionalHeader.SizeOfImage = 0;
 
     ret = pImageGetDigestStream(file, 0, accumulating_stream_output, &accum);
-    todo_wine
     ok(ret, "ImageGetDigestStream failed: %d\n", GetLastError());
-    todo_wine
     check_updates("flags = 0", &a1, &accum);
     free_updates(&accum);
     ret = pImageGetDigestStream(file, CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO,
      accumulating_stream_output, &accum);
-    todo_wine
     ok(ret, "ImageGetDigestStream failed: %d\n", GetLastError());
-    todo_wine
     check_updates("flags = CERT_PE_IMAGE_DIGEST_ALL_IMPORT_INFO", &a2, &accum);
     free_updates(&accum);
     CloseHandle(file);