github.com/wader/devd@v0.0.0-20221031103345-441c7e455249/responselogger.go (about)

     1  package devd
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"strconv"
     7  
     8  	"github.com/cortesi/termlog"
     9  	"github.com/dustin/go-humanize"
    10  	"github.com/fatih/color"
    11  	"github.com/wader/devd/timer"
    12  )
    13  
    14  // ResponseLogWriter is a ResponseWriter that logs
    15  type ResponseLogWriter struct {
    16  	Log         termlog.Logger
    17  	Resp        http.ResponseWriter
    18  	Flusher     http.Flusher
    19  	Timer       *timer.Timer
    20  	wroteHeader bool
    21  }
    22  
    23  func (rl *ResponseLogWriter) logCode(code int, status string) {
    24  	var codestr string
    25  	switch {
    26  	case code >= 200 && code < 300:
    27  		codestr = color.GreenString("%d %s", code, status)
    28  	case code >= 300 && code < 400:
    29  		codestr = color.BlueString("%d %s", code, status)
    30  	case code >= 400 && code < 500:
    31  		codestr = color.YellowString("%d %s", code, status)
    32  	case code >= 500 && code < 600:
    33  		codestr = color.RedString("%d %s", code, status)
    34  	default:
    35  		codestr = fmt.Sprintf("%d %s", code, status)
    36  	}
    37  	cl := rl.Header().Get("content-length")
    38  	clstr := ""
    39  	if cl != "" {
    40  		cli, err := strconv.Atoi(cl)
    41  		if err != nil {
    42  			rl.Log.Warn("Invalid content-length header")
    43  		} else if cli > 0 {
    44  			clstr = fmt.Sprintf("%s", humanize.Bytes(uint64(cli)))
    45  		}
    46  	}
    47  	rl.Log.Say("<- %s %s", codestr, clstr)
    48  }
    49  
    50  // Header returns the header map that will be sent by WriteHeader.
    51  // Changing the header after a call to WriteHeader (or Write) has
    52  // no effect.
    53  func (rl *ResponseLogWriter) Header() http.Header {
    54  	return rl.Resp.Header()
    55  }
    56  
    57  // Write writes the data to the connection as part of an HTTP reply.
    58  // If WriteHeader has not yet been called, Write calls WriteHeader(http.StatusOK)
    59  // before writing the data.  If the Header does not contain a
    60  // Content-Type line, Write adds a Content-Type set to the result of passing
    61  // the initial 512 bytes of written data to DetectContentType.
    62  func (rl *ResponseLogWriter) Write(data []byte) (int, error) {
    63  	if !rl.wroteHeader {
    64  		rl.WriteHeader(http.StatusOK)
    65  	}
    66  	ret, err := rl.Resp.Write(data)
    67  	rl.Timer.ResponseDone()
    68  	return ret, err
    69  }
    70  
    71  // WriteHeader sends an HTTP response header with status code.
    72  // If WriteHeader is not called explicitly, the first call to Write
    73  // will trigger an implicit WriteHeader(http.StatusOK).
    74  // Thus explicit calls to WriteHeader are mainly used to
    75  // send error codes.
    76  func (rl *ResponseLogWriter) WriteHeader(code int) {
    77  	rl.wroteHeader = true
    78  	rl.logCode(code, http.StatusText(code))
    79  	LogHeader(rl.Log, rl.Resp.Header())
    80  	rl.Timer.ResponseHeaders()
    81  	rl.Resp.WriteHeader(code)
    82  	rl.Timer.ResponseDone()
    83  }
    84  
    85  func (rl *ResponseLogWriter) Flush() {
    86  	if rl.Flusher != nil {
    87  		rl.Flusher.Flush()
    88  	}
    89  }