From b0744e1eab4367abf5cabc34875c7eb97a75b274 Mon Sep 17 00:00:00 2001
From: Alex Weibel <aweibel@amazon.com>
Date: Fri, 4 Oct 2024 14:47:06 -0700
Subject: [PATCH] Add ML-KEM Feature Test

---
 crypto/s2n_libcrypto.h                        |  1 +
 crypto/s2n_pq.c                               | 12 +++++++
 crypto/s2n_pq.h                               |  1 +
 tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c | 31 +++++++++++++++++
 .../S2N_LIBCRYPTO_SUPPORTS_MLKEM.flags        |  0
 tests/unit/s2n_mlkem_test.c                   | 34 +++++++++++++++++++
 6 files changed, 79 insertions(+)
 create mode 100644 tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c
 create mode 100644 tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.flags
 create mode 100644 tests/unit/s2n_mlkem_test.c

diff --git a/crypto/s2n_libcrypto.h b/crypto/s2n_libcrypto.h
index 7ec83557eaf..c2232b9bead 100644
--- a/crypto/s2n_libcrypto.h
+++ b/crypto/s2n_libcrypto.h
@@ -17,6 +17,7 @@
 
 #include "utils/s2n_result.h"
 
+uint64_t s2n_libcrypto_awslc_api_version(void);
 S2N_RESULT s2n_libcrypto_validate_runtime(void);
 const char *s2n_libcrypto_get_version_name(void);
 bool s2n_libcrypto_supports_flag_no_check_time();
diff --git a/crypto/s2n_pq.c b/crypto/s2n_pq.c
index ed902095adb..cb9795311a3 100644
--- a/crypto/s2n_pq.c
+++ b/crypto/s2n_pq.c
@@ -33,3 +33,15 @@ bool s2n_pq_is_enabled()
 {
     return s2n_libcrypto_supports_evp_kem();
 }
+
+bool s2n_libcrypto_supports_mlkem()
+{
+    /* S2N_LIBCRYPTO_SUPPORTS_MLKEM will be auto-detected and #defined if
+     * ./tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c successfully compiles
+     */
+#if defined(S2N_LIBCRYPTO_SUPPORTS_MLKEM)
+    return true;
+#else
+    return false;
+#endif
+}
diff --git a/crypto/s2n_pq.h b/crypto/s2n_pq.h
index 650f1c11344..089362569d1 100644
--- a/crypto/s2n_pq.h
+++ b/crypto/s2n_pq.h
@@ -23,3 +23,4 @@
 
 bool s2n_pq_is_enabled(void);
 bool s2n_libcrypto_supports_evp_kem(void);
+bool s2n_libcrypto_supports_mlkem(void);
diff --git a/tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c b/tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c
new file mode 100644
index 00000000000..6bfaf117af6
--- /dev/null
+++ b/tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.c
@@ -0,0 +1,31 @@
+/*
+* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+*
+* Licensed under the Apache License, Version 2.0 (the "License").
+* You may not use this file except in compliance with the License.
+* A copy of the License is located at
+*
+*  http://aws.amazon.com/apache2.0
+*
+* or in the "license" file accompanying this file. This file is distributed
+* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+* express or implied. See the License for the specific language governing
+* permissions and limitations under the License.
+*/
+
+#include <openssl/evp.h>
+#include <openssl/nid.h>
+
+int main()
+{
+    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_KEM, NULL);
+    if (ctx == NULL) {
+        return 1;
+    }
+    if (!EVP_PKEY_CTX_kem_set_params(ctx, NID_MLKEM768)) {
+        EVP_PKEY_CTX_free(ctx);
+        return 1;
+    }
+    EVP_PKEY_CTX_free(ctx);
+    return 0;
+}
diff --git a/tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.flags b/tests/features/S2N_LIBCRYPTO_SUPPORTS_MLKEM.flags
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/s2n_mlkem_test.c b/tests/unit/s2n_mlkem_test.c
new file mode 100644
index 00000000000..1bec449a960
--- /dev/null
+++ b/tests/unit/s2n_mlkem_test.c
@@ -0,0 +1,34 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+#include "api/s2n.h"
+#include "crypto/s2n_libcrypto.h"
+#include "crypto/s2n_openssl.h"
+#include "crypto/s2n_pq.h"
+#include "s2n_test.h"
+#include "testlib/s2n_testlib.h"
+
+int main()
+{
+    BEGIN_TEST();
+    /* MLKEM Support was added to AWSLC when AWSLC_API_VERSION == 29 */
+    if (s2n_libcrypto_is_awslc() && s2n_libcrypto_awslc_api_version() >= 30) {
+        EXPECT_TRUE(s2n_libcrypto_supports_mlkem());
+    } else if (s2n_libcrypto_is_awslc() && s2n_libcrypto_awslc_api_version() < 29) {
+        EXPECT_FALSE(s2n_libcrypto_supports_mlkem());
+    }
+
+    END_TEST();
+}