github.com/hhsnopek/up@v0.1.1/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 = !isText(w.header.Get("Content-Type"))
    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  // isText returns true if the content type represents textual data.
    82  func isText(kind string) bool {
    83  	// TODO: refactor textual mime type stuff
    84  	switch {
    85  	case strings.HasSuffix(kind, "svg+xml"):
    86  		return true
    87  	case strings.HasPrefix(kind, "text/"):
    88  		return true
    89  	case strings.HasPrefix(kind, "application/") && strings.HasSuffix(kind, "json"):
    90  		return true
    91  	case strings.HasPrefix(kind, "application/") && strings.HasSuffix(kind, "xml"):
    92  		return true
    93  	default:
    94  		return false
    95  	}
    96  }