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