1 files changed, 47 insertions, 15 deletions

diff --git a/src/main/java/com/lumaserv/bgp/protocol/message/BGPUpdate.java b/src/main/java/com/lumaserv/bgp/protocol/message/BGPUpdate.java
index fa23432..ad6c39b 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/message/BGPUpdate.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/message/BGPUpdate.java
@@ -1,63 +1,95 @@
 package com.lumaserv.bgp.protocol.message;
 
+import com.lumaserv.bgp.protocol.DataBuilder;
 import com.lumaserv.bgp.protocol.IPPrefix;
 import com.lumaserv.bgp.protocol.attribute.PathAttribute;
 import lombok.Getter;
+import lombok.NoArgsConstructor;
 import lombok.Setter;
 
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
 
+@NoArgsConstructor
 @Getter
 @Setter
-public class BGPUpdate {
+public class BGPUpdate implements DataBuilder {
 
     List<IPPrefix> withdrawnPrefixes = new ArrayList<>();
     List<PathAttribute> attributes = new ArrayList<>();
     List<IPPrefix> prefixes = new ArrayList<>();
 
-    public BGPUpdate(byte[] message) {
-        int routesLength = ((message[0] & 0xFF) << 8) | (message[1] & 0xFF);
+    public BGPUpdate(ByteBuffer message) {
+        int routesLength = message.getShort() & 0xFFFF;
         int offset = 2;
         int offsetOffset = 2;
         while ((offset - offsetOffset) < routesLength) {
             IPPrefix prefix = new IPPrefix()
-                    .setLength(message[offset])
+                    .setLength(message.get())
                     .setAddress(new byte[4]);
             offset++;
             int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
-            System.arraycopy(message, offset, prefix.getAddress(), 0, addressLen);
+            message.get(prefix.getAddress(), 0, addressLen);
             offset += addressLen;
             withdrawnPrefixes.add(prefix);
         }
-        int attributesLength = ((message[offset] & 0xFF) << 8) | (message[offset + 1] & 0xFF);
+        int attributesLength = message.getShort() & 0xFFFF;
         offset += 2;
         offsetOffset = offset;
         while ((offset - offsetOffset) < attributesLength) {
-            byte flags = message[offset];
-            byte typeCode = message[offset + 1];
-            int length = message[offset + 2] & 0xFF;
+            byte flags = message.get();
+            byte typeCode = message.get();
+            int length = message.get() & 0xFF;
             offset += 3;
             if((flags & 0b0001_0000) > 0) {
                 length <<= 8;
-                length |= message[offset] & 0xFF;
+                length |= message.get() & 0xFF;
                 offset++;
             }
             byte[] data = new byte[length];
-            System.arraycopy(message, offset, data, 0, length);
+            message.get(data);
             offset += length;
-            attributes.add(PathAttribute.from(typeCode, data));
+            PathAttribute attr = PathAttribute.from(typeCode, data);
+            attributes.add(attr);
         }
-        while (offset < message.length) {
+        while (message.hasRemaining()) {
             IPPrefix prefix = new IPPrefix()
-                    .setLength(message[offset])
+                    .setLength(message.get())
                     .setAddress(new byte[4]);
             offset++;
             int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
-            System.arraycopy(message, offset, prefix.getAddress(), 0, addressLen);
+            message.get(prefix.getAddress(), 0, addressLen);
             offset += addressLen;
             prefixes.add(prefix);
         }
     }
 
+    public void build(ByteBuffer message) {
+        int first = message.position();
+        message.putShort((short)0); // Withdrawn
+        message.putShort((short)0); // Attributes
+        // TODO
+        int pos = message.position();
+        for (PathAttribute attr: attributes) {
+// FIXME add header
+            byte flags = 0b0100_0000; // TODO
+            int start = message.position();
+            message.put(flags);
+            message.put(attr.getTypeCode());
+            message.put((byte)0); // Dummy length
+            attr.build(message);
+            int length = message.position() - (start + 3);
+            if (length > 255)
+                throw new RuntimeException("attribute length > 255");
+            message.put(start + 2, (byte)(length & 0xFF));
+        }
+        int attributesLength = message.position() - pos;
+        message.putShort(first + 2, (short)attributesLength);
+        for (IPPrefix prefix : prefixes) {
+            message.put((byte)prefix.getLength());
+            int addressLen = (int) Math.ceil(prefix.getLength() / 8d);
+            message.put(prefix.getAddress(), 0, addressLen);
+        }
+    }
 }