1// Copyright 2013 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	"bufio"
9	"bytes"
10	"net"
11	"net/http"
12	"reflect"
13	"strings"
14	"testing"
15)
16
17var subprotocolTests = []struct {
18	h         string
19	protocols []string
20}{
21	{"", nil},
22	{"foo", []string{"foo"}},
23	{"foo,bar", []string{"foo", "bar"}},
24	{"foo, bar", []string{"foo", "bar"}},
25	{" foo, bar", []string{"foo", "bar"}},
26	{" foo, bar ", []string{"foo", "bar"}},
27}
28
29func TestSubprotocols(t *testing.T) {
30	for _, st := range subprotocolTests {
31		r := http.Request{Header: http.Header{"Sec-Websocket-Protocol": {st.h}}}
32		protocols := Subprotocols(&r)
33		if !reflect.DeepEqual(st.protocols, protocols) {
34			t.Errorf("SubProtocols(%q) returned %#v, want %#v", st.h, protocols, st.protocols)
35		}
36	}
37}
38
39var isWebSocketUpgradeTests = []struct {
40	ok bool
41	h  http.Header
42}{
43	{false, http.Header{"Upgrade": {"websocket"}}},
44	{false, http.Header{"Connection": {"upgrade"}}},
45	{true, http.Header{"Connection": {"upgRade"}, "Upgrade": {"WebSocket"}}},
46}
47
48func TestIsWebSocketUpgrade(t *testing.T) {
49	for _, tt := range isWebSocketUpgradeTests {
50		ok := IsWebSocketUpgrade(&http.Request{Header: tt.h})
51		if tt.ok != ok {
52			t.Errorf("IsWebSocketUpgrade(%v) returned %v, want %v", tt.h, ok, tt.ok)
53		}
54	}
55}
56
57var checkSameOriginTests = []struct {
58	ok bool
59	r  *http.Request
60}{
61	{false, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://other.org"}}}},
62	{true, &http.Request{Host: "example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
63	{true, &http.Request{Host: "Example.org", Header: map[string][]string{"Origin": {"https://example.org"}}}},
64}
65
66func TestCheckSameOrigin(t *testing.T) {
67	for _, tt := range checkSameOriginTests {
68		ok := checkSameOrigin(tt.r)
69		if tt.ok != ok {
70			t.Errorf("checkSameOrigin(%+v) returned %v, want %v", tt.r, ok, tt.ok)
71		}
72	}
73}
74
75type reuseTestResponseWriter struct {
76	brw *bufio.ReadWriter
77	http.ResponseWriter
78}
79
80func (resp *reuseTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
81	return fakeNetConn{strings.NewReader(""), &bytes.Buffer{}}, resp.brw, nil
82}
83
84var bufioReuseTests = []struct {
85	n     int
86	reuse bool
87}{
88	{4096, true},
89	{128, false},
90}
91
92func TestBufioReuse(t *testing.T) {
93	for i, tt := range bufioReuseTests {
94		br := bufio.NewReaderSize(strings.NewReader(""), tt.n)
95		bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n)
96		resp := &reuseTestResponseWriter{
97			brw: bufio.NewReadWriter(br, bw),
98		}
99		upgrader := Upgrader{}
100		c, err := upgrader.Upgrade(resp, &http.Request{
101			Method: "GET",
102			Header: http.Header{
103				"Upgrade":               []string{"websocket"},
104				"Connection":            []string{"upgrade"},
105				"Sec-Websocket-Key":     []string{"dGhlIHNhbXBsZSBub25jZQ=="},
106				"Sec-Websocket-Version": []string{"13"},
107			}}, nil)
108		if err != nil {
109			t.Fatal(err)
110		}
111		if reuse := c.br == br; reuse != tt.reuse {
112			t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse)
113		}
114		writeBuf := bufioWriterBuffer(c.UnderlyingConn(), bw)
115		if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse {
116			t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse)
117		}
118	}
119}
120