packets: Handle flow control callbacks

Signed-off-by: Aris Adamantiadis <aris@0xbadc0de.be>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>

2 files changed, 45 insertions, 1 deletions

diff --git a/include/libssh/callbacks.h b/include/libssh/callbacks.h
index 20c540e5..6b07b087 100644
--- a/include/libssh/callbacks.h
+++ b/include/libssh/callbacks.h
@@ -694,6 +694,22 @@ typedef int (*ssh_channel_subsystem_request_callback) (ssh_session session,
                                             const char *subsystem,
                                             void *userdata);
 
+/**
+ * @brief SSH channel write will not block (flow control).
+ *
+ * @param channel the channel
+ *
+ * @param[in] bytes size of the remote window in bytes. Writing as much data
+ *            will not block.
+ *
+ * @param[in] userdata Userdata to be passed to the callback function.
+ *
+ * @returns 0 default return value (other return codes may be added in future).
+ */
+typedef int (*ssh_channel_write_wontblock_callback) (ssh_session session,
+                                                     ssh_channel channel,
+                                                     size_t bytes,
+                                                     void *userdata);
 
 struct ssh_channel_callbacks_struct {
   /** DON'T SET THIS use ssh_callbacks_init() instead. */
@@ -758,6 +774,10 @@ struct ssh_channel_callbacks_struct {
    * (like sftp).
    */
   ssh_channel_subsystem_request_callback channel_subsystem_request_function;
+  /** This function will be called when the channel write is guaranteed
+   * not to block.
+   */
+  ssh_channel_write_wontblock_callback channel_write_wontblock_function;
 };
 
 typedef struct ssh_channel_callbacks_struct *ssh_channel_callbacks;
diff --git a/src/packet.c b/src/packet.c
index 7d6957ca..6e8e3521 100644
--- a/src/packet.c
+++ b/src/packet.c
@@ -360,10 +360,34 @@ error:
     return processed;
 }
 
+static void ssh_packet_socket_controlflow_callback(int code, void *userdata)
+{
+    ssh_session session = userdata;
+    struct ssh_iterator *it;
+    ssh_channel channel;
+
+    if (code == SSH_SOCKET_FLOW_WRITEWONTBLOCK) {
+        /* the out pipe is empty so we can forward this to channels */
+        it = ssh_list_get_iterator(session->channels);
+        while (it != NULL) {
+            channel = ssh_iterator_value(ssh_channel, it);
+            if (ssh_callbacks_exists(channel->callbacks,
+                                     channel_write_wontblock_function)) {
+                SSH_LOG(SSH_LOG_TRACE, "Executing write_wontblock callback for channel");
+                channel->callbacks->channel_write_wontblock_function(session,
+                                                                     channel,
+                                                                     channel->remote_window,
+                                                                     channel->callbacks->userdata);
+            }
+            it = it->next;
+        }
+    }
+}
+
 void ssh_packet_register_socket_callback(ssh_session session, ssh_socket s){
 	session->socket_callbacks.data=ssh_packet_socket_callback;
 	session->socket_callbacks.connected=NULL;
-	session->socket_callbacks.controlflow=NULL;
+    session->socket_callbacks.controlflow = ssh_packet_socket_controlflow_callback;
 	session->socket_callbacks.exception=NULL;
 	session->socket_callbacks.userdata=session;
 	ssh_socket_set_callbacks(s,&session->socket_callbacks);