diff --git a/src/main/java/com/iota/iri/network/neighbor/Neighbor.java b/src/main/java/com/iota/iri/network/neighbor/Neighbor.java index 34f482ba7a..398ae0b277 100644 --- a/src/main/java/com/iota/iri/network/neighbor/Neighbor.java +++ b/src/main/java/com/iota/iri/network/neighbor/Neighbor.java @@ -2,6 +2,7 @@ import com.iota.iri.network.protocol.Handshake; import com.iota.iri.network.protocol.Heartbeat; +import com.iota.iri.network.protocol.MilestoneRequest; import java.io.IOException; import java.nio.ByteBuffer; @@ -47,6 +48,13 @@ public interface Neighbor { */ Heartbeat heartbeat() throws IOException; + /** + * Instructs the {@link Neighbor} to read from its source channel a {@link MilestoneRequest} packet. + * @return The {@link MilestoneRequest} object. + * @throws IOException thrown when reading from the source channel fails + */ + MilestoneRequest milestoneRequest() throws IOException; + /** * Instructs the {@link Neighbor} to send the given {@link ByteBuffer} to its destination channel. * diff --git a/src/main/java/com/iota/iri/network/neighbor/impl/NeighborImpl.java b/src/main/java/com/iota/iri/network/neighbor/impl/NeighborImpl.java index 28cb1cc6ce..591acf1a18 100644 --- a/src/main/java/com/iota/iri/network/neighbor/impl/NeighborImpl.java +++ b/src/main/java/com/iota/iri/network/neighbor/impl/NeighborImpl.java @@ -58,6 +58,7 @@ private enum ReadState { private MessageReader msgReader; private Handshake handshake = new Handshake(); private Heartbeat heartbeat = new Heartbeat(); + private MilestoneRequest milestoneRequest = new MilestoneRequest(); /** * Creates a new {@link NeighborImpl} using the given channel. @@ -92,6 +93,11 @@ public Heartbeat heartbeat() throws IOException { return heartbeat; } + public MilestoneRequest milestoneRequest() throws IOException{ + read(); + return milestoneRequest; + } + @Override public int read() throws IOException { int bytesRead = msgReader.readMessage(channel); @@ -166,6 +172,9 @@ private void handleMessage(ByteBuffer msg) { case HEARTBEAT: heartbeat = Heartbeat.fromByteBuffer(msg); break; + case MS_REQUEST: + milestoneRequest = MilestoneRequest.fromByteBuffer(msg); + break; default: // do nothing } diff --git a/src/main/java/com/iota/iri/network/protocol/MilestoneRequest.java b/src/main/java/com/iota/iri/network/protocol/MilestoneRequest.java new file mode 100644 index 0000000000..4eb7a6b5ca --- /dev/null +++ b/src/main/java/com/iota/iri/network/protocol/MilestoneRequest.java @@ -0,0 +1,41 @@ +package com.iota.iri.network.protocol; + +import java.nio.ByteBuffer; + +/** + * Defines the information contained in a milestone request + */ +public class MilestoneRequest { + + private int milestoneIndex; + + /** + * Parses the given message into a {@link MilestoneRequest} object. + * + * @param msg the buffer containing the milestone request info + * @return the {@link MilestoneRequest} object + */ + public static MilestoneRequest fromByteBuffer(ByteBuffer msg) { + MilestoneRequest milestoneRequest = new MilestoneRequest(); + milestoneRequest.setMilestoneIndex(msg.getInt()); + return milestoneRequest; + } + + /** + * Gets the milestone index of the milestone request + * + * @return The milestone index + */ + public int getMilestoneIndex() { + return milestoneIndex; + } + + /** + * Sets the milestone index of the milestone request + * + * @param milestoneIndex The milestone index + */ + public void setMilestoneIndex(int milestoneIndex) { + this.milestoneIndex = milestoneIndex; + } +} diff --git a/src/main/java/com/iota/iri/network/protocol/Protocol.java b/src/main/java/com/iota/iri/network/protocol/Protocol.java index 2e40e986ed..ac8ffed4ab 100644 --- a/src/main/java/com/iota/iri/network/protocol/Protocol.java +++ b/src/main/java/com/iota/iri/network/protocol/Protocol.java @@ -60,6 +60,11 @@ public class Protocol { */ public final static byte PROTOCOL_HEARTBEAT_BYTES_LENGTH = 8; + /** + * The amount of byes to store the index of the milestone to request + */ + public final static byte PROTOCOL_MS_REQUEST_BYTES_LENGTH = 4; + /** * Parses the given buffer into a {@link ProtocolHeader}. * @@ -123,6 +128,15 @@ public static ByteBuffer createHeartbeatPacket(Heartbeat heartbeat) { return buf; } + public static ByteBuffer createMilestoneRequestPacket(MilestoneRequest milestoneRequest){ + final short payloadLengthBytes = Protocol.PROTOCOL_MS_REQUEST_BYTES_LENGTH; + ByteBuffer buf = ByteBuffer.allocate(ProtocolMessage.HEADER.getMaxLength() + payloadLengthBytes); + addProtocolHeader(buf, ProtocolMessage.MS_REQUEST, payloadLengthBytes); + buf.putInt(milestoneRequest.getMilestoneIndex()); + buf.flip(); + return buf; + } + /** * Adds the protocol header to the given {@link ByteBuffer}. * diff --git a/src/main/java/com/iota/iri/network/protocol/ProtocolMessage.java b/src/main/java/com/iota/iri/network/protocol/ProtocolMessage.java index 9be1e23f28..29736856dc 100644 --- a/src/main/java/com/iota/iri/network/protocol/ProtocolMessage.java +++ b/src/main/java/com/iota/iri/network/protocol/ProtocolMessage.java @@ -28,7 +28,8 @@ public enum ProtocolMessage { TRANSACTION_GOSSIP((byte) 2, (short) (Protocol.GOSSIP_REQUESTED_TX_HASH_BYTES_LENGTH + TransactionTruncator.NON_SIG_TX_PART_BYTES_LENGTH + TransactionTruncator.SIG_DATA_MAX_BYTES_LENGTH), true), - HEARTBEAT((byte) 3, Protocol.PROTOCOL_HEARTBEAT_BYTES_LENGTH, true); + HEARTBEAT((byte) 3, Protocol.PROTOCOL_HEARTBEAT_BYTES_LENGTH, true), + MS_REQUEST((byte) 4, Protocol.PROTOCOL_MS_REQUEST_BYTES_LENGTH, true); private static final ProtocolMessage[] lookup = new ProtocolMessage[256]; @@ -47,6 +48,7 @@ public enum ProtocolMessage { lookup[1] = HANDSHAKE; lookup[2] = TRANSACTION_GOSSIP; lookup[3] = HEARTBEAT; + lookup[4] = MS_REQUEST; } /** diff --git a/src/main/java/com/iota/iri/network/protocol/message/MessageReaderFactory.java b/src/main/java/com/iota/iri/network/protocol/message/MessageReaderFactory.java index e9ac92dd6d..d68c80985a 100644 --- a/src/main/java/com/iota/iri/network/protocol/message/MessageReaderFactory.java +++ b/src/main/java/com/iota/iri/network/protocol/message/MessageReaderFactory.java @@ -22,6 +22,7 @@ public static MessageReader create(ProtocolMessage protoMsg) throws UnknownMessa case HANDSHAKE: case TRANSACTION_GOSSIP: case HEARTBEAT: + case MS_REQUEST: return create(protoMsg, protoMsg.getMaxLength()); // there might be message types in the future which need a separate message reader implementation default: diff --git a/src/test/java/com/iota/iri/network/neighbor/impl/NeighborImplTest.java b/src/test/java/com/iota/iri/network/neighbor/impl/NeighborImplTest.java index 22c013e408..5b70d33dc3 100644 --- a/src/test/java/com/iota/iri/network/neighbor/impl/NeighborImplTest.java +++ b/src/test/java/com/iota/iri/network/neighbor/impl/NeighborImplTest.java @@ -6,16 +6,13 @@ import com.iota.iri.network.neighbor.Neighbor; import com.iota.iri.network.neighbor.NeighborState; import com.iota.iri.network.pipeline.TransactionProcessingPipeline; -import com.iota.iri.network.protocol.Handshake; +import com.iota.iri.network.protocol.*; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.SelectionKey; import java.nio.channels.Selector; -import com.iota.iri.network.protocol.Heartbeat; -import com.iota.iri.network.protocol.Protocol; -import com.iota.iri.network.protocol.ProtocolMessage; import org.junit.Rule; import org.junit.Test; import org.mockito.Mock; @@ -267,4 +264,60 @@ public int read(ByteBuffer dst) { fail("didnt expect an exception: " + e.getMessage()); } } + + @Test + public void writeMilestoneRequest() { + MilestoneRequest milestoneRequest = new MilestoneRequest(); + milestoneRequest.setMilestoneIndex(1); + ByteBuffer milestoneRequestPacket = Protocol.createMilestoneRequestPacket(milestoneRequest); + + Neighbor neighbor = new NeighborImpl<>(selector, new FakeChannel() { + + @Override + public int write(ByteBuffer buf) { + int bytesWritten = 0; + while (buf.hasRemaining()) { + buf.get(); + bytesWritten++; + } + return bytesWritten; + } + }, localAddr, serverSocketPort, pipeline); + + neighbor.send(milestoneRequestPacket); + + try { + assertEquals("should have written the entire milestone request packet", milestoneRequestPacket.capacity(), neighbor.write()); + } catch (IOException e) { + fail("Didn't expect an exception: " + e.getMessage()); + } + } + + @Test + public void readMilestoneRequest() { + MilestoneRequest milestoneRequest = new MilestoneRequest(); + milestoneRequest.setMilestoneIndex(1); + ByteBuffer milestoneRequestPacket = Protocol.createMilestoneRequestPacket(milestoneRequest); + + Neighbor neighbor = new NeighborImpl<>(selector, new FakeChannel() { + // fake having a heartbeat packet in the socket + @Override + public int read(ByteBuffer dst) { + while (dst.hasRemaining()) { + dst.put(milestoneRequestPacket.get()); + } + return 0; + } + }, localAddr, serverSocketPort, pipeline); + + // set the neighbor as ready for other messages + neighbor.setState(NeighborState.READY_FOR_MESSAGES); + + try { + MilestoneRequest readMilestoneRequest = neighbor.milestoneRequest(); + assertEquals("milestone index of sent and read ms request should be equal", readMilestoneRequest.getMilestoneIndex(), milestoneRequest.getMilestoneIndex()); + } catch (IOException e) { + fail("Didn't expect an exception: " + e.getMessage()); + } + } } \ No newline at end of file diff --git a/src/test/java/com/iota/iri/network/protocol/ProtocolTest.java b/src/test/java/com/iota/iri/network/protocol/ProtocolTest.java index 5362d5f7a5..c78eaa1c26 100644 --- a/src/test/java/com/iota/iri/network/protocol/ProtocolTest.java +++ b/src/test/java/com/iota/iri/network/protocol/ProtocolTest.java @@ -128,4 +128,30 @@ public void parseHeartbeat() { fail("didn't expect any exceptions"); } } + + @Test + public void createMilestoneRequestPacket(){ + MilestoneRequest milestoneRequest = new MilestoneRequest(); + milestoneRequest.setMilestoneIndex(1); + + ByteBuffer buf = Protocol.createMilestoneRequestPacket(milestoneRequest); + final int expectedMessageSize = Protocol.PROTOCOL_MS_REQUEST_BYTES_LENGTH; + assertEquals("buffer should have the right capacity", + Protocol.PROTOCOL_HEADER_BYTES_LENGTH + expectedMessageSize, buf.capacity()); + assertEquals("should be of type milestone request message", ProtocolMessage.MS_REQUEST.getTypeID(), buf.get()); + assertEquals("should have correct message length", expectedMessageSize, buf.getShort()); + } + + @Test + public void parseMilestoneRequest() { + try { + ByteBuffer buf = ByteBuffer.allocate(Protocol.PROTOCOL_MS_REQUEST_BYTES_LENGTH); + buf.putInt(1); + buf.flip(); + MilestoneRequest milestoneRequest = MilestoneRequest.fromByteBuffer(buf); + assertEquals("should have correct milestone index", milestoneRequest.getMilestoneIndex(), 1); + } catch (Exception e) { + fail("Didn't expect any exceptions: " + e.getMessage()); + } + } } \ No newline at end of file