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 }