From 16dc5b53276956110834156a19e1306774d91d6c Mon Sep 17 00:00:00 2001
From: CalDescent <caldescent@protonmail.com>
Date: Sun, 20 Jun 2021 19:49:10 +0100
Subject: [PATCH] Include the data length in DataFileMessage, which is more
 similar to the approach used by ArbitraryDataMessage. These two message types
 are very similar, except arbitrary code currently has a requirement of one
 piece of data per signature, whereas DataFile code is independent and can
 support multiple files per transaction. Maybe the two can be combined at some
 point, but for now I'll keep them separate.

---
 .../network/message/DataFileMessage.java      | 34 ++++++++++++++++---
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/qortal/network/message/DataFileMessage.java b/src/main/java/org/qortal/network/message/DataFileMessage.java
index 06bace5f..e31dde67 100644
--- a/src/main/java/org/qortal/network/message/DataFileMessage.java
+++ b/src/main/java/org/qortal/network/message/DataFileMessage.java
@@ -1,7 +1,10 @@
 package org.qortal.network.message;
 
+import com.google.common.primitives.Ints;
 import org.qortal.storage.DataFile;
 
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.nio.ByteBuffer;
 
@@ -26,19 +29,40 @@ public class DataFileMessage extends Message {
 	}
 
 	public static Message fromByteBuffer(int id, ByteBuffer byteBuffer) throws UnsupportedEncodingException {
-		byte[] bytes = new byte[byteBuffer.remaining()];
-		byteBuffer.get(bytes);
-		DataFile dataFile = new DataFile(bytes);
+		int dataLength = byteBuffer.getInt();
+
+		if (byteBuffer.remaining() != dataLength)
+			return null;
+
+		byte[] data = new byte[dataLength];
+		byteBuffer.get(data);
+		DataFile dataFile = new DataFile(data);
 
 		return new DataFileMessage(id, dataFile);
 	}
 
 	@Override
 	protected byte[] toData() {
-		if (this.dataFile == null)
+		if (this.dataFile == null) {
+			return null;
+		}
+
+		byte[] data = this.dataFile.getBytes();
+		if (data == null) {
 			return null;
+		}
 
-		return this.dataFile.getBytes();
+		try {
+			ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+
+			bytes.write(Ints.toByteArray(data.length));
+
+			bytes.write(data);
+
+			return bytes.toByteArray();
+		} catch (IOException e) {
+			return null;
+		}
 	}
 
 	public DataFileMessage cloneWithNewId(int newId) {