Fixed a few issues

6 files changed, 24 insertions, 4 deletions

diff --git a/src/main/java/com/lumaserv/bgp/BGPServer.java b/src/main/java/com/lumaserv/bgp/BGPServer.java
index d06f341..b8f934a 100644
--- a/src/main/java/com/lumaserv/bgp/BGPServer.java
+++ b/src/main/java/com/lumaserv/bgp/BGPServer.java
@@ -17,7 +17,11 @@ public class BGPServer implements Runnable {
     final List<BGPSessionConfiguration> sessionConfigurations = new ArrayList<>();
 
     public BGPServer() throws IOException {
-        serverSocket = new ServerSocket(179);
+        this(179);
+    }
+
+    public BGPServer(int port) throws IOException {
+        serverSocket = new ServerSocket(port);
     }
 
     private static boolean checkEqual(byte[] a, byte[] b) {
@@ -48,6 +52,8 @@ public class BGPServer implements Runnable {
                         .orElse(null);
                 if(config == null)
                     continue;
+                BGPSession session = new BGPSession(socket, config);
+                config.getListener().onOpen(session);
                 BGPOpen response = new BGPOpen()
                         .setAsn(config.getLocalAs())
                         .setHoldTime(request.getHoldTime())
@@ -58,7 +64,6 @@ public class BGPServer implements Runnable {
                 } catch (IOException e) {
                     e.printStackTrace();
                 }
-                BGPSession session = new BGPSession(socket, config);
                 session.keepAlive();
                 new Thread(session).start();
             } catch (IOException ex) {
diff --git a/src/main/java/com/lumaserv/bgp/BGPSession.java b/src/main/java/com/lumaserv/bgp/BGPSession.java
index de1fc66..1d40f9d 100644
--- a/src/main/java/com/lumaserv/bgp/BGPSession.java
+++ b/src/main/java/com/lumaserv/bgp/BGPSession.java
@@ -15,6 +15,8 @@ public class BGPSession implements Runnable {
     final BGPSessionConfiguration configuration;
     final InputStream inputStream;
     final OutputStream outputStream;
+    @Getter
+    boolean closed;
 
     public BGPSession(Socket socket, BGPSessionConfiguration configuration) throws IOException {
         this.configuration = configuration;
@@ -32,6 +34,10 @@ public class BGPSession implements Runnable {
 
     private void handle(BGPPacket packet) {
         switch (packet.getType()) {
+            case NOTIFICATION:
+                closed = true;
+                configuration.getListener().onClose(this);
+                break;
             case KEEPALIVE:
                 keepAlive();
                 break;
@@ -44,7 +50,7 @@ public class BGPSession implements Runnable {
 
     public void run() {
         try {
-            while (true)
+            while (!closed)
                 handle(BGPPacket.read(inputStream));
         } catch (IOException ex) {
             ex.printStackTrace();
diff --git a/src/main/java/com/lumaserv/bgp/protocol/attribute/AS4PathAttribute.java b/src/main/java/com/lumaserv/bgp/protocol/attribute/AS4PathAttribute.java
index 415dd5f..796fbea 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/attribute/AS4PathAttribute.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/attribute/AS4PathAttribute.java
@@ -6,6 +6,8 @@ import lombok.Setter;
 import java.util.ArrayList;
 import java.util.List;
 
+@Getter
+@Setter
 public class AS4PathAttribute implements PathAttribute {
 
     List<Segment> segments = new ArrayList<>();
diff --git a/src/main/java/com/lumaserv/bgp/protocol/attribute/ASPathAttribute.java b/src/main/java/com/lumaserv/bgp/protocol/attribute/ASPathAttribute.java
index 64be45b..d6f14ea 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/attribute/ASPathAttribute.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/attribute/ASPathAttribute.java
@@ -6,6 +6,8 @@ import lombok.Setter;
 import java.util.ArrayList;
 import java.util.List;
 
+@Getter
+@Setter
 public class ASPathAttribute implements PathAttribute {
 
     List<Segment> segments = new ArrayList<>();
diff --git a/src/main/java/com/lumaserv/bgp/protocol/attribute/AggregatorAttribute.java b/src/main/java/com/lumaserv/bgp/protocol/attribute/AggregatorAttribute.java
index 8a3dfdb..25f3e70 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/attribute/AggregatorAttribute.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/attribute/AggregatorAttribute.java
@@ -7,7 +7,7 @@ import lombok.Setter;
 @Setter
 public class AggregatorAttribute implements PathAttribute {
 
-    int as;
+    final int as;
     byte[] origin;
 
     public AggregatorAttribute(byte typeCode, byte[] data) {
diff --git a/src/main/java/com/lumaserv/bgp/protocol/attribute/NextHopAttribute.java b/src/main/java/com/lumaserv/bgp/protocol/attribute/NextHopAttribute.java
index f60d236..14d5367 100644
--- a/src/main/java/com/lumaserv/bgp/protocol/attribute/NextHopAttribute.java
+++ b/src/main/java/com/lumaserv/bgp/protocol/attribute/NextHopAttribute.java
@@ -1,5 +1,10 @@
 package com.lumaserv.bgp.protocol.attribute;
 
+import lombok.Getter;
+import lombok.Setter;
+
+@Getter
+@Setter
 public class NextHopAttribute implements PathAttribute {
 
     byte[] address;