libgo/go/net/http/httputil/reverseproxy_test.go

   1 // Copyright 2011 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 // Reverse proxy tests.
   6
   7 package httputil
   8
   9 import (
  10         "bufio"
  11         "bytes"
  12         "context"
  13         "errors"
  14         "fmt"
  15         "io"
  16         "log"
  17         "net/http"
  18         "net/http/httptest"
  19         "net/url"
  20         "os"
  21         "reflect"
  22         "sort"
  23         "strconv"
  24         "strings"
  25         "sync"
  26         "testing"
  27         "time"
  28 )
  29
  30 const fakeHopHeader = "X-Fake-Hop-Header-For-Test"
  31
  32 func init() {
  33         inOurTests = true
  34         hopHeaders = append(hopHeaders, fakeHopHeader)
  35 }
  36
  37 func TestReverseProxy(t *testing.T) {
  38         const backendResponse = "I am the backend"
  39         const backendStatus = 404
  40         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  41                 if r.Method == "GET" && r.FormValue("mode") == "hangup" {
  42                         c, _, _ := w.(http.Hijacker).Hijack()
  43                         c.Close()
  44                         return
  45                 }
  46                 if len(r.TransferEncoding) > 0 {
  47                         t.Errorf("backend got unexpected TransferEncoding: %v", r.TransferEncoding)
  48                 }
  49                 if r.Header.Get("X-Forwarded-For") == "" {
  50                         t.Errorf("didn't get X-Forwarded-For header")
  51                 }
  52                 if c := r.Header.Get("Connection"); c != "" {
  53                         t.Errorf("handler got Connection header value %q", c)
  54                 }
  55                 if c := r.Header.Get("Te"); c != "trailers" {
  56                         t.Errorf("handler got Te header value %q; want 'trailers'", c)
  57                 }
  58                 if c := r.Header.Get("Upgrade"); c != "" {
  59                         t.Errorf("handler got Upgrade header value %q", c)
  60                 }
  61                 if c := r.Header.Get("Proxy-Connection"); c != "" {
  62                         t.Errorf("handler got Proxy-Connection header value %q", c)
  63                 }
  64                 if g, e := r.Host, "some-name"; g != e {
  65                         t.Errorf("backend got Host header %q, want %q", g, e)
  66                 }
  67                 w.Header().Set("Trailers", "not a special header field name")
  68                 w.Header().Set("Trailer", "X-Trailer")
  69                 w.Header().Set("X-Foo", "bar")
  70                 w.Header().Set("Upgrade", "foo")
  71                 w.Header().Set(fakeHopHeader, "foo")
  72                 w.Header().Add("X-Multi-Value", "foo")
  73                 w.Header().Add("X-Multi-Value", "bar")
  74                 http.SetCookie(w, &http.Cookie{Name: "flavor", Value: "chocolateChip"})
  75                 w.WriteHeader(backendStatus)
  76                 w.Write([]byte(backendResponse))
  77                 w.Header().Set("X-Trailer", "trailer_value")
  78                 w.Header().Set(http.TrailerPrefix+"X-Unannounced-Trailer", "unannounced_trailer_value")
  79         }))
  80         defer backend.Close()
  81         backendURL, err := url.Parse(backend.URL)
  82         if err != nil {
  83                 t.Fatal(err)
  84         }
  85         proxyHandler := NewSingleHostReverseProxy(backendURL)
  86         proxyHandler.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
  87         frontend := httptest.NewServer(proxyHandler)
  88         defer frontend.Close()
  89         frontendClient := frontend.Client()
  90
  91         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
  92         getReq.Host = "some-name"
  93         getReq.Header.Set("Connection", "close, TE")
  94         getReq.Header.Add("Te", "foo")
  95         getReq.Header.Add("Te", "bar, trailers")
  96         getReq.Header.Set("Proxy-Connection", "should be deleted")
  97         getReq.Header.Set("Upgrade", "foo")
  98         getReq.Close = true
  99         res, err := frontendClient.Do(getReq)
 100         if err != nil {
 101                 t.Fatalf("Get: %v", err)
 102         }
 103         if g, e := res.StatusCode, backendStatus; g != e {
 104                 t.Errorf("got res.StatusCode %d; expected %d", g, e)
 105         }
 106         if g, e := res.Header.Get("X-Foo"), "bar"; g != e {
 107                 t.Errorf("got X-Foo %q; expected %q", g, e)
 108         }
 109         if c := res.Header.Get(fakeHopHeader); c != "" {
 110                 t.Errorf("got %s header value %q", fakeHopHeader, c)
 111         }
 112         if g, e := res.Header.Get("Trailers"), "not a special header field name"; g != e {
 113                 t.Errorf("header Trailers = %q; want %q", g, e)
 114         }
 115         if g, e := len(res.Header["X-Multi-Value"]), 2; g != e {
 116                 t.Errorf("got %d X-Multi-Value header values; expected %d", g, e)
 117         }
 118         if g, e := len(res.Header["Set-Cookie"]), 1; g != e {
 119                 t.Fatalf("got %d SetCookies, want %d", g, e)
 120         }
 121         if g, e := res.Trailer, (http.Header{"X-Trailer": nil}); !reflect.DeepEqual(g, e) {
 122                 t.Errorf("before reading body, Trailer = %#v; want %#v", g, e)
 123         }
 124         if cookie := res.Cookies()[0]; cookie.Name != "flavor" {
 125                 t.Errorf("unexpected cookie %q", cookie.Name)
 126         }
 127         bodyBytes, _ := io.ReadAll(res.Body)
 128         if g, e := string(bodyBytes), backendResponse; g != e {
 129                 t.Errorf("got body %q; expected %q", g, e)
 130         }
 131         if g, e := res.Trailer.Get("X-Trailer"), "trailer_value"; g != e {
 132                 t.Errorf("Trailer(X-Trailer) = %q ; want %q", g, e)
 133         }
 134         if g, e := res.Trailer.Get("X-Unannounced-Trailer"), "unannounced_trailer_value"; g != e {
 135                 t.Errorf("Trailer(X-Unannounced-Trailer) = %q ; want %q", g, e)
 136         }
 137
 138         // Test that a backend failing to be reached or one which doesn't return
 139         // a response results in a StatusBadGateway.
 140         getReq, _ = http.NewRequest("GET", frontend.URL+"/?mode=hangup", nil)
 141         getReq.Close = true
 142         res, err = frontendClient.Do(getReq)
 143         if err != nil {
 144                 t.Fatal(err)
 145         }
 146         res.Body.Close()
 147         if res.StatusCode != http.StatusBadGateway {
 148                 t.Errorf("request to bad proxy = %v; want 502 StatusBadGateway", res.Status)
 149         }
 150
 151 }
 152
 153 // Issue 16875: remove any proxied headers mentioned in the "Connection"
 154 // header value.
 155 func TestReverseProxyStripHeadersPresentInConnection(t *testing.T) {
 156         const fakeConnectionToken = "X-Fake-Connection-Token"
 157         const backendResponse = "I am the backend"
 158
 159         // someConnHeader is some arbitrary header to be declared as a hop-by-hop header
 160         // in the Request's Connection header.
 161         const someConnHeader = "X-Some-Conn-Header"
 162
 163         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 164                 if c := r.Header.Get("Connection"); c != "" {
 165                         t.Errorf("handler got header %q = %q; want empty", "Connection", c)
 166                 }
 167                 if c := r.Header.Get(fakeConnectionToken); c != "" {
 168                         t.Errorf("handler got header %q = %q; want empty", fakeConnectionToken, c)
 169                 }
 170                 if c := r.Header.Get(someConnHeader); c != "" {
 171                         t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 172                 }
 173                 w.Header().Add("Connection", "Upgrade, "+fakeConnectionToken)
 174                 w.Header().Add("Connection", someConnHeader)
 175                 w.Header().Set(someConnHeader, "should be deleted")
 176                 w.Header().Set(fakeConnectionToken, "should be deleted")
 177                 io.WriteString(w, backendResponse)
 178         }))
 179         defer backend.Close()
 180         backendURL, err := url.Parse(backend.URL)
 181         if err != nil {
 182                 t.Fatal(err)
 183         }
 184         proxyHandler := NewSingleHostReverseProxy(backendURL)
 185         frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 186                 proxyHandler.ServeHTTP(w, r)
 187                 if c := r.Header.Get(someConnHeader); c != "should be deleted" {
 188                         t.Errorf("handler modified header %q = %q; want %q", someConnHeader, c, "should be deleted")
 189                 }
 190                 if c := r.Header.Get(fakeConnectionToken); c != "should be deleted" {
 191                         t.Errorf("handler modified header %q = %q; want %q", fakeConnectionToken, c, "should be deleted")
 192                 }
 193                 c := r.Header["Connection"]
 194                 var cf []string
 195                 for _, f := range c {
 196                         for _, sf := range strings.Split(f, ",") {
 197                                 if sf = strings.TrimSpace(sf); sf != "" {
 198                                         cf = append(cf, sf)
 199                                 }
 200                         }
 201                 }
 202                 sort.Strings(cf)
 203                 expectedValues := []string{"Upgrade", someConnHeader, fakeConnectionToken}
 204                 sort.Strings(expectedValues)
 205                 if !reflect.DeepEqual(cf, expectedValues) {
 206                         t.Errorf("handler modified header %q = %q; want %q", "Connection", cf, expectedValues)
 207                 }
 208         }))
 209         defer frontend.Close()
 210
 211         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 212         getReq.Header.Add("Connection", "Upgrade, "+fakeConnectionToken)
 213         getReq.Header.Add("Connection", someConnHeader)
 214         getReq.Header.Set(someConnHeader, "should be deleted")
 215         getReq.Header.Set(fakeConnectionToken, "should be deleted")
 216         res, err := frontend.Client().Do(getReq)
 217         if err != nil {
 218                 t.Fatalf("Get: %v", err)
 219         }
 220         defer res.Body.Close()
 221         bodyBytes, err := io.ReadAll(res.Body)
 222         if err != nil {
 223                 t.Fatalf("reading body: %v", err)
 224         }
 225         if got, want := string(bodyBytes), backendResponse; got != want {
 226                 t.Errorf("got body %q; want %q", got, want)
 227         }
 228         if c := res.Header.Get("Connection"); c != "" {
 229                 t.Errorf("handler got header %q = %q; want empty", "Connection", c)
 230         }
 231         if c := res.Header.Get(someConnHeader); c != "" {
 232                 t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 233         }
 234         if c := res.Header.Get(fakeConnectionToken); c != "" {
 235                 t.Errorf("handler got header %q = %q; want empty", fakeConnectionToken, c)
 236         }
 237 }
 238
 239 func TestReverseProxyStripEmptyConnection(t *testing.T) {
 240         // See Issue 46313.
 241         const backendResponse = "I am the backend"
 242
 243         // someConnHeader is some arbitrary header to be declared as a hop-by-hop header
 244         // in the Request's Connection header.
 245         const someConnHeader = "X-Some-Conn-Header"
 246
 247         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 248                 if c := r.Header.Values("Connection"); len(c) != 0 {
 249                         t.Errorf("handler got header %q = %v; want empty", "Connection", c)
 250                 }
 251                 if c := r.Header.Get(someConnHeader); c != "" {
 252                         t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 253                 }
 254                 w.Header().Add("Connection", "")
 255                 w.Header().Add("Connection", someConnHeader)
 256                 w.Header().Set(someConnHeader, "should be deleted")
 257                 io.WriteString(w, backendResponse)
 258         }))
 259         defer backend.Close()
 260         backendURL, err := url.Parse(backend.URL)
 261         if err != nil {
 262                 t.Fatal(err)
 263         }
 264         proxyHandler := NewSingleHostReverseProxy(backendURL)
 265         frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 266                 proxyHandler.ServeHTTP(w, r)
 267                 if c := r.Header.Get(someConnHeader); c != "should be deleted" {
 268                         t.Errorf("handler modified header %q = %q; want %q", someConnHeader, c, "should be deleted")
 269                 }
 270         }))
 271         defer frontend.Close()
 272
 273         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 274         getReq.Header.Add("Connection", "")
 275         getReq.Header.Add("Connection", someConnHeader)
 276         getReq.Header.Set(someConnHeader, "should be deleted")
 277         res, err := frontend.Client().Do(getReq)
 278         if err != nil {
 279                 t.Fatalf("Get: %v", err)
 280         }
 281         defer res.Body.Close()
 282         bodyBytes, err := io.ReadAll(res.Body)
 283         if err != nil {
 284                 t.Fatalf("reading body: %v", err)
 285         }
 286         if got, want := string(bodyBytes), backendResponse; got != want {
 287                 t.Errorf("got body %q; want %q", got, want)
 288         }
 289         if c := res.Header.Get("Connection"); c != "" {
 290                 t.Errorf("handler got header %q = %q; want empty", "Connection", c)
 291         }
 292         if c := res.Header.Get(someConnHeader); c != "" {
 293                 t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 294         }
 295 }
 296
 297 func TestXForwardedFor(t *testing.T) {
 298         const prevForwardedFor = "client ip"
 299         const backendResponse = "I am the backend"
 300         const backendStatus = 404
 301         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 302                 if r.Header.Get("X-Forwarded-For") == "" {
 303                         t.Errorf("didn't get X-Forwarded-For header")
 304                 }
 305                 if !strings.Contains(r.Header.Get("X-Forwarded-For"), prevForwardedFor) {
 306                         t.Errorf("X-Forwarded-For didn't contain prior data")
 307                 }
 308                 w.WriteHeader(backendStatus)
 309                 w.Write([]byte(backendResponse))
 310         }))
 311         defer backend.Close()
 312         backendURL, err := url.Parse(backend.URL)
 313         if err != nil {
 314                 t.Fatal(err)
 315         }
 316         proxyHandler := NewSingleHostReverseProxy(backendURL)
 317         frontend := httptest.NewServer(proxyHandler)
 318         defer frontend.Close()
 319
 320         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 321         getReq.Host = "some-name"
 322         getReq.Header.Set("Connection", "close")
 323         getReq.Header.Set("X-Forwarded-For", prevForwardedFor)
 324         getReq.Close = true
 325         res, err := frontend.Client().Do(getReq)
 326         if err != nil {
 327                 t.Fatalf("Get: %v", err)
 328         }
 329         if g, e := res.StatusCode, backendStatus; g != e {
 330                 t.Errorf("got res.StatusCode %d; expected %d", g, e)
 331         }
 332         bodyBytes, _ := io.ReadAll(res.Body)
 333         if g, e := string(bodyBytes), backendResponse; g != e {
 334                 t.Errorf("got body %q; expected %q", g, e)
 335         }
 336 }
 337
 338 // Issue 38079: don't append to X-Forwarded-For if it's present but nil
 339 func TestXForwardedFor_Omit(t *testing.T) {
 340         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 341                 if v := r.Header.Get("X-Forwarded-For"); v != "" {
 342                         t.Errorf("got X-Forwarded-For header: %q", v)
 343                 }
 344                 w.Write([]byte("hi"))
 345         }))
 346         defer backend.Close()
 347         backendURL, err := url.Parse(backend.URL)
 348         if err != nil {
 349                 t.Fatal(err)
 350         }
 351         proxyHandler := NewSingleHostReverseProxy(backendURL)
 352         frontend := httptest.NewServer(proxyHandler)
 353         defer frontend.Close()
 354
 355         oldDirector := proxyHandler.Director
 356         proxyHandler.Director = func(r *http.Request) {
 357                 r.Header["X-Forwarded-For"] = nil
 358                 oldDirector(r)
 359         }
 360
 361         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 362         getReq.Host = "some-name"
 363         getReq.Close = true
 364         res, err := frontend.Client().Do(getReq)
 365         if err != nil {
 366                 t.Fatalf("Get: %v", err)
 367         }
 368         res.Body.Close()
 369 }
 370
 371 var proxyQueryTests = []struct {
 372         baseSuffix string // suffix to add to backend URL
 373         reqSuffix  string // suffix to add to frontend's request URL
 374         want       string // what backend should see for final request URL (without ?)
 375 }{
 376         {"", "", ""},
 377         {"?sta=tic", "?us=er", "sta=tic&us=er"},
 378         {"", "?us=er", "us=er"},
 379         {"?sta=tic", "", "sta=tic"},
 380 }
 381
 382 func TestReverseProxyQuery(t *testing.T) {
 383         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 384                 w.Header().Set("X-Got-Query", r.URL.RawQuery)
 385                 w.Write([]byte("hi"))
 386         }))
 387         defer backend.Close()
 388
 389         for i, tt := range proxyQueryTests {
 390                 backendURL, err := url.Parse(backend.URL + tt.baseSuffix)
 391                 if err != nil {
 392                         t.Fatal(err)
 393                 }
 394                 frontend := httptest.NewServer(NewSingleHostReverseProxy(backendURL))
 395                 req, _ := http.NewRequest("GET", frontend.URL+tt.reqSuffix, nil)
 396                 req.Close = true
 397                 res, err := frontend.Client().Do(req)
 398                 if err != nil {
 399                         t.Fatalf("%d. Get: %v", i, err)
 400                 }
 401                 if g, e := res.Header.Get("X-Got-Query"), tt.want; g != e {
 402                         t.Errorf("%d. got query %q; expected %q", i, g, e)
 403                 }
 404                 res.Body.Close()
 405                 frontend.Close()
 406         }
 407 }
 408
 409 func TestReverseProxyFlushInterval(t *testing.T) {
 410         const expected = "hi"
 411         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 412                 w.Write([]byte(expected))
 413         }))
 414         defer backend.Close()
 415
 416         backendURL, err := url.Parse(backend.URL)
 417         if err != nil {
 418                 t.Fatal(err)
 419         }
 420
 421         proxyHandler := NewSingleHostReverseProxy(backendURL)
 422         proxyHandler.FlushInterval = time.Microsecond
 423
 424         frontend := httptest.NewServer(proxyHandler)
 425         defer frontend.Close()
 426
 427         req, _ := http.NewRequest("GET", frontend.URL, nil)
 428         req.Close = true
 429         res, err := frontend.Client().Do(req)
 430         if err != nil {
 431                 t.Fatalf("Get: %v", err)
 432         }
 433         defer res.Body.Close()
 434         if bodyBytes, _ := io.ReadAll(res.Body); string(bodyBytes) != expected {
 435                 t.Errorf("got body %q; expected %q", bodyBytes, expected)
 436         }
 437 }
 438
 439 func TestReverseProxyFlushIntervalHeaders(t *testing.T) {
 440         const expected = "hi"
 441         stopCh := make(chan struct{})
 442         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 443                 w.Header().Add("MyHeader", expected)
 444                 w.WriteHeader(200)
 445                 w.(http.Flusher).Flush()
 446                 <-stopCh
 447         }))
 448         defer backend.Close()
 449         defer close(stopCh)
 450
 451         backendURL, err := url.Parse(backend.URL)
 452         if err != nil {
 453                 t.Fatal(err)
 454         }
 455
 456         proxyHandler := NewSingleHostReverseProxy(backendURL)
 457         proxyHandler.FlushInterval = time.Microsecond
 458
 459         frontend := httptest.NewServer(proxyHandler)
 460         defer frontend.Close()
 461
 462         req, _ := http.NewRequest("GET", frontend.URL, nil)
 463         req.Close = true
 464
 465         ctx, cancel := context.WithTimeout(req.Context(), 10*time.Second)
 466         defer cancel()
 467         req = req.WithContext(ctx)
 468
 469         res, err := frontend.Client().Do(req)
 470         if err != nil {
 471                 t.Fatalf("Get: %v", err)
 472         }
 473         defer res.Body.Close()
 474
 475         if res.Header.Get("MyHeader") != expected {
 476                 t.Errorf("got header %q; expected %q", res.Header.Get("MyHeader"), expected)
 477         }
 478 }
 479
 480 func TestReverseProxyCancellation(t *testing.T) {
 481         const backendResponse = "I am the backend"
 482
 483         reqInFlight := make(chan struct{})
 484         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 485                 close(reqInFlight) // cause the client to cancel its request
 486
 487                 select {
 488                 case <-time.After(10 * time.Second):
 489                         // Note: this should only happen in broken implementations, and the
 490                         // closenotify case should be instantaneous.
 491                         t.Error("Handler never saw CloseNotify")
 492                         return
 493                 case <-w.(http.CloseNotifier).CloseNotify():
 494                 }
 495
 496                 w.WriteHeader(http.StatusOK)
 497                 w.Write([]byte(backendResponse))
 498         }))
 499
 500         defer backend.Close()
 501
 502         backend.Config.ErrorLog = log.New(io.Discard, "", 0)
 503
 504         backendURL, err := url.Parse(backend.URL)
 505         if err != nil {
 506                 t.Fatal(err)
 507         }
 508
 509         proxyHandler := NewSingleHostReverseProxy(backendURL)
 510
 511         // Discards errors of the form:
 512         // http: proxy error: read tcp 127.0.0.1:44643: use of closed network connection
 513         proxyHandler.ErrorLog = log.New(io.Discard, "", 0)
 514
 515         frontend := httptest.NewServer(proxyHandler)
 516         defer frontend.Close()
 517         frontendClient := frontend.Client()
 518
 519         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 520         go func() {
 521                 <-reqInFlight
 522                 frontendClient.Transport.(*http.Transport).CancelRequest(getReq)
 523         }()
 524         res, err := frontendClient.Do(getReq)
 525         if res != nil {
 526                 t.Errorf("got response %v; want nil", res.Status)
 527         }
 528         if err == nil {
 529                 // This should be an error like:
 530                 // Get "http://127.0.0.1:58079": read tcp 127.0.0.1:58079:
 531                 //    use of closed network connection
 532                 t.Error("Server.Client().Do() returned nil error; want non-nil error")
 533         }
 534 }
 535
 536 func req(t *testing.T, v string) *http.Request {
 537         req, err := http.ReadRequest(bufio.NewReader(strings.NewReader(v)))
 538         if err != nil {
 539                 t.Fatal(err)
 540         }
 541         return req
 542 }
 543
 544 // Issue 12344
 545 func TestNilBody(t *testing.T) {
 546         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 547                 w.Write([]byte("hi"))
 548         }))
 549         defer backend.Close()
 550
 551         frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
 552                 backURL, _ := url.Parse(backend.URL)
 553                 rp := NewSingleHostReverseProxy(backURL)
 554                 r := req(t, "GET / HTTP/1.0\r\n\r\n")
 555                 r.Body = nil // this accidentally worked in Go 1.4 and below, so keep it working
 556                 rp.ServeHTTP(w, r)
 557         }))
 558         defer frontend.Close()
 559
 560         res, err := http.Get(frontend.URL)
 561         if err != nil {
 562                 t.Fatal(err)
 563         }
 564         defer res.Body.Close()
 565         slurp, err := io.ReadAll(res.Body)
 566         if err != nil {
 567                 t.Fatal(err)
 568         }
 569         if string(slurp) != "hi" {
 570                 t.Errorf("Got %q; want %q", slurp, "hi")
 571         }
 572 }
 573
 574 // Issue 15524
 575 func TestUserAgentHeader(t *testing.T) {
 576         const explicitUA = "explicit UA"
 577         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 578                 if r.URL.Path == "/noua" {
 579                         if c := r.Header.Get("User-Agent"); c != "" {
 580                                 t.Errorf("handler got non-empty User-Agent header %q", c)
 581                         }
 582                         return
 583                 }
 584                 if c := r.Header.Get("User-Agent"); c != explicitUA {
 585                         t.Errorf("handler got unexpected User-Agent header %q", c)
 586                 }
 587         }))
 588         defer backend.Close()
 589         backendURL, err := url.Parse(backend.URL)
 590         if err != nil {
 591                 t.Fatal(err)
 592         }
 593         proxyHandler := NewSingleHostReverseProxy(backendURL)
 594         proxyHandler.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
 595         frontend := httptest.NewServer(proxyHandler)
 596         defer frontend.Close()
 597         frontendClient := frontend.Client()
 598
 599         getReq, _ := http.NewRequest("GET", frontend.URL, nil)
 600         getReq.Header.Set("User-Agent", explicitUA)
 601         getReq.Close = true
 602         res, err := frontendClient.Do(getReq)
 603         if err != nil {
 604                 t.Fatalf("Get: %v", err)
 605         }
 606         res.Body.Close()
 607
 608         getReq, _ = http.NewRequest("GET", frontend.URL+"/noua", nil)
 609         getReq.Header.Set("User-Agent", "")
 610         getReq.Close = true
 611         res, err = frontendClient.Do(getReq)
 612         if err != nil {
 613                 t.Fatalf("Get: %v", err)
 614         }
 615         res.Body.Close()
 616 }
 617
 618 type bufferPool struct {
 619         get func() []byte
 620         put func([]byte)
 621 }
 622
 623 func (bp bufferPool) Get() []byte  { return bp.get() }
 624 func (bp bufferPool) Put(v []byte) { bp.put(v) }
 625
 626 func TestReverseProxyGetPutBuffer(t *testing.T) {
 627         const msg = "hi"
 628         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 629                 io.WriteString(w, msg)
 630         }))
 631         defer backend.Close()
 632
 633         backendURL, err := url.Parse(backend.URL)
 634         if err != nil {
 635                 t.Fatal(err)
 636         }
 637
 638         var (
 639                 mu  sync.Mutex
 640                 log []string
 641         )
 642         addLog := func(event string) {
 643                 mu.Lock()
 644                 defer mu.Unlock()
 645                 log = append(log, event)
 646         }
 647         rp := NewSingleHostReverseProxy(backendURL)
 648         const size = 1234
 649         rp.BufferPool = bufferPool{
 650                 get: func() []byte {
 651                         addLog("getBuf")
 652                         return make([]byte, size)
 653                 },
 654                 put: func(p []byte) {
 655                         addLog("putBuf-" + strconv.Itoa(len(p)))
 656                 },
 657         }
 658         frontend := httptest.NewServer(rp)
 659         defer frontend.Close()
 660
 661         req, _ := http.NewRequest("GET", frontend.URL, nil)
 662         req.Close = true
 663         res, err := frontend.Client().Do(req)
 664         if err != nil {
 665                 t.Fatalf("Get: %v", err)
 666         }
 667         slurp, err := io.ReadAll(res.Body)
 668         res.Body.Close()
 669         if err != nil {
 670                 t.Fatalf("reading body: %v", err)
 671         }
 672         if string(slurp) != msg {
 673                 t.Errorf("msg = %q; want %q", slurp, msg)
 674         }
 675         wantLog := []string{"getBuf", "putBuf-" + strconv.Itoa(size)}
 676         mu.Lock()
 677         defer mu.Unlock()
 678         if !reflect.DeepEqual(log, wantLog) {
 679                 t.Errorf("Log events = %q; want %q", log, wantLog)
 680         }
 681 }
 682
 683 func TestReverseProxy_Post(t *testing.T) {
 684         const backendResponse = "I am the backend"
 685         const backendStatus = 200
 686         var requestBody = bytes.Repeat([]byte("a"), 1<<20)
 687         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 688                 slurp, err := io.ReadAll(r.Body)
 689                 if err != nil {
 690                         t.Errorf("Backend body read = %v", err)
 691                 }
 692                 if len(slurp) != len(requestBody) {
 693                         t.Errorf("Backend read %d request body bytes; want %d", len(slurp), len(requestBody))
 694                 }
 695                 if !bytes.Equal(slurp, requestBody) {
 696                         t.Error("Backend read wrong request body.") // 1MB; omitting details
 697                 }
 698                 w.Write([]byte(backendResponse))
 699         }))
 700         defer backend.Close()
 701         backendURL, err := url.Parse(backend.URL)
 702         if err != nil {
 703                 t.Fatal(err)
 704         }
 705         proxyHandler := NewSingleHostReverseProxy(backendURL)
 706         frontend := httptest.NewServer(proxyHandler)
 707         defer frontend.Close()
 708
 709         postReq, _ := http.NewRequest("POST", frontend.URL, bytes.NewReader(requestBody))
 710         res, err := frontend.Client().Do(postReq)
 711         if err != nil {
 712                 t.Fatalf("Do: %v", err)
 713         }
 714         if g, e := res.StatusCode, backendStatus; g != e {
 715                 t.Errorf("got res.StatusCode %d; expected %d", g, e)
 716         }
 717         bodyBytes, _ := io.ReadAll(res.Body)
 718         if g, e := string(bodyBytes), backendResponse; g != e {
 719                 t.Errorf("got body %q; expected %q", g, e)
 720         }
 721 }
 722
 723 type RoundTripperFunc func(*http.Request) (*http.Response, error)
 724
 725 func (fn RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
 726         return fn(req)
 727 }
 728
 729 // Issue 16036: send a Request with a nil Body when possible
 730 func TestReverseProxy_NilBody(t *testing.T) {
 731         backendURL, _ := url.Parse("http://fake.tld/")
 732         proxyHandler := NewSingleHostReverseProxy(backendURL)
 733         proxyHandler.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
 734         proxyHandler.Transport = RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
 735                 if req.Body != nil {
 736                         t.Error("Body != nil; want a nil Body")
 737                 }
 738                 return nil, errors.New("done testing the interesting part; so force a 502 Gateway error")
 739         })
 740         frontend := httptest.NewServer(proxyHandler)
 741         defer frontend.Close()
 742
 743         res, err := frontend.Client().Get(frontend.URL)
 744         if err != nil {
 745                 t.Fatal(err)
 746         }
 747         defer res.Body.Close()
 748         if res.StatusCode != 502 {
 749                 t.Errorf("status code = %v; want 502 (Gateway Error)", res.Status)
 750         }
 751 }
 752
 753 // Issue 33142: always allocate the request headers
 754 func TestReverseProxy_AllocatedHeader(t *testing.T) {
 755         proxyHandler := new(ReverseProxy)
 756         proxyHandler.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
 757         proxyHandler.Director = func(*http.Request) {}     // noop
 758         proxyHandler.Transport = RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
 759                 if req.Header == nil {
 760                         t.Error("Header == nil; want a non-nil Header")
 761                 }
 762                 return nil, errors.New("done testing the interesting part; so force a 502 Gateway error")
 763         })
 764
 765         proxyHandler.ServeHTTP(httptest.NewRecorder(), &http.Request{
 766                 Method:     "GET",
 767                 URL:        &url.URL{Scheme: "http", Host: "fake.tld", Path: "/"},
 768                 Proto:      "HTTP/1.0",
 769                 ProtoMajor: 1,
 770         })
 771 }
 772
 773 // Issue 14237. Test ModifyResponse and that an error from it
 774 // causes the proxy to return StatusBadGateway, or StatusOK otherwise.
 775 func TestReverseProxyModifyResponse(t *testing.T) {
 776         backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 777                 w.Header().Add("X-Hit-Mod", fmt.Sprintf("%v", r.URL.Path == "/mod"))
 778         }))
 779         defer backendServer.Close()
 780
 781         rpURL, _ := url.Parse(backendServer.URL)
 782         rproxy := NewSingleHostReverseProxy(rpURL)
 783         rproxy.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
 784         rproxy.ModifyResponse = func(resp *http.Response) error {
 785                 if resp.Header.Get("X-Hit-Mod") != "true" {
 786                         return fmt.Errorf("tried to by-pass proxy")
 787                 }
 788                 return nil
 789         }
 790
 791         frontendProxy := httptest.NewServer(rproxy)
 792         defer frontendProxy.Close()
 793
 794         tests := []struct {
 795                 url      string
 796                 wantCode int
 797         }{
 798                 {frontendProxy.URL + "/mod", http.StatusOK},
 799                 {frontendProxy.URL + "/schedule", http.StatusBadGateway},
 800         }
 801
 802         for i, tt := range tests {
 803                 resp, err := http.Get(tt.url)
 804                 if err != nil {
 805                         t.Fatalf("failed to reach proxy: %v", err)
 806                 }
 807                 if g, e := resp.StatusCode, tt.wantCode; g != e {
 808                         t.Errorf("#%d: got res.StatusCode %d; expected %d", i, g, e)
 809                 }
 810                 resp.Body.Close()
 811         }
 812 }
 813
 814 type failingRoundTripper struct{}
 815
 816 func (failingRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
 817         return nil, errors.New("some error")
 818 }
 819
 820 type staticResponseRoundTripper struct{ res *http.Response }
 821
 822 func (rt staticResponseRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
 823         return rt.res, nil
 824 }
 825
 826 func TestReverseProxyErrorHandler(t *testing.T) {
 827         tests := []struct {
 828                 name           string
 829                 wantCode       int
 830                 errorHandler   func(http.ResponseWriter, *http.Request, error)
 831                 transport      http.RoundTripper // defaults to failingRoundTripper
 832                 modifyResponse func(*http.Response) error
 833         }{
 834                 {
 835                         name:     "default",
 836                         wantCode: http.StatusBadGateway,
 837                 },
 838                 {
 839                         name:         "errorhandler",
 840                         wantCode:     http.StatusTeapot,
 841                         errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
 842                 },
 843                 {
 844                         name: "modifyresponse_noerr",
 845                         transport: staticResponseRoundTripper{
 846                                 &http.Response{StatusCode: 345, Body: http.NoBody},
 847                         },
 848                         modifyResponse: func(res *http.Response) error {
 849                                 res.StatusCode++
 850                                 return nil
 851                         },
 852                         errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
 853                         wantCode:     346,
 854                 },
 855                 {
 856                         name: "modifyresponse_err",
 857                         transport: staticResponseRoundTripper{
 858                                 &http.Response{StatusCode: 345, Body: http.NoBody},
 859                         },
 860                         modifyResponse: func(res *http.Response) error {
 861                                 res.StatusCode++
 862                                 return errors.New("some error to trigger errorHandler")
 863                         },
 864                         errorHandler: func(rw http.ResponseWriter, req *http.Request, err error) { rw.WriteHeader(http.StatusTeapot) },
 865                         wantCode:     http.StatusTeapot,
 866                 },
 867         }
 868
 869         for _, tt := range tests {
 870                 t.Run(tt.name, func(t *testing.T) {
 871                         target := &url.URL{
 872                                 Scheme: "http",
 873                                 Host:   "dummy.tld",
 874                                 Path:   "/",
 875                         }
 876                         rproxy := NewSingleHostReverseProxy(target)
 877                         rproxy.Transport = tt.transport
 878                         rproxy.ModifyResponse = tt.modifyResponse
 879                         if rproxy.Transport == nil {
 880                                 rproxy.Transport = failingRoundTripper{}
 881                         }
 882                         rproxy.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
 883                         if tt.errorHandler != nil {
 884                                 rproxy.ErrorHandler = tt.errorHandler
 885                         }
 886                         frontendProxy := httptest.NewServer(rproxy)
 887                         defer frontendProxy.Close()
 888
 889                         resp, err := http.Get(frontendProxy.URL + "/test")
 890                         if err != nil {
 891                                 t.Fatalf("failed to reach proxy: %v", err)
 892                         }
 893                         if g, e := resp.StatusCode, tt.wantCode; g != e {
 894                                 t.Errorf("got res.StatusCode %d; expected %d", g, e)
 895                         }
 896                         resp.Body.Close()
 897                 })
 898         }
 899 }
 900
 901 // Issue 16659: log errors from short read
 902 func TestReverseProxy_CopyBuffer(t *testing.T) {
 903         backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 904                 out := "this call was relayed by the reverse proxy"
 905                 // Coerce a wrong content length to induce io.UnexpectedEOF
 906                 w.Header().Set("Content-Length", fmt.Sprintf("%d", len(out)*2))
 907                 fmt.Fprintln(w, out)
 908         }))
 909         defer backendServer.Close()
 910
 911         rpURL, err := url.Parse(backendServer.URL)
 912         if err != nil {
 913                 t.Fatal(err)
 914         }
 915
 916         var proxyLog bytes.Buffer
 917         rproxy := NewSingleHostReverseProxy(rpURL)
 918         rproxy.ErrorLog = log.New(&proxyLog, "", log.Lshortfile)
 919         donec := make(chan bool, 1)
 920         frontendProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 921                 defer func() { donec <- true }()
 922                 rproxy.ServeHTTP(w, r)
 923         }))
 924         defer frontendProxy.Close()
 925
 926         if _, err = frontendProxy.Client().Get(frontendProxy.URL); err == nil {
 927                 t.Fatalf("want non-nil error")
 928         }
 929         // The race detector complains about the proxyLog usage in logf in copyBuffer
 930         // and our usage below with proxyLog.Bytes() so we're explicitly using a
 931         // channel to ensure that the ReverseProxy's ServeHTTP is done before we
 932         // continue after Get.
 933         <-donec
 934
 935         expected := []string{
 936                 "EOF",
 937                 "read",
 938         }
 939         for _, phrase := range expected {
 940                 if !bytes.Contains(proxyLog.Bytes(), []byte(phrase)) {
 941                         t.Errorf("expected log to contain phrase %q", phrase)
 942                 }
 943         }
 944 }
 945
 946 type staticTransport struct {
 947         res *http.Response
 948 }
 949
 950 func (t *staticTransport) RoundTrip(r *http.Request) (*http.Response, error) {
 951         return t.res, nil
 952 }
 953
 954 func BenchmarkServeHTTP(b *testing.B) {
 955         res := &http.Response{
 956                 StatusCode: 200,
 957                 Body:       io.NopCloser(strings.NewReader("")),
 958         }
 959         proxy := &ReverseProxy{
 960                 Director:  func(*http.Request) {},
 961                 Transport: &staticTransport{res},
 962         }
 963
 964         w := httptest.NewRecorder()
 965         r := httptest.NewRequest("GET", "/", nil)
 966
 967         b.ReportAllocs()
 968         for i := 0; i < b.N; i++ {
 969                 proxy.ServeHTTP(w, r)
 970         }
 971 }
 972
 973 func TestServeHTTPDeepCopy(t *testing.T) {
 974         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 975                 w.Write([]byte("Hello Gopher!"))
 976         }))
 977         defer backend.Close()
 978         backendURL, err := url.Parse(backend.URL)
 979         if err != nil {
 980                 t.Fatal(err)
 981         }
 982
 983         type result struct {
 984                 before, after string
 985         }
 986
 987         resultChan := make(chan result, 1)
 988         proxyHandler := NewSingleHostReverseProxy(backendURL)
 989         frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 990                 before := r.URL.String()
 991                 proxyHandler.ServeHTTP(w, r)
 992                 after := r.URL.String()
 993                 resultChan <- result{before: before, after: after}
 994         }))
 995         defer frontend.Close()
 996
 997         want := result{before: "/", after: "/"}
 998
 999         res, err := frontend.Client().Get(frontend.URL)
