1 files changed, 26 insertions, 4 deletions

diff --git a/kexkemc.c b/kexkemc.c
index 47f15c30c..13f36a116 100644
--- a/kexkemc.c
+++ b/kexkemc.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: kexkemc.c,v 1.1 2019/01/21 10:20:12 djm Exp $ */
+/* $OpenBSD: kexkemc.c,v 1.2 2019/01/21 10:24:09 djm Exp $ */
 /*
  * Copyright (c) 2019 Markus Friedl.  All rights reserved.
  *
@@ -47,7 +47,18 @@ kex_kem_client(struct ssh *ssh)
         struct kex *kex = ssh->kex;
         int r;
 
-        if ((r = kex_kem_sntrup4591761x25519_keypair(kex)) != 0)
+        switch (kex->kex_type) {
+        case KEX_C25519_SHA256:
+                r = kex_c25519_keypair(kex);
+                break;
+        case KEX_KEM_SNTRUP4591761X25519_SHA512:
+                r = kex_kem_sntrup4591761x25519_keypair(kex);
+                break;
+        default:
+                r = SSH_ERR_INVALID_ARGUMENT;
+                break;
+        }
+        if (r != 0)
                 return r;
         if ((r = sshpkt_start(ssh, SSH2_MSG_KEX_ECDH_INIT)) != 0 ||
             (r = sshpkt_put_stringb(ssh, kex->kem_client_pub)) != 0 ||
@@ -87,8 +98,19 @@ input_kex_kem_reply(int type, u_int32_t seq, struct ssh *ssh)
                 goto out;
 
         /* compute shared secret */
-        if ((r = kex_kem_sntrup4591761x25519_dec(kex, server_pubkey, pklen,
-            &shared_secret)) != 0)
+        switch (kex->kex_type) {
+        case KEX_C25519_SHA256:
+                r = kex_c25519_dec(kex, server_pubkey, pklen, &shared_secret);
+                break;
+        case KEX_KEM_SNTRUP4591761X25519_SHA512:
+                r = kex_kem_sntrup4591761x25519_dec(kex, server_pubkey, pklen,
+                    &shared_secret);
+                break;
+        default:
+                r = SSH_ERR_INVALID_ARGUMENT;
+                break;
+        }
+        if (r !=0 )
                 goto out;
 
         /* calc and verify H */