dh: Confine DH KEX keys handling into fewer functions

In preparation for deferring computation on DH secret material to
crypto library specific backends

Signed-off-by: Simo Sorce <simo@redhat.com>
Reviewed-by: Anderson Toshiyuki Sasaki <ansasaki@redhat.com>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>

include/libssh/dh.h

@@ -25,24 +25,18 @@
 
 #include "libssh/crypto.h"
 
-struct dh_keypair {
+struct dh_ctx;
-    bignum priv_key;
-    bignum pub_key;
-};
 
-struct dh_ctx {
+#define DH_CLIENT_KEYPAIR 0
-    struct dh_keypair client;
+#define DH_SERVER_KEYPAIR 1
-    struct dh_keypair server;
-    bignum generator;
-    bignum modulus;
-};
 
 int ssh_dh_init(void);
 void ssh_dh_finalize(void);
 
 int ssh_dh_import_next_pubkey_blob(ssh_session session, ssh_string pubkey_blob);
 
-int ssh_dh_build_k(ssh_session session);
+int ssh_dh_compute_shared_secret(struct dh_ctx *ctx, int local, int remote,
+                                 bignum *dest);
 int ssh_client_dh_init(ssh_session session);
 
 ssh_key ssh_dh_get_current_server_publickey(ssh_session session);
@@ -56,9 +50,18 @@ int ssh_dh_get_next_server_publickey_blob(ssh_session session,
 void ssh_server_dh_init(ssh_session session);
 #endif /* WITH_SERVER */
 
-int ssh_dh_init_common(ssh_session session);
+int ssh_dh_init_common(struct ssh_crypto_struct *crypto);
 void ssh_dh_cleanup(struct ssh_crypto_struct *crypto);
-int ssh_dh_generate_secret(ssh_session session, bignum dest);
 int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet);
 
+int ssh_dh_get_parameters(struct dh_ctx *ctx,
+                          const_bignum *modulus, const_bignum *generator);
+int ssh_dh_set_parameters(struct dh_ctx *ctx,
+                          bignum modulus, bignum generator);
+int ssh_dh_keypair_gen_keys(struct dh_ctx *ctx, int peer);
+int ssh_dh_keypair_get_keys(struct dh_ctx *ctx, int peer,
+                            const_bignum *priv, const_bignum *pub);
+int ssh_dh_keypair_set_keys(struct dh_ctx *ctx, int peer,
+                            bignum priv, bignum pub);
+
 #endif /* DH_H_ */

include/libssh/libcrypto.h

@@ -64,6 +64,7 @@ typedef void *EVPCTX;
 #define BROKEN_AES_CTR
 #endif
 typedef BIGNUM*  bignum;
+typedef const BIGNUM* const_bignum;
 typedef BN_CTX* bignum_CTX;
 
 #define bignum_new() BN_new()

include/libssh/libgcrypt.h

@@ -50,6 +50,7 @@ typedef gcry_md_hd_t EVPCTX;
 #define EVP_DIGEST_LEN EVP_MAX_MD_SIZE
 
 typedef gcry_mpi_t bignum;
+typedef const struct gcry_mpi *const_bignum;
 typedef void* bignum_CTX;
 
 /* Constants for curves.  */

include/libssh/libmbedcrypto.h

@@ -60,6 +60,7 @@ typedef mbedtls_md_context_t *EVPCTX;
 #define EVP_DIGEST_LEN EVP_MAX_MD_SIZE
 
 typedef mbedtls_mpi *bignum;
+typedef const mbedtls_mpi *const_bignum;
 typedef void* bignum_CTX;
 
 /* Constants for curves */

src/dh-gex.c

@@ -65,7 +65,7 @@ int ssh_client_dhgex_init(ssh_session session)
 {
     int rc;
 
-    rc = ssh_dh_init_common(session);
+    rc = ssh_dh_init_common(session->next_crypto);
     if (rc != SSH_OK){
         goto error;
     }
@@ -107,6 +107,8 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)
     int blen;
     bignum pmin1 = NULL, one = NULL;
     bignum_CTX ctx = bignum_ctx_new();
+    bignum modulus, generator;
+    const_bignum pubkey;
     (void) type;
     (void) user;
 
@@ -130,8 +132,8 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)
     }
     rc = ssh_buffer_unpack(packet,
                            "BB",
-                           &session->next_crypto->dh_ctx->modulus,
+                           &modulus,
-                           &session->next_crypto->dh_ctx->generator);
+                           &generator);
     if (rc != SSH_OK) {
         ssh_set_error(session, SSH_FATAL, "Invalid DH_GEX_GROUP packet");
         goto error;
@@ -141,7 +143,7 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)
     if (rc != 1) {
         goto error;
     }
