Add `--nul-output|-0` flag to separate element with NUL character (#1550)

This is to ensure solid parsing of complex data (with any binary
content except NUL chars) by separating the `yq` root collection
member's output with NUL char. As a safe-guard, an error will be cast
if trying to use NUL character with content that contains itself NUL
characters inside.

pkg/yqlib/printer.go

@@ -2,6 +2,7 @@ package yqlib
 
 import (
 	"bufio"
+	"bytes"
 	"container/list"
 	"fmt"
 	"io"
@@ -15,6 +16,7 @@ type Printer interface {
 	PrintedAnything() bool
 	//e.g. when given a front-matter doc, like jekyll
 	SetAppendix(reader io.Reader)
+	SetNulSepOutput(nulSepOutput bool)
 }
 
 type PrinterOutputFormat uint32
@@ -59,6 +61,7 @@ type resultsPrinter struct {
 	printedMatches    bool
 	treeNavigator     DataTreeNavigator
 	appendixReader    io.Reader
+	nulSepOutput      bool
 }
 
 func NewPrinter(encoder Encoder, printerWriter PrinterWriter) Printer {
@@ -67,9 +70,16 @@ func NewPrinter(encoder Encoder, printerWriter PrinterWriter) Printer {
 		printerWriter:     printerWriter,
 		firstTimePrinting: true,
 		treeNavigator:     NewDataTreeNavigator(),
+		nulSepOutput:      false,
 	}
 }
 
+func (p *resultsPrinter) SetNulSepOutput(nulSepOutput bool) {
+	log.Debug("Setting NUL separator output")
+
+	p.nulSepOutput = nulSepOutput
+}
+
 func (p *resultsPrinter) SetAppendix(reader io.Reader) {
 	p.appendixReader = reader
 }
@@ -84,6 +94,16 @@ func (p *resultsPrinter) printNode(node *yaml.Node, writer io.Writer) error {
 	return p.encoder.Encode(writer, node)
 }
 
+func removeLastEOL(b *bytes.Buffer) {
+	data := b.Bytes()
+	n := len(data)
+	if n >= 2 && data[n-2] == '\r' && data[n-1] == '\n' {
+		b.Truncate(n - 2)
+	} else if n >= 1 && (data[n-1] == '\r' || data[n-1] == '\n') {
+		b.Truncate(n - 1)
+	}
+}
+
 func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 	log.Debug("PrintResults for %v matches", matchingNodes.Len())
 
@@ -128,18 +148,40 @@ func (p *resultsPrinter) PrintResults(matchingNodes *list.List) error {
 			}
 		}
 
-		if err := p.encoder.PrintLeadingContent(writer, mappedDoc.LeadingContent); err != nil {
+		var destination io.Writer = writer
+		tempBuffer := bytes.NewBuffer(nil)
+		if p.nulSepOutput {
+			destination = tempBuffer
+		}
+
+		if err := p.encoder.PrintLeadingContent(destination, mappedDoc.LeadingContent); err != nil {
 			return err
 		}
 
-		if err := p.printNode(mappedDoc.Node, writer); err != nil {
+		if err := p.printNode(mappedDoc.Node, destination); err != nil {
 			return err
 		}
 
-		if err := p.encoder.PrintLeadingContent(writer, mappedDoc.TrailingContent); err != nil {
+		if err := p.encoder.PrintLeadingContent(destination, mappedDoc.TrailingContent); err != nil {
 			return err
 		}
 
+		if p.nulSepOutput {
+			removeLastEOL(tempBuffer)
+			tempBufferBytes := tempBuffer.Bytes()
+			if bytes.IndexByte(tempBufferBytes, 0) != -1 {
+				return fmt.Errorf(
+					"Can't serialize value because it contains NUL char and you are using NUL separated output",
+				)
+			}
+			if _, err := writer.Write(tempBufferBytes); err != nil {
+				return err
+			}
+			if _, err := writer.Write([]byte{0}); err != nil {
+				return err
+			}
+		}
+
 		p.previousDocIndex = mappedDoc.Document
 		if err := writer.Flush(); err != nil {
 			return err

pkg/yqlib/printer_test.go

@@ -340,3 +340,53 @@ func TestPrinterMultipleDocsJson(t *testing.T) {
 	writer.Flush()
 	test.AssertResult(t, expected, output.String())
 }
+
+func TestPrinterNulSeparator(t *testing.T) {
+	var output bytes.Buffer
+	var writer = bufio.NewWriter(&output)
+	printer := NewSimpleYamlPrinter(writer, YamlOutputFormat, true, false, 2, false)
+	printer.SetNulSepOutput(true)
+	node, err := getExpressionParser().ParseExpression(".a")
+	if err != nil {
+		panic(err)
+	}
+	streamEvaluator := NewStreamEvaluator()
+	_, err = streamEvaluator.Evaluate("sample", strings.NewReader(multiDocSample), node, printer, NewYamlDecoder(ConfiguredYamlPreferences))
+	if err != nil {
+		panic(err)
+	}
+
+	writer.Flush()
+	expected := "banana\x00apple\x00coconut\x00"
+	test.AssertResult(t, expected, output.String())
+}
+
+func TestPrinterNulSeparatorWithJson(t *testing.T) {
+	var output bytes.Buffer
+	var writer = bufio.NewWriter(&output)
+	// note printDocSeparators is true, it should still not print document separators
+	// when outputing JSON.
+	encoder := NewJSONEncoder(0, false, false)
+	if encoder == nil {
+		t.Skipf("no support for %s output format", "json")
+	}
+	printer := NewPrinter(encoder, NewSinglePrinterWriter(writer))
+	printer.SetNulSepOutput(true)
+
+	inputs, err := readDocuments(strings.NewReader(multiDocSample), "sample.yml", 0, NewYamlDecoder(ConfiguredYamlPreferences))
+	if err != nil {
+		panic(err)
+	}
+
+	inputs.Front().Value.(*CandidateNode).LeadingContent = "# ignore this\n"
+
+	err = printer.PrintResults(inputs)
+	if err != nil {
+		panic(err)
+	}
+
+	expected := `{"a":"banana"}` + "\x00" + `{"a":"apple"}` + "\x00" + `{"a":"coconut"}` + "\x00"
+
+	writer.Flush()
+	test.AssertResult(t, expected, output.String())
+}