From a61368a06a827fb7cf33f2c3a24a621acc56fcc6 Mon Sep 17 00:00:00 2001
From: Jakub Jelen <jjelen@redhat.com>
Date: Wed, 7 Nov 2018 14:49:43 +0100
Subject: [PATCH] packet: Prepare counters to handle rekeying limits

Signed-off-by: Jakub Jelen <jjelen@redhat.com>
Reviewed-by: Daiki Ueno <dueno@redhat.com>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>
---
 include/libssh/crypto.h  |  5 +++++
 include/libssh/session.h |  2 ++
 src/packet.c             | 15 ++++++++++++++-
 src/packet_cb.c          | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/include/libssh/crypto.h b/include/libssh/crypto.h
index 0160bafc..16f07325 100644
--- a/include/libssh/crypto.h
+++ b/include/libssh/crypto.h
@@ -164,6 +164,11 @@ struct ssh_cipher_struct {
     struct chacha20_poly1305_keysched *chacha20_schedule;
     unsigned int keysize; /* bytes of key used. != keylen */
     size_t tag_size; /* overhead required for tag */
+    /* Counters for rekeying initialization */
+    uint32_t packets;
+    uint64_t blocks;
+    /* Rekeying limit for the cipher or manually enforced */
+    uint64_t max_blocks;
     /* sets the new key for immediate use */
     int (*set_encrypt_key)(struct ssh_cipher_struct *cipher, void *key, void *IV);
     int (*set_decrypt_key)(struct ssh_cipher_struct *cipher, void *key, void *IV);
diff --git a/include/libssh/session.h b/include/libssh/session.h
index 109cb5f5..fb443b59 100644
--- a/include/libssh/session.h
+++ b/include/libssh/session.h
@@ -30,6 +30,7 @@
 #include "libssh/channels.h"
 #include "libssh/poll.h"
 #include "libssh/config.h"
+#include "libssh/misc.h"
 
 /* These are the different states a SSH session can be into its life */
 enum ssh_session_state_e {
@@ -113,6 +114,7 @@ struct ssh_session_struct {
     int openssh;
     uint32_t send_seq;
     uint32_t recv_seq;
+    struct ssh_timestamp last_rekey_time;
 
     int connected;
     /* !=0 when the user got a session handle */
diff --git a/src/packet.c b/src/packet.c
index 847a363d..d6162432 100644
--- a/src/packet.c
+++ b/src/packet.c
@@ -1151,6 +1151,13 @@ int ssh_packet_socket_callback(const void *data, size_t receivedlen, void *user)
 #endif /* WITH_ZLIB */
             payloadsize = ssh_buffer_get_len(session->in_buffer);
             session->recv_seq++;
+            if (session->current_crypto != NULL) {
+                struct ssh_cipher_struct *cipher = NULL;
+
+                cipher = session->current_crypto->in_cipher;
+                cipher->packets++;
+                cipher->blocks += payloadsize / cipher->blocksize;
+            }
             if (session->raw_counter != NULL) {
                 session->raw_counter->in_bytes += payloadsize;
                 session->raw_counter->in_packets++;
@@ -1486,6 +1493,13 @@ static int packet_send2(ssh_session session)
 
     rc = ssh_packet_write(session);
     session->send_seq++;
+    if (session->current_crypto != NULL) {
+        struct ssh_cipher_struct *cipher = NULL;
+
+        cipher = session->current_crypto->out_cipher;
+        cipher->packets++;
+        cipher->blocks += payloadsize / cipher->blocksize;
+    }
     if (session->raw_counter != NULL) {
         session->raw_counter->out_bytes += payloadsize;
         session->raw_counter->out_packets++;
@@ -1504,7 +1518,6 @@ error:
     return rc; /* SSH_OK, AGAIN or ERROR */
 }
 
-
 int ssh_packet_send(ssh_session session) {
     return packet_send2(session);
 }
diff --git a/src/packet_cb.c b/src/packet_cb.c
index 8e244cee..03b2b961 100644
--- a/src/packet_cb.c
+++ b/src/packet_cb.c
@@ -136,6 +136,33 @@ error:
   return SSH_PACKET_USED;
 }
 
+static void
+ssh_init_rekey_state(struct ssh_session_struct *session,
+                     struct ssh_cipher_struct *cipher)
+{
+    /* Reset the counters: should be NOOP */
+    cipher->packets = 0;
+    cipher->blocks = 0;
+
+    /* Default rekey limits for ciphers as specified in RFC4344, Section 3.2 */
+    if (cipher->blocksize >= 16) {
+        /* For larger block size (L bits) use maximum of 2**(L/4) blocks */
+        cipher->max_blocks = (uint64_t)1 << (cipher->blocksize*2);
+    } else {
+        /* For smaller blocks use limit of 1 GB as recommended in RFC4253 */
+        cipher->max_blocks = ((uint64_t)1 << 30) / cipher->blocksize;
+    }
+    /* If we have limit provided by user, use the smaller one */
+    if (session->opts.rekey_data != 0) {
+        cipher->max_blocks = MIN(cipher->max_blocks,
+                                 session->opts.rekey_data / cipher->blocksize);
+    }
+
+    SSH_LOG(SSH_LOG_PROTOCOL,
+            "Set rekey after %" PRIu64 " blocks",
+            cipher->max_blocks);
+}
+
 SSH_PACKET_CALLBACK(ssh_packet_newkeys){
   ssh_string sig_blob = NULL;
   ssh_signature sig = NULL;
@@ -234,6 +261,17 @@ SSH_PACKET_CALLBACK(ssh_packet_newkeys){
     /* FIXME later, include a function to change keys */
     session->current_crypto = session->next_crypto;
 
+    /* Initialize rekeying states */
+    ssh_init_rekey_state(session,
+                         session->current_crypto->out_cipher);
+    ssh_init_rekey_state(session,
+                         session->current_crypto->in_cipher);
+    if (session->opts.rekey_time != 0) {
+        ssh_timestamp_init(&session->last_rekey_time);
+        SSH_LOG(SSH_LOG_PROTOCOL, "Set rekey after %" PRIu32 " seconds",
+                session->opts.rekey_time/1000);
+    }
+
     session->next_crypto = crypto_new();
     if (session->next_crypto == NULL) {
       ssh_set_error_oom(session);