-    blen = bignum_num_bits(session->next_crypto->dh_ctx->modulus);
+    blen = bignum_num_bits(modulus);
     if (blen < DH_PMIN || blen > DH_PMAX) {
         ssh_set_error(session,
                 SSH_FATAL,
@@ -151,43 +153,46 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)
                 DH_PMAX);
         goto error;
     }
-    if (bignum_cmp(session->next_crypto->dh_ctx->modulus, one) <= 0) {
+    if (bignum_cmp(modulus, one) <= 0) {
         /* p must be positive and preferably bigger than one */
         ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter p");
     }
-    if (!bignum_is_bit_set(session->next_crypto->dh_ctx->modulus, 0)) {
+    if (!bignum_is_bit_set(modulus, 0)) {
         /* p must be a prime and therefore not divisible by 2 */
         ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter p");
         goto error;
     }
-    bignum_sub(pmin1, session->next_crypto->dh_ctx->modulus, one);
+    bignum_sub(pmin1, modulus, one);
-    if (bignum_cmp(session->next_crypto->dh_ctx->generator, one) <= 0 ||
+    if (bignum_cmp(generator, one) <= 0 ||
-        bignum_cmp(session->next_crypto->dh_ctx->generator, pmin1) > 0) {
+        bignum_cmp(generator, pmin1) > 0) {
         /* generator must be at least 2 and smaller than p-1*/
         ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter g");
         goto error;
     }
-    /* compute and send DH public parameter */
-    rc = ssh_dh_generate_secret(session, session->next_crypto->dh_ctx->client.priv_key);
-    if (rc == SSH_ERROR) {
-        goto error;
-    }
-    rc = bignum_mod_exp(session->next_crypto->dh_ctx->client.pub_key,
-                        session->next_crypto->dh_ctx->generator,
-                        session->next_crypto->dh_ctx->client.priv_key,
-                        session->next_crypto->dh_ctx->modulus,
-                        ctx);
-    if (rc != 1) {
-        goto error;
-    }
-
     bignum_ctx_free(ctx);
     ctx = NULL;
 
+    /* all checks passed, set parameters */
+    rc = ssh_dh_set_parameters(session->next_crypto->dh_ctx,
+                               modulus, generator);
+    if (rc != SSH_OK) {
+        bignum_safe_free(modulus);
+        bignum_safe_free(generator);
+        goto error;
+    }
+
+    /* compute and send DH public parameter */
+    rc = ssh_dh_keypair_gen_keys(session->next_crypto->dh_ctx,
+                                 DH_CLIENT_KEYPAIR);
+    if (rc == SSH_ERROR) {
+        goto error;
+    }
+    rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,
+                                 DH_CLIENT_KEYPAIR, NULL, &pubkey);
     rc = ssh_buffer_pack(session->out_buffer,
                          "bB",
                          SSH2_MSG_KEX_DH_GEX_INIT,
-                         session->next_crypto->dh_ctx->client.pub_key);
+                         pubkey);
     if (rc != SSH_OK) {
         goto error;
     }
@@ -199,6 +204,7 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)
     bignum_safe_free(one);
     bignum_safe_free(pmin1);
     return SSH_PACKET_USED;
+
 error:
     bignum_safe_free(one);
     bignum_safe_free(pmin1);
@@ -216,6 +222,7 @@ static SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_reply)
     struct ssh_crypto_struct *crypto=session->next_crypto;
     int rc;
     ssh_string pubkey_blob = NULL;
+    bignum server_pubkey = NULL;
     (void)type;
     (void)user;
     SSH_LOG(SSH_LOG_PROTOCOL, "SSH_MSG_KEX_DH_GEX_REPLY received");
@@ -223,20 +230,28 @@ static SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_reply)
     ssh_packet_remove_callbacks(session, &ssh_dhgex_client_callbacks);
     rc = ssh_buffer_unpack(packet,
                            "SBS",
-                           &pubkey_blob, &crypto->dh_ctx->server.pub_key,
+                           &pubkey_blob, &server_pubkey,
                            &crypto->dh_server_signature);
-
     if (rc == SSH_ERROR) {
         ssh_set_error(session, SSH_FATAL, "Invalid DH_GEX_REPLY packet");
         goto error;
     }