1000         if err != nil {
1001                 t.Fatalf("Do: %v", err)
1002         }
1003         res.Body.Close()
1004
1005         got := <-resultChan
1006         if got != want {
1007                 t.Errorf("got = %+v; want = %+v", got, want)
1008         }
1009 }
1010
1011 // Issue 18327: verify we always do a deep copy of the Request.Header map
1012 // before any mutations.
1013 func TestClonesRequestHeaders(t *testing.T) {
1014         log.SetOutput(io.Discard)
1015         defer log.SetOutput(os.Stderr)
1016         req, _ := http.NewRequest("GET", "http://foo.tld/", nil)
1017         req.RemoteAddr = "1.2.3.4:56789"
1018         rp := &ReverseProxy{
1019                 Director: func(req *http.Request) {
1020                         req.Header.Set("From-Director", "1")
1021                 },
1022                 Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
1023                         if v := req.Header.Get("From-Director"); v != "1" {
1024                                 t.Errorf("From-Directory value = %q; want 1", v)
1025                         }
1026                         return nil, io.EOF
1027                 }),
1028         }
1029         rp.ServeHTTP(httptest.NewRecorder(), req)
1030
1031         if req.Header.Get("From-Director") == "1" {
1032                 t.Error("Director header mutation modified caller's request")
1033         }
1034         if req.Header.Get("X-Forwarded-For") != "" {
1035                 t.Error("X-Forward-For header mutation modified caller's request")
1036         }
1037
1038 }
1039
1040 type roundTripperFunc func(req *http.Request) (*http.Response, error)
1041
1042 func (fn roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
1043         return fn(req)
1044 }
1045
1046 func TestModifyResponseClosesBody(t *testing.T) {
1047         req, _ := http.NewRequest("GET", "http://foo.tld/", nil)
1048         req.RemoteAddr = "1.2.3.4:56789"
1049         closeCheck := new(checkCloser)
1050         logBuf := new(bytes.Buffer)
1051         outErr := errors.New("ModifyResponse error")
1052         rp := &ReverseProxy{
1053                 Director: func(req *http.Request) {},
1054                 Transport: &staticTransport{&http.Response{
1055                         StatusCode: 200,
1056                         Body:       closeCheck,
1057                 }},
1058                 ErrorLog: log.New(logBuf, "", 0),
1059                 ModifyResponse: func(*http.Response) error {
1060                         return outErr
1061                 },
1062         }
1063         rec := httptest.NewRecorder()
1064         rp.ServeHTTP(rec, req)
1065         res := rec.Result()
1066         if g, e := res.StatusCode, http.StatusBadGateway; g != e {
1067                 t.Errorf("got res.StatusCode %d; expected %d", g, e)
1068         }
1069         if !closeCheck.closed {
1070                 t.Errorf("body should have been closed")
1071         }
1072         if g, e := logBuf.String(), outErr.Error(); !strings.Contains(g, e) {
1073                 t.Errorf("ErrorLog %q does not contain %q", g, e)
1074         }
1075 }
1076
1077 type checkCloser struct {
1078         closed bool
1079 }
1080
1081 func (cc *checkCloser) Close() error {
1082         cc.closed = true
1083         return nil
1084 }
1085
1086 func (cc *checkCloser) Read(b []byte) (int, error) {
1087         return len(b), nil
1088 }
1089
1090 // Issue 23643: panic on body copy error
1091 func TestReverseProxy_PanicBodyError(t *testing.T) {
1092         log.SetOutput(io.Discard)
1093         defer log.SetOutput(os.Stderr)
1094         backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1095                 out := "this call was relayed by the reverse proxy"
1096                 // Coerce a wrong content length to induce io.ErrUnexpectedEOF
1097                 w.Header().Set("Content-Length", fmt.Sprintf("%d", len(out)*2))
1098                 fmt.Fprintln(w, out)
1099         }))
1100         defer backendServer.Close()
1101
1102         rpURL, err := url.Parse(backendServer.URL)
1103         if err != nil {
1104                 t.Fatal(err)
1105         }
1106
1107         rproxy := NewSingleHostReverseProxy(rpURL)
1108
1109         // Ensure that the handler panics when the body read encounters an
1110         // io.ErrUnexpectedEOF
1111         defer func() {
1112                 err := recover()
1113                 if err == nil {
1114                         t.Fatal("handler should have panicked")
1115                 }
1116                 if err != http.ErrAbortHandler {
1117                         t.Fatal("expected ErrAbortHandler, got", err)
1118                 }
1119         }()
1120         req, _ := http.NewRequest("GET", "http://foo.tld/", nil)
1121         rproxy.ServeHTTP(httptest.NewRecorder(), req)
1122 }
1123
1124 func TestSelectFlushInterval(t *testing.T) {
1125         tests := []struct {
1126                 name string
1127                 p    *ReverseProxy
1128                 res  *http.Response
1129                 want time.Duration
1130         }{
1131                 {
1132                         name: "default",
1133                         res:  &http.Response{},
1134                         p:    &ReverseProxy{FlushInterval: 123},
1135                         want: 123,
1136                 },
1137                 {
1138                         name: "server-sent events overrides non-zero",
1139                         res: &http.Response{
1140                                 Header: http.Header{
1141                                         "Content-Type": {"text/event-stream"},
1142                                 },
1143                         },
1144                         p:    &ReverseProxy{FlushInterval: 123},
1145                         want: -1,
1146                 },
1147                 {
1148                         name: "server-sent events overrides zero",
1149                         res: &http.Response{
1150                                 Header: http.Header{
1151                                         "Content-Type": {"text/event-stream"},
1152                                 },
1153                         },
1154                         p:    &ReverseProxy{FlushInterval: 0},
1155                         want: -1,
1156                 },
1157                 {
1158                         name: "Content-Length: -1, overrides non-zero",
1159                         res: &http.Response{
1160                                 ContentLength: -1,
1161                         },
1162                         p:    &ReverseProxy{FlushInterval: 123},
1163                         want: -1,
1164                 },
1165                 {
1166                         name: "Content-Length: -1, overrides zero",
1167                         res: &http.Response{
1168                                 ContentLength: -1,
1169                         },
1170                         p:    &ReverseProxy{FlushInterval: 0},
1171                         want: -1,
1172                 },
1173         }
1174         for _, tt := range tests {
1175                 t.Run(tt.name, func(t *testing.T) {
1176                         got := tt.p.flushInterval(tt.res)
1177                         if got != tt.want {
1178                                 t.Errorf("flushLatency = %v; want %v", got, tt.want)
1179                         }
1180                 })
1181         }
1182 }
1183
1184 func TestReverseProxyWebSocket(t *testing.T) {
1185         backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1186                 if upgradeType(r.Header) != "websocket" {
1187                         t.Error("unexpected backend request")
1188                         http.Error(w, "unexpected request", 400)
1189                         return
1190                 }
1191                 c, _, err := w.(http.Hijacker).Hijack()
1192                 if err != nil {
1193                         t.Error(err)
1194                         return
1195                 }
1196                 defer c.Close()
1197                 io.WriteString(c, "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\n\r\n")
1198                 bs := bufio.NewScanner(c)
1199                 if !bs.Scan() {
1200                         t.Errorf("backend failed to read line from client: %v", bs.Err())
1201                         return
1202                 }
1203                 fmt.Fprintf(c, "backend got %q\n", bs.Text())
1204         }))
1205         defer backendServer.Close()
1206
1207         backURL, _ := url.Parse(backendServer.URL)
1208         rproxy := NewSingleHostReverseProxy(backURL)
1209         rproxy.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
1210         rproxy.ModifyResponse = func(res *http.Response) error {
1211                 res.Header.Add("X-Modified", "true")
1212                 return nil
1213         }
1214
1215         handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
1216                 rw.Header().Set("X-Header", "X-Value")
1217                 rproxy.ServeHTTP(rw, req)
1218                 if got, want := rw.Header().Get("X-Modified"), "true"; got != want {
1219                         t.Errorf("response writer X-Modified header = %q; want %q", got, want)
1220                 }
1221         })
1222
1223         frontendProxy := httptest.NewServer(handler)
1224         defer frontendProxy.Close()
1225
1226         req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
1227         req.Header.Set("Connection", "Upgrade")
1228         req.Header.Set("Upgrade", "websocket")
1229
1230         c := frontendProxy.Client()
1231         res, err := c.Do(req)
1232         if err != nil {
1233                 t.Fatal(err)
1234         }
1235         if res.StatusCode != 101 {
1236                 t.Fatalf("status = %v; want 101", res.Status)
1237         }
1238
1239         got := res.Header.Get("X-Header")
1240         want := "X-Value"
1241         if got != want {
1242                 t.Errorf("Header(XHeader) = %q; want %q", got, want)
1243         }
1244
1245         if upgradeType(res.Header) != "websocket" {
1246                 t.Fatalf("not websocket upgrade; got %#v", res.Header)
1247         }
1248         rwc, ok := res.Body.(io.ReadWriteCloser)
1249         if !ok {
1250                 t.Fatalf("response body is of type %T; does not implement ReadWriteCloser", res.Body)
1251         }
1252         defer rwc.Close()
1253
1254         if got, want := res.Header.Get("X-Modified"), "true"; got != want {
1255                 t.Errorf("response X-Modified header = %q; want %q", got, want)
1256         }
1257
1258         io.WriteString(rwc, "Hello\n")
1259         bs := bufio.NewScanner(rwc)
1260         if !bs.Scan() {
1261                 t.Fatalf("Scan: %v", bs.Err())
1262         }
1263         got = bs.Text()
1264         want = `backend got "Hello"`
1265         if got != want {
1266                 t.Errorf("got %#q, want %#q", got, want)
1267         }
1268 }
1269
1270 func TestReverseProxyWebSocketCancelation(t *testing.T) {
1271         n := 5
1272         triggerCancelCh := make(chan bool, n)
1273         nthResponse := func(i int) string {
1274                 return fmt.Sprintf("backend response #%d\n", i)
1275         }
1276         terminalMsg := "final message"
1277
1278         cst := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1279                 if g, ws := upgradeType(r.Header), "websocket"; g != ws {
1280                         t.Errorf("Unexpected upgrade type %q, want %q", g, ws)
1281                         http.Error(w, "Unexpected request", 400)
1282                         return
1283                 }
1284                 conn, bufrw, err := w.(http.Hijacker).Hijack()
1285                 if err != nil {
1286                         t.Error(err)
1287                         return
1288                 }
1289                 defer conn.Close()
1290
1291                 upgradeMsg := "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\n\r\n"
1292                 if _, err := io.WriteString(conn, upgradeMsg); err != nil {
1293                         t.Error(err)
1294                         return
1295                 }
1296                 if _, _, err := bufrw.ReadLine(); err != nil {
1297                         t.Errorf("Failed to read line from client: %v", err)
1298                         return
1299                 }
1300
1301                 for i := 0; i < n; i++ {
1302                         if _, err := bufrw.WriteString(nthResponse(i)); err != nil {
1303                                 select {
1304                                 case <-triggerCancelCh:
1305                                 default:
1306                                         t.Errorf("Writing response #%d failed: %v", i, err)
1307                                 }
1308                                 return
1309                         }
1310                         bufrw.Flush()
1311                         time.Sleep(time.Second)
1312                 }
1313                 if _, err := bufrw.WriteString(terminalMsg); err != nil {
1314                         select {
1315                         case <-triggerCancelCh:
1316                         default:
1317                                 t.Errorf("Failed to write terminal message: %v", err)
1318                         }
1319                 }
1320                 bufrw.Flush()
1321         }))
1322         defer cst.Close()
1323
1324         backendURL, _ := url.Parse(cst.URL)
1325         rproxy := NewSingleHostReverseProxy(backendURL)
1326         rproxy.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
1327         rproxy.ModifyResponse = func(res *http.Response) error {
1328                 res.Header.Add("X-Modified", "true")
1329                 return nil
1330         }
1331
1332         handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
1333                 rw.Header().Set("X-Header", "X-Value")
1334                 ctx, cancel := context.WithCancel(req.Context())
1335                 go func() {
1336                         <-triggerCancelCh
1337                         cancel()
1338                 }()
1339                 rproxy.ServeHTTP(rw, req.WithContext(ctx))
1340         })
1341
1342         frontendProxy := httptest.NewServer(handler)
1343         defer frontendProxy.Close()
1344
1345         req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
1346         req.Header.Set("Connection", "Upgrade")
1347         req.Header.Set("Upgrade", "websocket")
1348
1349         res, err := frontendProxy.Client().Do(req)
1350         if err != nil {
1351                 t.Fatalf("Dialing to frontend proxy: %v", err)
1352         }
1353         defer res.Body.Close()
1354         if g, w := res.StatusCode, 101; g != w {
1355                 t.Fatalf("Switching protocols failed, got: %d, want: %d", g, w)
1356         }
1357
1358         if g, w := res.Header.Get("X-Header"), "X-Value"; g != w {
1359                 t.Errorf("X-Header mismatch\n\tgot:  %q\n\twant: %q", g, w)
1360         }
1361
1362         if g, w := upgradeType(res.Header), "websocket"; g != w {
1363                 t.Fatalf("Upgrade header mismatch\n\tgot:  %q\n\twant: %q", g, w)
1364         }
1365
1366         rwc, ok := res.Body.(io.ReadWriteCloser)
1367         if !ok {
1368                 t.Fatalf("Response body type mismatch, got %T, want io.ReadWriteCloser", res.Body)
1369         }
1370
1371         if got, want := res.Header.Get("X-Modified"), "true"; got != want {
1372                 t.Errorf("response X-Modified header = %q; want %q", got, want)
1373         }
1374
1375         if _, err := io.WriteString(rwc, "Hello\n"); err != nil {
1376                 t.Fatalf("Failed to write first message: %v", err)
1377         }
1378
1379         // Read loop.
1380
1381         br := bufio.NewReader(rwc)
1382         for {
1383                 line, err := br.ReadString('\n')
1384                 switch {
1385                 case line == terminalMsg: // this case before "err == io.EOF"
1386                         t.Fatalf("The websocket request was not canceled, unfortunately!")
1387
1388                 case err == io.EOF:
1389                         return
1390
1391                 case err != nil:
1392                         t.Fatalf("Unexpected error: %v", err)
1393
1394                 case line == nthResponse(0): // We've gotten the first response back
1395                         // Let's trigger a cancel.
1396                         close(triggerCancelCh)
1397                 }
1398         }
1399 }
1400
1401 func TestUnannouncedTrailer(t *testing.T) {
1402         backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1403                 w.WriteHeader(http.StatusOK)
1404                 w.(http.Flusher).Flush()
1405                 w.Header().Set(http.TrailerPrefix+"X-Unannounced-Trailer", "unannounced_trailer_value")
1406         }))
1407         defer backend.Close()
1408         backendURL, err := url.Parse(backend.URL)
1409         if err != nil {
1410                 t.Fatal(err)
1411         }
1412         proxyHandler := NewSingleHostReverseProxy(backendURL)
1413         proxyHandler.ErrorLog = log.New(io.Discard, "", 0) // quiet for tests
1414         frontend := httptest.NewServer(proxyHandler)
1415         defer frontend.Close()
1416         frontendClient := frontend.Client()
1417
1418         res, err := frontendClient.Get(frontend.URL)
1419         if err != nil {
1420                 t.Fatalf("Get: %v", err)
1421         }
1422
1423         io.ReadAll(res.Body)
1424
1425         if g, w := res.Trailer.Get("X-Unannounced-Trailer"), "unannounced_trailer_value"; g != w {
1426                 t.Errorf("Trailer(X-Unannounced-Trailer) = %q; want %q", g, w)
1427         }
1428
1429 }
1430
1431 func TestSingleJoinSlash(t *testing.T) {
1432         tests := []struct {
1433                 slasha   string
1434                 slashb   string
1435                 expected string
1436         }{
1437                 {"https://www.google.com/", "/favicon.ico", "https://www.google.com/favicon.ico"},
1438                 {"https://www.google.com", "/favicon.ico", "https://www.google.com/favicon.ico"},
1439                 {"https://www.google.com", "favicon.ico", "https://www.google.com/favicon.ico"},
1440                 {"https://www.google.com", "", "https://www.google.com/"},
1441                 {"", "favicon.ico", "/favicon.ico"},
1442         }
1443         for _, tt := range tests {
1444                 if got := singleJoiningSlash(tt.slasha, tt.slashb); got != tt.expected {
1445                         t.Errorf("singleJoiningSlash(%q,%q) want %q got %q",
1446                                 tt.slasha,
1447                                 tt.slashb,
1448                                 tt.expected,
1449                                 got)
1450                 }
1451         }
1452 }
1453
1454 func TestJoinURLPath(t *testing.T) {
1455         tests := []struct {
1456                 a        *url.URL
1457                 b        *url.URL
1458                 wantPath string
1459                 wantRaw  string
1460         }{
1461                 {&url.URL{Path: "/a/b"}, &url.URL{Path: "/c"}, "/a/b/c", ""},
1462                 {&url.URL{Path: "/a/b", RawPath: "badpath"}, &url.URL{Path: "c"}, "/a/b/c", "/a/b/c"},
1463                 {&url.URL{Path: "/a/b", RawPath: "/a%2Fb"}, &url.URL{Path: "/c"}, "/a/b/c", "/a%2Fb/c"},
1464                 {&url.URL{Path: "/a/b", RawPath: "/a%2Fb"}, &url.URL{Path: "/c"}, "/a/b/c", "/a%2Fb/c"},
1465                 {&url.URL{Path: "/a/b/", RawPath: "/a%2Fb%2F"}, &url.URL{Path: "c"}, "/a/b//c", "/a%2Fb%2F/c"},
1466                 {&url.URL{Path: "/a/b/", RawPath: "/a%2Fb/"}, &url.URL{Path: "/c/d", RawPath: "/c%2Fd"}, "/a/b/c/d", "/a%2Fb/c%2Fd"},
1467         }
1468
1469         for _, tt := range tests {
1470                 p, rp := joinURLPath(tt.a, tt.b)
1471                 if p != tt.wantPath || rp != tt.wantRaw {
1472                         t.Errorf("joinURLPath(URL(%q,%q),URL(%q,%q)) want (%q,%q) got (%q,%q)",
1473                                 tt.a.Path, tt.a.RawPath,
1474                                 tt.b.Path, tt.b.RawPath,
1475                                 tt.wantPath, tt.wantRaw,
1476                                 p, rp)
1477                 }
1478         }
1479 }