From 49f5952a4e9bd3849fe450d9881322ef19194a20 Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Wed, 19 Feb 2020 21:11:54 +0530
Subject: [PATCH 1/6] Have json.Marshal function just for string type

---
 query/outputnode.go | 81 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 2 deletions(-)

diff --git a/query/outputnode.go b/query/outputnode.go
index 6b40be0c154..ee11e5c6659 100644
--- a/query/outputnode.go
+++ b/query/outputnode.go
@@ -18,12 +18,12 @@ package query
 
 import (
 	"bytes"
-	"encoding/json"
 	"fmt"
 	"sort"
 	"strconv"
 	"strings"
 	"time"
+	"unicode/utf8"
 
 	"github.com/golang/glog"
 	"github.com/pkg/errors"
@@ -135,10 +135,87 @@ var (
 	emptyString = []byte(`""`)
 )
 
+// stringJsonMarshal is replacement for json.Marshal() function only for string type.
+// This function is encodeState.string(string, escapeHTML) in "encoding/json/encode.go".
+func stringJsonMarshal(s string) []byte {
+	var hex = "0123456789abcdef"
+	escapeHTML := true
+	var e bytes.Buffer
+	e.WriteByte('"')
+	start := 0
+	for i := 0; i < len(s); {
+		if b := s[i]; b < utf8.RuneSelf {
+			if htmlSafeSet[b] || (!escapeHTML && safeSet[b]) {
+				i++
+				continue
+			}
+			if start < i {
+				e.WriteString(s[start:i])
+			}
+			e.WriteByte('\\')
+			switch b {
+			case '\\', '"':
+				e.WriteByte(b)
+			case '\n':
+				e.WriteByte('n')
+			case '\r':
+				e.WriteByte('r')
+			case '\t':
+				e.WriteByte('t')
+			default:
+				// This encodes bytes < 0x20 except for \t, \n and \r.
+				// If escapeHTML is set, it also escapes <, >, and &
+				// because they can lead to security holes when
+				// user-controlled strings are rendered into JSON
+				// and served to some browsers.
+				e.WriteString(`u00`)
+				e.WriteByte(hex[b>>4])
+				e.WriteByte(hex[b&0xF])
+			}
+			i++
+			start = i
+			continue
+		}
+		c, size := utf8.DecodeRuneInString(s[i:])
+		if c == utf8.RuneError && size == 1 {
+			if start < i {
+				e.WriteString(s[start:i])
+			}
+			e.WriteString(`\ufffd`)
+			i += size
+			start = i
+			continue
+		}
+		// U+2028 is LINE SEPARATOR.
+		// U+2029 is PARAGRAPH SEPARATOR.
+		// They are both technically valid characters in JSON strings,
+		// but don't work in JSONP, which has to be evaluated as JavaScript,
+		// and can lead to security holes there. It is valid JSON to
+		// escape them, so we do so unconditionally.
+		// See http://timelessrepo.com/json-isnt-a-javascript-subset for discussion.
+		if c == '\u2028' || c == '\u2029' {
+			if start < i {
+				e.WriteString(s[start:i])
+			}
+			e.WriteString(`\u202`)
+			e.WriteByte(hex[c&0xF])
+			i += size
+			start = i
+			continue
+		}
+		i += size
+	}
+	if start < len(s) {
+		e.WriteString(s[start:])
+	}
+	e.WriteByte('"')
+	return e.Bytes()
+}
+
 func valToBytes(v types.Val) ([]byte, error) {
 	switch v.Tid {
 	case types.StringID, types.DefaultID:
-		return json.Marshal(v.Value)
+		return stringJsonMarshal(v.Value.(string)), nil
 	case types.BinaryID:
 		return []byte(fmt.Sprintf("%q", v.Value)), nil
 	case types.IntID:

From a9c37ceb2d8b5e036c937c8afffad15d95a03982 Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Thu, 19 Mar 2020 15:53:30 +0530
Subject: [PATCH 2/6] Add table.go

---
 query/table.go | 217 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 217 insertions(+)
 create mode 100644 query/table.go

diff --git a/query/table.go b/query/table.go
new file mode 100644
index 00000000000..da4ad23dd20
--- /dev/null
+++ b/query/table.go
@@ -0,0 +1,217 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+package query
+
+import "unicode/utf8"
+
+// safeSet holds the value true if the ASCII character with the given array
+// position can be represented inside a JSON string without any further
+// escaping.
+//
+// All values are true except for the ASCII control characters (0-31), the
+// double quote ("), and the backslash character ("\").
+var safeSet = [utf8.RuneSelf]bool{
+	' ':      true,
+	'!':      true,
+	'"':      false,
+	'#':      true,
+	'$':      true,
+	'%':      true,
+	'&':      true,
+	'\'':     true,
+	'(':      true,
+	')':      true,
+	'*':      true,
+	'+':      true,
+	',':      true,
+	'-':      true,
+	'.':      true,
+	'/':      true,
+	'0':      true,
+	'1':      true,
+	'2':      true,
+	'3':      true,
+	'4':      true,
+	'5':      true,
+	'6':      true,
+	'7':      true,
+	'8':      true,
+	'9':      true,
+	':':      true,
+	';':      true,
+	'<':      true,
+	'=':      true,
+	'>':      true,
+	'?':      true,
+	'@':      true,
+	'A':      true,
+	'B':      true,
+	'C':      true,
+	'D':      true,
+	'E':      true,
+	'F':      true,
+	'G':      true,
+	'H':      true,
+	'I':      true,
+	'J':      true,
+	'K':      true,
+	'L':      true,
+	'M':      true,
+	'N':      true,
+	'O':      true,
+	'P':      true,
+	'Q':      true,
+	'R':      true,
+	'S':      true,
+	'T':      true,
+	'U':      true,
+	'V':      true,
+	'W':      true,
+	'X':      true,
+	'Y':      true,
+	'Z':      true,
+	'[':      true,
+	'\\':     false,
+	']':      true,
+	'^':      true,
+	'_':      true,
+	'`':      true,
+	'a':      true,
+	'b':      true,
+	'c':      true,
+	'd':      true,
+	'e':      true,
+	'f':      true,
+	'g':      true,
+	'h':      true,
+	'i':      true,
+	'j':      true,
+	'k':      true,
+	'l':      true,
+	'm':      true,
+	'n':      true,
+	'o':      true,
+	'p':      true,
+	'q':      true,
+	'r':      true,
+	's':      true,
+	't':      true,
+	'u':      true,
+	'v':      true,
+	'w':      true,
+	'x':      true,
+	'y':      true,
+	'z':      true,
+	'{':      true,
+	'|':      true,
+	'}':      true,
+	'~':      true,
+	'\u007f': true,
+}
+
+// htmlSafeSet holds the value true if the ASCII character with the given
+// array position can be safely represented inside a JSON string, embedded
+// inside of HTML <script> tags, without any additional escaping.
+//
+// All values are true except for the ASCII control characters (0-31), the
+// double quote ("), the backslash character ("\"), HTML opening and closing
+// tags ("<" and ">"), and the ampersand ("&").
+var htmlSafeSet = [utf8.RuneSelf]bool{
+	' ':      true,
+	'!':      true,
+	'"':      false,
+	'#':      true,
+	'$':      true,
+	'%':      true,
+	'&':      false,
+	'\'':     true,
+	'(':      true,
+	')':      true,
+	'*':      true,
+	'+':      true,
+	',':      true,
+	'-':      true,
+	'.':      true,
+	'/':      true,
+	'0':      true,
+	'1':      true,
+	'2':      true,
+	'3':      true,
+	'4':      true,
+	'5':      true,
+	'6':      true,
+	'7':      true,
+	'8':      true,
+	'9':      true,
+	':':      true,
+	';':      true,
+	'<':      false,
+	'=':      true,
+	'>':      false,
+	'?':      true,
+	'@':      true,
+	'A':      true,
+	'B':      true,
+	'C':      true,
+	'D':      true,
+	'E':      true,
+	'F':      true,
+	'G':      true,
+	'H':      true,
+	'I':      true,
+	'J':      true,
+	'K':      true,
+	'L':      true,
+	'M':      true,
+	'N':      true,
+	'O':      true,
+	'P':      true,
+	'Q':      true,
+	'R':      true,
+	'S':      true,
+	'T':      true,
+	'U':      true,
+	'V':      true,
+	'W':      true,
+	'X':      true,
+	'Y':      true,
+	'Z':      true,
+	'[':      true,
+	'\\':     false,
+	']':      true,
+	'^':      true,
+	'_':      true,
+	'`':      true,
+	'a':      true,
+	'b':      true,
+	'c':      true,
+	'd':      true,
+	'e':      true,
+	'f':      true,
+	'g':      true,
+	'h':      true,
+	'i':      true,
+	'j':      true,
+	'k':      true,
+	'l':      true,
+	'm':      true,
+	'n':      true,
+	'o':      true,
+	'p':      true,
+	'q':      true,
+	'r':      true,
+	's':      true,
+	't':      true,
+	'u':      true,
+	'v':      true,
+	'w':      true,
+	'x':      true,
+	'y':      true,
+	'z':      true,
+	'{':      true,
+	'|':      true,
+	'}':      true,
+	'~':      true,
+	'\u007f': true,
+}

From 3d272cb16a5026be015207d43c0e6acbf1cbcecb Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Thu, 19 Mar 2020 22:06:42 +0530
Subject: [PATCH 3/6] Add tests

---
 query/outputnode.go      |  6 ++++--
 query/outputnode_test.go | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/query/outputnode.go b/query/outputnode.go
index ee11e5c6659..c662a9932c0 100644
--- a/query/outputnode.go
+++ b/query/outputnode.go
@@ -133,13 +133,15 @@ var (
 	boolTrue    = []byte("true")
 	boolFalse   = []byte("false")
 	emptyString = []byte(`""`)
+
+	// Used in stringJsonMarshal function.
+	hex        = "0123456789abcdef"
+	escapeHTML = true
 )
 
 // stringJsonMarshal is replacement for json.Marshal() function only for string type.
 // This function is encodeState.string(string, escapeHTML) in "encoding/json/encode.go".
 func stringJsonMarshal(s string) []byte {
-	var hex = "0123456789abcdef"
-	escapeHTML := true
 	var e bytes.Buffer
 	e.WriteByte('"')
 	start := 0
diff --git a/query/outputnode_test.go b/query/outputnode_test.go
index bfa23a08c0a..5da76fbea4d 100644
--- a/query/outputnode_test.go
+++ b/query/outputnode_test.go
@@ -18,6 +18,7 @@ package query
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"runtime"
 	"sync"
@@ -94,3 +95,24 @@ func TestNormalizeJSONLimit(t *testing.T) {
 	_, err := n.normalize()
 	require.Error(t, err, "Couldn't evaluate @normalize directive - too many results")
 }
+
+func BenchmarkJsonMarshal(b *testing.B) {
+	a := "abcasfasfsadflkjlkjsalfdjsalfjlsajf;llksafs"
+	var result []byte
+
+	b.Run("stringJsonMarshal", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			result = stringJsonMarshal(a)
+		}
+
+		_ = result
+	})
+
+	b.Run("StdJsonMarshal", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			result, _ = json.Marshal(a)
+		}
+
+		_ = result
+	})
+}

