imagehlp: Implement ImageAddCertificate.

dlls/imagehlp/integrity.c

@@ -167,6 +167,67 @@ static BOOL IMAGEHLP_GetSecurityDirOffset( HANDLE handle,
 }
 
 /***********************************************************************
+ * IMAGEHLP_SetSecurityDirOffset (INTERNAL)
+ *
+ * Read a file's PE header, and update the offset and size of the
+ *  security directory.
+ */
+static BOOL IMAGEHLP_SetSecurityDirOffset(HANDLE handle,
+                                          DWORD dwOfs, DWORD dwSize)
+{
+    IMAGE_NT_HEADERS32 nt_hdr32;
+    IMAGE_NT_HEADERS64 nt_hdr64;
+    IMAGE_DATA_DIRECTORY *sd;
+    int ret, nt_hdr_size = 0;
+    DWORD pe_offset;
+    void *nt_hdr;
+    DWORD count;
+    BOOL r;
+
+    ret = IMAGEHLP_GetNTHeaders(handle, &pe_offset, &nt_hdr32, &nt_hdr64);
+
+    if (ret == HDR_NT32)
+    {
+        sd = &nt_hdr32.OptionalHeader.DataDirectory[IMAGE_FILE_SECURITY_DIRECTORY];
+
+        nt_hdr = &nt_hdr32;
+        nt_hdr_size = sizeof(IMAGE_NT_HEADERS32);
+    }
+    else if (ret == HDR_NT64)
+    {
+        sd = &nt_hdr64.OptionalHeader.DataDirectory[IMAGE_FILE_SECURITY_DIRECTORY];
+
+        nt_hdr = &nt_hdr64;
+        nt_hdr_size = sizeof(IMAGE_NT_HEADERS64);
+    }
+    else
+        return FALSE;
+
+    sd->Size = dwSize;
+    sd->VirtualAddress = dwOfs;
+
+    TRACE("size = %x addr = %x\n", sd->Size, sd->VirtualAddress);
+
+    /* write the header back again */
+    count = SetFilePointer(handle, pe_offset, NULL, FILE_BEGIN);
+
+    if (count == INVALID_SET_FILE_POINTER)
+        return FALSE;
+
+    count = 0;
+
+    r = WriteFile(handle, nt_hdr, nt_hdr_size, &count, NULL);
+
+    if (!r)
+        return FALSE;
+
+    if (count != nt_hdr_size)
+        return FALSE;
+
+    return TRUE;
+}
+
+/***********************************************************************
  * IMAGEHLP_GetCertificateOffset (INTERNAL)
  *
  * Read a file's PE header, and return the offset and size of the 
@@ -227,16 +288,111 @@ static BOOL IMAGEHLP_GetCertificateOffset( HANDLE handle, DWORD num,
 
 /***********************************************************************
  *		ImageAddCertificate (IMAGEHLP.@)
+ *
+ * Adds the specified certificate to the security directory of
+ * open PE file.
  */
 
 BOOL WINAPI ImageAddCertificate(
   HANDLE FileHandle, LPWIN_CERTIFICATE Certificate, PDWORD Index)
 {
-  FIXME("(%p, %p, %p): stub\n",
-    FileHandle, Certificate, Index
-  );
-  SetLastError(ERROR_CALL_NOT_IMPLEMENTED);
-  return FALSE;
+    DWORD size = 0, count = 0, offset = 0, sd_VirtualAddr = 0, index = 0;
+    WIN_CERTIFICATE hdr;
+    const size_t cert_hdr_size = sizeof hdr - sizeof hdr.bCertificate;
+    BOOL r;
+
+    TRACE("(%p, %p, %p)\n", FileHandle, Certificate, Index);
+
+    r = IMAGEHLP_GetSecurityDirOffset(FileHandle, &sd_VirtualAddr, &size);
+
+    /* If we've already got a security directory, find the end of it */
+    if ((r) && (sd_VirtualAddr != 0))
+    {
+        offset = 0;
+        index = 0;
+        count = 0;
+
+        /* Check if the security directory is at the end of the file.
+           If not, we should probably relocate it. */
+        if (GetFileSize(FileHandle, NULL) != sd_VirtualAddr + size)
+        {
+            FIXME("Security directory already present but not located at EOF, not adding certificate\n");
+
+            SetLastError(ERROR_NOT_SUPPORTED);
+            return FALSE;
+        }
+
+        while (offset < size)
+        {
+            /* read the length of the current certificate */
+            count = SetFilePointer (FileHandle, sd_VirtualAddr + offset,
+                                     NULL, FILE_BEGIN);
+
+            if (count == INVALID_SET_FILE_POINTER)
+                return FALSE;
+
+            r = ReadFile(FileHandle, &hdr, cert_hdr_size, &count, NULL);
+
+            if (!r)
+                return FALSE;
+
+            if (count != cert_hdr_size)
+                return FALSE;
+
+            /* check the certificate is not too big or too small */
+            if (hdr.dwLength < cert_hdr_size)
+                return FALSE;
+
+            if (hdr.dwLength > (size-offset))
+                return FALSE;
+
+            /* next certificate */
+            offset += hdr.dwLength;
+
+            /* padded out to the nearest 8-byte boundary */
+            if (hdr.dwLength % 8)
+                offset += 8 - (hdr.dwLength % 8);
+
+            index++;
+        }
+
+        count = SetFilePointer (FileHandle, sd_VirtualAddr + offset, NULL, FILE_BEGIN);
+
+        if (count == INVALID_SET_FILE_POINTER)
+            return FALSE;
+    }
+    else
+    {
+        sd_VirtualAddr = SetFilePointer(FileHandle, 0, NULL, FILE_END);
+
+        if (sd_VirtualAddr == INVALID_SET_FILE_POINTER)
+            return FALSE;
+    }
+
+    /* Write the certificate to the file */
+    r = WriteFile(FileHandle, Certificate, Certificate->dwLength, &count, NULL);
+
+    if (!r)
+        return FALSE;
+
+    /* Pad out if necessary */
+    if (Certificate->dwLength % 8)
+    {
+        char null[8];
+
+        ZeroMemory(null, 8);
+        WriteFile(FileHandle, null, 8 - (Certificate->dwLength % 8), NULL, NULL);
+
+        size += 8 - (Certificate->dwLength % 8);
+    }
+
+    size += Certificate->dwLength;
+
+    /* Update the security directory offset and size */
+    if (!IMAGEHLP_SetSecurityDirOffset(FileHandle, sd_VirtualAddr, size))
+        return FALSE;
+
+    return TRUE;
 }
 
 /***********************************************************************

dlls/imagehlp/tests/integrity.c

@@ -146,7 +146,7 @@ static void test_add_certificate(void)
     cert->wCertificateType = WIN_CERT_TYPE_PKCS_SIGNED_DATA;
     CopyMemory(cert->bCertificate, test_cert_data, sizeof(test_cert_data));
 
-    todo_wine ok(pImageAddCertificate(hFile, cert, &index), "Unable to add certificate to image, error %x\n", GetLastError());
+    ok(pImageAddCertificate(hFile, cert, &index), "Unable to add certificate to image, error %x\n", GetLastError());
 
     HeapFree(GetProcessHeap(), 0, cert);
     CloseHandle(hFile);
@@ -170,7 +170,7 @@ static void test_get_certificate(void)
     ret = pImageGetCertificateData(hFile, 0, NULL, &cert_len);
     err = GetLastError();
 
-    todo_wine ok ((ret == FALSE) && (err == ERROR_INSUFFICIENT_BUFFER), "ImageGetCertificateData gave unexpected result; ret=%d / err=%x\n", ret, err);
+    ok ((ret == FALSE) && (err == ERROR_INSUFFICIENT_BUFFER), "ImageGetCertificateData gave unexpected result; ret=%d / err=%x\n", ret, err);
 
     cert = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, cert_len);
 
@@ -181,8 +181,8 @@ static void test_get_certificate(void)
         return;
     }
 
-    todo_wine ok(ret = pImageGetCertificateData(hFile, 0, cert, &cert_len), "Unable to retrieve certificate; err=%x\n", GetLastError());
-    todo_wine ok(memcmp(cert->bCertificate, test_cert_data, cert_len - sizeof(WIN_CERTIFICATE)) == 0, "Certificate retrieved did not match original\n");
+    ok(ret = pImageGetCertificateData(hFile, 0, cert, &cert_len), "Unable to retrieve certificate; err=%x\n", GetLastError());
+    ok(memcmp(cert->bCertificate, test_cert_data, cert_len - sizeof(WIN_CERTIFICATE)) == 0, "Certificate retrieved did not match original\n");
 
     HeapFree(GetProcessHeap(), 0, cert);
     CloseHandle(hFile);
@@ -204,7 +204,7 @@ static void test_remove_certificate(void)
     todo_wine ok (pImageRemoveCertificate(hFile, 0), "Unable to remove certificate from file; err=%x\n", GetLastError());
 
     /* Test to see if the certificate has actually been removed */
-    ok(pImageGetCertificateHeader(hFile, 0, &cert) == FALSE, "Certificate header retrieval succeeded when it should have failed\n");
+    todo_wine ok(pImageGetCertificateHeader(hFile, 0, &cert) == FALSE, "Certificate header retrieval succeeded when it should have failed\n");
 
     CloseHandle(hFile);
 }