+    rc = ssh_dh_keypair_set_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR,
+                                 NULL, server_pubkey);
+    if (rc != SSH_OK) {
+        bignum_safe_free(server_pubkey);
+        goto error;
+    }
+
     rc = ssh_dh_import_next_pubkey_blob(session, pubkey_blob);
     ssh_string_free(pubkey_blob);
     if (rc != 0) {
         goto error;
     }
 
-    rc = ssh_dh_build_k(session);
+    rc = ssh_dh_compute_shared_secret(session->next_crypto->dh_ctx,
+                                      DH_CLIENT_KEYPAIR, DH_SERVER_KEYPAIR,
+                                      &session->next_crypto->shared_secret);
     if (rc == SSH_ERROR) {
         ssh_set_error(session, SSH_FATAL, "Could not generate shared secret");
         goto error;
@@ -526,12 +541,13 @@ static struct ssh_packet_callbacks_struct ssh_dhgex_server_callbacks = {
 void ssh_server_dhgex_init(ssh_session session){
     /* register the packet callbacks */
     ssh_packet_set_callbacks(session, &ssh_dhgex_server_callbacks);
-    ssh_dh_init_common(session);
+    ssh_dh_init_common(session->next_crypto);
     session->dh_handshake_state = DH_STATE_INIT;
 }
 
 static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request)
 {
+    bignum modulus = NULL, generator = NULL;
     uint32_t pmin, pn, pmax;
     size_t size = 0;
     int rc;
@@ -578,8 +594,8 @@ static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request)
                               pn,
                               pmax,
                               &size,
-                              &session->next_crypto->dh_ctx->modulus,
+                              &modulus,
-                              &session->next_crypto->dh_ctx->generator);
+                              &generator);
     if (rc == SSH_ERROR) {
         ssh_set_error(session,
                       SSH_FATAL,
@@ -589,11 +605,18 @@ static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request)
                       pmax);
         goto error;
     }
+    rc = ssh_dh_set_parameters(session->next_crypto->dh_ctx,
+                               modulus, generator);
+    if (rc != SSH_OK) {
+        bignum_safe_free(generator);
+        bignum_safe_free(modulus);
+        goto error;
+    }
     rc = ssh_buffer_pack(session->out_buffer,
                          "bBB",
                          SSH2_MSG_KEX_DH_GEX_GROUP,
-                         session->next_crypto->dh_ctx->modulus,
+                         modulus,
-                         session->next_crypto->dh_ctx->generator);
+                         generator);
     if (rc != SSH_OK) {
         ssh_set_error_invalid(session);
         goto error;
@@ -605,8 +628,8 @@ static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request)
     if (rc == SSH_ERROR) {
         goto error;
     }
-error:
 
+error:
     return SSH_PACKET_USED;
 }
 

src/dh.c

@@ -226,6 +226,18 @@ static bignum p_group16;
 static bignum p_group18;
 static int dh_crypto_initialized;
 
