sdk/go/arvadostest/proxy.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package arvadostest
   6
   7 import (
   8         "crypto/tls"
   9         "net"
  10         "net/http"
  11         "net/http/httptest"
  12         "net/http/httputil"
  13         "net/url"
  14         "sync"
  15         "time"
  16
  17         "git.arvados.org/arvados.git/sdk/go/arvados"
  18         "gopkg.in/check.v1"
  19 )
  20
  21 type Proxy struct {
  22         *httptest.Server
  23
  24         // URL where the proxy is listening. Same as Server.URL, but
  25         // with parsing already done for you.
  26         URL *url.URL
  27
  28         // A dump of each request that has been proxied.
  29         RequestDumps [][]byte
  30
  31         // If non-nil, func will be called on each incoming request
  32         // before proxying it.
  33         Director func(*http.Request)
  34
  35         wg sync.WaitGroup
  36 }
  37
  38 // NewProxy returns a new Proxy that saves a dump of each reqeust
  39 // before forwarding to the indicated service.
  40 func NewProxy(c *check.C, svc arvados.Service) *Proxy {
  41         var target url.URL
  42         c.Assert(svc.InternalURLs, check.HasLen, 1)
  43         for u := range svc.InternalURLs {
  44                 target = url.URL(u)
  45                 break
  46         }
  47         rp := httputil.NewSingleHostReverseProxy(&target)
  48         rp.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
  49                 dump, _ := httputil.DumpRequest(r, false)
  50                 c.Logf("arvadostest.Proxy ErrorHandler(%s): %s\n%s", r.URL, err, dump)
  51                 http.Error(w, err.Error(), http.StatusBadGateway)
  52         }
  53         rp.Transport = &http.Transport{
  54                 DialContext: (&net.Dialer{
  55                         Timeout:   30 * time.Second,
  56                         KeepAlive: 30 * time.Second,
  57                         DualStack: true,
  58                 }).DialContext,
  59                 MaxIdleConns:          100,
  60                 IdleConnTimeout:       90 * time.Second,
  61                 TLSHandshakeTimeout:   10 * time.Second,
  62                 ExpectContinueTimeout: 1 * time.Second,
  63                 TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
  64         }
  65         srv := httptest.NewServer(rp)
  66         u, err := url.Parse(srv.URL)
  67         c.Assert(err, check.IsNil)
  68         proxy := &Proxy{
  69                 Server: srv,
  70                 URL:    u,
  71         }
  72         var mtx sync.Mutex
  73         rp.Director = func(r *http.Request) {
  74                 proxy.wg.Add(1)
  75                 defer proxy.wg.Done()
  76                 if proxy.Director != nil {
  77                         proxy.Director(r)
  78                 }
  79                 dump, _ := httputil.DumpRequest(r, true)
  80                 mtx.Lock()
  81                 proxy.RequestDumps = append(proxy.RequestDumps, dump)
  82                 mtx.Unlock()
  83                 r.URL.Scheme = target.Scheme
  84                 r.URL.Host = target.Host
  85         }
  86         return proxy
  87 }
  88
  89 // Wait waits until all of the proxied requests that have been sent to
  90 // Director() have also been recorded in RequestDumps.
  91 func (proxy *Proxy) Wait() {
  92         proxy.wg.Wait()
  93 }