From a30fd2239cbc5ee6413232b8a97e08bd1e0374da Mon Sep 17 00:00:00 2001
From: Allen <leloucharcher@163.com>
Date: Sat, 25 Apr 2020 09:38:04 +0800
Subject: [PATCH] Escape adjacent quotation marks marshaling in multiline
 string (#365)

---
 marshal_test.go   | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 tomltree_write.go | 15 +++++++++++++--
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/marshal_test.go b/marshal_test.go
index 5d66652..5d3038d 100644
--- a/marshal_test.go
+++ b/marshal_test.go
@@ -1467,6 +1467,55 @@ func TestMarshalCustomMultiline(t *testing.T) {
 	}
 }
 
+func TestMultilineWithAdjacentQuotationMarks(t *testing.T) {
+	type testStruct struct {
+		Str string `multiline:"true"`
+	}
+	type testCase struct {
+		expected []byte
+		data     testStruct
+	}
+
+	testCases := []testCase{
+		{
+			expected: []byte(`Str = """
+hello\""""
+`),
+			data: testStruct{
+				Str: "hello\"",
+			},
+		},
+		{
+			expected: []byte(`Str = """
+""\"""\"""\""""
+`),
+			data: testStruct{
+				Str: "\"\"\"\"\"\"\"\"\"",
+			},
+		},
+	}
+	for i := range testCases {
+		result, err := Marshal(testCases[i].data)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !bytes.Equal(result, testCases[i].expected) {
+			t.Errorf("Bad marshal: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n",
+				testCases[i].expected, result)
+		} else {
+			var data testStruct
+			if err = Unmarshal(result, &data); err != nil {
+				t.Fatal(err)
+			}
+			if data.Str != testCases[i].data.Str {
+				t.Errorf("Round trip test fail: expected\n-----\n%s\n-----\ngot\n-----\n%s\n-----\n",
+					testCases[i].data.Str, data.Str)
+			}
+		}
+	}
+}
+
 func TestMarshalEmbedTree(t *testing.T) {
 	expected := []byte(`OuterField1 = "Out"
 OuterField2 = 1024
diff --git a/tomltree_write.go b/tomltree_write.go
index 16c1986..9acc2f3 100644
--- a/tomltree_write.go
+++ b/tomltree_write.go
@@ -30,9 +30,15 @@ type sortNode struct {
 // are preserved. Quotation marks and backslashes are also not escaped.
 func encodeMultilineTomlString(value string, commented string) string {
 	var b bytes.Buffer
+	adjacentQuoteCount := 0
 
 	b.WriteString(commented)
-	for _, rr := range value {
+	for i, rr := range value {
+		if rr != '"' {
+			adjacentQuoteCount = 0
+		} else {
+			adjacentQuoteCount++
+		}
 		switch rr {
 		case '\b':
 			b.WriteString(`\b`)
@@ -45,7 +51,12 @@ func encodeMultilineTomlString(value string, commented string) string {
 		case '\r':
 			b.WriteString("\r")
 		case '"':
-			b.WriteString(`"`)
+			if adjacentQuoteCount >= 3 || i == len(value)-1 {
+				adjacentQuoteCount = 0
+				b.WriteString(`\"`)
+			} else {
+				b.WriteString(`"`)
+			}
 		case '\\':
 			b.WriteString(`\`)
 		default: