github.com/orteth01/up@v0.2.0/internal/proxy/response.go (about)

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