Add JoinMessages

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
+}

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)
+			}
+		}
+	}
+}