github.com/llimllib/devd@v0.0.0-20230426145215-4d29fc25f909/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/llimllib/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 = 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 }