From 0404d45c295c3e9d89a1c803c559b774cc4fcedc Mon Sep 17 00:00:00 2001
From: Aris Adamantiadis <aris@0xbadc0de.be>
Date: Sun, 5 Jan 2014 22:33:45 +0100
Subject: poll: fix poll_handles ownerships

---
 src/poll.c | 32 +++++++++++++++++++++++++++-----
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/poll.c b/src/poll.c
index 4d92e3f3..f79e1659 100644
--- a/src/poll.c
+++ b/src/poll.c
@@ -64,6 +64,7 @@
 
 struct ssh_poll_handle_struct {
   ssh_poll_ctx ctx;
+  ssh_session session;
   union {
     socket_t fd;
     size_t idx;
@@ -790,6 +791,10 @@ int ssh_event_add_session(ssh_event event, ssh_session session) {
         p = session->default_poll_ctx->pollptrs[i];
         ssh_poll_ctx_remove(session->default_poll_ctx, p);
         ssh_poll_ctx_add(event->ctx, p);
+        /* associate the pollhandler with a session so we can put it back
+         * at ssh_event_free()
+         */
+        p->session = session;
     }
 #ifdef WITH_SERVER
     iterator = ssh_list_get_iterator(event->sessions);
@@ -852,6 +857,10 @@ int ssh_event_remove_fd(ssh_event event, socket_t fd) {
     for (i = 0; i < used; i++) {
         if(fd == event->ctx->pollfds[i].fd) {
             ssh_poll_handle p = event->ctx->pollptrs[i];
+            if (p->session != NULL){
+            	/* we cannot free that handle, it's owned by its session */
+            	continue;
+            }
             if (p->cb == ssh_event_fd_wrapper_callback) {
                 struct ssh_event_fd_wrapper *pw = p->cb_data;
                 SAFE_FREE(pw);
@@ -886,7 +895,6 @@ int ssh_event_remove_session(ssh_event event, ssh_session session) {
     ssh_poll_handle p;
     register size_t i, used;
     int rc = SSH_ERROR;
-    socket_t session_fd;
 #ifdef WITH_SERVER
     struct ssh_iterator *iterator;
 #endif
@@ -895,14 +903,15 @@ int ssh_event_remove_session(ssh_event event, ssh_session session) {
         return SSH_ERROR;
     }
 
-    session_fd = ssh_get_fd(session);
     used = event->ctx->polls_used;
     for(i = 0; i < used; i++) {
-        if(session_fd == event->ctx->pollfds[i].fd) {
-            p = event->ctx->pollptrs[i];
+    	p = event->ctx->pollptrs[i];
+    	if(p->session == session){
             ssh_poll_ctx_remove(event->ctx, p);
+            p->session = NULL;
             ssh_poll_ctx_add(session->default_poll_ctx, p);
             rc = SSH_OK;
+            used = 0;
         }
     }
 #ifdef WITH_SERVER
@@ -929,10 +938,23 @@ int ssh_event_remove_session(ssh_event event, ssh_session session) {
  *
  */
 void ssh_event_free(ssh_event event) {
-    if(event == NULL) {
+	int used, i;
+	ssh_poll_handle p;
+	if(event == NULL) {
         return;
     }
     if(event->ctx != NULL) {
+        used = event->ctx->polls_used;
+        for(i = 0; i < used; i++) {
+        	p = event->ctx->pollptrs[i];
+        	if(p->session != NULL){
+                ssh_poll_ctx_remove(event->ctx, p);
+                ssh_poll_ctx_add(p->session->default_poll_ctx, p);
+                p->session = NULL;
+                used = 0;
+            }
+        }
+
         ssh_poll_ctx_free(event->ctx);
     }
 #ifdef WITH_SERVER
-- 
cgit v1.2.3