vendor/google.golang.org/grpc/proxy_test.go

   1 // +build !race
   2
   3 /*
   4  *
   5  * Copyright 2017 gRPC authors.
   6  *
   7  * Licensed under the Apache License, Version 2.0 (the "License");
   8  * you may not use this file except in compliance with the License.
   9  * You may obtain a copy of the License at
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  13  * Unless required by applicable law or agreed to in writing, software
  14  * distributed under the License is distributed on an "AS IS" BASIS,
  15  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16  * See the License for the specific language governing permissions and
  17  * limitations under the License.
  18  *
  19  */
  20
  21 package grpc
  22
  23 import (
  24         "bufio"
  25         "io"
  26         "net"
  27         "net/http"
  28         "net/url"
  29         "testing"
  30         "time"
  31
  32         "golang.org/x/net/context"
  33         "google.golang.org/grpc/test/leakcheck"
  34 )
  35
  36 const (
  37         envTestAddr  = "1.2.3.4:8080"
  38         envProxyAddr = "2.3.4.5:7687"
  39 )
  40
  41 // overwriteAndRestore overwrite function httpProxyFromEnvironment and
  42 // returns a function to restore the default values.
  43 func overwrite(hpfe func(req *http.Request) (*url.URL, error)) func() {
  44         backHPFE := httpProxyFromEnvironment
  45         httpProxyFromEnvironment = hpfe
  46         return func() {
  47                 httpProxyFromEnvironment = backHPFE
  48         }
  49 }
  50
  51 type proxyServer struct {
  52         t   *testing.T
  53         lis net.Listener
  54         in  net.Conn
  55         out net.Conn
  56 }
  57
  58 func (p *proxyServer) run() {
  59         in, err := p.lis.Accept()
  60         if err != nil {
  61                 return
  62         }
  63         p.in = in
  64
  65         req, err := http.ReadRequest(bufio.NewReader(in))
  66         if err != nil {
  67                 p.t.Errorf("failed to read CONNECT req: %v", err)
  68                 return
  69         }
  70         if req.Method != http.MethodConnect || req.UserAgent() != grpcUA {
  71                 resp := http.Response{StatusCode: http.StatusMethodNotAllowed}
  72                 resp.Write(p.in)
  73                 p.in.Close()
  74                 p.t.Errorf("get wrong CONNECT req: %+v", req)
  75                 return
  76         }
  77
  78         out, err := net.Dial("tcp", req.URL.Host)
  79         if err != nil {
  80                 p.t.Errorf("failed to dial to server: %v", err)
  81                 return
  82         }
  83         resp := http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.0"}
  84         resp.Write(p.in)
  85         p.out = out
  86         go io.Copy(p.in, p.out)
  87         go io.Copy(p.out, p.in)
  88 }
  89
  90 func (p *proxyServer) stop() {
  91         p.lis.Close()
  92         if p.in != nil {
  93                 p.in.Close()
  94         }
  95         if p.out != nil {
  96                 p.out.Close()
  97         }
  98 }
  99
 100 func TestHTTPConnect(t *testing.T) {
 101         defer leakcheck.Check(t)
 102         plis, err := net.Listen("tcp", "localhost:0")
 103         if err != nil {
 104                 t.Fatalf("failed to listen: %v", err)
 105         }
 106         p := &proxyServer{t: t, lis: plis}
 107         go p.run()
 108         defer p.stop()
 109
 110         blis, err := net.Listen("tcp", "localhost:0")
 111         if err != nil {
 112                 t.Fatalf("failed to listen: %v", err)
 113         }
 114
 115         msg := []byte{4, 3, 5, 2}
 116         recvBuf := make([]byte, len(msg), len(msg))
 117         done := make(chan struct{})
 118         go func() {
 119                 in, err := blis.Accept()
 120                 if err != nil {
 121                         t.Errorf("failed to accept: %v", err)
 122                         return
 123                 }
 124                 defer in.Close()
 125                 in.Read(recvBuf)
 126                 close(done)
 127         }()
 128
 129         // Overwrite the function in the test and restore them in defer.
 130         hpfe := func(req *http.Request) (*url.URL, error) {
 131                 return &url.URL{Host: plis.Addr().String()}, nil
 132         }
 133         defer overwrite(hpfe)()
 134
 135         // Dial to proxy server.
 136         dialer := newProxyDialer(func(ctx context.Context, addr string) (net.Conn, error) {
 137                 if deadline, ok := ctx.Deadline(); ok {
 138                         return net.DialTimeout("tcp", addr, deadline.Sub(time.Now()))
 139                 }
 140                 return net.Dial("tcp", addr)
 141         })
 142         ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 143         defer cancel()
 144         c, err := dialer(ctx, blis.Addr().String())
 145         if err != nil {
 146                 t.Fatalf("http connect Dial failed: %v", err)
 147         }
 148         defer c.Close()
 149
 150         // Send msg on the connection.
 151         c.Write(msg)
 152         <-done
 153
 154         // Check received msg.
 155         if string(recvBuf) != string(msg) {
 156                 t.Fatalf("received msg: %v, want %v", recvBuf, msg)
 157         }
 158 }
 159
 160 func TestMapAddressEnv(t *testing.T) {
 161         defer leakcheck.Check(t)
 162         // Overwrite the function in the test and restore them in defer.
 163         hpfe := func(req *http.Request) (*url.URL, error) {
 164                 if req.URL.Host == envTestAddr {
 165                         return &url.URL{
 166                                 Scheme: "https",
 167                                 Host:   envProxyAddr,
 168                         }, nil
 169                 }
 170                 return nil, nil
 171         }
 172         defer overwrite(hpfe)()
 173
 174         // envTestAddr should be handled by ProxyFromEnvironment.
 175         got, err := mapAddress(context.Background(), envTestAddr)
 176         if err != nil {
 177                 t.Error(err)
 178         }
 179         if got != envProxyAddr {
 180                 t.Errorf("want %v, got %v", envProxyAddr, got)
 181         }
 182 }