From 79a790e6b577467a64d030a2d4c240cec2ea734a Mon Sep 17 00:00:00 2001
From: Miroslav Lichvar <mlichvar@redhat.com>
Date: Tue, 1 Oct 2024 11:44:16 +0200
Subject: [PATCH] test: improve nts_ke_client unit test

---
 test/unit/nts_ke_client.c | 53 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/test/unit/nts_ke_client.c b/test/unit/nts_ke_client.c
index e78cdc7..587efc2 100644
--- a/test/unit/nts_ke_client.c
+++ b/test/unit/nts_ke_client.c
@@ -23,8 +23,24 @@
 
 #ifdef FEAT_NTS
 
-#include <nts_ke_client.c>
 #include <local.h>
+#include <nts_ke_session.h>
+#include <util.h>
+
+#define NKSN_GetKeys get_keys
+
+static int
+get_keys(NKSN_Instance session, SIV_Algorithm algorithm, SIV_Algorithm exporter_algorithm,
+         int next_protocol, NKE_Key *c2s, NKE_Key *s2c)
+{
+  c2s->length = SIV_GetKeyLength(algorithm);
+  UTI_GetRandomBytes(c2s->key, c2s->length);
+  s2c->length = SIV_GetKeyLength(algorithm);
+  UTI_GetRandomBytes(s2c->key, s2c->length);
+  return 1;
+}
+
+#include <nts_ke_client.c>
 
 static void
 prepare_response(NKSN_Instance session, int valid)
@@ -95,8 +111,10 @@ prepare_response(NKSN_Instance session, int valid)
     TEST_CHECK(NKSN_AddRecord(session, 0, NKE_RECORD_COMPLIANT_128GCM_EXPORT, NULL, 0));
 
   if (index != 8) {
-    for (i = 0; i < NKE_MAX_COOKIES; i++) {
+    for (i = random() % NKE_MAX_COOKIES; i >= 0; i--) {
       length = (random() % sizeof (data) + 1) / 4 * 4;
+      if (i == 0 && length == 0)
+        length = 4;
       if (index == 9)
         length += (length < sizeof (data) ? 1 : -1) * (random() % 3 + 1);
       TEST_CHECK(NKSN_AddRecord(session, 0, NKE_RECORD_COOKIE, data, length));
@@ -106,12 +124,16 @@ prepare_response(NKSN_Instance session, int valid)
   TEST_CHECK(NKSN_EndMessage(session));
 }
 
+#define MAX_COOKIES (2 * NKE_MAX_COOKIES)
+
 void
 test_unit(void)
 {
+  NKE_Context context, alt_context;
+  NKE_Cookie cookies[MAX_COOKIES];
+  int i, j, r, valid, num_cookies;
   NKC_Instance inst;
   IPSockAddr addr;
-  int i, r, valid;
 
   char conf[][100] = {
     "nosystemcert",
@@ -130,10 +152,33 @@ test_unit(void)
   TEST_CHECK(inst);
 
   for (i = 0; i < 10000; i++) {
+    inst->got_response = 0;
     valid = random() % 2;
     prepare_response(inst->session, valid);
-    r = process_response(inst);
+    r = handle_message(inst);
     TEST_CHECK(r == valid);
+
+    memset(&context, 0, sizeof (context));
+    memset(&alt_context, 0, sizeof (alt_context));
+    num_cookies = 0;
+
+    r = NKC_GetNtsData(inst, &context, &alt_context,
+                       cookies, &num_cookies, random() % MAX_COOKIES + 1, &addr);
+    TEST_CHECK(r == valid);
+    if (r) {
+      TEST_CHECK(context.algorithm != AEAD_SIV_INVALID);
+      TEST_CHECK(context.c2s.length > 0);
+      TEST_CHECK(context.c2s.length == SIV_GetKeyLength(context.algorithm));
+      TEST_CHECK(context.s2c.length == SIV_GetKeyLength(context.algorithm));
+      if (alt_context.algorithm != AEAD_SIV_INVALID) {
+        TEST_CHECK(context.c2s.length > 0);
+        TEST_CHECK(alt_context.c2s.length == SIV_GetKeyLength(alt_context.algorithm));
+        TEST_CHECK(alt_context.s2c.length == SIV_GetKeyLength(alt_context.algorithm));
+      }
+      TEST_CHECK(num_cookies > 0 && num_cookies <= NKE_MAX_COOKIES);
+      for (j = 0; j < num_cookies; j++)
+        TEST_CHECK(cookies[j].length > 0 && cookies[j].length % 4 == 0);
+    }
   }
 
   NKC_DestroyInstance(inst);