From 1f6d24565cc3a4ea75747bc8c836a884253602aa Mon Sep 17 00:00:00 2001 From: Owen Rudge Date: Thu, 3 Dec 2009 13:51:02 -0600 Subject: [PATCH] imagehlp: Implement ImageAddCertificate. --- dlls/imagehlp/integrity.c | 166 ++++++++++++++++++++++++++++++++++++++-- dlls/imagehlp/tests/integrity.c | 10 +-- 2 files changed, 166 insertions(+), 10 deletions(-) diff --git a/dlls/imagehlp/integrity.c b/dlls/imagehlp/integrity.c index a46c853d95f..6a13c53bb23 100644 --- a/dlls/imagehlp/integrity.c +++ b/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; } /*********************************************************************** diff --git a/dlls/imagehlp/tests/integrity.c b/dlls/imagehlp/tests/integrity.c index 01fd4e9fdf0..e89167c4fd4 100644 --- a/dlls/imagehlp/tests/integrity.c +++ b/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); } -- 2.11.4.GIT