Fixed a few issues

1 files changed, 7 insertions, 1 deletions

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();