From a277dd92776c5c18d4e08a555e169c46cf75ff94 Mon Sep 17 00:00:00 2001 From: Audrius Butkevicius Date: Wed, 12 Feb 2014 16:53:18 +0000 Subject: [PATCH] Add session/channel byte/packet counters Signed-off-by: Audrius Butkevicius --- include/libssh/channels.h | 2 ++ include/libssh/libssh.h | 11 +++++++++++ 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, 106 insertions(+) diff --git a/include/libssh/channels.h b/include/libssh/channels.h index 4c726453..b05c4c02 100644 --- a/include/libssh/channels.h +++ b/include/libssh/channels.h @@ -75,6 +75,8 @@ struct ssh_channel_struct { 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 index b10dca28..ea63feeb 100644 --- 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; @@ -395,6 +402,8 @@ LIBSSH_API int ssh_channel_send_eof(ssh_channel channel); LIBSSH_API int ssh_channel_select(ssh_channel *readchans, ssh_channel *writechans, ssh_channel *exceptchans, struct timeval * timeout); LIBSSH_API void ssh_channel_set_blocking(ssh_channel channel, int blocking); +LIBSSH_API void ssh_channel_set_counter(ssh_channel channel, + ssh_counter counter); 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); @@ -575,6 +584,8 @@ LIBSSH_API int ssh_select(ssh_channel *channels, ssh_channel *outchannels, socke LIBSSH_API int ssh_service_request(ssh_session session, const char *service); LIBSSH_API int ssh_set_agent_channel(ssh_session session, ssh_channel channel); LIBSSH_API void ssh_set_blocking(ssh_session session, int blocking); +LIBSSH_API void ssh_set_counters(ssh_session session, ssh_counter scounter, + ssh_counter rcounter); LIBSSH_API void ssh_set_fd_except(ssh_session session); LIBSSH_API void ssh_set_fd_toread(ssh_session session); LIBSSH_API void ssh_set_fd_towrite(ssh_session session); diff --git a/include/libssh/session.h b/include/libssh/session.h index c360a70c..8480135d 100644 --- a/include/libssh/session.h +++ b/include/libssh/session.h @@ -188,6 +188,9 @@ struct ssh_session_struct { 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 index b6f59f5f..7a72cfe4 100644 --- a/src/channels.c +++ b/src/channels.c @@ -555,6 +555,9 @@ SSH_PACKET_CALLBACK(channel_rcv_data){ 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 @@ static int channel_write_common(ssh_channel channel, 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 @@ int ssh_channel_read_timeout(ssh_channel channel, 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) { @@ -3269,6 +3278,31 @@ int ssh_channel_select(ssh_channel *readchans, ssh_channel *writechans, return 0; } +/** + * @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_channel_set_counter(channel, &counter); + * @endcode + * + * @param[in] channel The SSH channel. + * + * @param[in] counter Counter for bytes handled by the channel. + */ +void ssh_channel_set_counter(ssh_channel channel, + ssh_counter counter) { + if (channel != NULL) { + channel->counter = counter; + } +} + #if WITH_SERVER /** * @brief Blocking write on a channel stderr. diff --git a/src/channels1.c b/src/channels1.c index 2b307408..4d82c636 100644 --- a/src/channels1.c +++ b/src/channels1.c @@ -367,6 +367,9 @@ int channel_write1(ssh_channel channel, const void *data, int len) { 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 index 4aed2fac..71ba40d1 100644 --- a/src/packet.c +++ b/src/packet.c @@ -311,6 +311,10 @@ int ssh_packet_socket_callback(const void *data, size_t receivedlen, void *user) #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 @@ static int packet_send2(ssh_session session) { 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 index 93c047f9..3f22f111 100644 --- a/src/session.c +++ b/src/session.c @@ -831,6 +831,45 @@ error: 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_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_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 index 8f4c9b22..ab47b96a 100644 --- a/src/socket.c +++ b/src/socket.c @@ -280,6 +280,9 @@ int ssh_socket_pollcallback(struct ssh_poll_handle_struct *p, socket_t fd, int r } } 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 @@ int ssh_socket_nonblocking_flush(ssh_socket s) { 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? */