github.com/blend/go-sdk@v1.20220411.3/reverseproxy/upstream.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package reverseproxy
     9  
    10  import (
    11  	"net/http"
    12  	"net/http/httputil"
    13  	"net/url"
    14  	"time"
    15  
    16  	"golang.org/x/net/http2"
    17  
    18  	"github.com/blend/go-sdk/ex"
    19  	"github.com/blend/go-sdk/logger"
    20  	"github.com/blend/go-sdk/webutil"
    21  )
    22  
    23  // NewUpstream returns a new upstram.
    24  func NewUpstream(target *url.URL, opts ...UpstreamOption) *Upstream {
    25  	rp := httputil.NewSingleHostReverseProxy(target)
    26  	u := &Upstream{
    27  		URL:          target,
    28  		ReverseProxy: rp,
    29  	}
    30  	// NOTE: This creates a reference cycle `u -> rp -> u`.
    31  	rp.ErrorHandler = u.errorHandler
    32  	return u
    33  }
    34  
    35  // Upstream represents a proxyable server.
    36  type Upstream struct {
    37  	// Name is the name of the upstream.
    38  	Name string
    39  	// Log is a logger agent.
    40  	Log logger.Log
    41  	// URL represents the target of the upstream.
    42  	URL *url.URL
    43  	// ReverseProxy is what actually forwards requests.
    44  	ReverseProxy *httputil.ReverseProxy
    45  }
    46  
    47  // UseHTTP2 sets the upstream to use http2.
    48  func (u *Upstream) UseHTTP2() error {
    49  	if u.ReverseProxy.Transport == nil {
    50  		u.ReverseProxy.Transport = &http.Transport{}
    51  	}
    52  	if typed, ok := u.ReverseProxy.Transport.(*http.Transport); ok {
    53  		if err := http2.ConfigureTransport(typed); err != nil {
    54  			return ex.New(err)
    55  		}
    56  	}
    57  	return nil
    58  }
    59  
    60  // ServeHTTP
    61  func (u *Upstream) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
    62  	w := webutil.NewStatusResponseWriter(rw)
    63  
    64  	if u.Log != nil {
    65  		start := time.Now()
    66  		defer func() {
    67  			wre := webutil.NewHTTPRequestEvent(req,
    68  				webutil.OptHTTPRequestStatusCode(w.StatusCode()),
    69  				webutil.OptHTTPRequestContentLength(w.ContentLength()),
    70  				webutil.OptHTTPRequestElapsed(time.Since(start)),
    71  			)
    72  			if value := w.Header().Get("Content-Type"); len(value) > 0 {
    73  				wre.ContentType = value
    74  			}
    75  			if value := w.Header().Get("Content-Encoding"); len(value) > 0 {
    76  				wre.ContentEncoding = value
    77  			}
    78  
    79  			u.Log.TriggerContext(req.Context(), wre)
    80  		}()
    81  	}
    82  	u.ReverseProxy.ServeHTTP(w, req)
    83  }
    84  
    85  // errorHandler is intended to be used as an `(net/http/httputil).ReverseProxy.ErrorHandler`
    86  // This implementation is based on:
    87  // https://github.com/golang/go/blob/go1.13.6/src/net/http/httputil/reverseproxy.go#L151-L154
    88  func (u *Upstream) errorHandler(rw http.ResponseWriter, req *http.Request, err error) {
    89  	logger.MaybeErrorfContext(req.Context(), u.Log, "http: proxy error: %v", err)
    90  	rw.WriteHeader(http.StatusBadGateway)
    91  }