From c500ed2187d20809edf60b572fc64f242b07f784 Mon Sep 17 00:00:00 2001
From: Ratin Gao <ratin@knsoft.org>
Date: Mon, 23 Sep 2024 00:18:46 +0800
Subject: [PATCH] [STRSAFE] Implement `StrSafe_CchCopy(A/W)`

---
 Source/Include/KNSoft/NDK/Extension/StrSafe.h | 100 +++++++++++++++---
 Source/Test/StrSafe.c                         |   7 ++
 2 files changed, 95 insertions(+), 12 deletions(-)

diff --git a/Source/Include/KNSoft/NDK/Extension/StrSafe.h b/Source/Include/KNSoft/NDK/Extension/StrSafe.h
index b3f65f4..d8378cd 100644
--- a/Source/Include/KNSoft/NDK/Extension/StrSafe.h
+++ b/Source/Include/KNSoft/NDK/Extension/StrSafe.h
@@ -1,18 +1,24 @@
-#pragma once
+/*
+ * KNSoft.NDK StrSafe.h licensed under the MIT license.
+ *
+ * Make C string routines a bit more safe and easy.
+ * Not equals to the strsafe.h in Windows SDK.
+ *
+ * When output to a buffer, usually:
+ *     Return == 0: Error or no data
+ *     Return < BufferCount: Success, returns the number of characters or bytes written, not including null-terminator
+ *     Return >= BufferCount: Truncated, returns required size in character or byte, not including null-terminator
+ *
+ * Define `_NO_CRT_STDIO_INLINE` to use `legacy_stdio_definitions.lib`.
+ */
+
+#pragma once
 
 #include <stdio.h>
 #include <stdarg.h>
 
 #pragma region String PrintF
 
-/*
- * StrSafe_Cch[V]Printf(A/W)
- *
- * Return == 0: Error or no data
- * Return < BufferCount: Success
- * Return >= BufferCount: Truncated, returns required size in character, not including null-terminator
- */
-
 _Success_(
     return > 0 && return < BufferCount
 )
@@ -110,7 +116,8 @@ __cdecl
 StrSafe_CchPrintfA(
     _Out_writes_opt_(BufferCount) _Always_(_Post_z_) char* const Buffer,
     _In_ size_t const BufferCount,
-    _In_z_ _Printf_format_string_ const char* Format, ...)
+    _In_z_ _Printf_format_string_ const char* Format,
+    ...)
 {
     va_list ArgList;
 
@@ -127,7 +134,8 @@ __cdecl
 StrSafe_CchPrintfW(
     _Out_writes_opt_(BufferCount) _Always_(_Post_z_) wchar_t* const Buffer,
     _In_ size_t const BufferCount,
-    _In_z_ _Printf_format_string_ const wchar_t* Format, ...)
+    _In_z_ _Printf_format_string_ const wchar_t* Format,
+    ...)
 {
     va_list ArgList;
 
@@ -135,4 +143,72 @@ StrSafe_CchPrintfW(
     return StrSafe_CchVPrintfW(Buffer, BufferCount, Format, ArgList);
 }
 
-#pragma endregion
+#pragma endregion StrSafe_Cch[V]Printf(A/W)
+
+#pragma region String Copy
+
+_Success_(
+    return > 0 && return < BufferCount
+)
+__inline
+size_t
+__cdecl
+StrSafe_CchCopyA(
+    _Out_writes_opt_(BufferCount) _When_(BufferCount > 0, _Notnull_) _Always_(_Post_z_) char* const Buffer,
+    _In_ size_t const BufferCount,
+    _In_z_ const char* Source)
+{
+    size_t i;
+
+    for (i = 0; i < BufferCount; i++)
+    {
+        if ((Buffer[i] = Source[i]) == '\0')
+        {
+            return i;
+        }
+    }
+
+    if (BufferCount > 0)
+    {
+        Buffer[BufferCount - 1] = '\0';
+    }
+    while (Source[i] != '\0')
+    {
+        i++;
+    }
+    return i;
+}
+
+_Success_(
+    return > 0 && return < BufferCount
+)
+__inline
+size_t
+__cdecl
+StrSafe_CchCopyW(
+    _Out_writes_opt_(BufferCount) _When_(BufferCount > 0, _Notnull_) _Always_(_Post_z_) wchar_t* const Buffer,
+    _In_ size_t const BufferCount,
+    _In_z_ const wchar_t* Source)
+{
+    size_t i;
+
+    for (i = 0; i < BufferCount; i++)
+    {
+        if ((Buffer[i] = Source[i]) == L'\0')
+        {
+            return i;
+        }
+    }
+
+    if (BufferCount > 0)
+    {
+        Buffer[BufferCount - 1] = L'\0';
+    }
+    while (Source[i] != L'\0')
+    {
+        i++;
+    }
+    return i;
+}
+
+#pragma endregion StrSafe_CchCopy(A/W)
diff --git a/Source/Test/StrSafe.c b/Source/Test/StrSafe.c
index eb1b925..d60f865 100644
--- a/Source/Test/StrSafe.c
+++ b/Source/Test/StrSafe.c
@@ -33,4 +33,11 @@ TEST_FUNC(StrSafeFunc)
     TEST_OK(wcscmp(szTempW, TEST_STRING1_PART1_W) == 0);
     TEST_OK(StrSafe_CchPrintfW(szTempW, ARRAYSIZE(szTempW), L"%ls$", TEST_STRING1_W) == _STR_CCH_LEN(TEST_STRING1_W) + 1);
     TEST_OK(wcscmp(szTempW, TEST_STRING1_W) == 0);
+
+    TEST_OK(StrSafe_CchCopyA(NULL, 0, TEST_STRING1_A) == _STR_CCH_LEN(TEST_STRING1_A));
+    TEST_OK(StrSafe_CchCopyA(szTempA, ARRAYSIZE(szTempA), TEST_STRING1_A) == _STR_CCH_LEN(TEST_STRING1_A));
+    TEST_OK(strcmp(szTempA, TEST_STRING1_A) == 0);
+    TEST_OK(StrSafe_CchCopyW(NULL, 0, TEST_STRING1_W) == _STR_CCH_LEN(TEST_STRING1_W));
+    TEST_OK(StrSafe_CchCopyW(szTempW, ARRAYSIZE(szTempW), TEST_STRING1_W) == _STR_CCH_LEN(TEST_STRING1_W));
+    TEST_OK(wcscmp(szTempW, TEST_STRING1_W) == 0);
 }