From 85937d6e26234363bec48812c78672fc93b99e76 Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Fri, 20 Mar 2020 18:21:10 +0530
Subject: [PATCH 4/6] Add tests and benchmarks

---
 query/outputnode.go      |  7 +++--
 query/outputnode_test.go | 55 ++++++++++++++++++++++++++++++----------
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/query/outputnode.go b/query/outputnode.go
index c662a9932c0..0dab1606901 100644
--- a/query/outputnode.go
+++ b/query/outputnode.go
@@ -133,15 +133,14 @@ var (
 	boolTrue    = []byte("true")
 	boolFalse   = []byte("false")
 	emptyString = []byte(`""`)
-
-	// Used in stringJsonMarshal function.
-	hex        = "0123456789abcdef"
-	escapeHTML = true
 )
 
 // stringJsonMarshal is replacement for json.Marshal() function only for string type.
 // This function is encodeState.string(string, escapeHTML) in "encoding/json/encode.go".
+// It should be in sync with encodeState.string function.
 func stringJsonMarshal(s string) []byte {
+	hex := "0123456789abcdef"
+	escapeHTML := true
 	var e bytes.Buffer
 	e.WriteByte('"')
 	start := 0
diff --git a/query/outputnode_test.go b/query/outputnode_test.go
index 5da76fbea4d..ab40b4975a8 100644
--- a/query/outputnode_test.go
+++ b/query/outputnode_test.go
@@ -97,22 +97,51 @@ func TestNormalizeJSONLimit(t *testing.T) {
 }
 
 func BenchmarkJsonMarshal(b *testing.B) {
-	a := "abcasfasfsadflkjlkjsalfdjsalfjlsajf;llksafs"
+	inputStrings := [][]string{
+		[]string{"smallstring", "abcdef"},
+		[]string{"largestring", string(bytes.Repeat([]byte("a"), 1000))},
+		[]string{"specialchars", "<><>^)(*&(%*&%&^$*&%)(*&)^)"},
+	}
+
 	var result []byte
 
-	b.Run("stringJsonMarshal", func(b *testing.B) {
-		for i := 0; i < b.N; i++ {
-			result = stringJsonMarshal(a)
-		}
+	for _, input := range inputStrings {
+		b.Run(fmt.Sprintf("stringJsonMarshal-%s", input[0]), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				result = stringJsonMarshal(input[1])
+			}
 
-		_ = result
-	})
+			_ = result
+		})
 
-	b.Run("StdJsonMarshal", func(b *testing.B) {
-		for i := 0; i < b.N; i++ {
-			result, _ = json.Marshal(a)
-		}
+		b.Run(fmt.Sprintf("STDJsonMarshal-%s", input[0]), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				result, _ = json.Marshal(input[1])
+			}
+
+			_ = result
+		})
+	}
+}
+
+func TestStringJsonMarshal(t *testing.T) {
+	inputs := []string{
+		"",
+		"0",
+		"true",
+		"1.909045927350",
+		"nil",
+		"null",
+		"<&>",
+		`quoted"str"ing`,
+	}
 
-		_ = result
-	})
+	for _, input := range inputs {
+		gm, err := json.Marshal(input)
+		require.NoError(t, err)
+
+		sm := stringJsonMarshal(input)
+
+		require.Equal(t, gm, sm)
+	}
 }

