From 175375bc0995d552223d3c09f1e17fdd63cc8be8 Mon Sep 17 00:00:00 2001
From: Jon Simons <jon@jonsimons.org>
Date: Wed, 24 Apr 2019 09:50:17 -0700
Subject: [PATCH] tests/pkd: support --rekey to set rekey data limit

Support an optional --rekey input to set the SSH session rekey data
limit using SSH_OPTIONS_REKEY_DATA.  This flag can be used together
with --buffer to test out server rekeying.

Signed-off-by: Jon Simons <jon@jonsimons.org>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>
---
 tests/pkd/pkd_daemon.c | 7 +++++++
 tests/pkd/pkd_daemon.h | 2 ++
 tests/pkd/pkd_hello.c  | 5 +++++
 3 files changed, 14 insertions(+)

diff --git a/tests/pkd/pkd_daemon.c b/tests/pkd/pkd_daemon.c
index 29948d88..6cecbedc 100644
--- a/tests/pkd/pkd_daemon.c
+++ b/tests/pkd/pkd_daemon.c
@@ -249,6 +249,7 @@ static int pkd_exec_hello(int fd, struct pkd_daemon_args *args)
     const char *default_kex = NULL;
     char *all_kex = NULL;
     size_t kex_len = 0;
+    const uint64_t rekey_data_limit = args->rekey_data_limit;
 
     pkd_state.eof_received = 0;
     pkd_state.close_received  = 0;
@@ -311,6 +312,12 @@ static int pkd_exec_hello(int fd, struct pkd_daemon_args *args)
         goto outclose;
     }
 
+    rc = ssh_options_set(s, SSH_OPTIONS_REKEY_DATA, &rekey_data_limit);
+    if (rc != 0) {
+        pkderr("ssh_options_set rekey data: %s\n", ssh_get_error(s));
+        goto outclose;
+    }
+
     /*
      * ssh_bind_accept loads host key as side-effect.  If this
      * succeeds, the given 'fd' will be closed upon 'ssh_free(s)'.
diff --git a/tests/pkd/pkd_daemon.h b/tests/pkd/pkd_daemon.h
index ffc36c6c..493326c1 100644
--- a/tests/pkd/pkd_daemon.h
+++ b/tests/pkd/pkd_daemon.h
@@ -28,6 +28,8 @@ struct pkd_daemon_args {
         size_t len;
     } payload;
 
+    uint64_t rekey_data_limit;
+
     struct {
         int list;
 
diff --git a/tests/pkd/pkd_hello.c b/tests/pkd/pkd_hello.c
index f5588bd6..c255ef1c 100644
--- a/tests/pkd/pkd_hello.c
+++ b/tests/pkd/pkd_hello.c
@@ -64,6 +64,8 @@ static struct argp_option options[] = {
       "Run in socket-wrapper mode using the given mkdtemp directory template", 0 },
     { "stdout", 'o', NULL, 0,
       "Emit pkd stdout messages", 0 },
+    { "rekey", 'r', "limit", 0,
+      "Set the given rekey data limit, in bytes, using SSH_OPTIONS_REKEY_DATA", 0 },
     { "test", 't', "testname", 0,
       "Run tests matching the given testname", 0 },
     { "verbose", 'v', NULL, 0,
@@ -96,6 +98,9 @@ static error_t parse_opt(int key, char *arg, struct argp_state *state) {
     case 'o':
         pkd_dargs.opts.log_stdout = 1;
         break;
+    case 'r':
+        pkd_dargs.rekey_data_limit = atoi(arg);
+        break;
     case 't':
         pkd_dargs.opts.testname = arg;
         break;