[PATCH V2] Add session/channel byte/packet counters

[Audrius Butkevicius <audrius.butkevicius@xxxxxxxxxxxxxxxx>]

 include/libssh/channels.h |   2 ++
 include/libssh/libssh.h   |  12 ++++++++++++
 include/libssh/session.h  |   3 +++
 src/channels.c            |  34 ++++++++++++++++++++++++++++++++++
 src/channels1.c           |   3 +++
 src/packet.c              |   8 ++++++++
 src/session.c             |  39 +++++++++++++++++++++++++++++++++++++++
 src/socket.c              |   6 ++++++
 8 files changed, 107 insertions(+), 0 deletions(-)


# HG changeset patch
# User Audrius Butkevicius <audrius.butkevicius@xxxxxxxxxxxxxxxx>
# Date 1392141424 0
#      Tue Feb 11 17:57:04 2014 +0000
# Node ID 795795a61c967a9df169870fd98a776d330e7260
# Parent  441875df1ed307a82bd7ef59a9a788014c3d7742
Add session/channel byte/packet counters

diff --git a/include/libssh/channels.h b/include/libssh/channels.h
--- a/include/libssh/channels.h
+++ b/include/libssh/channels.h
@@ -75,6 +75,8 @@
     int exit_status;
     enum ssh_channel_request_state_e request_state;
     ssh_channel_callbacks callbacks;
+    /* counters */
+    ssh_counter counter;
 };
 
 SSH_PACKET_CALLBACK(ssh_packet_channel_open_conf);
diff --git a/include/libssh/libssh.h b/include/libssh/libssh.h
--- a/include/libssh/libssh.h
+++ b/include/libssh/libssh.h
@@ -104,6 +104,13 @@
 extern "C" {
 #endif
 
+struct ssh_counter_struct {
+    uint64_t in_bytes;
+    uint64_t out_bytes;
+    uint64_t in_packets;
+    uint64_t out_packets;
+};
+typedef struct ssh_counter_struct *ssh_counter;
 
 typedef struct ssh_agent_struct* ssh_agent;
 typedef struct ssh_buffer_struct* ssh_buffer;
@@ -397,6 +404,8 @@
 LIBSSH_API void ssh_channel_set_blocking(ssh_channel channel, int blocking);
 LIBSSH_API int ssh_channel_write(ssh_channel channel, const void *data, uint32_t len);
 LIBSSH_API uint32_t ssh_channel_window_size(ssh_channel channel);
+LIBSSH_API void ssh_set_channel_counter(ssh_channel channel,
+                                        ssh_counter counter);
 
 LIBSSH_API char *ssh_basename (const char *path);
 LIBSSH_API void ssh_clean_pubkey_hash(unsigned char **hash);
@@ -485,6 +494,9 @@
 LIBSSH_API int ssh_options_get(ssh_session session, enum ssh_options_e type,
     char **value);
 LIBSSH_API int ssh_options_get_port(ssh_session session, unsigned int * port_target);
+LIBSSH_API void ssh_set_session_counters(ssh_session session,
+                                         ssh_counter scounter,
+                                         ssh_counter rcounter);
 LIBSSH_API int ssh_pcap_file_close(ssh_pcap_file pcap);
 LIBSSH_API void ssh_pcap_file_free(ssh_pcap_file pcap);
 LIBSSH_API ssh_pcap_file ssh_pcap_file_new(void);
diff --git a/include/libssh/session.h b/include/libssh/session.h
--- a/include/libssh/session.h
+++ b/include/libssh/session.h
@@ -188,6 +188,9 @@
         char *gss_client_identity;
         int gss_delegate_creds;
     } opts;
+    /* counters */
+    ssh_counter socket_counter;
+    ssh_counter raw_counter;
 };
 
 /** @internal
diff --git a/src/channels.c b/src/channels.c
--- a/src/channels.c
+++ b/src/channels.c
@@ -555,6 +555,9 @@
                                                 is_stderr,
                                                 channel->callbacks->userdata);
       if(rest > 0) {
+        if (channel->counter != NULL) {
+            channel->counter->in_bytes += rest;
+        }
         buffer_pass_bytes(buf, rest);
       }
       if (channel->local_window + buffer_get_rest_len(buf) < WINDOWLIMIT) {
@@ -1406,6 +1409,9 @@
     channel->remote_window -= effectivelen;
     len -= effectivelen;
     data = ((uint8_t*)data + effectivelen);
+    if (channel->counter != NULL) {
+        channel->counter->out_bytes += effectivelen;
+    }
   }
 
   /* it's a good idea to flush the socket now */
@@ -2843,6 +2849,9 @@
   len = (len > count ? count : len);
   memcpy(dest, buffer_get_rest(stdbuf), len);
   buffer_pass_bytes(stdbuf,len);
