github.com/xmidt-org/webpa-common@v1.11.9/bookkeeping/bookkeeper.go (about)

     1  package bookkeeping
     2  
     3  import (
     4  	"bytes"
     5  	"github.com/xmidt-org/webpa-common/logging"
     6  	"github.com/xmidt-org/webpa-common/xhttp/xcontext"
     7  	"net/http"
     8  	"time"
     9  )
    10  
    11  type CapturedResponse struct {
    12  	Code    int
    13  	Payload []byte
    14  	Header  http.Header
    15  }
    16  
    17  // RequestFunc takes the request and returns key value pairs from the request
    18  type RequestFunc func(request *http.Request) []interface{}
    19  
    20  // ResponseFunc takes the ResponseWriter and returns key value pairs from the request
    21  type ResponseFunc func(response CapturedResponse) []interface{}
    22  
    23  // handler is for logging of the request and response
    24  type handler struct {
    25  	next   http.Handler
    26  	before []RequestFunc
    27  	after  []ResponseFunc
    28  }
    29  
    30  // Option provides a single configuration option for a bookkeeping handler
    31  type Option func(h *handler)
    32  
    33  func New(options ...Option) func(next http.Handler) http.Handler {
    34  	return func(next http.Handler) http.Handler {
    35  		if next == nil {
    36  			panic("next can't be nil")
    37  		}
    38  		h := &handler{
    39  			next: next,
    40  		}
    41  		for _, o := range options {
    42  			o(h)
    43  		}
    44  		return h
    45  	}
    46  }
    47  
    48  // WithRequests take a one or many RequestFuncs to build an Option for logging key value pairs from the
    49  // RequestFun
    50  func WithRequests(requestFuncs ...RequestFunc) Option {
    51  	return func(h *handler) {
    52  		h.before = append(h.before, requestFuncs...)
    53  	}
    54  }
    55  
    56  // WithResponses take a one or many ResponseFunc to build an Option for logging key value pairs from the
    57  // ResponseFunc
    58  func WithResponses(responseFuncs ...ResponseFunc) Option {
    59  	return func(h *handler) {
    60  		h.after = append(h.after, responseFuncs...)
    61  	}
    62  }
    63  
    64  // ServeHTTP handles the bookkeeping given
    65  func (h *handler) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
    66  	kv := []interface{}{logging.MessageKey(), "Bookkeeping transactor"}
    67  	for _, before := range h.before {
    68  		kv = append(kv, before(request)...)
    69  	}
    70  	responseWriter, request = xcontext.WithContext(responseWriter, request, request.Context())
    71  	w := &writerInterceptor{ResponseWriter: responseWriter, ContextAware: responseWriter.(xcontext.ContextAware)}
    72  
    73  	start := time.Now()
    74  	defer func() {
    75  		ctx := xcontext.Context(responseWriter, request)
    76  		duration := time.Since(start)
    77  
    78  		kv = append(kv, "duration", duration)
    79  
    80  		response := CapturedResponse{
    81  			Code:    w.code,
    82  			Payload: w.buffer.Bytes(),
    83  			Header:  w.Header(),
    84  		}
    85  
    86  		for _, after := range h.after {
    87  			kv = append(kv, after(response)...)
    88  		}
    89  
    90  		logging.Info(logging.GetLogger(ctx)).Log(kv...)
    91  	}()
    92  
    93  	h.next.ServeHTTP(w, request)
    94  }
    95  
    96  type writerInterceptor struct {
    97  	xcontext.ContextAware
    98  	http.ResponseWriter
    99  	code   int
   100  	buffer bytes.Buffer
   101  }
   102  
   103  func (w *writerInterceptor) Write(data []byte) (int, error) {
   104  	w.buffer.Write(data)
   105  	return w.ResponseWriter.Write(data)
   106  
   107  }
   108  
   109  func (w *writerInterceptor) WriteHeader(code int) {
   110  	w.code = code
   111  	w.ResponseWriter.WriteHeader(code)
   112  }