common/io: Don't generate corrupted escape sequences in WriteData

Author: foxcpp

common/io.go

@@ -136,11 +136,13 @@ func min(a, b int) int {
 // Note: Error may occur even after some data is written so it's better
 // to just CAN transaction after WriteData error.
 func (p *Pipe) WriteData(input []byte) error {
-	encoded := []byte(escapeParameters(string(input)))
-	chunkLen := MaxLineLen - 3 // 3 is for 'D ' and line feed.
-	for i := 0; i < len(encoded); i += chunkLen {
-		chunk := encoded[i:min(i+chunkLen, len(encoded))]
-		chunk = append([]byte{'D', ' '}, chunk...)
+	// Encoding of all bytes increases stream size 3 times.
+	// So by decreasing "chunk" size to 1/3 of it we
+	// will never break this limit.
+	chunkLen := MaxLineLen/3 - 3 // -3 is for 'D ' and line feed.
+	for i := 0; i < len(input); i += chunkLen {
+		chunk := input[i:min(i+chunkLen, len(input))]
+		chunk = append([]byte{'D', ' '}, []byte(escapeParameters(string(chunk)))...)
 		chunk = append(chunk, '\n')
 
 		if _, err := p.w.Write(chunk); err != nil {
@@ -153,7 +155,10 @@ func (p *Pipe) WriteData(input []byte) error {
 // WriteDataReader is similar to WriteData but sends data from input Reader
 // until EOF.
 func (p *Pipe) WriteDataReader(input io.Reader) error {
-	chunkLen := MaxLineLen - 3 // 3 is for 'D ' and line feed.
+	// Encoding of all bytes increases stream size 3 times.
+	// So by decreasing "chunk" size to 1/3 of it we
+	// will never break this limit.
+	chunkLen := MaxLineLen/3 - 3 // -3 is for 'D ' and line feed.
 	buf := make([]byte, chunkLen)
 
 	for {

common/io_test.go

@@ -124,6 +124,26 @@ func TestPipe_WriteData(t *testing.T) {
 			}
 		}
 	})
+	t.Run("line wrapping + escaping", func(t *testing.T) {
+		buf := bytes.Buffer{}
+		pipe := common.NewPipe(nil, &buf)
+		defer pipe.Close()
+
+		data := []byte(strings.Repeat("\n", common.MaxLineLen*7))
+
+		err := pipe.WriteData(data)
+
+		if err != nil {
+			t.Error("Unexpected error on pipe.WriteData:", err)
+		}
+		splitten := strings.Split(buf.String(), "\n")
+		for _, part := range splitten {
+			if len(part)+1 > common.MaxLineLen {
+				t.Error("pipe.WriteData wrote line bigger than MaxLineLen")
+				t.FailNow()
+			}
+		}
+	})
 	t.Run("from io.Reader", func(t *testing.T) {
 		buf := bytes.Buffer{}
 		pipe := common.NewPipe(nil, &buf)
@@ -141,7 +161,27 @@ func TestPipe_WriteData(t *testing.T) {
 			t.Errorf("pipe.WriteData wrote wrong line: '%s'", buf.String())
 		}
 	})
-	// TODO: Test that wrapping is done correctly and no data is corrupted.
+	t.Run("line wrapping + escaping (from io.Reader)", func(t *testing.T) {
+		buf := bytes.Buffer{}
+		pipe := common.NewPipe(nil, &buf)
+		defer pipe.Close()
+
+		data := []byte(strings.Repeat("\n", common.MaxLineLen*7))
+
+		err := pipe.WriteDataReader(bytes.NewReader(data))
+
+		if err != nil {
+			t.Error("Unexpected error on pipe.WriteData:", err)
+			t.FailNow()
+		}
+		splitten := strings.Split(buf.String(), "\n")
+		for _, part := range splitten {
+			if len(part)+1 > common.MaxLineLen {
+				t.Error("pipe.WriteData wrote line bigger than MaxLineLen")
+				t.FailNow()
+			}
+		}
+	})
 }
 
 func TestPipe_ReadData(t *testing.T) {