From 033891a75306a3f821cee568136350be64199a7c Mon Sep 17 00:00:00 2001
From: Sam Clark <3758302+goatgoose@users.noreply.github.com>
Date: Wed, 20 Sep 2023 10:22:24 -0400
Subject: [PATCH] test wip

---
 tests/unit/s2n_client_hello_test.c | 50 +++++++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 8 deletions(-)

diff --git a/tests/unit/s2n_client_hello_test.c b/tests/unit/s2n_client_hello_test.c
index 9fb4c756173..228d799d9fe 100644
--- a/tests/unit/s2n_client_hello_test.c
+++ b/tests/unit/s2n_client_hello_test.c
@@ -1634,8 +1634,12 @@ int main(int argc, char **argv)
 
     /* s2n_client_hello_get_supported_groups tests */
     {
-        /* 2 length bytes + space for the maximum number of supported groups */
-#define TEST_SUPPORTED_GROUPS_EXTENSION_SIZE (2 + (S2N_RECEIVED_SUPPORTED_GROUPS_MAX * 2))
+
+        /* Each supported group is 2 bytes. */
+#define TEST_SUPPORTED_GROUPS_LIST_SIZE (S2N_RECEIVED_SUPPORTED_GROUPS_MAX * 2)
+
+        /* 2 length bytes + space for the maximum number of supported groups. */
+#define TEST_SUPPORTED_GROUPS_EXTENSION_SIZE (2 + TEST_SUPPORTED_GROUPS_LIST_SIZE)
 
         /* Safety */
         {
@@ -1677,8 +1681,7 @@ int main(int argc, char **argv)
             struct s2n_stuffer extension_stuffer = { 0 };
             EXPECT_SUCCESS(s2n_stuffer_init(&extension_stuffer, &extension_blob));
 
-            uint16_t maximum_groups_size = S2N_RECEIVED_SUPPORTED_GROUPS_MAX * 2;
-            EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, maximum_groups_size));
+            EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, TEST_SUPPORTED_GROUPS_LIST_SIZE));
 
             uint16_t supported_groups[S2N_RECEIVED_SUPPORTED_GROUPS_MAX] = { 0 };
             uint16_t supported_groups_count = 0;
@@ -1772,8 +1775,7 @@ int main(int argc, char **argv)
                 struct s2n_stuffer extension_stuffer = { 0 };
                 EXPECT_SUCCESS(s2n_stuffer_init(&extension_stuffer, &extension_blob));
 
-                uint16_t maximum_groups_size = S2N_RECEIVED_SUPPORTED_GROUPS_MAX * 2;
-                EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, maximum_groups_size));
+                EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, TEST_SUPPORTED_GROUPS_LIST_SIZE));
 
                 uint16_t supported_groups[S2N_RECEIVED_SUPPORTED_GROUPS_MAX] = { 0 };
                 uint16_t supported_groups_count = 0;
@@ -1793,7 +1795,7 @@ int main(int argc, char **argv)
                 struct s2n_stuffer extension_stuffer = { 0 };
                 EXPECT_SUCCESS(s2n_stuffer_init(&extension_stuffer, &extension_blob));
 
-                uint16_t too_many_groups_size = (S2N_RECEIVED_SUPPORTED_GROUPS_MAX + 1) * 2;
+                uint16_t too_many_groups_size = TEST_SUPPORTED_GROUPS_LIST_SIZE + 2;
                 EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, too_many_groups_size));
 
                 uint16_t supported_groups[S2N_RECEIVED_SUPPORTED_GROUPS_MAX] = { 0 };
@@ -1829,10 +1831,42 @@ int main(int argc, char **argv)
         }
 
         /* Ensure that s2n_client_hello_get_supported_groups writes what is contained in the parsed
-         * extension in the client hello.
+         * supported groups extension in the client hello.
          */
         {
             struct s2n_client_hello client_hello = { 0 };
+
+            s2n_extension_type_id supported_groups_id = 0;
+            EXPECT_SUCCESS(s2n_extension_supported_iana_value_to_id(S2N_EXTENSION_SUPPORTED_GROUPS, &supported_groups_id));
+
+            s2n_parsed_extension *supported_groups_extension = &client_hello.extensions.parsed_extensions[supported_groups_id];
+            supported_groups_extension->extension_type = S2N_EXTENSION_SUPPORTED_GROUPS;
+
+            for (uint16_t groups_count = 0; groups_count < S2N_RECEIVED_SUPPORTED_GROUPS_MAX; groups_count++) {
+                uint16_t groups_list_size = groups_count * 2;
+
+                uint8_t test_groups_list_data[TEST_SUPPORTED_GROUPS_LIST_SIZE] = { 0 };
+                struct s2n_blob test_groups_list_data_blob = { 0 };
+                EXPECT_SUCCESS(s2n_blob_init(&test_groups_list_data_blob, test_groups_list_data, groups_list_size));
+                EXPECT_OK(s2n_get_public_random_data(&test_groups_list_data_blob));
+
+                uint8_t extension_data[TEST_SUPPORTED_GROUPS_EXTENSION_SIZE] = { 0 };
+                struct s2n_blob extension_blob = { 0 };
+                EXPECT_SUCCESS(s2n_blob_init(&extension_blob, extension_data, sizeof(extension_data)));
+                supported_groups_extension->extension = extension_blob;
+
+                struct s2n_stuffer extension_stuffer = { 0 };
+                EXPECT_SUCCESS(s2n_stuffer_init(&extension_stuffer, &extension_blob));
+                EXPECT_SUCCESS(s2n_stuffer_write_uint16(&extension_stuffer, groups_list_size));
+                EXPECT_SUCCESS(s2n_stuffer_write(&extension_stuffer, &test_groups_list_data_blob));
+
+                uint16_t supported_groups[S2N_RECEIVED_SUPPORTED_GROUPS_MAX] = { 0 };
+                uint16_t supported_groups_count = 0;
+                EXPECT_SUCCESS(s2n_client_hello_get_supported_groups(&client_hello, supported_groups,
+                        &supported_groups_count, S2N_RECEIVED_SUPPORTED_GROUPS_MAX));
+
+                EXPECT_EQUAL(supported_groups_count)
+            }
         }
 
         /* Self-talk: Ensure that the retrieved supported groups match what was sent by the client.