+  if (channel->counter != NULL) {
+      channel->counter->in_bytes += len;
+  }
   /* Authorize some buffering while userapp is busy */
   if (channel->local_window < WINDOWLIMIT) {
     if (grow_window(session, channel, 0) < 0) {
@@ -3575,6 +3584,31 @@
   return rc;
 }
 
+/**
+ * @brief Set the channel data counter.
+ *
+ * @code
+ * struct ssh_counter_struct counter = {
+ *     .in_bytes = 0,
+ *     .out_bytes = 0,
+ *     .in_packets = 0,
+ *     .out_packets = 0
+ * };
+ *
+ * ssh_set_channel_counter(channel, &counter);
+ * @endcode
+ *
+ * @param[in] channel The SSH channel.
+ *
+ * @param[in] counter Counter for bytes handled by the channel.
+ */
+void ssh_set_channel_counter(ssh_channel channel,
+                             ssh_counter counter) {
+    if (channel != NULL) {
+        channel->counter = counter;
+    }
+}
+
 #endif
 
 /* @} */
diff --git a/src/channels1.c b/src/channels1.c
--- a/src/channels1.c
+++ b/src/channels1.c
@@ -367,6 +367,9 @@
       return -1;
     }
     ssh_handle_packets(session, SSH_TIMEOUT_NONBLOCKING);
+    if (channel->counter != NULL) {
+        channel->counter->out_bytes += effectivelen;
+    }
   }
   if (ssh_blocking_flush(session,SSH_TIMEOUT_USER) == SSH_ERROR)
       return -1;
diff --git a/src/packet.c b/src/packet.c
--- a/src/packet.c
+++ b/src/packet.c
@@ -311,6 +311,10 @@
 #endif /* WITH_ZLIB */
             payloadsize = buffer_get_rest_len(session->in_buffer);
             session->recv_seq++;
+            if (session->raw_counter != NULL) {
+                session->raw_counter->in_bytes += payloadsize;
+                session->raw_counter->in_packets++;
+            }
 
             /*
              * We don't want to rewrite a new packet while still executing the
@@ -560,6 +564,10 @@
 
   rc = ssh_packet_write(session);
   session->send_seq++;
+  if (session->raw_counter != NULL) {
+      session->raw_counter->out_bytes += payloadsize;
+      session->raw_counter->out_packets++;
+  }
 
   SSH_LOG(SSH_LOG_PACKET,
           "packet: wrote [len=%d,padding=%hhd,comp=%d,payload=%d]",
diff --git a/src/session.c b/src/session.c
--- a/src/session.c
+++ b/src/session.c
@@ -831,6 +831,45 @@
     return SSH_ERROR;
 }
 
+ /**
+ * @brief Set the session data counters.
+ *
+ * This functions sets the counter structures to be used to calculate data
+ * which comes in and goes out through the session at different levels.
+ *
+ * @code
+ * struct ssh_counter_struct scounter = {
+ *     .in_bytes = 0,
+ *     .out_bytes = 0,
+ *     .in_packets = 0,
+ *     .out_packets = 0
+ * };
+ *
+ * struct ssh_counter_struct rcounter = {
+ *     .in_bytes = 0,
+ *     .out_bytes = 0,
+ *     .in_packets = 0,
+ *     .out_packets = 0
+ * };
+ *
+ * ssh_set_session_counters(session, &scounter, &rcounter);
+ * @endcode
+ *
+ * @param[in] session   The SSH session.
+ *
+ * @param[in] scounter  Counter for byte data handled by the session sockets.
+ *
+ * @param[in] rcounter  Counter for byte and packet data handled by the session,
+ *                      prior compression and SSH overhead.
+ */
+void ssh_set_session_counters(ssh_session session, ssh_counter scounter,
+                              ssh_counter rcounter) {
+    if (session != NULL) {
+        session->socket_counter = scounter;
+        session->raw_counter = rcounter;
+    }
+}
+
 /** @} */
 
 /* vim: set ts=4 sw=4 et cindent: */
diff --git a/src/socket.c b/src/socket.c
--- a/src/socket.c
+++ b/src/socket.c
@@ -280,6 +280,9 @@
 			}
 		}
 		if(r>0){
+            if (s->session->socket_counter != NULL) {
+                s->session->socket_counter->in_bytes += r;
+            }
 			/* Bufferize the data and then call the callback */
             r = ssh_buffer_add_data(s->in_buffer,buffer,r);
             if (r < 0) {
@@ -659,6 +662,9 @@
       return SSH_ERROR;
     }
     buffer_pass_bytes(s->out_buffer, w);
+    if (s->session->socket_counter != NULL) {
+        s->session->socket_counter->out_bytes += w;
+    }
   }
 
   /* Is there some data pending? */