1/*
2 * Copyright 2016 gRPC authors.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17// Package test contains tests.
18package test
19
20import (
21	"bytes"
22	"errors"
23	"io"
24	"strings"
25	"testing"
26	"time"
27
28	"golang.org/x/net/http2"
29	"golang.org/x/net/http2/hpack"
30)
31
32// This is a subset of http2's serverTester type.
33//
34// serverTester wraps a io.ReadWriter (acting like the underlying
35// network connection) and provides utility methods to read and write
36// http2 frames.
37//
38// NOTE(bradfitz): this could eventually be exported somewhere. Others
39// have asked for it too. For now I'm still experimenting with the
40// API and don't feel like maintaining a stable testing API.
41
42type serverTester struct {
43	cc io.ReadWriteCloser // client conn
44	t  testing.TB
45	fr *http2.Framer
46
47	// writing headers:
48	headerBuf bytes.Buffer
49	hpackEnc  *hpack.Encoder
50
51	// reading frames:
52	frc    chan http2.Frame
53	frErrc chan error
54}
55
56func newServerTesterFromConn(t testing.TB, cc io.ReadWriteCloser) *serverTester {
57	st := &serverTester{
58		t:      t,
59		cc:     cc,
60		frc:    make(chan http2.Frame, 1),
61		frErrc: make(chan error, 1),
62	}
63	st.hpackEnc = hpack.NewEncoder(&st.headerBuf)
64	st.fr = http2.NewFramer(cc, cc)
65	st.fr.ReadMetaHeaders = hpack.NewDecoder(4096 /*initialHeaderTableSize*/, nil)
66
67	return st
68}
69
70func (st *serverTester) readFrame() (http2.Frame, error) {
71	go func() {
72		fr, err := st.fr.ReadFrame()
73		if err != nil {
74			st.frErrc <- err
75		} else {
76			st.frc <- fr
77		}
78	}()
79	t := time.NewTimer(2 * time.Second)
80	defer t.Stop()
81	select {
82	case f := <-st.frc:
83		return f, nil
84	case err := <-st.frErrc:
85		return nil, err
86	case <-t.C:
87		return nil, errors.New("timeout waiting for frame")
88	}
89}
90
91// greet initiates the client's HTTP/2 connection into a state where
92// frames may be sent.
93func (st *serverTester) greet() {
94	st.writePreface()
95	st.writeInitialSettings()
96	st.wantSettings()
97	st.writeSettingsAck()
98	for {
99		f, err := st.readFrame()
100		if err != nil {
101			st.t.Fatal(err)
102		}
103		switch f := f.(type) {
104		case *http2.WindowUpdateFrame:
105			// grpc's transport/http2_server sends this
106			// before the settings ack. The Go http2
107			// server uses a setting instead.
108		case *http2.SettingsFrame:
109			if f.IsAck() {
110				return
111			}
112			st.t.Fatalf("during greet, got non-ACK settings frame")
113		default:
114			st.t.Fatalf("during greet, unexpected frame type %T", f)
115		}
116	}
117}
118
119func (st *serverTester) writePreface() {
120	n, err := st.cc.Write([]byte(http2.ClientPreface))
121	if err != nil {
122		st.t.Fatalf("Error writing client preface: %v", err)
123	}
124	if n != len(http2.ClientPreface) {
125		st.t.Fatalf("Writing client preface, wrote %d bytes; want %d", n, len(http2.ClientPreface))
126	}
127}
128
129func (st *serverTester) writeInitialSettings() {
130	if err := st.fr.WriteSettings(); err != nil {
131		st.t.Fatalf("Error writing initial SETTINGS frame from client to server: %v", err)
132	}
133}
134
135func (st *serverTester) writeSettingsAck() {
136	if err := st.fr.WriteSettingsAck(); err != nil {
137		st.t.Fatalf("Error writing ACK of server's SETTINGS: %v", err)
138	}
139}
140
141func (st *serverTester) wantSettings() *http2.SettingsFrame {
142	f, err := st.readFrame()
143	if err != nil {
144		st.t.Fatalf("Error while expecting a SETTINGS frame: %v", err)
145	}
146	sf, ok := f.(*http2.SettingsFrame)
147	if !ok {
148		st.t.Fatalf("got a %T; want *SettingsFrame", f)
149	}
150	return sf
151}
152
153// wait for any activity from the server
154func (st *serverTester) wantAnyFrame() http2.Frame {
155	f, err := st.fr.ReadFrame()
156	if err != nil {
157		st.t.Fatal(err)
158	}
159	return f
160}
161
162func (st *serverTester) encodeHeaderField(k, v string) {
163	err := st.hpackEnc.WriteField(hpack.HeaderField{Name: k, Value: v})
164	if err != nil {
165		st.t.Fatalf("HPACK encoding error for %q/%q: %v", k, v, err)
166	}
167}
168
169// encodeHeader encodes headers and returns their HPACK bytes. headers
170// must contain an even number of key/value pairs.  There may be
171// multiple pairs for keys (e.g. "cookie").  The :method, :path, and
172// :scheme headers default to GET, / and https.
173func (st *serverTester) encodeHeader(headers ...string) []byte {
174	if len(headers)%2 == 1 {
175		panic("odd number of kv args")
176	}
177
178	st.headerBuf.Reset()
179
180	if len(headers) == 0 {
181		// Fast path, mostly for benchmarks, so test code doesn't pollute
182		// profiles when we're looking to improve server allocations.
183		st.encodeHeaderField(":method", "GET")
184		st.encodeHeaderField(":path", "/")
185		st.encodeHeaderField(":scheme", "https")
186		return st.headerBuf.Bytes()
187	}
188
189	if len(headers) == 2 && headers[0] == ":method" {
190		// Another fast path for benchmarks.
191		st.encodeHeaderField(":method", headers[1])
192		st.encodeHeaderField(":path", "/")
193		st.encodeHeaderField(":scheme", "https")
194		return st.headerBuf.Bytes()
195	}
196
197	pseudoCount := map[string]int{}
198	keys := []string{":method", ":path", ":scheme"}
199	vals := map[string][]string{
200		":method": {"GET"},
201		":path":   {"/"},
202		":scheme": {"https"},
203	}
204	for len(headers) > 0 {
205		k, v := headers[0], headers[1]
206		headers = headers[2:]
207		if _, ok := vals[k]; !ok {
208			keys = append(keys, k)
209		}
210		if strings.HasPrefix(k, ":") {
211			pseudoCount[k]++
212			if pseudoCount[k] == 1 {
213				vals[k] = []string{v}
214			} else {
215				// Allows testing of invalid headers w/ dup pseudo fields.
216				vals[k] = append(vals[k], v)
217			}
218		} else {
219			vals[k] = append(vals[k], v)
220		}
221	}
222	for _, k := range keys {
223		for _, v := range vals[k] {
224			st.encodeHeaderField(k, v)
225		}
226	}
227	return st.headerBuf.Bytes()
228}
229
230func (st *serverTester) writeHeadersGRPC(streamID uint32, path string) {
231	st.writeHeaders(http2.HeadersFrameParam{
232		StreamID: streamID,
233		BlockFragment: st.encodeHeader(
234			":method", "POST",
235			":path", path,
236			"content-type", "application/grpc",
237			"te", "trailers",
238		),
239		EndStream:  false,
240		EndHeaders: true,
241	})
242}
243
244func (st *serverTester) writeHeaders(p http2.HeadersFrameParam) {
245	if err := st.fr.WriteHeaders(p); err != nil {
246		st.t.Fatalf("Error writing HEADERS: %v", err)
247	}
248}
249
250func (st *serverTester) writeData(streamID uint32, endStream bool, data []byte) {
251	if err := st.fr.WriteData(streamID, endStream, data); err != nil {
252		st.t.Fatalf("Error writing DATA: %v", err)
253	}
254}
255
256func (st *serverTester) writeRSTStream(streamID uint32, code http2.ErrCode) {
257	if err := st.fr.WriteRSTStream(streamID, code); err != nil {
258		st.t.Fatalf("Error writing RST_STREAM: %v", err)
259	}
260}
261