From 8eb4dda26fac43623bcd939ac5ce84950af1b7d9 Mon Sep 17 00:00:00 2001 From: Jahan Khan Date: Wed, 10 Jan 2024 17:06:29 -0500 Subject: [PATCH 1/3] add test and copy goji 1.8+ over --- hlog/hlog_test.go | 28 +++++ hlog/internal/mutil/writer_proxy.go | 3 + hlog/internal/mutil/writer_proxy1_8.go | 137 +++++++++++++++++++++++++ 3 files changed, 168 insertions(+) create mode 100644 hlog/internal/mutil/writer_proxy1_8.go diff --git a/hlog/hlog_test.go b/hlog/hlog_test.go index 445d6b6f..2d9f5f54 100644 --- a/hlog/hlog_test.go +++ b/hlog/hlog_test.go @@ -13,6 +13,7 @@ import ( "net/url" "reflect" "testing" + "time" "github.com/rs/xid" "github.com/rs/zerolog" @@ -273,6 +274,33 @@ func TestResponseHeaderHandler(t *testing.T) { } } +func TestAccessHandler(t *testing.T) { + out := &bytes.Buffer{} + w := httptest.NewRecorder() + r := &http.Request{} + + var respSize int + f := func(r *http.Request, status, size int, duration time.Duration) { + respSize = size + } + + h := AccessHandler(f)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte{'a'}) + w.WriteHeader(http.StatusOK) + })) + + h2 := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h.ServeHTTP(w, r) + l := FromRequest(r) + l.Log().Int("size", respSize).Msg("") + }) + h3 := NewHandler(zerolog.New(out))(h2) + h3.ServeHTTP(w, r) + if want, got := `{"size":1}`+"\n", decodeIfBinary(out); want != got { + t.Errorf("Invalid log output, got: %s, want: %s", got, want) + } +} + func TestProtoHandler(t *testing.T) { out := &bytes.Buffer{} r := &http.Request{ diff --git a/hlog/internal/mutil/writer_proxy.go b/hlog/internal/mutil/writer_proxy.go index 9d427390..d4c338aa 100644 --- a/hlog/internal/mutil/writer_proxy.go +++ b/hlog/internal/mutil/writer_proxy.go @@ -1,3 +1,6 @@ +//go:build !go1.8 +// +build !go1.8 + package mutil import ( diff --git a/hlog/internal/mutil/writer_proxy1_8.go b/hlog/internal/mutil/writer_proxy1_8.go new file mode 100644 index 00000000..381cad65 --- /dev/null +++ b/hlog/internal/mutil/writer_proxy1_8.go @@ -0,0 +1,137 @@ +//go:build go1.8 +// +build go1.8 + +package mutil + +import ( + "bufio" + "io" + "net" + "net/http" +) + +// WriterProxy is a proxy around an http.ResponseWriter that allows you to hook +// into various parts of the response process. +type WriterProxy interface { + http.ResponseWriter + // Status returns the HTTP status of the request, or 0 if one has not + // yet been sent. + Status() int + // BytesWritten returns the total number of bytes sent to the client. + BytesWritten() int + // Tee causes the response body to be written to the given io.Writer in + // addition to proxying the writes through. Only one io.Writer can be + // tee'd to at once: setting a second one will overwrite the first. + // Writes will be sent to the proxy before being written to this + // io.Writer. It is illegal for the tee'd writer to be modified + // concurrently with writes. + Tee(io.Writer) + // Unwrap returns the original proxied target. + Unwrap() http.ResponseWriter +} + +// WrapWriter wraps an http.ResponseWriter, returning a proxy that allows you to +// hook into various parts of the response process. +func WrapWriter(w http.ResponseWriter) WriterProxy { + _, fl := w.(http.Flusher) + _, hj := w.(http.Hijacker) + _, rf := w.(io.ReaderFrom) + + bw := basicWriter{ResponseWriter: w} + if fl && hj && rf { + return &fancyWriter{bw} + } + return &bw +} + +// basicWriter wraps a http.ResponseWriter that implements the minimal +// http.ResponseWriter interface. +type basicWriter struct { + http.ResponseWriter + wroteHeader bool + code int + bytes int + tee io.Writer +} + +func (b *basicWriter) WriteHeader(code int) { + if !b.wroteHeader { + b.code = code + b.wroteHeader = true + b.ResponseWriter.WriteHeader(code) + } +} + +func (b *basicWriter) Write(buf []byte) (int, error) { + b.WriteHeader(http.StatusOK) + n, err := b.ResponseWriter.Write(buf) + if b.tee != nil { + _, err2 := b.tee.Write(buf[:n]) + // Prefer errors generated by the proxied writer. + if err == nil { + err = err2 + } + } + b.bytes += n + return n, err +} + +func (b *basicWriter) maybeWriteHeader() { + if !b.wroteHeader { + b.WriteHeader(http.StatusOK) + } +} + +func (b *basicWriter) Status() int { + return b.code +} + +func (b *basicWriter) BytesWritten() int { + return b.bytes +} + +func (b *basicWriter) Tee(w io.Writer) { + b.tee = w +} + +func (b *basicWriter) Unwrap() http.ResponseWriter { + return b.ResponseWriter +} + +// fancyWriter is a writer that additionally satisfies http.Pusher, +// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case +// of wrapping the http.ResponseWriter that package http gives you, in order to +// make the proxied object support the full method set of the proxied object. +type fancyWriter struct { + basicWriter +} + +func (f *fancyWriter) Push(target string, opts *http.PushOptions) error { + return f.basicWriter.ResponseWriter.(http.Pusher).Push(target, opts) +} + +func (f *fancyWriter) Flush() { + fl := f.basicWriter.ResponseWriter.(http.Flusher) + fl.Flush() +} + +func (f *fancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + hj := f.basicWriter.ResponseWriter.(http.Hijacker) + return hj.Hijack() +} + +func (f *fancyWriter) ReadFrom(r io.Reader) (int64, error) { + if f.basicWriter.tee != nil { + return io.Copy(&f.basicWriter, r) + } + rf := f.basicWriter.ResponseWriter.(io.ReaderFrom) + f.basicWriter.maybeWriteHeader() + return rf.ReadFrom(r) +} + +var ( + _ http.Pusher = &fancyWriter{} + _ http.Flusher = &fancyWriter{} + _ http.Hijacker = &fancyWriter{} + _ io.ReaderFrom = &fancyWriter{} +) From b32b2addbd70754082bd7a5cc0a74224ba760c0b Mon Sep 17 00:00:00 2001 From: Jahan Khan Date: Wed, 10 Jan 2024 17:13:35 -0500 Subject: [PATCH 2/3] remove pusher and extraneous flusher check --- hlog/internal/mutil/writer_proxy.go | 1 - hlog/internal/mutil/writer_proxy1_8.go | 5 ----- 2 files changed, 6 deletions(-) diff --git a/hlog/internal/mutil/writer_proxy.go b/hlog/internal/mutil/writer_proxy.go index d4c338aa..9e749ece 100644 --- a/hlog/internal/mutil/writer_proxy.go +++ b/hlog/internal/mutil/writer_proxy.go @@ -153,5 +153,4 @@ var ( _ http.Flusher = &fancyWriter{} _ http.Hijacker = &fancyWriter{} _ io.ReaderFrom = &fancyWriter{} - _ http.Flusher = &flushWriter{} ) diff --git a/hlog/internal/mutil/writer_proxy1_8.go b/hlog/internal/mutil/writer_proxy1_8.go index 381cad65..d4072659 100644 --- a/hlog/internal/mutil/writer_proxy1_8.go +++ b/hlog/internal/mutil/writer_proxy1_8.go @@ -106,10 +106,6 @@ type fancyWriter struct { basicWriter } -func (f *fancyWriter) Push(target string, opts *http.PushOptions) error { - return f.basicWriter.ResponseWriter.(http.Pusher).Push(target, opts) -} - func (f *fancyWriter) Flush() { fl := f.basicWriter.ResponseWriter.(http.Flusher) fl.Flush() @@ -130,7 +126,6 @@ func (f *fancyWriter) ReadFrom(r io.Reader) (int64, error) { } var ( - _ http.Pusher = &fancyWriter{} _ http.Flusher = &fancyWriter{} _ http.Hijacker = &fancyWriter{} _ io.ReaderFrom = &fancyWriter{} From 3fb7b707ce87f0f6127aab73aa3cc26c9d24153b Mon Sep 17 00:00:00 2001 From: Jahan Khan Date: Wed, 10 Jan 2024 17:22:13 -0500 Subject: [PATCH 3/3] remove w.Write check --- hlog/hlog_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hlog/hlog_test.go b/hlog/hlog_test.go index 2d9f5f54..a6ed987c 100644 --- a/hlog/hlog_test.go +++ b/hlog/hlog_test.go @@ -285,7 +285,7 @@ func TestAccessHandler(t *testing.T) { } h := AccessHandler(f)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte{'a'}) + w.Write([]byte{'a'}) w.WriteHeader(http.StatusOK) }))