From cc33a6e5afd5f68b6969e0cffbf6bb762f2dbc4f Mon Sep 17 00:00:00 2001
From: "Roger A. Light" <roger@atchoo.org>
Date: Mon, 23 Jun 2014 18:18:29 +0100
Subject: [PATCH] Do proper cleanup, including websockets.

---
 src/loop.c       | 7 +++++++
 src/websockets.c | 6 +++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/loop.c b/src/loop.c
index f32391e0..39ade3ac 100644
--- a/src/loop.c
+++ b/src/loop.c
@@ -130,7 +130,11 @@ int mosquitto_main_loop(struct mosquitto_db *db, int *listensock, int listensock
 			}
 			context->pollfd_index = -1;
 
+#ifdef WITH_WEBSOCKETS
+			if(context->sock != INVALID_SOCKET || context->wsi){
+#else
 			if(context->sock != INVALID_SOCKET){
+#endif
 #ifdef WITH_BRIDGE
 				if(context->bridge){
 					_mosquitto_check_keepalive(db, context);
@@ -348,6 +352,9 @@ static void do_disconnect(struct mosquitto_db *db, struct mosquitto *context)
 		}
 	}
 	mqtt3_context_disconnect(db, context);
+	if(context->clean_session){
+		mqtt3_context_cleanup(db, context, true);
+	}
 }
 
 /* Error ocurred, probably an fd has been closed. 
diff --git a/src/websockets.c b/src/websockets.c
index 159d623e..d31ed5ed 100644
--- a/src/websockets.c
+++ b/src/websockets.c
@@ -129,7 +129,11 @@ static int callback_mqtt(struct libwebsocket_context *context,
 
 		case LWS_CALLBACK_CLOSED:
 			mosq = u->mosq;
-			mqtt3_context_cleanup(db, mosq, true);
+			mqtt3_context_disconnect(db, mosq);
+			if(context->clean_session){
+				mqtt3_context_cleanup(db, mosq, true);
+			}
+			mosq->wsi = NULL;
 			break;
 
 		case LWS_CALLBACK_SERVER_WRITEABLE: