| /* |
| * |
| * Copyright 2017, Google Inc. |
| * All rights reserved. |
| * |
| * Redistribution and use in source and binary forms, with or without |
| * modification, are permitted provided that the following conditions are |
| * met: |
| * |
| * * Redistributions of source code must retain the above copyright |
| * notice, this list of conditions and the following disclaimer. |
| * * Redistributions in binary form must reproduce the above |
| * copyright notice, this list of conditions and the following disclaimer |
| * in the documentation and/or other materials provided with the |
| * distribution. |
| * * Neither the name of Google Inc. nor the names of its |
| * contributors may be used to endorse or promote products derived from |
| * this software without specific prior written permission. |
| * |
| * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| * |
| */ |
| |
| package grpc |
| |
| import ( |
| "bufio" |
| "io" |
| "net" |
| "net/http" |
| "net/url" |
| "testing" |
| "time" |
| |
| "golang.org/x/net/context" |
| ) |
| |
| const ( |
| envTestAddr = "1.2.3.4:8080" |
| envProxyAddr = "2.3.4.5:7687" |
| ) |
| |
| // overwriteAndRestore overwrite function httpProxyFromEnvironment and |
| // returns a function to restore the default values. |
| func overwrite(hpfe func(req *http.Request) (*url.URL, error)) func() { |
| backHPFE := httpProxyFromEnvironment |
| httpProxyFromEnvironment = hpfe |
| return func() { |
| httpProxyFromEnvironment = backHPFE |
| } |
| } |
| |
| func TestMapAddressEnv(t *testing.T) { |
| // Overwrite the function in the test and restore them in defer. |
| hpfe := func(req *http.Request) (*url.URL, error) { |
| if req.URL.Host == envTestAddr { |
| return &url.URL{ |
| Scheme: "https", |
| Host: envProxyAddr, |
| }, nil |
| } |
| return nil, nil |
| } |
| defer overwrite(hpfe)() |
| |
| // envTestAddr should be handled by ProxyFromEnvironment. |
| got, err := mapAddress(context.Background(), envTestAddr) |
| if err != nil { |
| t.Error(err) |
| } |
| if got != envProxyAddr { |
| t.Errorf("want %v, got %v", envProxyAddr, got) |
| } |
| } |
| |
| type proxyServer struct { |
| t *testing.T |
| lis net.Listener |
| in net.Conn |
| out net.Conn |
| } |
| |
| func (p *proxyServer) run() { |
| in, err := p.lis.Accept() |
| if err != nil { |
| return |
| } |
| p.in = in |
| |
| req, err := http.ReadRequest(bufio.NewReader(in)) |
| if err != nil { |
| p.t.Errorf("failed to read CONNECT req: %v", err) |
| return |
| } |
| if req.Method != http.MethodConnect || req.UserAgent() != grpcUA { |
| resp := http.Response{StatusCode: http.StatusMethodNotAllowed} |
| resp.Write(p.in) |
| p.in.Close() |
| p.t.Errorf("get wrong CONNECT req: %+v", req) |
| return |
| } |
| |
| out, err := net.Dial("tcp", req.URL.Host) |
| if err != nil { |
| p.t.Errorf("failed to dial to server: %v", err) |
| return |
| } |
| resp := http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.0"} |
| resp.Write(p.in) |
| p.out = out |
| go io.Copy(p.in, p.out) |
| go io.Copy(p.out, p.in) |
| } |
| |
| func (p *proxyServer) stop() { |
| p.lis.Close() |
| if p.in != nil { |
| p.in.Close() |
| } |
| if p.out != nil { |
| p.out.Close() |
| } |
| } |
| |
| func TestHTTPConnect(t *testing.T) { |
| plis, err := net.Listen("tcp", "localhost:0") |
| if err != nil { |
| t.Fatalf("failed to listen: %v", err) |
| } |
| p := &proxyServer{t: t, lis: plis} |
| go p.run() |
| defer p.stop() |
| |
| blis, err := net.Listen("tcp", "localhost:0") |
| if err != nil { |
| t.Fatalf("failed to listen: %v", err) |
| } |
| |
| msg := []byte{4, 3, 5, 2} |
| recvBuf := make([]byte, len(msg), len(msg)) |
| done := make(chan struct{}) |
| go func() { |
| in, err := blis.Accept() |
| if err != nil { |
| t.Errorf("failed to accept: %v", err) |
| return |
| } |
| defer in.Close() |
| in.Read(recvBuf) |
| close(done) |
| }() |
| |
| // Overwrite the function in the test and restore them in defer. |
| hpfe := func(req *http.Request) (*url.URL, error) { |
| return &url.URL{Host: plis.Addr().String()}, nil |
| } |
| defer overwrite(hpfe)() |
| |
| // Dial to proxy server. |
| dialer := newProxyDialer(func(ctx context.Context, addr string) (net.Conn, error) { |
| if deadline, ok := ctx.Deadline(); ok { |
| return net.DialTimeout("tcp", addr, deadline.Sub(time.Now())) |
| } |
| return net.Dial("tcp", addr) |
| }) |
| ctx, cancel := context.WithTimeout(context.Background(), time.Second) |
| defer cancel() |
| c, err := dialer(ctx, blis.Addr().String()) |
| if err != nil { |
| t.Fatalf("http connect Dial failed: %v", err) |
| } |
| defer c.Close() |
| |
| // Send msg on the connection. |
| c.Write(msg) |
| <-done |
| |
| // Check received msg. |
| if string(recvBuf) != string(msg) { |
| t.Fatalf("received msg: %v, want %v", recvBuf, msg) |
| } |
| } |