From 7b282481026e53c8dcc8938fa398d60b3a6ebda6 Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Mon, 23 Mar 2020 19:09:59 +0530
Subject: [PATCH 5/6] Minor change

---
 query/outputnode.go      | 21 +++++++++++++++++----
 query/outputnode_test.go | 17 ++++++++---------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/query/outputnode.go b/query/outputnode.go
index 0dab1606901..0daec898833 100644
--- a/query/outputnode.go
+++ b/query/outputnode.go
@@ -22,6 +22,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 	"unicode/utf8"
 
@@ -133,15 +134,25 @@ var (
 	boolTrue    = []byte("true")
 	boolFalse   = []byte("false")
 	emptyString = []byte(`""`)
+
+	// Below variables are used in stringJsonMarshal function.
+	bufferPool = sync.Pool{
+		New: func() interface{} {
+			return new(bytes.Buffer)
+		},
+	}
+
+	hex        = "0123456789abcdef"
+	escapeHTML = true
 )
 
 // stringJsonMarshal is replacement for json.Marshal() function only for string type.
 // This function is encodeState.string(string, escapeHTML) in "encoding/json/encode.go".
 // It should be in sync with encodeState.string function.
 func stringJsonMarshal(s string) []byte {
-	hex := "0123456789abcdef"
-	escapeHTML := true
-	var e bytes.Buffer
+	e := bufferPool.Get().(*bytes.Buffer)
+	e.Reset()
+
 	e.WriteByte('"')
 	start := 0
 	for i := 0; i < len(s); {
@@ -210,7 +221,9 @@ func stringJsonMarshal(s string) []byte {
 		e.WriteString(s[start:])
 	}
 	e.WriteByte('"')
-	return e.Bytes()
+	buf := append([]byte(nil), e.Bytes()...)
+	bufferPool.Put(e)
+	return buf
 }
 
 func valToBytes(v types.Val) ([]byte, error) {
diff --git a/query/outputnode_test.go b/query/outputnode_test.go
index ab40b4975a8..c05ecf10a02 100644
--- a/query/outputnode_test.go
+++ b/query/outputnode_test.go
@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"runtime"
+	"strings"
 	"sync"
 	"testing"
 
@@ -98,30 +99,28 @@ func TestNormalizeJSONLimit(t *testing.T) {
 
 func BenchmarkJsonMarshal(b *testing.B) {
 	inputStrings := [][]string{
+		[]string{"largestring", strings.Repeat("a", 1024)},
 		[]string{"smallstring", "abcdef"},
-		[]string{"largestring", string(bytes.Repeat([]byte("a"), 1000))},
 		[]string{"specialchars", "<><>^)(*&(%*&%&^$*&%)(*&)^)"},
 	}
 
 	var result []byte
 
 	for _, input := range inputStrings {
-		b.Run(fmt.Sprintf("stringJsonMarshal-%s", input[0]), func(b *testing.B) {
+		b.Run(fmt.Sprintf("STDJsonMarshal-%s", input[0]), func(b *testing.B) {
 			for i := 0; i < b.N; i++ {
-				result = stringJsonMarshal(input[1])
+				result, _ = json.Marshal(input[1])
 			}
-
-			_ = result
 		})
 
-		b.Run(fmt.Sprintf("STDJsonMarshal-%s", input[0]), func(b *testing.B) {
+		b.Run(fmt.Sprintf("stringJsonMarshal-%s", input[0]), func(b *testing.B) {
 			for i := 0; i < b.N; i++ {
-				result, _ = json.Marshal(input[1])
+				result = stringJsonMarshal(input[1])
 			}
-
-			_ = result
 		})
 	}
+
+	_ = result
 }
 
 func TestStringJsonMarshal(t *testing.T) {

From e336177f933d25818d97afd8222df534cbc08a68 Mon Sep 17 00:00:00 2001
From: Ashish Goswami <ashish@dgraph.io>
Date: Mon, 23 Mar 2020 21:22:29 +0530
Subject: [PATCH 6/6] Address review comments

---
 query/outputnode.go           | 8 +++++++-
 query/{table.go => tables.go} | 3 +++
 2 files changed, 10 insertions(+), 1 deletion(-)
 rename query/{table.go => tables.go} (96%)

diff --git a/query/outputnode.go b/query/outputnode.go
index 0daec898833..628710ea2a1 100644
--- a/query/outputnode.go
+++ b/query/outputnode.go
@@ -18,6 +18,7 @@ package query
 
 import (
 	"bytes"
+	"encoding/json"
 	"fmt"
 	"sort"
 	"strconv"
@@ -229,7 +230,12 @@ func stringJsonMarshal(s string) []byte {
 func valToBytes(v types.Val) ([]byte, error) {
 	switch v.Tid {
 	case types.StringID, types.DefaultID:
-		return stringJsonMarshal(v.Value.(string)), nil
+		switch str := v.Value.(type) {
+		case string:
+			return stringJsonMarshal(str), nil
+		default:
+			return json.Marshal(str)
+		}
 	case types.BinaryID:
 		return []byte(fmt.Sprintf("%q", v.Value)), nil
 	case types.IntID:
diff --git a/query/table.go b/query/tables.go
similarity index 96%
rename from query/table.go
rename to query/tables.go
index da4ad23dd20..e09adb83484 100644
--- a/query/table.go
+++ b/query/tables.go
@@ -5,6 +5,9 @@ package query
 
 import "unicode/utf8"
 
+// This file has been taken from go std lib(encoding/json/tables.go).
+// All variable declared here are used in stringJsonMarshal().
+
 // safeSet holds the value true if the ASCII character with the given array
 // position can be represented inside a JSON string without any further
 // escaping.