github.com/simpleiot/simpleiot@v0.18.3/api/http-logger.go (about)

     1  package api
     2  
     3  // this code based on https://github.com/unrolled/logger, but expanded
     4  // to optionally dump the req/resp body
     5  
     6  import (
     7  	"bufio"
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"log"
    12  	"net"
    13  	"net/http"
    14  	"os"
    15  )
    16  
    17  // HTTPLogger can be used to log http requests
    18  type HTTPLogger struct {
    19  	prefix string
    20  	*log.Logger
    21  }
    22  
    23  // NewHTTPLogger returns a http logger
    24  func NewHTTPLogger(prefix string) *HTTPLogger {
    25  	return &HTTPLogger{
    26  		prefix: prefix,
    27  		Logger: log.New(os.Stdout, prefix, 0),
    28  	}
    29  }
    30  
    31  // Handler wraps an HTTP handler and logs the request as necessary.
    32  func (l *HTTPLogger) Handler(next http.Handler) http.Handler {
    33  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    34  
    35  		var rdr io.ReadCloser
    36  		buf, err := io.ReadAll(r.Body)
    37  		if err == nil {
    38  			rdr = io.NopCloser(bytes.NewBuffer(buf))
    39  			rdr2 := io.NopCloser(bytes.NewBuffer(buf))
    40  			r.Body = rdr2
    41  		}
    42  
    43  		crw := newCustomResponseWriter(w)
    44  		next.ServeHTTP(crw, r)
    45  
    46  		addr := r.RemoteAddr
    47  		if err == nil {
    48  			rBuf := bytes.Buffer{}
    49  			_, _ = rBuf.ReadFrom(rdr)
    50  			l.Printf("(%s) \"%s %s\" %d -> %v -> %v", addr, r.Method, r.RequestURI,
    51  				crw.status, rBuf.String(), crw.buf.String())
    52  		} else {
    53  			l.Printf("(%s) \"%s %s\" %d", addr, r.Method, r.RequestURI, crw.status)
    54  		}
    55  
    56  	})
    57  }
    58  
    59  type customResponseWriter struct {
    60  	http.ResponseWriter
    61  	status int
    62  	size   int
    63  	buf    bytes.Buffer
    64  }
    65  
    66  func (c *customResponseWriter) WriteHeader(status int) {
    67  	c.status = status
    68  	c.ResponseWriter.WriteHeader(status)
    69  }
    70  
    71  func (c *customResponseWriter) Write(b []byte) (int, error) {
    72  	size, err := c.ResponseWriter.Write(b)
    73  	c.buf.Write(b)
    74  	c.size += size
    75  	return size, err
    76  }
    77  
    78  func (c *customResponseWriter) Flush() {
    79  	if f, ok := c.ResponseWriter.(http.Flusher); ok {
    80  		f.Flush()
    81  	}
    82  }
    83  
    84  func (c *customResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    85  	if hj, ok := c.ResponseWriter.(http.Hijacker); ok {
    86  		return hj.Hijack()
    87  	}
    88  	return nil, nil, fmt.Errorf("ResponseWriter does not implement the Hijacker interface")
    89  }
    90  
    91  func newCustomResponseWriter(w http.ResponseWriter) *customResponseWriter {
    92  	// When WriteHeader is not called, it's safe to assume the status will be 200.
    93  	return &customResponseWriter{
    94  		ResponseWriter: w,
    95  		status:         200,
    96  	}
    97  }