1 files changed, 41 insertions, 21 deletions
diff --git a/src/main/java/com/lumaserv/bgp/protocol/BGPPacket.java b/src/main/java/com/lumaserv/bgp/protocol/BGPPacket.java
index e7cf66b..e0aa3d6 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/BGPPacket.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/BGPPacket.java
@@ -1,5 +1,9 @@
 package com.lumaserv.bgp.protocol;
 
+import com.lumaserv.bgp.protocol.DataBuilder;
+import com.lumaserv.bgp.protocol.message.BGPOpen;
+import com.lumaserv.bgp.protocol.message.BGPUpdate;
+
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
@@ -7,6 +11,8 @@ import lombok.Setter;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.Arrays;
 
 @Setter
@@ -15,44 +21,58 @@ import java.util.Arrays;
 public class BGPPacket {
 
     Type type;
-    byte[] message;
+    DataBuilder message;
 
-    public BGPPacket(byte[] packet) {
-        type = Type.fromValue(packet[18]);
-        message = new byte[packet.length-19];
-        System.arraycopy(packet, 19, message, 0, message.length);
+    public BGPPacket(ByteBuffer packet) {
+        type = Type.fromValue(packet.get());
+        switch (type) {
+            case OPEN:
+                message = new BGPOpen(packet);
+                break;
+            case UPDATE:
+                message = new BGPUpdate(packet);
+                break;
+            default:
+                System.out.println("type: " + type);
+                break;
+        }
     }
 
     public byte[] build() {
-        byte[] packet = new byte[message.length + 19];
-        Arrays.fill(packet, 0, 16, (byte) 0xFF);
-        packet[16] = (byte)(packet.length >> 8);
-        packet[17] = (byte)(packet.length & 0xFF);
-        packet[18] = type.getValue();
-        System.arraycopy(message, 0, packet, 19, message.length);
+        ByteBuffer buf = ByteBuffer.allocate(4096).order(ByteOrder.BIG_ENDIAN); // Maximum BGP message size
+        for (int i=0; i<16; i++)
+            buf.put((byte) 0xFF);
+        buf.putShort((short)19); // Length without message payload
+        buf.put(type.getValue());
+        if (message != null) {
+            message.build(buf);
+            buf.putShort(16, (short)buf.position()); // Set length
+        }
+        byte[] packet = new byte[buf.position()];
+        buf.position(0);
+        buf.get(packet); // TODO is working?
         return packet;
     }
 
     public static BGPPacket read(InputStream stream) throws IOException {
-        byte[] packet = new byte[18];
+        ByteBuffer buf = ByteBuffer.allocate(4096).order(ByteOrder.BIG_ENDIAN);
         int value;
-        for(int i=0; i<packet.length; i++) {
+        for(int i=0; i<18; i++) {
             value = stream.read();
             if(value == -1)
                 throw new IOException("Unexpected end of stream");
-            packet[i] = (byte) value;
+            buf.put((byte) value);
         }
-        int length = ((packet[16] & 0xFF) << 8) | (packet[17] & 0xFF);
-        byte[] newPacket = new byte[length];
-        System.arraycopy(packet, 0, newPacket, 0, packet.length);
-        packet = newPacket;
-        for(int i=18; i<packet.length; i++) {
+        int length = buf.getShort(16);
+        for(int i=18; i<length; i++) {
             value = stream.read();
             if(value == -1)
                 throw new IOException("Unexpected end of stream");
-            packet[i] = (byte) value;
+            buf.put((byte) value);
         }
-        return new BGPPacket(packet);
+        buf.limit(buf.position());
+        buf.position(18);
+        return new BGPPacket(buf);
     }
 
     @AllArgsConstructor