vendor/google.golang.org/grpc/test/servertester.go

   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
  18
  19 import (
  20         "bytes"
  21         "errors"
  22         "io"
  23         "strings"
  24         "testing"
  25         "time"
  26
  27         "golang.org/x/net/http2"
  28         "golang.org/x/net/http2/hpack"
  29 )
  30
  31 // This is a subset of http2's serverTester type.
  32 //
  33 // serverTester wraps a io.ReadWriter (acting like the underlying
  34 // network connection) and provides utility methods to read and write
  35 // http2 frames.
  36 //
  37 // NOTE(bradfitz): this could eventually be exported somewhere. Others
  38 // have asked for it too. For now I'm still experimenting with the
  39 // API and don't feel like maintaining a stable testing API.
  40
  41 type serverTester struct {
  42         cc io.ReadWriteCloser // client conn
  43         t  testing.TB
  44         fr *http2.Framer
  45
  46         // writing headers:
  47         headerBuf bytes.Buffer
  48         hpackEnc  *hpack.Encoder
  49
  50         // reading frames:
  51         frc       chan http2.Frame
  52         frErrc    chan error
  53         readTimer *time.Timer
  54 }
  55
  56 func 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
  70 func (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.
  93 func (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
 119 func (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
 129 func (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
 135 func (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
 141 func (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 func (st *serverTester) wantSettingsAck() {
 154         f, err := st.readFrame()
 155         if err != nil {
 156                 st.t.Fatal(err)
 157         }
 158         sf, ok := f.(*http2.SettingsFrame)
 159         if !ok {
 160                 st.t.Fatalf("Wanting a settings ACK, received a %T", f)
 161         }
 162         if !sf.IsAck() {
 163                 st.t.Fatal("Settings Frame didn't have ACK set")
 164         }
 165 }
 166
 167 // wait for any activity from the server
 168 func (st *serverTester) wantAnyFrame() http2.Frame {
 169         f, err := st.fr.ReadFrame()
 170         if err != nil {
 171                 st.t.Fatal(err)
 172         }
 173         return f
 174 }
 175
 176 func (st *serverTester) encodeHeaderField(k, v string) {
 177         err := st.hpackEnc.WriteField(hpack.HeaderField{Name: k, Value: v})
 178         if err != nil {
 179                 st.t.Fatalf("HPACK encoding error for %q/%q: %v", k, v, err)
 180         }
 181 }
 182
 183 // encodeHeader encodes headers and returns their HPACK bytes. headers
 184 // must contain an even number of key/value pairs.  There may be
 185 // multiple pairs for keys (e.g. "cookie").  The :method, :path, and
 186 // :scheme headers default to GET, / and https.
 187 func (st *serverTester) encodeHeader(headers ...string) []byte {
 188         if len(headers)%2 == 1 {
 189                 panic("odd number of kv args")
 190         }
 191
 192         st.headerBuf.Reset()
 193
 194         if len(headers) == 0 {
 195                 // Fast path, mostly for benchmarks, so test code doesn't pollute
 196                 // profiles when we're looking to improve server allocations.
 197                 st.encodeHeaderField(":method", "GET")
 198                 st.encodeHeaderField(":path", "/")
 199                 st.encodeHeaderField(":scheme", "https")
 200                 return st.headerBuf.Bytes()
 201         }
 202
 203         if len(headers) == 2 && headers[0] == ":method" {
 204                 // Another fast path for benchmarks.
 205                 st.encodeHeaderField(":method", headers[1])
 206                 st.encodeHeaderField(":path", "/")
 207                 st.encodeHeaderField(":scheme", "https")
 208                 return st.headerBuf.Bytes()
 209         }
 210
 211         pseudoCount := map[string]int{}
 212         keys := []string{":method", ":path", ":scheme"}
 213         vals := map[string][]string{
 214                 ":method": {"GET"},
 215                 ":path":   {"/"},
 216                 ":scheme": {"https"},
 217         }
 218         for len(headers) > 0 {
 219                 k, v := headers[0], headers[1]
 220                 headers = headers[2:]
 221                 if _, ok := vals[k]; !ok {
 222                         keys = append(keys, k)
 223                 }
 224                 if strings.HasPrefix(k, ":") {
 225                         pseudoCount[k]++
 226                         if pseudoCount[k] == 1 {
 227                                 vals[k] = []string{v}
 228                         } else {
 229                                 // Allows testing of invalid headers w/ dup pseudo fields.
 230                                 vals[k] = append(vals[k], v)
 231                         }
 232                 } else {
 233                         vals[k] = append(vals[k], v)
 234                 }
 235         }
 236         for _, k := range keys {
 237                 for _, v := range vals[k] {
 238                         st.encodeHeaderField(k, v)
 239                 }
 240         }
 241         return st.headerBuf.Bytes()
 242 }
 243
 244 func (st *serverTester) writeHeadersGRPC(streamID uint32, path string) {
 245         st.writeHeaders(http2.HeadersFrameParam{
 246                 StreamID: streamID,
 247                 BlockFragment: st.encodeHeader(
 248                         ":method", "POST",
 249                         ":path", path,
 250                         "content-type", "application/grpc",
 251                         "te", "trailers",
 252                 ),
 253                 EndStream:  false,
 254                 EndHeaders: true,
 255         })
 256 }
 257
 258 func (st *serverTester) writeHeaders(p http2.HeadersFrameParam) {
 259         if err := st.fr.WriteHeaders(p); err != nil {
 260                 st.t.Fatalf("Error writing HEADERS: %v", err)
 261         }
 262 }
 263
 264 func (st *serverTester) writeData(streamID uint32, endStream bool, data []byte) {
 265         if err := st.fr.WriteData(streamID, endStream, data); err != nil {
 266                 st.t.Fatalf("Error writing DATA: %v", err)
 267         }
 268 }
 269
 270 func (st *serverTester) writeRSTStream(streamID uint32, code http2.ErrCode) {
 271         if err := st.fr.WriteRSTStream(streamID, code); err != nil {
 272                 st.t.Fatalf("Error writing RST_STREAM: %v", err)
 273         }
 274 }
 275
 276 func (st *serverTester) writeDataPadded(streamID uint32, endStream bool, data, padding []byte) {
 277         if err := st.fr.WriteDataPadded(streamID, endStream, data, padding); err != nil {
 278                 st.t.Fatalf("Error writing DATA with padding: %v", err)
 279         }
 280 }