1 files changed, 93 insertions, 6 deletions

diff --git a/src/main/java/com/lumaserv/bgp/protocol/attribute/TunnelEncapsAttribute.java b/src/main/java/com/lumaserv/bgp/protocol/attribute/TunnelEncapsAttribute.java
index 0cfd9bc..0a29820 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/attribute/TunnelEncapsAttribute.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/attribute/TunnelEncapsAttribute.java
@@ -149,16 +149,64 @@ public class TunnelEncapsAttribute implements PathAttribute {
 
     @Getter
     public static class WireGuard implements SubTlv {
+        public final static int DEFAULT_PROTOCOL_VERSION = 1;
+
+        public final static int FLAG_TYPE_2 = 0x100;
+        public final static int FLAG_PERSISTENT_KEEPALIVE = 0x80;
+        public final static int FLAG_PROTOCOL_VERSION = 0x40;
+        public final static int FLAG_PRESHARED_KEY = 0x20;
+
+        int flags;
         byte[] publicKey;
+        int persistentKeepalive;
+        byte[] presharedKey;
+        int protocolVersion;
 
         public static final int TYPE = Encapsulation.TYPE;
 
         WireGuard(byte[] src, int offset, int length) {
-            if (length != 32) {
+            if (length == 32) { // Type 1
+                flags = 0;
+                publicKey = new byte[32];
+                System.arraycopy(src, offset, publicKey, 0, 32);
+            } else if (length >= 4 + 32) { // Type 2
+                flags = src[offset] | FLAG_TYPE_2;
+
+                int neededLength = 4 + 32
+                    + ((flags & FLAG_PROTOCOL_VERSION) != 0 ? 4 : 0)
+                    + ((flags & FLAG_PRESHARED_KEY) != 0 ? 32 : 0);
+
+                if (length != neededLength) {
+                    throw new RuntimeException("WireGuard: bad length, got "
+                                               + length + " needed " + neededLength);
+                }
+
+                offset += 2;
+
+                if ((flags & FLAG_PERSISTENT_KEEPALIVE) != 0) {
+                    persistentKeepalive = ((int)src[offset] & 0xff) << 8
+                        | ((int)src[offset + 1] & 0xff);
+                }
+                offset += 2;
+
+                System.arraycopy(src, offset, publicKey, 0, 32);
+                offset += 32;
+
+                if ((flags & FLAG_PROTOCOL_VERSION) != 0) {
+                    protocolVersion = ((int)src[offset] & 0xff) << 24
+                        | ((int)src[offset + 1] & 0xff) << 16
+                        | ((int)src[offset + 2] & 0xff) << 8
+                        | ((int)src[offset + 3] & 0xff);
+                    offset += 4;
+                }
+
+                if ((flags & FLAG_PRESHARED_KEY) != 0) {
+                    System.arraycopy(src, offset, presharedKey, 0, 32);
+                    offset += 32;
+                }
+            } else {
                 throw new RuntimeException("WireGuard: bad length");
             }
-            publicKey = new byte[32];
-            System.arraycopy(src, offset, publicKey, 0, 32);
         }
 
         @Override
@@ -166,13 +214,52 @@ public class TunnelEncapsAttribute implements PathAttribute {
 
         @Override
         public int getValue(byte[] dst, int offset) {
-            System.arraycopy(publicKey, 0, dst, offset, 32);
-            return 32;
+            if ((flags & FLAG_TYPE_2) != 0) {
+                int length = 0;
+
+                dst[offset++] = (byte)(flags & 0xff);
+                dst[offset++] = 0;
+                if ((flags & FLAG_PERSISTENT_KEEPALIVE) != 0) {
+                    dst[offset++] = (byte)(persistentKeepalive >> 8 & 0xff);
+                    dst[offset++] = (byte)(persistentKeepalive & 0xff);
+                } else {
+                    dst[offset++] = 0;
+                    dst[offset++] = 0;
+                }
+                length += 4;
+
+                System.arraycopy(publicKey, 0, dst, offset, 32);
+                offset += 32;
+                length += 32;
+
+                if ((flags & FLAG_PROTOCOL_VERSION) != 0) {
+                    dst[offset++] = (byte)(protocolVersion >> 24 & 0xff);
+                    dst[offset++] = (byte)(protocolVersion >> 16 & 0xff);
+                    dst[offset++] = (byte)(protocolVersion >> 8 & 0xff);
+                    dst[offset++] = (byte)(protocolVersion & 0xff);
+                    length += 4;
+                }
+
+                if ((flags & FLAG_PRESHARED_KEY) != 0) {
+                    System.arraycopy(presharedKey, 0, dst, offset, 32);
+                    offset += 32;
+                    length += 32;
+                }
+
+                return length;
+            } else {
+                System.arraycopy(publicKey, 0, dst, offset, 32);
+                return 32;
+            }
         }
 
         @Override
         public String toString() {
-            return "WireGuard:" + Base64.getEncoder().encodeToString(publicKey);
+            return "WireGuard:"
+                + ((flags & FLAG_PERSISTENT_KEEPALIVE) != 0 ? " persistent keepalive:" + persistentKeepalive : "")
+                + " public key:" + Base64.getEncoder().encodeToString(publicKey)
+                + ((flags & FLAG_PROTOCOL_VERSION) != 0 ? " protocol version:" + protocolVersion : "")
+                + ((flags & FLAG_PRESHARED_KEY) != 0 ? " preshared key:" + Base64.getEncoder().encodeToString(presharedKey) : "");
         }
     }