Ensure data is fully written to PE files

ebourg · Mar 5, 2024 · 16824e3 · 16824e3
1 parent 5259dba
commit 16824e3
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 23 deletions.
diff --git a/jsign-core/src/main/java/net/jsign/pe/DataDirectory.java b/jsign-core/src/main/java/net/jsign/pe/DataDirectory.java
@@ -68,7 +68,7 @@ void check() throws IOException {
      * 
      * @since 2.0
      */
-    public void erase() {
+    public void erase() throws IOException {
         peFile.write(getVirtualAddress(), new byte[getSize()]);
     }
 
@@ -83,11 +83,12 @@ public boolean isTrailing() throws IOException {
         return getVirtualAddress() + getSize() == peFile.channel.size();
     }
 
-    public void write(long virtualAddress, int size) {
+    public void write(long virtualAddress, int size) throws IOException {
         ByteBuffer buffer = ByteBuffer.allocate(8);
         buffer.order(ByteOrder.LITTLE_ENDIAN);
         buffer.putInt((int) virtualAddress);
         buffer.putInt(size);
-        peFile.write(peFile.getDataDirectoryOffset() + index * 8, buffer.array());
+        buffer.flip();
+        peFile.write(peFile.getDataDirectoryOffset() + index * 8, buffer);
     }
 }
diff --git a/jsign-core/src/main/java/net/jsign/pe/PEFile.java b/jsign-core/src/main/java/net/jsign/pe/PEFile.java
@@ -199,12 +199,14 @@ synchronized long readQWord(long base, int offset) {
         return valueBuffer.getLong();
     }
 
-    synchronized void write(long base, byte[] data) {
-        try {
-            channel.position(base);
-            channel.write(ByteBuffer.wrap(data));
-        } catch (IOException e) {
-            throw new RuntimeException(e);
+    synchronized void write(long base, byte[] data) throws IOException {
+        write(base, ByteBuffer.wrap(data));
+    }
+
+    synchronized void write(long base, ByteBuffer data) throws IOException {
+        channel.position(base);
+        while (data.hasRemaining()) {
+            channel.write(data);
         }
     }
 
@@ -527,8 +529,7 @@ public synchronized void updateChecksum() {
         buffer.flip();
 
         try {
-            channel.position(peHeaderOffset + 88);
-            channel.write(buffer);
+            write(peHeaderOffset + 88, buffer);
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
@@ -658,32 +659,28 @@ public synchronized void writeDataDirectory(DataDirectoryType type, byte[] data)
             // append the data directory at the end of the file
             long offset = channel.size();
 
-            channel.position(offset);
-            channel.write(ByteBuffer.wrap(data));
+            write(offset, data);
 
             // update the entry in the data directory table
             directory.write(offset, data.length);
 
         } else {
             if (data.length == directory.getSize()) {
                 // same size as before, just overwrite
-                channel.position(directory.getVirtualAddress());
-                channel.write(ByteBuffer.wrap(data));
+                write(directory.getVirtualAddress(), data);
 
             } else if (data.length < directory.getSize() && type != DataDirectoryType.CERTIFICATE_TABLE) {
                 // the new data is smaller, erase and rewrite in-place
                 // this doesn't work with the certificate table since it changes the file digest
                 directory.erase();
-                channel.position(directory.getVirtualAddress());
-                channel.write(ByteBuffer.wrap(data));
+                write(directory.getVirtualAddress(), data);
 
                 // update the size in the data directory table
                 directory.write(directory.getVirtualAddress(), data.length);
 
             } else if (directory.isTrailing()) {
                 // the data is at the end of the file, overwrite it
-                channel.position(directory.getVirtualAddress());
-                channel.write(ByteBuffer.wrap(data));
+                write(directory.getVirtualAddress(), data);
                 channel.truncate(directory.getVirtualAddress() + data.length); // trim the file if the data shrunk
 
                 // update the size in the data directory table
@@ -699,8 +696,7 @@ public synchronized void writeDataDirectory(DataDirectoryType type, byte[] data)
 
                 long offset = channel.size();
 
-                channel.position(offset);
-                channel.write(ByteBuffer.wrap(data));
+                write(offset, data);
 
                 // update the entry in the data directory table
                 directory.write(offset, data.length);
@@ -965,7 +961,6 @@ public ASN1Object createIndirectData(DigestAlgorithm digestAlgorithm) throws IOE
      */
     public synchronized void pad(int multiple) throws IOException {
         long padding = (multiple - channel.size() % multiple) % multiple;
-        channel.position(channel.size());
-        channel.write(ByteBuffer.allocate((int) padding));
+        write(channel.size(), ByteBuffer.allocate((int) padding));
     }
 }