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) {