Merge pull request #1306 from dcantah/cp-unicodestr-fix

[release/0.8] Backport 'Bugfix for UnicodeString constructor'

internal/winapi/utils.go

@@ -20,36 +20,41 @@ func Uint16BufferToSlice(buffer *uint16, bufferLength int) (result []uint16) {
 	return
 }
 
+// UnicodeString corresponds to UNICODE_STRING win32 struct defined here
+// https://docs.microsoft.com/en-us/windows/win32/api/ntdef/ns-ntdef-_unicode_string
 type UnicodeString struct {
 	Length        uint16
 	MaximumLength uint16
 	Buffer        *uint16
 }
 
+// NTSTRSAFE_UNICODE_STRING_MAX_CCH is a constant defined in ntstrsafe.h. This value
+// 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
 func (uni UnicodeString) String() string {
 	// UnicodeString is not guaranteed to be null terminated, therefore
 	// use the UnicodeString's Length field
-	return syscall.UTF16ToString(Uint16BufferToSlice(uni.Buffer, int(uni.Length/2)))
+	return windows.UTF16ToString(Uint16BufferToSlice(uni.Buffer, int(uni.Length/2)))
 }
 
 // NewUnicodeString allocates a new UnicodeString and copies `s` into
 // the buffer of the new UnicodeString.
 func NewUnicodeString(s string) (*UnicodeString, error) {
-	// Get length of original `s` to use in the UnicodeString since the `buf`
-	// created later will have an additional trailing null character
-	length := len(s)
-	if length > 32767 {
-		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{
-		Length:        uint16(length * 2),
-		MaximumLength: uint16(length * 2),
+		// The length is in bytes and should not include the trailing null character.
+		Length:        uint16((len(buf) - 1) * 2),
+		MaximumLength: uint16((len(buf) - 1) * 2),
 		Buffer:        &buf[0],
 	}
 	return uni, nil

internal/winapi/winapi_test.go

@@ -1,6 +1,7 @@
 package winapi
 
 import (
+	"strings"
 	"testing"
 	"unicode/utf16"
 )
@@ -19,10 +20,11 @@ func wideStringsEqual(target, actual []uint16) bool {
 }
 
 func TestNewUnicodeString(t *testing.T) {
-	targetStrings := []string{"abcde", "abcd\n", "C:\\Test", "\\&_Test"}
+	// include UTF8 chars which take more than 8 bits to encode
+	targetStrings := []string{"abcde", "abcd\n", "C:\\Test", "\\&_Test", "Äb", "\u8483\u119A2\u0041"}
 	for _, target := range targetStrings {
-		targetLength := uint16(len(target) * 2)
 		targetWideString := utf16.Encode(([]rune)(target))
+		targetLength := uint16(len(targetWideString) * 2)
 
 		uni, err := NewUnicodeString(target)
 		if err != nil {
@@ -36,7 +38,7 @@ func TestNewUnicodeString(t *testing.T) {
 			t.Fatalf("Expected new Unicode String maximum length to be %d for target string %s, got %d instead", targetLength, target, uni.MaximumLength)
 		}
 
-		uniBufferStringAsSlice := Uint16BufferToSlice(uni.Buffer, len(target))
+		uniBufferStringAsSlice := Uint16BufferToSlice(uni.Buffer, int(targetLength/2))
 
 		if !wideStringsEqual(targetWideString, uniBufferStringAsSlice) {
 			t.Fatalf("Expected wide string %v, got %v instead", targetWideString, uniBufferStringAsSlice)
@@ -45,7 +47,7 @@ func TestNewUnicodeString(t *testing.T) {
 }
 
 func TestUnicodeToString(t *testing.T) {
-	targetStrings := []string{"abcde", "abcd\n", "C:\\Test", "\\&_Test"}
+	targetStrings := []string{"abcde", "abcd\n", "C:\\Test", "\\&_Test", "Äb", "\u8483\u119A2\u0041"}
 	for _, target := range targetStrings {
 		uni, err := NewUnicodeString(target)
 		if err != nil {
@@ -58,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)
+		}
+	}
+}