Skip to content

Commit

Permalink
Merge pull request #35 from swismer/resources-rebuild-support
Browse files Browse the repository at this point in the history 
Resources rebuild support
  • Loading branch information
@swismer
swismer committed Aug 24, 2023
2 parents 454e1de + 0961a00 commit f33f55c
Show file tree
Hide file tree
Showing 21 changed files with 572 additions and 84 deletions.
137 changes: 136 additions & 1 deletion src/main/java/com/kichik/pecoff4j/PE.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,20 @@
package com.kichik.pecoff4j;

import com.kichik.pecoff4j.constant.ImageDataDirectoryType;
import com.kichik.pecoff4j.constant.SectionFlag;
import com.kichik.pecoff4j.io.DataEntry;
import com.kichik.pecoff4j.io.DataReader;
import com.kichik.pecoff4j.io.DataWriter;
import com.kichik.pecoff4j.io.IDataReader;
import com.kichik.pecoff4j.io.IDataWriter;
import com.kichik.pecoff4j.util.PaddingType;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

public class PE {
import static com.kichik.pecoff4j.util.Alignment.align;

public class PE implements WritableStructure {
private DOSHeader dosHeader;
private DOSStub stub;
private PESignature signature;
Expand Down Expand Up @@ -158,6 +165,134 @@ private static void readDebugRawData(PE pe, DataEntry entry, IDataReader dr)
id.setDebugRawData(b);
}

/**
* Rebuild the resources section and update the various size fields and the PE checksum.
*/
public int rebuild(PaddingType paddingType) {
// rebuild resource table data
ResourceDirectory resourceTable = imageData.getResourceTable();
if (resourceTable != null) {
ImageDataDirectory resTable = optionalHeader.getDataDirectory(ImageDataDirectoryType.RESOURCE_TABLE);
int resBaseAddress = resTable.getVirtualAddress();

int resSize = resourceTable.rebuild(resBaseAddress);
byte[] resData = resourceTable.toByteArray(paddingType, optionalHeader.getFileAlignment());

// update resource data section
for (int i = 0; i < sectionTable.getNumberOfSections(); i++) {
SectionHeader header = sectionTable.getHeader(i);
SectionData section = sectionTable.getSection(i);
if (header.getVirtualAddress() == resTable.getVirtualAddress()) {
if (header.getVirtualSize() == resTable.getSize()) {
header.setVirtualSize(resSize);
header.setSizeOfRawData(align(resSize, optionalHeader.getFileAlignment()));
resTable.setSize(resSize);
section.setData(resData);
break;
} else {
throw new UnsupportedOperationException("Partial update of section is not supported");
}
}
}

// update virtual address and raw data pointer
int virtualAddress = 0;
int fileAddress = 0;
for (SectionHeader header : sectionTable.getHeadersPointerSorted()) {
if (header.getVirtualAddress() > resTable.getVirtualAddress()) {
for (ImageDataDirectory dataDirectory : optionalHeader.getDataDirectories()) {
if (dataDirectory.getVirtualAddress() == header.getVirtualAddress()) {
dataDirectory.setVirtualAddress(virtualAddress);
break;
}
}
header.setVirtualAddress(virtualAddress);
header.setPointerToRawData(fileAddress);
}
virtualAddress = align(header.getVirtualAddress() + header.getVirtualSize(), optionalHeader.getSectionAlignment());
fileAddress = align(header.getPointerToRawData() + header.getSizeOfRawData(), optionalHeader.getFileAlignment());
}
}

// update size fields
optionalHeader.setSizeOfInitializedData(
calculateSizeOfInitializedData(sectionTable, optionalHeader.getFileAlignment()));
optionalHeader.setSizeOfImage(calculateSizeOfImage(sectionTable, optionalHeader.getSectionAlignment()));

int length = updateChecksum(paddingType);

return length;
}

private int calculateSizeOfInitializedData(SectionTable sectionTable, int fileAlignment) {
int sum = 0;
for (int i = 0; i < sectionTable.getNumberOfSections(); i++) {
SectionHeader header = sectionTable.getHeader(i);
if ((header.getCharacteristics() & SectionFlag.IMAGE_SCN_CNT_INITIALIZED_DATA) != 0) {
sum += Math.max(header.getSizeOfRawData(), align(header.getVirtualSize(), fileAlignment));
}
}
return sum;
}

private int calculateSizeOfImage(SectionTable sectionTable, int sectionAlignment) {
// find the highest address used in any section (virtual address space)
int max = 0;
for (int i = 0; i < sectionTable.getNumberOfSections(); i++) {
SectionHeader header = sectionTable.getHeader(i);
max = Math.max(max, header.getVirtualAddress() + header.getVirtualSize());
}
return align(max, sectionAlignment);
}

/**
* Update Checksum and return the length of the executable file.
*/
private int updateChecksum(PaddingType paddingType) {
try {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
DataWriter dataWriter = new DataWriter(baos);
dataWriter.setPaddingMode(paddingType);
write(dataWriter);
dataWriter.close();

int length = baos.size();
// align to DWORD
baos.write(new byte[(4 - (baos.size() % 4)) % 4]);
byte[] data = baos.toByteArray();

long checksum = 0;

try (DataReader reader = new DataReader(data)) {
while (reader.hasMore()) {
// skip data where checksum is stored
if (reader.getPosition() == (dosHeader.getAddressOfNewExeHeader() + 0x58)) {
reader.readDoubleWord();
continue;
}
long doubleWord = ((long) reader.readDoubleWord()) & 0xffffffffL;
checksum = (checksum & 0xffffffffL) + (doubleWord & 0xffffffffL);

if (checksum >= 1L << 32) {
checksum = (checksum & 0xffffffffL) + (checksum >> 32);
}
}
}

checksum = (checksum & 0xffff) + (checksum >> 16);
checksum = checksum + (checksum >> 16);
checksum = checksum & 0xffff;
checksum += length; // must be original data length

optionalHeader.setCheckSum((int) checksum);
return length;
} catch (IOException e) {
// cannot happen due to in-memory implementation of writer and reader
return 0;
}
}

@Override
public void write(IDataWriter dw) throws IOException {
getDosHeader().write(dw);
getStub().write(dw);
Expand Down
27 changes: 0 additions & 27 deletions src/main/java/com/kichik/pecoff4j/RebuildableStructure.java
View file

This file was deleted.

68 changes: 64 additions & 4 deletions src/main/java/com/kichik/pecoff4j/ResourceDirectory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,13 @@
import com.kichik.pecoff4j.io.IDataReader;
import com.kichik.pecoff4j.io.IDataWriter;
import com.kichik.pecoff4j.util.DataObject;
import com.kichik.pecoff4j.util.PaddingType;

public class ResourceDirectory extends DataObject {
import static com.kichik.pecoff4j.util.Alignment.align;

public class ResourceDirectory extends DataObject implements WritableStructure {
private ResourceDirectoryTable table;
private List<ResourceEntry> entries = new ArrayList();
private final List<ResourceEntry> entries = new ArrayList<>();

public static ResourceDirectory read(IDataReader dr, int baseAddress) throws IOException {
ResourceDirectory d = new ResourceDirectory();
Expand All @@ -37,10 +40,56 @@ public static ResourceDirectory read(IDataReader dr, int baseAddress) throws IOE
return d;
}

public int rebuild(int baseAddress) {
// update count of entries
updateCount();

// re-calculate offsets of resource entries
List<ResourceDirectory> dirs = new ArrayList<>();
List<ResourceEntry> flatEntries = new ArrayList<>();
flattenInto(dirs, flatEntries, true);
int pos = 16 + 8 * entries.size();
for (ResourceEntry re : flatEntries) {
if (re.getDirectory() != null) {
re.setOffset(pos | 0x80000000);
pos += 16 + 8 * re.getDirectory().size();
} else {
re.setOffset(pos);
pos += 16;
}
}

// re-calculate relative virtual address of resource entries' data
dirs.clear();
flatEntries.clear();
flattenInto(dirs, flatEntries, false);
for (ResourceEntry re : flatEntries) {
pos = align(pos, 4);
re.setDataRVA(baseAddress + pos);
pos += re.getData().length;
}
return align(pos, 4);
}

private void updateCount() {
int numNameEntries = 0;
for (ResourceEntry re : entries) {
if (re.getDirectory() != null) {
re.getDirectory().updateCount();
}
if (re.getName() != null) {
numNameEntries++;
}
}
table.setNumNameEntries(numNameEntries);
table.setNumIdEntries(entries.size() - numNameEntries);
}

@Override
public void write(IDataWriter dw) throws IOException {
List<ResourceDirectory> dirs = new ArrayList<>();
List<ResourceEntry> entries = new ArrayList<>();
flattenInto(dirs, entries);
flattenInto(dirs, entries, false);

for (ResourceDirectory dir : dirs) {
dir.getTable().write(dw);
Expand Down Expand Up @@ -70,7 +119,8 @@ public void write(IDataWriter dw) throws IOException {
/**
* Traverse tree in breath-first order and fill directories and entries into given lists.
*/
private void flattenInto(List<ResourceDirectory> dirs, List<ResourceEntry> entries) {
private void flattenInto(List<ResourceDirectory> dirs, List<ResourceEntry> entries,
boolean includeIntermediateEntries) {
Queue<ResourceDirectory> toProcess = new ArrayDeque<>();
toProcess.add(this);
dirs.add(this);
Expand All @@ -82,6 +132,9 @@ private void flattenInto(List<ResourceDirectory> dirs, List<ResourceEntry> entri
if (entry.getDirectory() != null) {
toProcess.add(entry.getDirectory());
dirs.add(entry.getDirectory());
if (includeIntermediateEntries) {
entries.add(entry);
}
} else {
entries.add(entry);
}
Expand All @@ -97,15 +150,22 @@ public void setTable(ResourceDirectoryTable table) {
this.table = table;
}

@Deprecated
public void add(ResourceEntry entry) {
this.entries.add(entry);
}

@Deprecated
public ResourceEntry get(int index) {
return entries.get(index);
}

@Deprecated
public int size() {
return entries.size();
}

public List<ResourceEntry> getEntries() {
return entries;
}
}
48 changes: 48 additions & 0 deletions src/main/java/com/kichik/pecoff4j/WritableStructure.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package com.kichik.pecoff4j;

import com.kichik.pecoff4j.io.DataWriter;
import com.kichik.pecoff4j.io.IDataWriter;
import com.kichik.pecoff4j.util.PaddingType;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
* A structure that can be written using a data writer.
*/
public interface WritableStructure {
/**
* Write this structure to a data writer.
*
* @param dw the writer
*/
void write(IDataWriter dw) throws IOException;

/**
* Write this structure to a byte array using {@link PaddingType#PATTERN} and no alignment.
* @return the byte array
*/
default byte[] toByteArray() {
return toByteArray(PaddingType.PATTERN, 1);
}
/**
* Write this structure to a byte array.
* @param paddingType the padding mode
* @param alignment the alignment to use after writing this structure
* @return the byte array
*/
default byte[] toByteArray(PaddingType paddingType, int alignment) {
try {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
DataWriter writer = new DataWriter(baos);
writer.setPaddingMode(paddingType);
write(writer);
writer.align(alignment);
writer.close();
return baos.toByteArray();
} catch (IOException e) {
// cannot happen due to in-memory implementation of writer
throw new RuntimeException("DataWriter must not throw IOException when using ByteArrayOutputStream");
}
}
}
12 changes: 11 additions & 1 deletion src/main/java/com/kichik/pecoff4j/io/DataWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
*******************************************************************************/
package com.kichik.pecoff4j.io;

import com.kichik.pecoff4j.util.PaddingType;

import java.io.BufferedOutputStream;
import java.io.EOFException;
import java.io.File;
Expand All @@ -20,6 +22,7 @@
public class DataWriter implements IDataWriter {
private BufferedOutputStream out;
private int position;
private PaddingType paddingType = PaddingType.PATTERN;

public DataWriter(File output) throws FileNotFoundException {
this(new FileOutputStream(output));
Expand All @@ -29,6 +32,10 @@ public DataWriter(OutputStream out) {
this.out = new BufferedOutputStream(out);
}

public void setPaddingMode(PaddingType paddingType) {
this.paddingType = paddingType;
}

@Override
public void writeByte(int b) throws IOException {
out.write(b);
Expand Down Expand Up @@ -129,7 +136,10 @@ public void writeUnicode(String s, int len) throws IOException {
public int align(int alignment) throws IOException {
int off = (alignment - (getPosition() % alignment)) % alignment;
try {
writeByte(0, off);
byte[] padding = paddingType.getData();
for (int i = 0; i < off; i++) {
writeByte(padding[i % padding.length]);
}
} catch (EOFException ignored) {
// no need to align when it's at the end of its data
}
Expand Down
Loading

0 comments on commit f33f55c

Please sign in to comment.