diff --git a/join.go b/join.go new file mode 100644 index 0000000..c64f8c8 --- /dev/null +++ b/join.go @@ -0,0 +1,42 @@ +// Copyright 2019 The Gorilla WebSocket 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 websocket + +import ( + "io" + "strings" +) + +// JoinMessages concatenates received messages to create a single io.Reader. +// The string term is appended to each message. The returned reader does not +// support concurrent calls to the Read method. +func JoinMessages(c *Conn, term string) io.Reader { + return &joinReader{c: c, term: term} +} + +type joinReader struct { + c *Conn + term string + r io.Reader +} + +func (r *joinReader) Read(p []byte) (int, error) { + if r.r == nil { + var err error + _, r.r, err = r.c.NextReader() + if err != nil { + return 0, err + } + if r.term != "" { + r.r = io.MultiReader(r.r, strings.NewReader(r.term)) + } + } + n, err := r.r.Read(p) + if err == io.EOF { + err = nil + r.r = nil + } + return n, err +} diff --git a/join_test.go b/join_test.go new file mode 100644 index 0000000..961ac04 --- /dev/null +++ b/join_test.go @@ -0,0 +1,36 @@ +// Copyright 2019 The Gorilla WebSocket 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 websocket + +import ( + "bytes" + "io" + "strings" + "testing" +) + +func TestJoinMessages(t *testing.T) { + messages := []string{"a", "bc", "def", "ghij", "klmno", "0", "12", "345", "6789"} + for _, readChunk := range []int{1, 2, 3, 4, 5, 6, 7} { + for _, term := range []string{"", ","} { + var connBuf bytes.Buffer + wc := newTestConn(nil, &connBuf, true) + rc := newTestConn(&connBuf, nil, false) + for _, m := range messages { + wc.WriteMessage(BinaryMessage, []byte(m)) + } + + var result bytes.Buffer + _, err := io.CopyBuffer(&result, JoinMessages(rc, term), make([]byte, readChunk)) + if IsUnexpectedCloseError(err, CloseAbnormalClosure) { + t.Errorf("readChunk=%d, term=%q: unexpected error %v", readChunk, term, err) + } + want := strings.Join(messages, term) + term + if result.String() != want { + t.Errorf("readChunk=%d, term=%q, got %q, want %q", readChunk, term, result.String(), want) + } + } + } +}