2 files changed, 46 insertions, 0 deletions

diff --git a/libssh/poll.c b/libssh/poll.c
index 4a84cbe3..7d9e3e5a 100644
--- a/libssh/poll.c
+++ b/libssh/poll.c
@@ -33,6 +33,7 @@
 #include "libssh/priv.h"
 #include "libssh/libssh.h"
 #include "libssh/poll.h"
+#include "libssh/socket.h"
 
 #ifndef SSH_POLL_CTX_CHUNK
 #define SSH_POLL_CTX_CHUNK			5
@@ -310,6 +311,21 @@ void ssh_poll_set_events(ssh_poll_handle p, short events) {
 }
 
 /**
+ * @brief  Set the file descriptor of a poll object. The FD will also be propagated
+ *         to an associated poll context.
+ *
+ * @param  p            Pointer to an already allocated poll object.
+ * @param  fd       New file descriptor.
+ */
+void ssh_poll_set_fd(ssh_poll_handle p, socket_t fd) {
+  if (p->ctx != NULL) {
+    p->ctx->pollfds[p->x.idx].fd = fd;
+  } else {
+  	p->x.fd = fd;
+  }
+}
+
+/**
  * @brief  Add extra events to a poll object. Duplicates are ignored.
  *         The events will also be propagated to an associated poll context.
  *
@@ -475,6 +491,22 @@ int ssh_poll_ctx_add(ssh_poll_ctx ctx, ssh_poll_handle p) {
 }
 
 /**
+ * @brief  Add a socket object to a poll context.
+ *
+ * @param  ctx          Pointer to an already allocated poll context.
+ * @param  s            A SSH socket handle
+ *
+ * @return              0 on success, < 0 on error
+ */
+int ssh_poll_ctx_add_socket (ssh_poll_ctx ctx, struct socket *s) {
+  ssh_poll_handle p=ssh_socket_get_poll_handle(s);
+  if(p==NULL)
+  	return -1;
+  return ssh_poll_ctx_add(ctx,p);
+}
+
+
+/**
  * @brief  Remove a poll object from a poll context.
  *
  * @param  ctx          Pointer to an already allocated poll context.
diff --git a/libssh/socket.c b/libssh/socket.c
index 8acc76ce..a250f2ea 100644
--- a/libssh/socket.c
+++ b/libssh/socket.c
@@ -181,6 +181,18 @@ void ssh_socket_register_pollcallback(struct socket *s, ssh_poll_handle p){
 	s->poll=p;
 }
 
+/** @internal
+ * @brief returns the poll handle corresponding to the socket,
+ * creates it if it does not exist.
+ * @returns allocated and initialized ssh_poll_handle object
+ */
+ssh_poll_handle ssh_socket_get_poll_handle(struct socket *s){
+	if(s->poll)
+		return s->poll;
+	s->poll=ssh_poll_new(s->fd,0,ssh_socket_pollcallback,s);
+	return s->poll;
+}
+
 /** \internal
  * \brief Deletes a socket object
  */
@@ -191,6 +203,8 @@ void ssh_socket_free(struct socket *s){
   ssh_socket_close(s);
   buffer_free(s->in_buffer);
   buffer_free(s->out_buffer);
+  if(s->poll)
+  	ssh_poll_free(s->poll);
   SAFE_FREE(s);
 }