1// Copyright 2017 The Gorilla WebSocket Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package websocket
6
7import (
8	"bytes"
9	"compress/flate"
10	"math/rand"
11	"testing"
12)
13
14var preparedMessageTests = []struct {
15	messageType            int
16	isServer               bool
17	enableWriteCompression bool
18	compressionLevel       int
19}{
20	// Server
21	{TextMessage, true, false, flate.BestSpeed},
22	{TextMessage, true, true, flate.BestSpeed},
23	{TextMessage, true, true, flate.BestCompression},
24	{PingMessage, true, false, flate.BestSpeed},
25	{PingMessage, true, true, flate.BestSpeed},
26
27	// Client
28	{TextMessage, false, false, flate.BestSpeed},
29	{TextMessage, false, true, flate.BestSpeed},
30	{TextMessage, false, true, flate.BestCompression},
31	{PingMessage, false, false, flate.BestSpeed},
32	{PingMessage, false, true, flate.BestSpeed},
33}
34
35func TestPreparedMessage(t *testing.T) {
36	for _, tt := range preparedMessageTests {
37		var data = []byte("this is a test")
38		var buf bytes.Buffer
39		c := newTestConn(nil, &buf, tt.isServer)
40		if tt.enableWriteCompression {
41			c.newCompressionWriter = compressNoContextTakeover
42		}
43		c.SetCompressionLevel(tt.compressionLevel)
44
45		// Seed random number generator for consistent frame mask.
46		rand.Seed(1234)
47
48		if err := c.WriteMessage(tt.messageType, data); err != nil {
49			t.Fatal(err)
50		}
51		want := buf.String()
52
53		pm, err := NewPreparedMessage(tt.messageType, data)
54		if err != nil {
55			t.Fatal(err)
56		}
57
58		// Scribble on data to ensure that NewPreparedMessage takes a snapshot.
59		copy(data, "hello world")
60
61		// Seed random number generator for consistent frame mask.
62		rand.Seed(1234)
63
64		buf.Reset()
65		if err := c.WritePreparedMessage(pm); err != nil {
66			t.Fatal(err)
67		}
68		got := buf.String()
69
70		if got != want {
71			t.Errorf("write message != prepared message for %+v", tt)
72		}
73	}
74}
75