+struct dh_keypair {
+    bignum priv_key;
+    bignum pub_key;
+};
+
+struct dh_ctx {
+    /* 0 is client, 1 is server */
+    struct dh_keypair keypair[2];
+    bignum generator;
+    bignum modulus;
+};
+
 /**
  * @internal
  * @brief Initialize global constants used in DH key agreement
@@ -294,6 +306,24 @@ void ssh_dh_finalize(void)
     dh_crypto_initialized = 0;
 }
 
+static void ssh_dh_free_modulus(struct dh_ctx *ctx)
+{
+    if ((ctx->modulus != p_group1) &&
+        (ctx->modulus != p_group14) &&
+        (ctx->modulus != p_group16) &&
+        (ctx->modulus != p_group18)) {
+        bignum_safe_free(ctx->modulus);
+    }
+    ctx->modulus = NULL;
+}
+
+static void ssh_dh_free_generator(struct dh_ctx *ctx)
+{
+    if (ctx->generator != g) {
+        bignum_safe_free(ctx->generator);
+    }
+}
+
 static void ssh_dh_free_dh_keypair(struct dh_keypair *keypair)
 {
     bignum_safe_free(keypair->priv_key);
@@ -323,13 +353,111 @@ done:
     return rc;
 }
 
+int ssh_dh_keypair_get_keys(struct dh_ctx *ctx, int peer,
+                            const_bignum *priv, const_bignum *pub)
+{
+    if (((peer != DH_CLIENT_KEYPAIR) && (peer != DH_SERVER_KEYPAIR)) ||
+        ((priv == NULL) && (pub == NULL)) || (ctx == NULL)) {
+        return SSH_ERROR;
+    }
+
+    if (priv) {
+        /* check that we have something in it */
+        if (bignum_num_bits(ctx->keypair[peer].priv_key)) {
+            *priv = ctx->keypair[peer].priv_key;
+        } else {
+            return SSH_ERROR;
+        }
+    }
+
+    if (pub) {
+        /* check that we have something in it */
+        if (bignum_num_bits(ctx->keypair[peer].pub_key)) {
+            *pub = ctx->keypair[peer].pub_key;
+        } else {
+            return SSH_ERROR;
+        }
+    }
+
+    return SSH_OK;
+}
+
+int ssh_dh_keypair_set_keys(struct dh_ctx *ctx, int peer,
+                            bignum priv, bignum pub)
+{
+    if (((peer != DH_CLIENT_KEYPAIR) && (peer != DH_SERVER_KEYPAIR)) ||
+        ((priv == NULL) && (pub == NULL)) || (ctx == NULL)) {
+        return SSH_ERROR;
+    }
+
+    if (priv) {
+        bignum_safe_free(ctx->keypair[peer].priv_key);
+        ctx->keypair[peer].priv_key = priv;
+    }
+    if (pub) {
+        bignum_safe_free(ctx->keypair[peer].pub_key);
+        ctx->keypair[peer].pub_key = pub;
+    }
+    return SSH_OK;
+}
+
+int ssh_dh_get_parameters(struct dh_ctx *ctx,
+                          const_bignum *modulus, const_bignum *generator)
+{
+    if (ctx == NULL) {
+        return SSH_ERROR;
+    }
+    if (modulus) {
+        *modulus = ctx->modulus;
+    }
+    if (generator) {
+        *generator = ctx->generator;
+    }
+
+    return SSH_OK;
+}
+
+int ssh_dh_set_parameters(struct dh_ctx *ctx,
+                          bignum modulus, bignum generator)
+{
+    int rc;
+
+    if ((ctx == NULL) || ((modulus == NULL) && (generator == NULL))) {
+        return SSH_ERROR;
+    }
+    /* when setting modulus or generator,
+     * make sure to invalidate existing keys */
+    ssh_dh_free_dh_keypair(&ctx->keypair[DH_CLIENT_KEYPAIR]);
+    ssh_dh_free_dh_keypair(&ctx->keypair[DH_SERVER_KEYPAIR]);
+
+    rc = ssh_dh_init_dh_keypair(&ctx->keypair[DH_CLIENT_KEYPAIR]);
+    if (rc != SSH_OK) {
+        goto done;
+    }
+    rc = ssh_dh_init_dh_keypair(&ctx->keypair[DH_SERVER_KEYPAIR]);
+    if (rc != SSH_OK) {
+        goto done;
+    }
+
+    if (modulus) {
+        ssh_dh_free_modulus(ctx);
+        ctx->modulus = modulus;
+    }
+    if (generator) {
+        ssh_dh_free_generator(ctx);
+        ctx->generator = generator;
+    }
+
+done:
+    return rc;
+}
+
 /**
  * @internal
  * @brief allocate and initialize ephemeral values used in dh kex
  */
-int ssh_dh_init_common(ssh_session session)
+int ssh_dh_init_common(struct ssh_crypto_struct *crypto)
 {
-    struct ssh_crypto_struct *crypto = session->next_crypto;
     struct dh_ctx *ctx = NULL;
     int rc;
 
@@ -340,48 +468,26 @@ int ssh_dh_init_common(ssh_session session)
 
     switch (crypto->kex_type) {
     case SSH_KEX_DH_GROUP1_SHA1:
-        ctx->generator = g;
+        rc = ssh_dh_set_parameters(ctx, p_group1, g);
-        ctx->modulus = p_group1;
         break;
     case SSH_KEX_DH_GROUP14_SHA1:
-        ctx->generator = g;
+        rc = ssh_dh_set_parameters(ctx, p_group14, g);
-        ctx->modulus = p_group14;
         break;
     case SSH_KEX_DH_GROUP16_SHA512:
-        ctx->generator = g;
+        rc = ssh_dh_set_parameters(ctx, p_group16, g);
-        ctx->modulus = p_group16;
         break;
     case SSH_KEX_DH_GROUP18_SHA512:
-        ctx->generator = g;
+        rc = ssh_dh_set_parameters(ctx, p_group18, g);
-        ctx->modulus = p_group18;
         break;
     default:
+        rc = SSH_OK;
         break;
     }
 
-    rc = ssh_dh_init_dh_keypair(&ctx->client);
-    if (rc != SSH_OK) {
-        goto done;
-    }
-
-    rc = ssh_dh_init_dh_keypair(&ctx->server);
-    if (rc != SSH_OK) {
-        goto done;
-    }
-
-    crypto->shared_secret = bignum_new();
-    if (crypto->shared_secret == NULL) {
-        goto done;
-    }
-
     crypto->dh_ctx = ctx;
-    rc = SSH_OK;
+
-done:
     if (rc != SSH_OK) {
-        ssh_dh_free_dh_keypair(&ctx->client);
+        ssh_dh_cleanup(crypto);
-        ssh_dh_free_dh_keypair(&ctx->server);
-        free(ctx);
-        ssh_set_error_oom(session);
     }
     return rc;
 }
