Do not escape CONNECT and CONNECTED frame headers

frame/writer.go

@@ -3,6 +3,7 @@ package frame
 import (
 	"bufio"
 	"io"
+	"strings"
 )
 
 // slices used to write frames
@@ -27,6 +28,33 @@ func NewWriterSize(writer io.Writer, bufferSize int) *Writer {
 	return &Writer{writer: bufio.NewWriterSize(writer, bufferSize)}
 }
 
+func headersShouldBeEscaped(command string) bool {
+	// STOMP 1.2 excludes CONNECT and CONNECTED frames from header escaping.
+	return command != CONNECT && command != CONNECTED
+}
+
+func validateUnescapedHeaders(header *Header) error {
+	if header == nil {
+		return nil
+	}
+
+	for i := 0; i < header.Len(); i++ {
+		key, value := header.GetAt(i)
+		if key == "" || strings.ContainsAny(key, ":\r\n\x00") || strings.ContainsAny(value, "\r\n\x00") {
+			return ErrInvalidFrameFormat
+		}
+	}
+
+	return nil
+}
+
+func (w *Writer) writeHeaderString(value string, escape bool) (int, error) {
+	if escape {
+		return replacerForEncodeValue.WriteString(w.writer, value)
+	}
+	return w.writer.WriteString(value)
+}
+
 // Write the contents of a frame to the underlying io.Writer.
 func (w *Writer) Write(f *Frame) error {
 	var err error
@@ -38,6 +66,13 @@ func (w *Writer) Write(f *Frame) error {
 			return err
 		}
 	} else {
+		needEscapeHeaders := headersShouldBeEscaped(f.Command)
+		if !needEscapeHeaders {
+			if err = validateUnescapedHeaders(f.Header); err != nil {
+				return err
+			}
+		}
+
 		_, err = w.writer.Write([]byte(f.Command))
 		if err != nil {
 			return err
@@ -53,15 +88,15 @@ func (w *Writer) Write(f *Frame) error {
 			for i := 0; i < f.Header.Len(); i++ {
 				key, value := f.Header.GetAt(i)
 				//println("   ", key, ":", value)
-				_, err = replacerForEncodeValue.WriteString(w.writer, key)
+				_, err = w.writeHeaderString(key, needEscapeHeaders)
 				if err != nil {
 					return err
 				}
 				_, err = w.writer.Write(colonSlice)
 				if err != nil {
 					return err
 				}
-				_, err = replacerForEncodeValue.WriteString(w.writer, value)
+				_, err = w.writeHeaderString(value, needEscapeHeaders)
 				if err != nil {
 					return err
 				}

frame/writer_test.go

@@ -28,6 +28,78 @@ func TestWriterWrites(t *testing.T) {
 	}
 }
 
+func TestWriterConnectHeadersAreNotEscaped(t *testing.T) {
+	tests := []struct {
+		name     string
+		frame    *Frame
+		expected string
+	}{
+		{
+			name: "CONNECT",
+			frame: New(CONNECT,
+				"login", `user\name`,
+				"passcode", `=x/xx-test:password:xx+xxx`,
+			),
+			expected: "CONNECT\n" +
+				"login:user\\name\n" +
+				"passcode:=x/xx-test:password:xx+xxx\n\n\x00",
+		},
+		{
+			name:     "CONNECTED",
+			frame:    New(CONNECTED, "server", `broker:1\main`),
+			expected: "CONNECTED\nserver:broker:1\\main\n\n\x00",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var buffer bytes.Buffer
+			err := NewWriter(&buffer).Write(test.frame)
+			require.NoError(t, err)
+			require.Equal(t, test.expected, buffer.String())
+		})
+	}
+}
+
+func TestWriterOtherFrameHeadersAreEscaped(t *testing.T) {
+	for _, command := range []string{STOMP, SEND} {
+		t.Run(command, func(t *testing.T) {
+			var buffer bytes.Buffer
+			err := NewWriter(&buffer).Write(New(command, "key", `abc:def\ghi`))
+			require.NoError(t, err)
+			require.Equal(t, command+"\nkey:abc\\cdef\\\\ghi\n\n\x00", buffer.String())
+		})
+	}
+}
+
+func TestWriterInvalidUnescapedHeadersAreRejected(t *testing.T) {
+	tests := []struct {
+		name  string
+		key   string
+		value string
+	}{
+		{name: "empty key", key: "", value: "value"},
+		{name: "colon in key", key: "bad:key", value: "value"},
+		{name: "carriage return in key", key: "bad\rkey", value: "value"},
+		{name: "line feed in key", key: "bad\nkey", value: "value"},
+		{name: "NUL in key", key: "bad\x00key", value: "value"},
+		{name: "carriage return in value", key: "key", value: "bad\rvalue"},
+		{name: "line feed in value", key: "key", value: "bad\nvalue"},
+		{name: "NUL in value", key: "key", value: "bad\x00value"},
+	}
+
+	for _, command := range []string{CONNECT, CONNECTED} {
+		for _, test := range tests {
+			t.Run(command+"/"+test.name, func(t *testing.T) {
+				var buffer bytes.Buffer
+				err := NewWriter(&buffer).Write(New(command, test.key, test.value))
+				require.ErrorIs(t, err, ErrInvalidFrameFormat)
+				require.Empty(t, buffer.String())
+			})
+		}
+	}
+}
+
 func writeToBufferAndCheck(t *testing.T, frameText string) {
 	reader := NewReader(strings.NewReader(frameText))