diff --git a/frame/writer.go b/frame/writer.go index 7c6e83e..a246415 100644 --- a/frame/writer.go +++ b/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,7 +88,7 @@ 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 } @@ -61,7 +96,7 @@ func (w *Writer) Write(f *Frame) error { if err != nil { return err } - _, err = replacerForEncodeValue.WriteString(w.writer, value) + _, err = w.writeHeaderString(value, needEscapeHeaders) if err != nil { return err } diff --git a/frame/writer_test.go b/frame/writer_test.go index 4eeed6b..f9d41ea 100644 --- a/frame/writer_test.go +++ b/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))