@@ -394,18 +500,11 @@ void ssh_dh_cleanup(struct ssh_crypto_struct *crypto)
         return;
     }
 
-    ssh_dh_free_dh_keypair(&ctx->client);
+    ssh_dh_free_dh_keypair(&ctx->keypair[DH_CLIENT_KEYPAIR]);
-    ssh_dh_free_dh_keypair(&ctx->server);
+    ssh_dh_free_dh_keypair(&ctx->keypair[DH_SERVER_KEYPAIR]);
 
-    if ((ctx->modulus != p_group1) &&
+    ssh_dh_free_modulus(ctx);
-        (ctx->modulus != p_group14) &&
+    ssh_dh_free_generator(ctx);
-        (ctx->modulus != p_group16) &&
-        (ctx->modulus != p_group18)) {
-        bignum_safe_free(ctx->modulus);
-    }
-    if (ctx->generator != g) {
-        bignum_safe_free(ctx->generator);
-    }
     free(ctx);
     crypto->dh_ctx = NULL;
 }
@@ -420,12 +519,16 @@ int ssh_dh_import_next_pubkey_blob(ssh_session session, ssh_string pubkey_blob)
 #ifdef DEBUG_CRYPTO
 static void ssh_dh_debug(ssh_session session)
 {
-    ssh_print_bignum("p", session->next_crypto->dh_ctx->modulus);
+    struct ssh_crypto_struct *crypto = session->next_crypto;
-    ssh_print_bignum("g", session->next_crypto->dh_ctx->generator);
+    const_bignum x, y, e, f;
-    ssh_print_bignum("x", session->next_crypto->dh_ctx->client.priv_key);
+    ssh_dh_keypair_get_keys(crypto->dh_ctx, DH_CLIENT_KEYPAIR, &x, &e);
-    ssh_print_bignum("y", session->next_crypto->dh_ctx->server.priv_key);
+    ssh_dh_keypair_get_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR, &y, &f);
-    ssh_print_bignum("e", session->next_crypto->dh_ctx->client.pub_key);
+    ssh_print_bignum("p", crypto->dh_ctx->modulus);
-    ssh_print_bignum("f", session->next_crypto->dh_ctx->server.pub_key);
+    ssh_print_bignum("g", crypto->dh_ctx->generator);
+    ssh_print_bignum("x", x);
+    ssh_print_bignum("y", y);
+    ssh_print_bignum("e", e);
+    ssh_print_bignum("f", f);
  
     ssh_print_hexa("Session server cookie",
                    session->next_crypto->server_kex.cookie, 16);
@@ -439,11 +542,12 @@ static void ssh_dh_debug(ssh_session session)
 
 /** @internal
  * @brief generates a secret DH parameter of at least DH_SECURITY_BITS
- *        security.
+ *        security as well as the corresponding public key.
- * @param[out] dest preallocated bignum where to store the parameter
+ * @param[out] parms a dh_kex paramters structure with preallocated bignum
+ *             where to store the parameters
  * @return SSH_OK on success, SSH_ERROR on error
  */
-int ssh_dh_generate_secret(ssh_session session, bignum dest)
+int ssh_dh_keypair_gen_keys(struct dh_ctx *dh_ctx, int peer)
 {
     bignum tmp = NULL;
     bignum_CTX ctx = NULL;
@@ -459,7 +563,7 @@ int ssh_dh_generate_secret(ssh_session session, bignum dest)
     if (tmp == NULL) {
         goto error;
     }
-    p_bits = bignum_num_bits(session->next_crypto->dh_ctx->modulus);
+    p_bits = bignum_num_bits(dh_ctx->modulus);
     /* we need at most DH_SECURITY_BITS */
     bits = MIN(DH_SECURITY_BITS * 2, p_bits);
     /* ensure we're not too close of p so rnd()%p stays uniform */
@@ -470,7 +574,13 @@ int ssh_dh_generate_secret(ssh_session session, bignum dest)
     if (rc != 1) {
         goto error;
     }
-    rc = bignum_mod(dest, tmp, session->next_crypto->dh_ctx->modulus, ctx);
+    rc = bignum_mod(dh_ctx->keypair[peer].priv_key, tmp, dh_ctx->modulus, ctx);
+    if (rc != 1) {
+        goto error;
+    }
+    /* Now compute the corresponding public key */
+    rc = bignum_mod_exp(dh_ctx->keypair[peer].pub_key, dh_ctx->generator,
+                        dh_ctx->keypair[peer].priv_key, dh_ctx->modulus, ctx);
     if (rc != 1) {
         goto error;
     }
@@ -478,13 +588,20 @@ int ssh_dh_generate_secret(ssh_session session, bignum dest)
     bignum_ctx_free(ctx);
     return SSH_OK;
 error:
-    ssh_set_error_oom(session);
     bignum_safe_free(tmp);
     bignum_ctx_free(ctx);
     return SSH_ERROR;
 }
 
