github.com/google/martian/v3@v3.3.3/h2/testing/fixture.go (about) 1 // Copyright 2021 Google Inc. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 // Package testing contains a test fixture for working with gRPC over HTTP/2. 16 package testing 17 18 import ( 19 "crypto/tls" 20 "fmt" 21 "io/ioutil" 22 "net" 23 "net/http" 24 "os" 25 "strconv" 26 "sync" 27 "time" 28 29 "github.com/google/martian/v3" 30 "github.com/google/martian/v3/h2" 31 "github.com/google/martian/v3/mitm" 32 "google.golang.org/grpc" 33 "google.golang.org/grpc/credentials" 34 35 tspb "github.com/google/martian/v3/h2/testservice" 36 ) 37 38 var ( 39 // proxyPort is a global variable that stores the listener used by the proxy. This value is 40 // shared globally because golang http transport code caches the environment variable values, in 41 // particular HTTPS_PROXY. 42 proxyPort int 43 ) 44 45 // Fixture encapsulates the TestService gRPC server, a proxy and a gRPC client. 46 type Fixture struct { 47 // TestServiceClient is a client pointing at the service and redirected through the proxy. 48 tspb.TestServiceClient 49 50 wg sync.WaitGroup 51 server *grpc.Server 52 // serverErr is any error returned by invoking `Serve` on the gRPC server. 53 serverErr error 54 55 proxyListener net.Listener 56 proxy *martian.Proxy 57 58 conn *grpc.ClientConn 59 } 60 61 // New creates a new instance of the Fixture. It is not possible for there to be more than one 62 // instance concurrently because clients decide whether to use the proxy based on the global 63 // HTTPS_PROXY environment variable. 64 func New(spf []h2.StreamProcessorFactory) (*Fixture, error) { 65 f := &Fixture{} 66 67 // Starts the gRPC server. 68 f.server = grpc.NewServer(grpc.Creds(credentials.NewServerTLSFromCert(Localhost))) 69 tspb.RegisterTestServiceServer(f.server, &Server{}) 70 71 lis, err := net.Listen("tcp", ":0") 72 if err != nil { 73 return nil, fmt.Errorf("creating listener for gRPC service: %w", err) 74 } 75 76 f.wg.Add(1) 77 go func() { 78 defer f.wg.Done() 79 f.serverErr = f.server.Serve(lis) 80 }() 81 82 hostname, err := os.Hostname() 83 if err != nil { 84 return nil, fmt.Errorf("getting hostname: %w", err) 85 } 86 87 // Creates a listener for the proxy, obtaining a new port if needed. 88 if proxyPort == 0 { 89 // Attempts a query to port server first, falling back if it is unavailable. Ports that are 90 // provided by listening on ":0" can be recyled by the OS leading to flakiness in certain 91 // environments since we need the same port to be available across multiple instances of the 92 // test fixture. 93 proxyPort = queryPortServer() 94 if proxyPort == 0 { 95 var err error 96 f.proxyListener, err = net.Listen("tcp", ":0") 97 if err != nil { 98 return nil, fmt.Errorf("creating listener for proxy; %w", err) 99 } 100 proxyPort = f.proxyListener.Addr().(*net.TCPAddr).Port 101 } 102 proxyTarget := hostname + ":" + strconv.Itoa(proxyPort) 103 // Sets the HTTPS_PROXY environment variable so that http requests will go through the proxy. 104 os.Setenv("HTTPS_PROXY", fmt.Sprintf("http://%s", proxyTarget)) 105 fmt.Printf("proxy at %s\n", proxyTarget) 106 } 107 if f.proxyListener == nil { 108 var err error 109 f.proxyListener, err = net.Listen("tcp", fmt.Sprintf(":%d", proxyPort)) 110 if err != nil { 111 return nil, fmt.Errorf("creating listener for proxy; %w", err) 112 } 113 } 114 115 // Starts the proxy. 116 f.proxy, err = newProxy(spf) 117 if err != nil { 118 return nil, fmt.Errorf("creating proxy: %w", err) 119 } 120 go func() { 121 f.proxy.Serve(f.proxyListener) 122 }() 123 124 port := lis.Addr().(*net.TCPAddr).Port 125 target := hostname + ":" + strconv.Itoa(port) 126 127 fmt.Printf("server at %s\n", target) 128 129 // Connects a gRPC client with the service via the proxy. 130 f.conn, err = grpc.Dial(target, grpc.WithTransportCredentials(ClientTLS)) 131 if err != nil { 132 return nil, fmt.Errorf("error dialing %s: %w", target, err) 133 } 134 f.TestServiceClient = tspb.NewTestServiceClient(f.conn) 135 136 return f, nil 137 } 138 139 // Close cleans up the servers and connections. 140 func (f *Fixture) Close() error { 141 f.conn.Close() 142 f.server.Stop() 143 f.proxy.Close() 144 f.wg.Wait() 145 146 if err := f.proxyListener.Close(); err != nil { 147 return fmt.Errorf("closing proxy listener: %w", err) 148 } 149 return f.serverErr 150 } 151 152 func newProxy(spf []h2.StreamProcessorFactory) (*martian.Proxy, error) { 153 p := martian.NewProxy() 154 mc, err := mitm.NewConfig(CA, CAKey) 155 if err != nil { 156 return nil, fmt.Errorf("creating mitm config: %w", err) 157 } 158 mc.SetValidity(time.Hour) 159 mc.SetOrganization("Martian Proxy") 160 mc.SetH2Config(&h2.Config{ 161 AllowedHostsFilter: func(_ string) bool { return true }, 162 RootCAs: RootCAs, 163 StreamProcessorFactories: spf, 164 EnableDebugLogs: true, 165 }) 166 167 p.SetMITM(mc) 168 169 tr := &http.Transport{ 170 TLSClientConfig: &tls.Config{ 171 RootCAs: RootCAs, 172 }, 173 } 174 p.SetRoundTripper(tr) 175 176 return p, nil 177 } 178 179 func queryPortServer() int { 180 // portpicker isn't available in third_party. 181 if portServer := os.Getenv("PORTSERVER_ADDRESS"); portServer != "" { 182 c, err := net.Dial("unix", portServer) 183 if err != nil { 184 // failed connection to portServer; this is normal in many circumstances. 185 return 0 186 } 187 defer c.Close() 188 if _, err := fmt.Fprintf(c, "%d\n", os.Getpid()); err != nil { 189 return 0 190 } 191 buf, err := ioutil.ReadAll(c) 192 if err != nil || len(buf) == 0 { 193 return 0 194 } 195 buf = buf[:len(buf)-1] // remove newline char 196 port, err := strconv.Atoi(string(buf)) 197 if err != nil { 198 return 0 199 } 200 fmt.Printf("got port %d\n", port) 201 return port 202 } 203 return 0 204 }