From e9e8265313614e4d34fd9a9b0aa1c2e10625b615 Mon Sep 17 00:00:00 2001
From: Allen <leloucharcher@163.com>
Date: Sat, 25 Apr 2020 09:41:25 +0800
Subject: [PATCH] Add support for tab in basic string value and quoted key
 (#364)

---
 lexer.go        | 11 ++++------
 lexer_test.go   | 18 +++++++++++++++
 marshal_test.go | 58 ++++++++++++++++++++++++++++++++++++-------------
 3 files changed, 65 insertions(+), 22 deletions(-)

diff --git a/lexer.go b/lexer.go
index 54b0969..0bde0a1 100644
--- a/lexer.go
+++ b/lexer.go
@@ -313,7 +313,7 @@ func (l *tomlLexer) lexKey() tomlLexStateFn {
 	for r := l.peek(); isKeyChar(r) || r == '\n' || r == '\r'; r = l.peek() {
 		if r == '"' {
 			l.next()
-			str, err := l.lexStringAsString(`"`, false, true, false)
+			str, err := l.lexStringAsString(`"`, false, true)
 			if err != nil {
 				return l.errorf(err.Error())
 			}
@@ -419,7 +419,7 @@ func (l *tomlLexer) lexLiteralString() tomlLexStateFn {
 // Lex a string and return the results as a string.
 // Terminator is the substring indicating the end of the token.
 // The resulting string does not include the terminator.
-func (l *tomlLexer) lexStringAsString(terminator string, discardLeadingNewLine, acceptNewLines bool, acceptTab bool) (string, error) {
+func (l *tomlLexer) lexStringAsString(terminator string, discardLeadingNewLine, acceptNewLines bool) (string, error) {
 	growingString := ""
 
 	if discardLeadingNewLine {
@@ -512,8 +512,7 @@ func (l *tomlLexer) lexStringAsString(terminator string, discardLeadingNewLine,
 		} else {
 			r := l.peek()
 
-			if 0x00 <= r && r <= 0x1F && !(acceptNewLines && (r == '\n' || r == '\r')) &&
-				!(acceptTab && r == '\t') {
+			if 0x00 <= r && r <= 0x1F && r != '\t' && !(acceptNewLines && (r == '\n' || r == '\r')) {
 				return "", fmt.Errorf("unescaped control character %U", r)
 			}
 			l.next()
@@ -535,17 +534,15 @@ func (l *tomlLexer) lexString() tomlLexStateFn {
 	terminator := `"`
 	discardLeadingNewLine := false
 	acceptNewLines := false
-	acceptTab := false
 	if l.follow(`""`) {
 		l.skip()
 		l.skip()
 		terminator = `"""`
 		discardLeadingNewLine = true
 		acceptNewLines = true
-		acceptTab = true
 	}
 
-	str, err := l.lexStringAsString(terminator, discardLeadingNewLine, acceptNewLines, acceptTab)
+	str, err := l.lexStringAsString(terminator, discardLeadingNewLine, acceptNewLines)
 
 	if err != nil {
 		return l.errorf(err.Error())
diff --git a/lexer_test.go b/lexer_test.go
index 2d06f4e..ff826e9 100644
--- a/lexer_test.go
+++ b/lexer_test.go
@@ -707,6 +707,15 @@ func TestEscapeInString(t *testing.T) {
 	})
 }
 
+func TestTabInString(t *testing.T) {
+	testFlow(t, `foo = "hello	world"`, []token{
+		{Position{1, 1}, tokenKey, "foo"},
+		{Position{1, 5}, tokenEqual, "="},
+		{Position{1, 8}, tokenString, "hello\tworld"},
+		{Position{1, 20}, tokenEOF, ""},
+	})
+}
+
 func TestKeyGroupArray(t *testing.T) {
 	testFlow(t, "[[foo]]", []token{
 		{Position{1, 1}, tokenDoubleLeftBracket, "[["},
@@ -725,6 +734,15 @@ func TestQuotedKey(t *testing.T) {
 	})
 }
 
+func TestQuotedKeyTab(t *testing.T) {
+	testFlow(t, "\"num\tber\" = 123", []token{
+		{Position{1, 1}, tokenKey, "\"num\tber\""},
+		{Position{1, 11}, tokenEqual, "="},
+		{Position{1, 13}, tokenInteger, "123"},
+		{Position{1, 16}, tokenEOF, ""},
+	})
+}
+
 func TestKeyNewline(t *testing.T) {
 	testFlow(t, "a\n= 4", []token{
 		{Position{1, 1}, tokenError, "keys cannot contain new lines"},
diff --git a/marshal_test.go b/marshal_test.go
index 5d3038d..7ac4522 100644
--- a/marshal_test.go
+++ b/marshal_test.go
@@ -1419,26 +1419,54 @@ func TestMarshalDirectMultilineString(t *testing.T) {
 	}
 }
 
-//issue 354
-func TestUnmarshalMultilineStringWithTab(t *testing.T) {
-	input := []byte(`
-Field = """
-hello	world"""
-`)
-
+func TestUnmarshalTabInStringAndQuotedKey(t *testing.T) {
 	type Test struct {
-		Field string
+		Field1 string `toml:"Fie	ld1"`
+		Field2 string
 	}
 
-	expected := Test{"hello\tworld"}
-	result := Test{}
-	err := Unmarshal(input, &result)
-	if err != nil {
-		t.Fatal("unmarshal should not error:", err)
+	type TestCase struct {
+		desc     string
+		input    []byte
+		expected Test
 	}
 
-	if !reflect.DeepEqual(result, expected) {
-		t.Errorf("Bad unmarshal: expected\n-----\n%+v\n-----\ngot\n-----\n%+v\n-----\n", expected, result)
+	testCases := []TestCase{
+		{
+			desc:  "multiline string with tab",
+			input: []byte("Field2 = \"\"\"\nhello\tworld\"\"\""),
+			expected: Test{
+				Field2: "hello\tworld",
+			},
+		},
+		{
+			desc:  "quoted key with tab",
+			input: []byte("\"Fie\tld1\" = \"key with tab\""),
+			expected: Test{
+				Field1: "key with tab",
+			},
+		},
+		{
+			desc:  "basic string tab",
+			input: []byte("Field2 = \"hello\tworld\""),
+			expected: Test{
+				Field2: "hello\tworld",
+			},
+		},
+	}
+
+	for i := range testCases {
+		result := Test{}
+		err := Unmarshal(testCases[i].input, &result)
+		if err != nil {
+			t.Errorf("%s test error:%v", testCases[i].desc, err)
+			continue
+		}
+
+		if !reflect.DeepEqual(result, testCases[i].expected) {
+			t.Errorf("%s test error: expected\n-----\n%+v\n-----\ngot\n-----\n%+v\n-----\n",
+				testCases[i].desc, testCases[i].expected, result)
+		}
 	}
 }