From 587be85a4cc3ae962665adcf20be1bfe511ab741 Mon Sep 17 00:00:00 2001
From: Amit Barve <ambarve@microsoft.com>
Date: Fri, 27 Aug 2021 16:12:50 -0700
Subject: [PATCH] Add test for UnicodeString Limit

Signed-off-by: Amit Barve <ambarve@microsoft.com>
---
 internal/winapi/utils.go       | 10 +++++-----
 internal/winapi/winapi_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/internal/winapi/utils.go b/internal/winapi/utils.go
index 68c3c8141f..859b753c24 100644
--- a/internal/winapi/utils.go
+++ b/internal/winapi/utils.go
@@ -29,7 +29,7 @@ type UnicodeString struct {
 }
 
 // NTSTRSAFE_UNICODE_STRING_MAX_CCH is a constant defined in ntstrsafe.h. This value
-// denotes the maximum number of bytes a path can have.
+// denotes the maximum number of wide chars a path can have.
 const NTSTRSAFE_UNICODE_STRING_MAX_CCH = 32767
 
 //String converts a UnicodeString to a golang string
@@ -42,15 +42,15 @@ func (uni UnicodeString) String() string {
 // NewUnicodeString allocates a new UnicodeString and copies `s` into
 // the buffer of the new UnicodeString.
 func NewUnicodeString(s string) (*UnicodeString, error) {
-	if len(s) > NTSTRSAFE_UNICODE_STRING_MAX_CCH {
-		return nil, syscall.ENAMETOOLONG
-	}
-
 	buf, err := windows.UTF16FromString(s)
 	if err != nil {
 		return nil, err
 	}
 
+	if len(buf) > NTSTRSAFE_UNICODE_STRING_MAX_CCH {
+		return nil, syscall.ENAMETOOLONG
+	}
+
 	uni := &UnicodeString{
 		// The length is in bytes and should not include the trailing null character.
 		Length:        uint16((len(buf) - 1) * 2),
diff --git a/internal/winapi/winapi_test.go b/internal/winapi/winapi_test.go
index 234ca3f521..80759edf8a 100644
--- a/internal/winapi/winapi_test.go
+++ b/internal/winapi/winapi_test.go
@@ -1,6 +1,7 @@
 package winapi
 
 import (
+	"strings"
 	"testing"
 	"unicode/utf16"
 )
@@ -59,3 +60,35 @@ func TestUnicodeToString(t *testing.T) {
 		}
 	}
 }
+
+func TestUnicodeStringLimit(t *testing.T) {
+	var sb strings.Builder
+
+	// limit in bytes of how long the input string can be
+	// -1 to account for null character.
+	limit := NTSTRSAFE_UNICODE_STRING_MAX_CCH - 1
+
+	lengths := []int{limit - 1, limit, limit + 1}
+	testStrings := []string{}
+	for _, len := range lengths {
+		sb.Reset()
+		for i := 0; i < len; i++ {
+			// We are deliberately writing byte 41 here as it takes only 8
+			// bits in UTF-8 encoding.  If we use non-ASCII chars the limit
+			// calculations used above won't work.
+			if err := sb.WriteByte(41); err != nil {
+				t.Fatalf("string creation failed: %s", err)
+			}
+		}
+		testStrings = append(testStrings, sb.String())
+	}
+
+	for i, testStr := range testStrings {
+		_, err := NewUnicodeString(testStr)
+		if lengths[i] > limit && err == nil {
+			t.Fatalf("input string of length %d should throw ENAMETOOLONG error", lengths[i])
+		} else if lengths[i] <= limit && err != nil {
+			t.Fatalf("unexpected error for length %d: %s", lengths[i], err)
+		}
+	}
+}