From 736b4835dfee77e34c09faa2623612691f7c1489 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20M=C3=A4rkl?= <info@florianmaerkl.de>
Date: Sun, 12 Jul 2020 19:28:36 +0200
Subject: [PATCH] Add Key State Tracker

---
 lib/include/chiaki/gkcrypt.h |  7 +++++
 lib/src/gkcrypt.c            | 16 ++++++++++
 test/CMakeLists.txt          |  1 +
 test/keystate.c              | 57 ++++++++++++++++++++++++++++++++++++
 test/main.c                  |  8 +++++
 5 files changed, 89 insertions(+)
 create mode 100644 test/keystate.c

diff --git a/lib/include/chiaki/gkcrypt.h b/lib/include/chiaki/gkcrypt.h
index a2c76e7..4573810 100644
--- a/lib/include/chiaki/gkcrypt.h
+++ b/lib/include/chiaki/gkcrypt.h
@@ -96,6 +96,13 @@ static inline void chiaki_gkcrypt_free(ChiakiGKCrypt *gkcrypt)
 	free(gkcrypt);
 }
 
+typedef struct chiaki_key_state_t {
+	uint64_t prev;
+} ChiakiKeyState;
+
+CHIAKI_EXPORT void chiaki_key_state_init(ChiakiKeyState *state);
+CHIAKI_EXPORT uint64_t chiaki_key_state_request_pos(ChiakiKeyState *state, uint32_t low);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/lib/src/gkcrypt.c b/lib/src/gkcrypt.c
index 0a2bc15..b398a28 100644
--- a/lib/src/gkcrypt.c
+++ b/lib/src/gkcrypt.c
@@ -542,3 +542,19 @@ static void *gkcrypt_thread_func(void *user)
 	chiaki_mutex_unlock(&gkcrypt->key_buf_mutex);
 	return NULL;
 }
+
+CHIAKI_EXPORT void chiaki_key_state_init(ChiakiKeyState *state)
+{
+	state->prev = 0;
+}
+
+CHIAKI_EXPORT uint64_t chiaki_key_state_request_pos(ChiakiKeyState *state, uint32_t low)
+{
+	uint32_t prev_low = (uint32_t)state->prev;
+	uint32_t high = (uint32_t)(state->prev >> 32);
+	if(chiaki_seq_num_32_gt(low, prev_low) && low < prev_low)
+		high++;
+	else if(chiaki_seq_num_32_lt(low, prev_low) && low > prev_low && high)
+		high--;
+	return state->prev = (((uint64_t)high) << 32) | ((uint64_t)low);
+}
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 52bc34a..827f4a7 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -9,6 +9,7 @@ add_executable(chiaki-unit
 		gkcrypt.c
 		takion.c
 		seqnum.c
+		keystate.c
 		reorderqueue.c
 		fec.c
 		test_log.c
diff --git a/test/keystate.c b/test/keystate.c
new file mode 100644
index 0000000..97bb6f2
--- /dev/null
+++ b/test/keystate.c
@@ -0,0 +1,57 @@
+/*
+ * This file is part of Chiaki.
+ *
+ * Chiaki is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Chiaki is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Chiaki.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include <munit.h>
+
+#include <chiaki/gkcrypt.h>
+
+static MunitResult test_key_state(const MunitParameter params[], void *user)
+{
+	ChiakiKeyState state;
+	chiaki_key_state_init(&state);
+
+	uint64_t pos = chiaki_key_state_request_pos(&state, 0);
+	munit_assert_uint64(pos, ==, 0);
+	pos = chiaki_key_state_request_pos(&state, 0x1337);
+	munit_assert_uint64(pos, ==, 0x1337);
+	pos = chiaki_key_state_request_pos(&state, 0xffff0000);
+	munit_assert_uint64(pos, ==, 0xffff0000);
+	pos = chiaki_key_state_request_pos(&state, 0x1337);
+	munit_assert_uint64(pos, ==, 0x100001337);
+	pos = chiaki_key_state_request_pos(&state, 0xffff1337);
+	munit_assert_uint64(pos, ==, 0xffff1337);
+	pos = chiaki_key_state_request_pos(&state, 0x50000000);
+	munit_assert_uint64(pos, ==, 0x150000000);
+	pos = chiaki_key_state_request_pos(&state, 0xb0000000);
+	munit_assert_uint64(pos, ==, 0x1b0000000);
+	pos = chiaki_key_state_request_pos(&state, 0x00000000);
+	munit_assert_uint64(pos, ==, 0x200000000);
+
+	return MUNIT_OK;
+}
+
+MunitTest tests_key_state[] = {
+	{
+		"/key_state",
+		test_key_state,
+		NULL,
+		NULL,
+		MUNIT_TEST_OPTION_NONE,
+		NULL
+	},
+	{ NULL, NULL, NULL, NULL, MUNIT_TEST_OPTION_NONE, NULL }
+};
diff --git a/test/main.c b/test/main.c
index 299382f..6a43def 100644
--- a/test/main.c
+++ b/test/main.c
@@ -18,6 +18,7 @@
 #include <munit.h>
 
 extern MunitTest tests_seq_num[];
+extern MunitTest tests_key_state[];
 extern MunitTest tests_reorder_queue[];
 extern MunitTest tests_http[];
 extern MunitTest tests_rpcrypt[];
@@ -34,6 +35,13 @@ static MunitSuite suites[] = {
 		1,
 		MUNIT_SUITE_OPTION_NONE
 	},
+	{
+		"/key_state",
+		tests_key_state,
+		NULL,
+		1,
+		MUNIT_SUITE_OPTION_NONE
+	},
 	{
 		"/reorder_queue",
 		tests_reorder_queue,