-int ssh_dh_build_k(ssh_session session)
+/** @internal
+ * @brief generates a shared secret between the local peer and the remote peer
+ * @param[in] local peer identifier
+ * @param[in] remote peer identifier
+ * @param[out] dest a preallocated bignum where to store parameter
+ * @return SSH_OK on success, SSH_ERROR on error
+ */
+int ssh_dh_compute_shared_secret(struct dh_ctx *dh_ctx, int local, int remote,
+                                 bignum *dest)
 {
     int rc;
     bignum_CTX ctx = bignum_ctx_new();
@@ -492,21 +609,19 @@ int ssh_dh_build_k(ssh_session session)
         return -1;
     }
 
-    /* the server and clients don't use the same numbers */
+    if (*dest == NULL) {
-    if (session->client) {
+        *dest = bignum_new();
-        rc = bignum_mod_exp(session->next_crypto->shared_secret,
+        if (*dest == NULL) {
-                            session->next_crypto->dh_ctx->server.pub_key,
+            rc = 0;
-                            session->next_crypto->dh_ctx->client.priv_key,
+            goto done;
-                            session->next_crypto->dh_ctx->modulus,
+        }
-                            ctx);
-    } else {
-        rc = bignum_mod_exp(session->next_crypto->shared_secret,
-                            session->next_crypto->dh_ctx->client.pub_key,
-                            session->next_crypto->dh_ctx->server.priv_key,
-                            session->next_crypto->dh_ctx->modulus,
-                            ctx);
     }
 
+    rc = bignum_mod_exp(*dest, dh_ctx->keypair[remote].pub_key,
+                        dh_ctx->keypair[local].priv_key,
+                        dh_ctx->modulus, ctx);
+
+done:
     bignum_ctx_free(ctx);
     ssh_dh_debug(session);
     if (rc != 1) {
@@ -532,29 +647,25 @@ static struct ssh_packet_callbacks_struct ssh_dh_client_callbacks = {
  * @brief Starts diffie-hellman-group1 key exchange
  */
 int ssh_client_dh_init(ssh_session session){
-  bignum_CTX ctx = bignum_ctx_new();
+  struct ssh_crypto_struct *crypto = session->next_crypto;
+  const_bignum pubkey;
   int rc;
 
-  if (bignum_ctx_invalid(ctx)) {
+  rc = ssh_dh_init_common(crypto);
-    goto error;
-  }
-  rc = ssh_dh_init_common(session);
   if (rc == SSH_ERROR) {
     goto error;
   }
 
-  rc = ssh_dh_generate_secret(session, session->next_crypto->dh_ctx->client.priv_key);
+  rc = ssh_dh_keypair_gen_keys(crypto->dh_ctx, DH_CLIENT_KEYPAIR);
   if (rc == SSH_ERROR){
       goto error;
   }
-  rc = bignum_mod_exp(session->next_crypto->dh_ctx->client.pub_key, session->next_crypto->dh_ctx->generator, session->next_crypto->dh_ctx->client.priv_key,
+  rc = ssh_dh_keypair_get_keys(crypto->dh_ctx, DH_CLIENT_KEYPAIR,
-      session->next_crypto->dh_ctx->modulus, ctx);
+                               NULL, &pubkey);
-
+  if (rc != SSH_OK) {
-  bignum_ctx_free(ctx);
-  if (rc != 1) {
     goto error;
   }
-  rc = ssh_buffer_pack(session->out_buffer, "bB", SSH2_MSG_KEXDH_INIT, session->next_crypto->dh_ctx->client.pub_key);
+  rc = ssh_buffer_pack(session->out_buffer, "bB", SSH2_MSG_KEXDH_INIT, pubkey);
   if (rc != SSH_OK) {
     goto error;
   }
@@ -566,13 +677,14 @@ int ssh_client_dh_init(ssh_session session){
   rc = ssh_packet_send(session);
   return rc;
 error:
-  ssh_dh_cleanup(session->next_crypto);
+  ssh_dh_cleanup(crypto);
   return SSH_ERROR;
 }
 
 SSH_PACKET_CALLBACK(ssh_packet_client_dh_reply){
   struct ssh_crypto_struct *crypto=session->next_crypto;
   ssh_string pubkey_blob = NULL;
+  bignum server_pubkey;
   int rc;
 
   (void)type;
@@ -580,18 +692,26 @@ SSH_PACKET_CALLBACK(ssh_packet_client_dh_reply){
 
   ssh_packet_remove_callbacks(session, &ssh_dh_client_callbacks);
 
-  rc = ssh_buffer_unpack(packet, "SBS", &pubkey_blob, &crypto->dh_ctx->server.pub_key,
+  rc = ssh_buffer_unpack(packet, "SBS", &pubkey_blob, &server_pubkey,
           &crypto->dh_server_signature);
   if (rc == SSH_ERROR) {
       goto error;
   }
+  rc = ssh_dh_keypair_set_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR,
+                               NULL, server_pubkey);
+  if (rc != SSH_OK) {
+      bignum_safe_free(server_pubkey);
+      goto error;
+  }
   rc = ssh_dh_import_next_pubkey_blob(session, pubkey_blob);
   ssh_string_free(pubkey_blob);
   if (rc != 0) {
       goto error;
   }
   
-  rc = ssh_dh_build_k(session);
+  rc = ssh_dh_compute_shared_secret(session->next_crypto->dh_ctx,
+                                    DH_CLIENT_KEYPAIR, DH_SERVER_KEYPAIR,
+                                    &session->next_crypto->shared_secret);
   if (rc == SSH_ERROR){
     ssh_set_error(session, SSH_FATAL, "Could not generate shared secret");
     goto error;
@@ -638,7 +758,7 @@ void ssh_server_dh_init(ssh_session session){
     /* register the packet callbacks */
     ssh_packet_set_callbacks(session, &ssh_dh_server_callbacks);
 
-    ssh_dh_init_common(session);
+    ssh_dh_init_common(session->next_crypto);
 }
 
 /** @internal
@@ -647,34 +767,30 @@ void ssh_server_dh_init(ssh_session session){
  */
 int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
 {
+    struct ssh_crypto_struct *crypto = session->next_crypto;
     ssh_key privkey = NULL;
     ssh_string sig_blob = NULL;
     ssh_string pubkey_blob = NULL;
-    bignum_CTX ctx = bignum_ctx_new();
+    bignum client_pubkey;
+    const_bignum server_pubkey;
     int packet_type;
     int rc;
 
-    if (bignum_ctx_invalid(ctx)) {
+    rc = ssh_buffer_unpack(packet, "B", &client_pubkey);
-        goto error;
-    }
-    rc = ssh_buffer_unpack(packet, "B", &session->next_crypto->dh_ctx->client.pub_key);
     if (rc == SSH_ERROR) {
         ssh_set_error(session, SSH_FATAL, "No e number in client request");
         goto error;
     }
 
-    rc = ssh_dh_generate_secret(session, session->next_crypto->dh_ctx->server.priv_key);
+    rc = ssh_dh_keypair_set_keys(crypto->dh_ctx, DH_CLIENT_KEYPAIR,
-    if (rc == SSH_ERROR){
+                                 NULL, client_pubkey);
+    if (rc != SSH_OK) {
+        bignum_safe_free(client_pubkey);
         goto error;
     }
 
-    rc = bignum_mod_exp(session->next_crypto->dh_ctx->server.pub_key,
+    rc = ssh_dh_keypair_gen_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR);
-                        session->next_crypto->dh_ctx->generator,
+    if (rc == SSH_ERROR) {
-                        session->next_crypto->dh_ctx->server.priv_key,
-                        session->next_crypto->dh_ctx->modulus, ctx);
-    bignum_ctx_free(ctx);
-    ctx = NULL;
-    if (rc != 1) {
         goto error;
     }
 
@@ -682,7 +798,9 @@ int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
     if (rc != SSH_OK) {
         goto error;
     }
-    rc = ssh_dh_build_k(session);
+    rc = ssh_dh_compute_shared_secret(crypto->dh_ctx,
+                                      DH_SERVER_KEYPAIR, DH_CLIENT_KEYPAIR,
+                                      &crypto->shared_secret);
     if (rc == SSH_ERROR) {
         ssh_set_error(session, SSH_FATAL, "Could not generate shared secret");
         goto error;
@@ -697,7 +815,7 @@ int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
         ssh_set_error(session, SSH_FATAL, "Could not sign the session id");
         goto error;
     }
-    switch (session->next_crypto->kex_type){
+    switch (crypto->kex_type){
     case SSH_KEX_DH_GROUP1_SHA1:
     case SSH_KEX_DH_GROUP14_SHA1:
     case SSH_KEX_DH_GROUP16_SHA512:
@@ -714,7 +832,11 @@ int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
         ssh_set_error(session, SSH_FATAL, "Invalid kex type");
         goto error;
     }
-   
+    rc = ssh_dh_keypair_get_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR,
+                                 NULL, &server_pubkey);
+    if (rc != SSH_OK){
+        goto error;
+    }
     rc = ssh_dh_get_next_server_publickey_blob(session, &pubkey_blob);
     if (rc != SSH_OK){
         ssh_set_error_oom(session);
@@ -724,7 +846,7 @@ int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
                          "bSBS",
                          packet_type,
                          pubkey_blob,
-                         session->next_crypto->dh_ctx->server.pub_key,
+                         server_pubkey,
                          sig_blob);
     SSH_STRING_FREE(sig_blob);
     SSH_STRING_FREE(pubkey_blob);
@@ -751,9 +873,6 @@ int ssh_server_dh_process_init(ssh_session session, ssh_buffer packet)
 
     return SSH_OK;
 error:
-    if (!bignum_ctx_invalid(ctx)) {
-        bignum_ctx_free(ctx);
-    }
     SSH_STRING_FREE(sig_blob);
     SSH_STRING_FREE(pubkey_blob);
 

src/kex.c

@@ -1047,6 +1047,10 @@ int ssh_make_sessionid(ssh_session session)
     ssh_buffer client_hash = NULL;
     ssh_buffer buf = NULL;
     ssh_string server_pubkey_blob = NULL;
+    const_bignum client_pubkey, server_pubkey;
+#ifdef WITH_GEX
+    const_bignum modulus, generator;
+#endif
     int rc = SSH_ERROR;
 
     buf = ssh_buffer_new();
@@ -1121,10 +1125,20 @@ int ssh_make_sessionid(ssh_session session)
     case SSH_KEX_DH_GROUP14_SHA1:
     case SSH_KEX_DH_GROUP16_SHA512:
     case SSH_KEX_DH_GROUP18_SHA512:
+        rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,
+                                     DH_CLIENT_KEYPAIR, NULL, &client_pubkey);
+        if (rc != SSH_OK) {
+            goto error;
+        }
+        rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,
+                                     DH_SERVER_KEYPAIR, NULL, &server_pubkey);
+        if (rc != SSH_OK) {
+            goto error;
+        }
         rc = ssh_buffer_pack(buf,
                              "BB",
-                             session->next_crypto->dh_ctx->client.pub_key,
+                             client_pubkey,
-                             session->next_crypto->dh_ctx->server.pub_key);
+                             server_pubkey);
         if (rc != SSH_OK) {
             goto error;
         }
@@ -1132,15 +1146,30 @@ int ssh_make_sessionid(ssh_session session)
 #ifdef WITH_GEX
     case SSH_KEX_DH_GEX_SHA1:
     case SSH_KEX_DH_GEX_SHA256:
+        rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,
+                                     DH_CLIENT_KEYPAIR, NULL, &client_pubkey);
+        if (rc != SSH_OK) {
+            goto error;
+        }
+        rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,
+                                     DH_SERVER_KEYPAIR, NULL, &server_pubkey);
+        if (rc != SSH_OK) {
+            goto error;
+        }
+        rc = ssh_dh_get_parameters(session->next_crypto->dh_ctx,
+                                   &modulus, &generator);
+        if (rc != SSH_OK) {
+            goto error;
+        }
         rc = ssh_buffer_pack(buf,
                     "dddBBBB",
                     session->next_crypto->dh_pmin,
                     session->next_crypto->dh_pn,
                     session->next_crypto->dh_pmax,
-                    session->next_crypto->dh_ctx->modulus,
+                    modulus,
-                    session->next_crypto->dh_ctx->generator,
+                    generator,
-                    session->next_crypto->dh_ctx->client.pub_key,
+                    client_pubkey,
-                    session->next_crypto->dh_ctx->server.pub_key);
+                    server_pubkey);
         if (rc != SSH_OK) {
             goto error;
         }