Fix Windows codepage transformer to handle streaming correctly

dlevy-msft-sql · dlevy-msft-sql · commit 4ddfd1af5ce9 · 2026-01-25T18:37:23.000-06:00
Address Copilot review: The windowsDecoder and windowsEncoder now properly
handle the atEOF parameter and buffer incomplete sequences between Transform
calls. This ensures correct behavior when transform.Reader/Writer splits
multibyte sequences across chunks.

Changes:
- Add buffer fields to windowsDecoder and windowsEncoder structs
- Use MB_ERR_INVALID_CHARS to detect incomplete sequences in decoder
- Use utf8.Valid to detect incomplete UTF-8 sequences in encoder
- Return transform.ErrShortSrc when more input is needed
- Return error for incomplete sequences at EOF
- Add TestWindowsEncodingStreaming test for streaming behavior
diff --git a/pkg/sqlcmd/codepage_test.go b/pkg/sqlcmd/codepage_test.go
@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/text/transform"
 )
 
 func TestParseCodePage(t *testing.T) {
@@ -316,3 +317,74 @@ func TestGetEncodingWindowsFallback(t *testing.T) {
 	assert.Error(t, err, "invalid codepage should fail on all platforms")
 	assert.Contains(t, err.Error(), "codepage")
 }
+
+func TestWindowsEncodingStreaming(t *testing.T) {
+	// This test verifies that the Windows API fallback handles streaming correctly
+	// by properly buffering incomplete multibyte sequences
+
+	// Japanese EBCDIC (20290) is a good test case as it's only available via Windows API
+	cp := 20290 // IBM EBCDIC Japanese Katakana Extended
+
+	enc, err := GetEncoding(cp)
+	if err != nil {
+		t.Skip("Codepage 20290 not available on this platform")
+	}
+
+	// Test decoder streaming with transform.Reader
+	t.Run("decoder streaming", func(t *testing.T) {
+		// Create a simple EBCDIC encoded string: "ABC" = 0xC1 0xC2 0xC3
+		ebcdicData := []byte{0xC1, 0xC2, 0xC3}
+
+		decoder := enc.NewDecoder()
+
+		// Simulate streaming by processing one byte at a time
+		var result []byte
+		for i := 0; i < len(ebcdicData); i++ {
+			decoder.Reset() // Reset between chunks for clean state
+			dst := make([]byte, 32)
+			nDst, _, err := decoder.Transform(dst, ebcdicData[i:i+1], i == len(ebcdicData)-1)
+			if err != nil && err != transform.ErrShortSrc {
+				t.Fatalf("Transform failed at byte %d: %v", i, err)
+			}
+			result = append(result, dst[:nDst]...)
+		}
+		assert.Equal(t, "ABC", string(result), "streaming decode should produce 'ABC'")
+	})
+
+	// Test encoder streaming
+	t.Run("encoder streaming", func(t *testing.T) {
+		// Test encoding "ABC" one character at a time
+		input := "ABC"
+		encoder := enc.NewEncoder()
+
+		var result []byte
+		for i := 0; i < len(input); i++ {
+			encoder.Reset() // Reset between chunks for clean state
+			dst := make([]byte, 32)
+			nDst, _, err := encoder.Transform(dst, []byte(input[i:i+1]), i == len(input)-1)
+			if err != nil && err != transform.ErrShortSrc {
+				t.Fatalf("Transform failed at char %d: %v", i, err)
+			}
+			result = append(result, dst[:nDst]...)
+		}
+		expected := []byte{0xC1, 0xC2, 0xC3} // "ABC" in EBCDIC
+		assert.Equal(t, expected, result, "streaming encode should produce EBCDIC ABC")
+	})
+
+	// Test encoder handles incomplete UTF-8 correctly
+	t.Run("encoder incomplete UTF-8", func(t *testing.T) {
+		encoder := enc.NewEncoder()
+		dst := make([]byte, 32)
+
+		// Send first byte of a 2-byte UTF-8 sequence (é = 0xC3 0xA9)
+		incompleteUTF8 := []byte{0xC3} // First byte of é
+		_, _, err := encoder.Transform(dst, incompleteUTF8, false)
+		// Should return ErrShortSrc because the sequence is incomplete
+		assert.Equal(t, transform.ErrShortSrc, err, "incomplete UTF-8 should return ErrShortSrc when not at EOF")
+
+		// At EOF, incomplete sequence should be an error
+		encoder.Reset()
+		_, _, err = encoder.Transform(dst, incompleteUTF8, true)
+		assert.Error(t, err, "incomplete UTF-8 at EOF should return error")
+	})
+}
diff --git a/pkg/sqlcmd/codepage_windows.go b/pkg/sqlcmd/codepage_windows.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"strconv"
 	"unicode/utf16"
+	"unicode/utf8"
 	"unsafe"
 
 	"github.com/microsoft/go-sqlcmd/internal/localizer"
@@ -17,6 +18,15 @@ import (
 	"golang.org/x/text/transform"
 )
 
+const (
+	// MB_ERR_INVALID_CHARS causes MultiByteToWideChar to fail if it encounters
+	// an invalid character in the source string (including incomplete sequences)
+	mbErrInvalidChars = 0x00000008
+	// Maximum bytes that might form a single character in any Windows codepage
+	// (most DBCS codepages use 2 bytes, but we use 4 for safety)
+	maxMultibyteCharLen = 4
+)
+
 var (
 	kernel32                = windows.NewLazySystemDLL("kernel32.dll")
 	procMultiByteToWideChar = kernel32.NewProc("MultiByteToWideChar")
@@ -36,43 +46,113 @@ func (e *windowsCodePageEncoding) NewEncoder() *encoding.Encoder {
 	return &encoding.Encoder{Transformer: &windowsEncoder{codepage: e.codepage}}
 }
 
-// windowsDecoder converts from a Windows codepage to UTF-8
+// windowsDecoder converts from a Windows codepage to UTF-8.
+// It buffers incomplete multibyte sequences between Transform calls.
 type windowsDecoder struct {
 	codepage uint32
+	buf      [maxMultibyteCharLen]byte // buffer for incomplete sequences
+	bufLen   int                       // number of bytes in buffer
 }
 
-func (d *windowsDecoder) Reset() {}
+func (d *windowsDecoder) Reset() {
+	d.bufLen = 0
+}
 
 func (d *windowsDecoder) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
-	if len(src) == 0 {
+	// Prepend any buffered bytes from previous call
+	var input []byte
+	if d.bufLen > 0 {
+		input = make([]byte, d.bufLen+len(src))
+		copy(input, d.buf[:d.bufLen])
+		copy(input[d.bufLen:], src)
+	} else {
+		input = src
+	}
+
+	if len(input) == 0 {
 		return 0, 0, nil
 	}
 
-	// First call to get required buffer size for wide chars
+	// Try to convert with MB_ERR_INVALID_CHARS to detect incomplete sequences
 	n, _, errno := procMultiByteToWideChar.Call(
 		uintptr(d.codepage),
-		0,
-		uintptr(unsafe.Pointer(&src[0])),
-		uintptr(len(src)),
+		mbErrInvalidChars,
+		uintptr(unsafe.Pointer(&input[0])),
+		uintptr(len(input)),
 		0,
 		0,
 	)
-	if n == 0 {
+
+	// If conversion failed, it might be due to incomplete trailing sequence
+	if n == 0 && errno == windows.ERROR_NO_UNICODE_TRANSLATION {
+		if atEOF {
+			// At EOF with incomplete sequence - this is an error
+			d.bufLen = 0
+			return 0, len(src), errors.New("incomplete multibyte sequence at end of input")
+		}
+
+		// Not at EOF - try removing bytes from the end until conversion succeeds
+		// This finds the incomplete trailing sequence
+		for trimLen := 1; trimLen <= len(input) && trimLen <= maxMultibyteCharLen; trimLen++ {
+			tryLen := len(input) - trimLen
+			if tryLen <= 0 {
+				// Need more input - buffer what we have
+				if len(input) <= maxMultibyteCharLen {
+					copy(d.buf[:], input)
+					d.bufLen = len(input)
+					return 0, len(src), transform.ErrShortSrc
+				}
+				break
+			}
+
+			n, _, errno = procMultiByteToWideChar.Call(
+				uintptr(d.codepage),
+				mbErrInvalidChars,
+				uintptr(unsafe.Pointer(&input[0])),
+				uintptr(tryLen),
+				0,
+				0,
+			)
+			if n > 0 || errno != windows.ERROR_NO_UNICODE_TRANSLATION {
+				// Found a valid prefix - buffer the trailing bytes
+				trailingBytes := input[tryLen:]
+				copy(d.buf[:], trailingBytes)
+				d.bufLen = len(trailingBytes)
+				input = input[:tryLen]
+				break
+			}
+		}
+
+		// If still failing, buffer everything and wait for more
+		if n == 0 {
+			if len(input) <= maxMultibyteCharLen {
+				copy(d.buf[:], input)
+				d.bufLen = len(input)
+				return 0, len(src), transform.ErrShortSrc
+			}
+			// Input is larger than max char length but still invalid - real error
+			d.bufLen = 0
+			return 0, len(src), errors.New("invalid multibyte sequence")
+		}
+	} else if n == 0 {
 		if errno != windows.ERROR_SUCCESS {
+			d.bufLen = 0
 			return 0, 0, errno
 		}
+		d.bufLen = 0
 		return 0, 0, errors.New("MultiByteToWideChar failed")
+	} else {
+		// Success - clear buffer since we'll consume all input
+		d.bufLen = 0
 	}
 
-	// Allocate wide char buffer
+	// Allocate wide char buffer and do the actual conversion
 	wideChars := make([]uint16, n)
-
-	// Convert to wide chars
 	n, _, errno = procMultiByteToWideChar.Call(
 		uintptr(d.codepage),
-		0,
-		uintptr(unsafe.Pointer(&src[0])),
-		uintptr(len(src)),
+		0, // Don't use MB_ERR_INVALID_CHARS here - we already validated
+		uintptr(unsafe.Pointer(&input[0])),
+		uintptr(len(input)),
 		uintptr(unsafe.Pointer(&wideChars[0])),
 		uintptr(len(wideChars)),
 	)
@@ -92,23 +172,71 @@ func (d *windowsDecoder) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int,
 	}
 
 	copy(dst, utf8Bytes)
-	return len(utf8Bytes), len(src), nil
+	return len(utf8Bytes), len(src), err
 }
 
-// windowsEncoder converts from UTF-8 to a Windows codepage
+// windowsEncoder converts from UTF-8 to a Windows codepage.
+// It buffers incomplete UTF-8 sequences between Transform calls.
 type windowsEncoder struct {
 	codepage uint32
+	buf      [utf8.UTFMax]byte // buffer for incomplete UTF-8 sequences
+	bufLen   int               // number of bytes in buffer
 }
 
-func (e *windowsEncoder) Reset() {}
+func (e *windowsEncoder) Reset() {
+	e.bufLen = 0
+}
 
 func (e *windowsEncoder) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int, err error) {
-	if len(src) == 0 {
+	// Prepend any buffered bytes from previous call
+	var input []byte
+	if e.bufLen > 0 {
+		input = make([]byte, e.bufLen+len(src))
+		copy(input, e.buf[:e.bufLen])
+		copy(input[e.bufLen:], src)
+	} else {
+		input = src
+	}
+
+	if len(input) == 0 {
 		return 0, 0, nil
 	}
 
+	// Find the last complete UTF-8 sequence
+	validLen := len(input)
+	for validLen > 0 && !utf8.Valid(input[:validLen]) {
+		validLen--
+	}
+
+	// Check for incomplete trailing sequence
+	if validLen < len(input) {
+		trailingBytes := input[validLen:]
+		if atEOF {
+			// At EOF with incomplete UTF-8 - this is an error
+			e.bufLen = 0
+			return 0, len(src), errors.New("incomplete UTF-8 sequence at end of input")
+		}
+		// Buffer the incomplete trailing bytes for next call
+		if len(trailingBytes) <= utf8.UTFMax {
+			copy(e.buf[:], trailingBytes)
+			e.bufLen = len(trailingBytes)
+		} else {
+			// Shouldn't happen with valid partial UTF-8, but handle it
+			e.bufLen = 0
+			return 0, len(src), errors.New("invalid UTF-8 sequence")
+		}
+		input = input[:validLen]
+	} else {
+		e.bufLen = 0
+	}
+
+	if len(input) == 0 {
+		// Only incomplete sequence - need more input
+		return 0, len(src), transform.ErrShortSrc
+	}
+
 	// Convert UTF-8 to UTF-16
-	runes := []rune(string(src))
+	runes := []rune(string(input))
 	wideChars := utf16.Encode(runes)
 
 	if len(wideChars) == 0 {
@@ -155,7 +283,7 @@ func (e *windowsEncoder) Transform(dst, src []byte, atEOF bool) (nDst, nSrc int,
 		return 0, 0, errors.New("WideCharToMultiByte failed")
 	}
 
-	return int(n), len(src), nil
+	return int(n), len(src), err
 }
 
 // isCodePageValid checks if a codepage is valid/installed on Windows
