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