lib: Improve hybrid KEX support

The algorithm was inferred from the key length. That will not always
work as SecP256r1MLKEM768 private key is 2432 bytes, exactly like an
X25519MLKEM768 one.

Raw OAP kex payloads now lead with the algorithm NID in network byte
order, so a peer reads the algorithm from the wire instead of guessing
it from the payload length.

Test if the KEX is hybrid KEM with kex_nid_is_hybrid() based on the
NID range.

The configured algorithm is passed to the raw key loaders. The public
key loader imports the key to validate it, so a corrupt or mismatched
file is reported at load time.

Signed-off-by: Dimitri Staessens <dimitri@ouroboros.rocks>
Signed-off-by: Sander Vrijders <sander@ouroboros.rocks>

6 files changed, 231 insertions, 88 deletions

diff --git a/src/irmd/oap/cli.c b/src/irmd/oap/cli.c
index 2203596f..02ad2005 100644
--- a/src/irmd/oap/cli.c
+++ b/src/irmd/oap/cli.c
@@ -51,7 +51,7 @@ struct oap_cli_ctx {
         uint8_t           __id[OAP_ID_SIZE];
         buffer_t          id;
         char              peer[NAME_SIZE + 1]; /* expected server name */
-        uint8_t           kex_buf[CRYPT_KEY_BUFSZ];
+        uint8_t           kex_buf[OAP_KEX_NIDSZ + CRYPT_KEY_BUFSZ];
         uint8_t           req_hash[MAX_HASH_SIZE];
         size_t            req_hash_len;
         int               req_md_nid;
@@ -113,13 +113,13 @@ int load_server_kem_pk(const char *        name,
         assert(cfg != NULL);
         assert(pk != NULL);
 
-        ext = IS_HYBRID_KEM(cfg->x.str) ? "raw" : "pem";
+        ext = IS_HYBRID_KEM_NID(cfg->x.nid) ? "raw" : "pem";
 
         snprintf(path, sizeof(path),
                  OUROBOROS_CLI_CRT_DIR "/%s/kex.srv.pub.%s", name, ext);
 
-        if (IS_HYBRID_KEM(cfg->x.str)) {
-                if (crypt_load_pubkey_raw_file(path, pk) < 0) {
+        if (IS_HYBRID_KEM_NID(cfg->x.nid)) {
+                if (crypt_load_pubkey_raw_file(path, cfg->x.str, pk) < 0) {
                         log_err("Failed to load %s pubkey from %s.", ext, path);
                         return -1;
                 }
@@ -173,12 +173,14 @@ static int do_client_kex_prepare_kem_encap(const char *         server_name,
                 return -ECRYPT;
         }
 
-        if (IS_HYBRID_KEM(scfg->x.str))
-                len = kex_kem_encap_raw(server_pk, kex->data,
+        if (IS_HYBRID_KEM_NID(scfg->x.nid)) {
+                len = kex_kem_encap_raw(scfg->x.str, server_pk,
+                                        kex->data + OAP_KEX_NIDSZ,
                                         scfg->k.nid, key_buf);
-        else
-                len = kex_kem_encap(server_pk, kex->data,
-                                    scfg->k.nid, key_buf);
+                len = oap_kex_tag_nid(kex->data, scfg->x.nid, len);
+        } else {
+                len = kex_kem_encap(server_pk, kex->data, scfg->k.nid, key_buf);
+        }
 
         freebuf(server_pk);
 
@@ -210,7 +212,13 @@ static int do_client_kex_prepare_kem_decap(struct oap_cli_ctx * s)
         ssize_t             len;
 
         /* Server encaps: generate keypair, send PK */
-        len = kex_pkp_create(scfg, &s->pkp, kex->data);
+        if (IS_HYBRID_KEM_NID(scfg->x.nid)) {
+                len = kex_pkp_create(scfg, &s->pkp, kex->data + OAP_KEX_NIDSZ);
+                len = oap_kex_tag_nid(kex->data, scfg->x.nid, len);
+        } else {
+                len = kex_pkp_create(scfg, &s->pkp, kex->data);
+        }
+
         if (len < 0) {
                 log_err_id(id, "Failed to generate KEM keypair.");
                 return -ECRYPT;
@@ -421,6 +429,13 @@ static int do_client_kex_complete_kem(struct oap_cli_ctx *   s,
         ct.data = peer_hdr->kex.data;
         ct.len  = peer_hdr->kex.len;
 
+        if (IS_HYBRID_KEM_NID(scfg->x.nid)) {
+                if (oap_kex_strip_nid(&ct, scfg->x.nid) < 0) {
+                        log_err_id(id, "KEX algo mismatch in CT.");
+                        return -ECRYPT;
+                }
+        }
+
         if (kex_kem_decap(s->pkp, ct, scfg->k.nid, key_buf) < 0) {
                 log_err_id(id, "Failed to decapsulate KEM.");
                 return -ECRYPT;
diff --git a/src/irmd/oap/hdr.c b/src/irmd/oap/hdr.c
index 0cff345c..d037d7b8 100644
--- a/src/irmd/oap/hdr.c
+++ b/src/irmd/oap/hdr.c
@@ -211,6 +211,56 @@ void oap_hdr_fini(struct oap_hdr * oap_hdr)
         memset(oap_hdr, 0, sizeof(*oap_hdr));
 }
 
+uint16_t oap_kex_nid(buffer_t kex)
+{
+        uint16_t nid;
+
+        if (kex.len <= OAP_KEX_NIDSZ)
+                return NID_undef;
+
+        memcpy(&nid, kex.data, sizeof(nid));
+
+        return ntoh16(nid);
+}
+
+void oap_kex_set_nid(uint8_t * buf,
+                     uint16_t  nid)
+{
+        uint16_t v;
+
+        assert(buf != NULL);
+
+        v = hton16(nid);
+
+        memcpy(buf, &v, sizeof(v));
+}
+
+int oap_kex_strip_nid(buffer_t * kex,
+                      uint16_t   nid)
+{
+        assert(kex != NULL);
+
+        if (oap_kex_nid(*kex) != nid)
+                return -1;
+
+        kex->data += OAP_KEX_NIDSZ;
+        kex->len  -= OAP_KEX_NIDSZ;
+
+        return 0;
+}
+
+ssize_t oap_kex_tag_nid(uint8_t * buf,
+                        uint16_t  nid,
+                        ssize_t   len)
+{
+        if (len < 0)
+                return len;
+
+        oap_kex_set_nid(buf, nid);
+
+        return len + OAP_KEX_NIDSZ;
+}
+
 int oap_hdr_copy_data(const struct oap_hdr * hdr,
                       buffer_t *             out)
 {
@@ -286,7 +336,7 @@ static void write_oap_fixed(uint8_t *                 buf,
 
         kex_len = (uint16_t) hdr->kex.len;
         if (hdr->kex.len > 0 && IS_KEM_ALGORITHM(scfg->x.str)) {
-                if (IS_HYBRID_KEM(scfg->x.str))
+                if (IS_HYBRID_KEM_NID(scfg->x.nid))
                         kex_len |= OAP_KEX_FMT_BIT;
                 if (scfg->x.mode == KEM_MODE_CLIENT_ENCAP)
                         kex_len |= OAP_KEX_ROLE_BIT;
diff --git a/src/irmd/oap/hdr.h b/src/irmd/oap/hdr.h
index 66fbac9a..4c5f885a 100644
--- a/src/irmd/oap/hdr.h
+++ b/src/irmd/oap/hdr.h
@@ -32,10 +32,13 @@
 #define OAP_ID_SIZE      (16)
 #define OAP_HDR_MIN_SIZE (OAP_ID_SIZE + sizeof(uint64_t) + 6 * sizeof(uint16_t))
 
-#define OAP_KEX_FMT_BIT  0x8000 /* bit 15: 0=X.509 DER, 1=Raw               */
+#define OAP_KEX_FMT_BIT  0x8000 /* bit 15: 0=X.509 DER, 1=Raw: NID + bytes  */
 #define OAP_KEX_ROLE_BIT 0x4000 /* bit 14: 0=Server encaps, 1=Client encaps */
 #define OAP_KEX_LEN_MASK 0x3FFF /* bits 0-13: Length (0-16383 bytes)        */
 
+/* Raw format kex payloads lead with the algorithm NID */
+#define OAP_KEX_NIDSZ    sizeof(uint16_t)
+
 #define OAP_KEX_ROLE(hdr) (hdr->kex_flags.role)
 #define OAP_KEX_FMT(hdr) (hdr->kex_flags.fmt)
 
@@ -160,40 +163,55 @@ struct oap_hdr {
 };
 
 
-void oap_hdr_init(struct oap_hdr * hdr,
-                  buffer_t         id,
-                  uint8_t *        kex_buf,
-                  buffer_t         data,
-                  uint16_t         nid);
+void     oap_hdr_init(struct oap_hdr * hdr,
+                      buffer_t         id,
+                      uint8_t *        kex_buf,
+                      buffer_t         data,
+                      uint16_t         nid);
+
+void     oap_hdr_fini(struct oap_hdr * oap_hdr);
+
+/* NID leading a raw format kex payload, NID_undef if too short */
+uint16_t oap_kex_nid(buffer_t kex);
+
+void     oap_kex_set_nid(uint8_t * buf,
+                         uint16_t  nid);
+
+/* Verify the NID leading a raw kex payload, then strip it */
+int      oap_kex_strip_nid(buffer_t * kex,
+                           uint16_t   nid);
 
-void oap_hdr_fini(struct oap_hdr * oap_hdr);
+/* Tag a raw payload written at buf + OAP_KEX_NIDSZ, len passthrough */
+ssize_t  oap_kex_tag_nid(uint8_t * buf,
+                         uint16_t  nid,
+                         ssize_t   len);
 
 /* oap_hdr_encode option flags */
 #define OAP_ENC_REKEY (1U << 0) /* signed, cert-less re-key packet */
 
-int  oap_hdr_encode(struct oap_hdr *    hdr,
-                    void *              pkp,
-                    void *              crt,
-                    struct sec_config * scfg,
-                    buffer_t            rsp_tag,
-                    int                 req_md_nid,
-                    const uint8_t *     seal_key,
-                    int                 flags);
+int      oap_hdr_encode(struct oap_hdr *    hdr,
+                        void *              pkp,
+                        void *              crt,
+                        struct sec_config * scfg,
+                        buffer_t            rsp_tag,
+                        int                 req_md_nid,
+                        const uint8_t *     seal_key,
+                        int                 flags);
 
-int  oap_hdr_decode(struct oap_hdr * hdr,
-                    buffer_t         buf,
-                    int              req_md_nid,
-                    bool             rekey);
+int      oap_hdr_decode(struct oap_hdr * hdr,
+                        buffer_t         buf,
+                        int              req_md_nid,
+                        bool             rekey);
 
 /* Decrypt a sealed response identity block; fills data, crt and sig. */
-int  oap_hdr_unseal(struct oap_hdr * hdr,
-                    const uint8_t *  key);
+int      oap_hdr_unseal(struct oap_hdr * hdr,
+                        const uint8_t *  key);
 
-void debug_oap_hdr_rcv(const struct oap_hdr * hdr);
+void     debug_oap_hdr_rcv(const struct oap_hdr * hdr);
 
-void debug_oap_hdr_snd(const struct oap_hdr * hdr);
+void     debug_oap_hdr_snd(const struct oap_hdr * hdr);
 
-int  oap_hdr_copy_data(const struct oap_hdr * hdr,
-                       buffer_t *             out);
+int      oap_hdr_copy_data(const struct oap_hdr * hdr,
+                           buffer_t *             out);
 
 #endif /* OUROBOROS_IRMD_OAP_HDR_H */
diff --git a/src/irmd/oap/srv.c b/src/irmd/oap/srv.c
index f78e0874..0ceba7b6 100644
--- a/src/irmd/oap/srv.c
+++ b/src/irmd/oap/srv.c
@@ -51,9 +51,9 @@ extern int load_srv_credentials(const struct name_info * info,
                                 void **                  crt);
 extern int load_srv_sec_config(const struct name_info * info,
                                struct sec_config *      cfg);
-extern int load_server_kem_keypair(const char * name,
-                                   bool         raw_fmt,
-                                   void **      pkp);
+extern int load_server_kem_keypair(const char *        name,
+                                   struct sec_config * cfg,
+                                   void **             pkp);
 #else
 
 int load_srv_credentials(const struct name_info * info,
@@ -79,23 +79,26 @@ int load_srv_sec_config(const struct name_info * info,
         return load_sec_config(info->name, info->s.sec, cfg);
 }
 
-int load_server_kem_keypair(const char * name,
-                            bool         raw_fmt,
-                            void **      pkp)
+int load_server_kem_keypair(const char *        name,
+                            struct sec_config * cfg,
+                            void **             pkp)
 {
         char         path[PATH_MAX];
         const char * ext;
+        bool         raw_fmt;
 
         assert(name != NULL);
+        assert(cfg != NULL);
         assert(pkp != NULL);
 
+        raw_fmt = IS_HYBRID_KEM_NID(cfg->x.nid);
         ext = raw_fmt ? "raw" : "pem";
 
         snprintf(path, sizeof(path),
                  OUROBOROS_SRV_CRT_DIR "/%s/kex.key.%s", name, ext);
 
         if (raw_fmt) {
-                if (crypt_load_privkey_raw_file(path, pkp) < 0) {
+                if (crypt_load_privkey_raw_file(path, cfg->x.str, pkp) < 0) {
                         log_err("Failed to load %s keypair from %s.",
                                 ext, path);
                         return -ECRYPT;
@@ -109,6 +112,7 @@ int load_server_kem_keypair(const char * name,
         }
 
         log_dbg("Loaded server KEM keypair from %s.", path);
+
         return 0;
 }
 
@@ -117,18 +121,19 @@ int load_server_kem_keypair(const char * name,
 static int get_algo_from_peer_key(const struct oap_hdr * peer_hdr,
                                   char *                 algo_buf)
 {
-        uint8_t * id = peer_hdr->id.data;
-        int       ret;
+        uint8_t *    id = peer_hdr->id.data;
+        const char * name;
 
         if (OAP_KEX_IS_RAW_FMT(peer_hdr)) {
-                ret = kex_get_algo_from_pk_raw(peer_hdr->kex, algo_buf);
-                if (ret < 0) {
-                        log_err_id(id, "Failed to get algo from raw key.");
+                name = kex_nid_to_str(oap_kex_nid(peer_hdr->kex));
+                if (name == NULL || !IS_HYBRID_KEM(name)) {
+                        log_err_id(id, "Bad algo NID in raw key.");
                         return -ECRYPT;
                 }
+
+                strcpy(algo_buf, name);
         } else {
-                ret = kex_get_algo_from_pk_der(peer_hdr->kex, algo_buf);
-                if (ret < 0) {
+                if (kex_get_algo_from_pk_der(peer_hdr->kex, algo_buf) < 0) {
                         log_err_id(id, "Failed to get algo from DER key.");
                         return -ECRYPT;
                 }
@@ -216,15 +221,20 @@ static int do_server_kem_decap(const struct name_info * info,
         int       ret;
         uint8_t * id = peer_hdr->id.data;
 
-        ret = load_server_kem_keypair(info->name,
-                                      peer_hdr->kex_flags.fmt,
-                                      &server_pkp);
-        if (ret < 0)
-                return ret;
-
         ct.data = peer_hdr->kex.data;
         ct.len  = peer_hdr->kex.len;
 
+        if (IS_HYBRID_KEM_NID(scfg->x.nid)) {
+                if (oap_kex_strip_nid(&ct, scfg->x.nid) < 0) {
+                        log_err_id(id, "KEX algo mismatch in CT.");
+                        return -ECRYPT;
+                }
+        }
+
+        ret = load_server_kem_keypair(info->name, scfg, &server_pkp);
+        if (ret < 0)
+                return ret;
+
         ret = kex_kem_decap(server_pkp, ct, scfg->k.nid, sk->key);
 
         crypt_free_key(server_pkp);
@@ -251,12 +261,20 @@ static int do_server_kem_encap(const struct oap_hdr * peer_hdr,
         client_pk.data = peer_hdr->kex.data;
         client_pk.len  = peer_hdr->kex.len;
 
-        if (IS_HYBRID_KEM(scfg->x.str))
-                ct_len = kex_kem_encap_raw(client_pk, kex->data,
+        if (IS_HYBRID_KEM_NID(scfg->x.nid)) {
+                if (oap_kex_strip_nid(&client_pk, scfg->x.nid) < 0) {
+                        log_err_id(id, "KEX algo mismatch in PK.");
+                        return -ECRYPT;
+                }
+
+                ct_len = kex_kem_encap_raw(scfg->x.str, client_pk,
+                                           kex->data + OAP_KEX_NIDSZ,
                                            scfg->k.nid, sk->key);
-        else
+                ct_len = oap_kex_tag_nid(kex->data, scfg->x.nid, ct_len);
+        } else {
                 ct_len = kex_kem_encap(client_pk, kex->data,
                                        scfg->k.nid, sk->key);
+        }
 
         if (ct_len < 0) {
                 log_err_id(id, "Failed to encapsulate KEM.");
@@ -385,27 +403,27 @@ int oap_srv_process(const struct name_info * info,
                     const buffer_t *         cached_crt,
                     buffer_t *               peer_crt)
 {
-        struct oap_hdr        peer_hdr;
-        struct oap_hdr        local_hdr;
-        struct sec_config     scfg;
-        uint8_t               kex_buf[CRYPT_KEY_BUFSZ];
-        uint8_t               hash_buf[MAX_HASH_SIZE];
-        uint8_t               kc_buf[MAX_HASH_SIZE];
-        uint8_t               resp_hash_buf[MAX_HASH_SIZE];
-        uint8_t               hs_key[SYMMKEYSZ];
-        const uint8_t *       seal_key = NULL;
-        buffer_t              req_hash  = BUF_INIT;
-        buffer_t              resp_hash = BUF_INIT;
-        buffer_t              crt_der   = BUF_INIT;
-        buffer_t              rsp_tag   = BUF_INIT;
-        ssize_t               hash_ret;
-        char                  cli_name[NAME_SIZE + 1];
-        uint8_t *             id;
-        void *                pkp = NULL;
-        void *                crt = NULL;
-        int                   req_md_nid;
-        int                   enc_flags = 0;
-        int                   ret;
+        struct oap_hdr    peer_hdr;
+        struct oap_hdr    local_hdr;
+        struct sec_config scfg;
+        uint8_t           kex_buf[OAP_KEX_NIDSZ + CRYPT_KEY_BUFSZ];
+        uint8_t           hash_buf[MAX_HASH_SIZE];
+        uint8_t           kc_buf[MAX_HASH_SIZE];
+        uint8_t           resp_hash_buf[MAX_HASH_SIZE];
+        uint8_t           hs_key[SYMMKEYSZ];
+        const uint8_t *   seal_key = NULL;
+        buffer_t          req_hash  = BUF_INIT;
+        buffer_t          resp_hash = BUF_INIT;
+        buffer_t          crt_der   = BUF_INIT;
+        buffer_t          rsp_tag   = BUF_INIT;
+        ssize_t           hash_ret;
+        char              cli_name[NAME_SIZE + 1];
+        uint8_t *         id;
+        void *            pkp = NULL;
+        void *            crt = NULL;
+        int               req_md_nid;
+        int               enc_flags = 0;
+        int               ret;
 
         assert(info != NULL);
         assert(rsp_buf != NULL);
diff --git a/src/irmd/oap/tests/oap_test.c b/src/irmd/oap/tests/oap_test.c
index b24bb786..3e2bae56 100644
--- a/src/irmd/oap/tests/oap_test.c
+++ b/src/irmd/oap/tests/oap_test.c
@@ -81,9 +81,9 @@ int mock_load_credentials(void ** pkp,
 }
 
 /* Stub KEM functions - ECDSA tests don't use KEM */
-int load_server_kem_keypair(__attribute__((unused)) const char * name,
-                            __attribute__((unused)) bool         raw_fmt,
-                            __attribute__((unused)) void **      pkp)
+int load_server_kem_keypair(__attribute__((unused)) const char *        name,
+                            __attribute__((unused)) struct sec_config * cfg,
+                            __attribute__((unused)) void **             pkp)
 {
         return -1;
 }
diff --git a/src/irmd/oap/tests/oap_test_ml_dsa.c b/src/irmd/oap/tests/oap_test_ml_dsa.c
index 477b712b..b9132b13 100644
--- a/src/irmd/oap/tests/oap_test_ml_dsa.c
+++ b/src/irmd/oap/tests/oap_test_ml_dsa.c
@@ -37,6 +37,7 @@
 
 #include <test/certs/ml_dsa.h>
 
+#include "oap/hdr.h"
 #include "oap.h"
 #include "common.h"
 
@@ -97,16 +98,16 @@ int mock_load_credentials(void ** pkp,
         return 0;
 }
 
-int load_server_kem_keypair(const char * name,
-                            bool         raw_fmt,
-                            void **      pkp)
+int load_server_kem_keypair(const char *        name,
+                            struct sec_config * cfg,
+                            void **             pkp)
 {
 #ifdef HAVE_OPENSSL
         struct sec_config local_cfg;
         ssize_t           pk_len;
 
         (void) name;
-        (void) raw_fmt;
+        (void) cfg;
 
         /*
         * Uses reference counting. The caller will call
@@ -150,7 +151,7 @@ int load_server_kem_keypair(const char * name,
 
 #else
         (void) name;
-        (void) raw_fmt;
+        (void) cfg;
         (void) pkp;
         return -1;
 #endif
@@ -720,6 +721,44 @@ static int test_oap_kem_srv_uncfg_all(void)
         return ret;
 }
 
+/* Server must reject a raw kex payload tagged with a bad NID */
+static int test_oap_kem_bad_nid_tag(uint16_t bad_nid)
+{
+        struct oap_test_ctx ctx;
+
+        test_cfg_init(NID_X25519MLKEM768, NID_aes_256_gcm,
+                      get_random_kdf(), SRV_ENCAP, NO_AUTH);
+
+        TEST_START("(%u)", bad_nid);
+
+        if (oap_test_setup_kem(&ctx, root_ca_crt_ml, im_ca_crt_ml) < 0)
+                goto fail;
+
+        if (oap_cli_prepare_ctx(&ctx) < 0) {
+                printf("Client prepare failed.\n");
+                goto fail_cleanup;
+        }
+
+        /* NO_AUTH request: raw kex starts after the fixed header */
+        oap_kex_set_nid(ctx.req_hdr.data + OAP_HDR_MIN_SIZE, bad_nid);
+
+        if (oap_srv_process_ctx(&ctx) == 0) {
+                printf("Server accepted bad NID tag %u.\n", bad_nid);
+                goto fail_cleanup;
+        }
+
+        oap_test_teardown_kem(&ctx);
+
+        TEST_SUCCESS("(%u)", bad_nid);
+
+        return TEST_RC_SUCCESS;
+ fail_cleanup:
+        oap_test_teardown_kem(&ctx);
+ fail:
+        TEST_FAIL("(%u)", bad_nid);
+        return TEST_RC_FAIL;
+}
+
 /* Client encap must be rejected when the server has no KEX config */
 static int test_oap_kem_cli_encap_srv_uncfg(void)
 {
@@ -782,6 +821,8 @@ int oap_test_ml_dsa(int    argc,
 
         ret |= test_oap_kem_srv_uncfg_all();
         ret |= test_oap_kem_cli_encap_srv_uncfg();
+        ret |= test_oap_kem_bad_nid_tag(NID_X25519); /* known, not hybrid */
+        ret |= test_oap_kem_bad_nid_tag(0xFFFF);     /* unknown           */
 
         ret |= test_oap_corrupted_request();
         ret |= test_oap_corrupted_response();
@@ -812,6 +853,7 @@ int oap_test_ml_dsa(int    argc,
         (void) test_oap_kem_srv_uncfg;
         (void) test_oap_kem_srv_uncfg_all;
         (void) test_oap_kem_cli_encap_srv_uncfg;
+        (void) test_oap_kem_bad_nid_tag;
         (void) test_oap_corrupted_request;
         (void) test_oap_corrupted_response;
         (void) test_oap_truncated_request;