github.com/apex/up@v1.7.1/internal/proxy/response.go (about)

     1  package proxy
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/base64"
     6  	"mime"
     7  	"net/http"
     8  	"strings"
     9  
    10  	"github.com/apex/up/internal/util"
    11  )
    12  
    13  // ResponseWriter implements the http.ResponseWriter interface
    14  // in order to support the API Gateway Lambda HTTP "protocol".
    15  type ResponseWriter struct {
    16  	out         Output
    17  	buf         bytes.Buffer
    18  	header      http.Header
    19  	wroteHeader bool
    20  }
    21  
    22  // NewResponse returns a new response writer to capture http output.
    23  func NewResponse() *ResponseWriter {
    24  	return &ResponseWriter{}
    25  }
    26  
    27  // Header implementation.
    28  func (w *ResponseWriter) Header() http.Header {
    29  	if w.header == nil {
    30  		w.header = make(http.Header)
    31  	}
    32  
    33  	return w.header
    34  }
    35  
    36  // Write implementation.
    37  func (w *ResponseWriter) Write(b []byte) (int, error) {
    38  	if !w.wroteHeader {
    39  		w.WriteHeader(http.StatusOK)
    40  	}
    41  
    42  	// TODO: HEAD? ignore
    43  
    44  	return w.buf.Write(b)
    45  }
    46  
    47  // WriteHeader implementation.
    48  func (w *ResponseWriter) WriteHeader(status int) {
    49  	if w.wroteHeader {
    50  		return
    51  	}
    52  
    53  	if w.Header().Get("Content-Type") == "" {
    54  		w.Header().Set("Content-Type", "text/plain; charset=utf8")
    55  	}
    56  
    57  	w.out.StatusCode = status
    58  
    59  	h := make(map[string]string)
    60  
    61  	// API Gateway does not support multiple set-cookie fields
    62  	// so we have to stagger the casing in order to support this.
    63  	util.FixMultipleSetCookie(w.Header())
    64  
    65  	for k, v := range w.Header() {
    66  		if len(v) > 0 {
    67  			h[k] = v[len(v)-1]
    68  		}
    69  	}
    70  
    71  	w.out.Headers = h
    72  	w.wroteHeader = true
    73  }
    74  
    75  // End the request.
    76  func (w *ResponseWriter) End() Output {
    77  	w.out.IsBase64Encoded = isBinary(w.header)
    78  
    79  	if w.out.IsBase64Encoded {
    80  		w.out.Body = base64.StdEncoding.EncodeToString(w.buf.Bytes())
    81  	} else {
    82  		w.out.Body = w.buf.String()
    83  	}
    84  
    85  	return w.out
    86  }
    87  
    88  // isBinary returns true if the response reprensents binary.
    89  func isBinary(h http.Header) bool {
    90  	if !isTextMime(h.Get("Content-Type")) {
    91  		return true
    92  	}
    93  
    94  	if h.Get("Content-Encoding") == "gzip" {
    95  		return true
    96  	}
    97  
    98  	return false
    99  }
   100  
   101  // isTextMime returns true if the content type represents textual data.
   102  func isTextMime(kind string) bool {
   103  	mt, _, err := mime.ParseMediaType(kind)
   104  	if err != nil {
   105  		return false
   106  	}
   107  
   108  	if strings.HasPrefix(mt, "text/") {
   109  		return true
   110  	}
   111  
   112  	switch mt {
   113  	case "image/svg+xml":
   114  		return true
   115  	case "application/json":
   116  		return true
   117  	case "application/xml":
   118  		return true
   119  	default:
   120  		return false
   121  	}
   122  }