From 93cf5d4f5f78058ef0e5fe5fbf597de2084c64fd Mon Sep 17 00:00:00 2001 From: Igor Drozdov Date: Thu, 28 Mar 2024 14:42:34 +0100 Subject: [PATCH] Implement Unwrap() for responseWriterDelegator (#1480) If the ResponseWriter implements any of the following methods, the ResponseController will call them as appropriate: Flush() FlushError() error // alternative Flush returning an error Hijack() (net.Conn, *bufio.ReadWriter, error) SetReadDeadline(deadline time.Time) error SetWriteDeadline(deadline time.Time) error EnableFullDuplex() error If the ResponseWriter doesn't implement the methods, the ResponseController will call Unwrap() method until it finds a ResponseWriter in the chain This commit implements Unwrap() method to simply return the wrapped ResponseWriter Signed-off-by: Igor Drozdov --- prometheus/promhttp/delegator.go | 6 +++ prometheus/promhttp/delegator_test.go | 78 +++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 prometheus/promhttp/delegator_test.go diff --git a/prometheus/promhttp/delegator.go b/prometheus/promhttp/delegator.go index 9819917..315eab5 100644 --- a/prometheus/promhttp/delegator.go +++ b/prometheus/promhttp/delegator.go @@ -76,6 +76,12 @@ func (r *responseWriterDelegator) Write(b []byte) (int, error) { return n, err } +// Unwrap lets http.ResponseController get the underlying http.ResponseWriter, +// by implementing the [rwUnwrapper](https://cs.opensource.google/go/go/+/refs/tags/go1.21.4:src/net/http/responsecontroller.go;l=42-44) interface. +func (r *responseWriterDelegator) Unwrap() http.ResponseWriter { + return r.ResponseWriter +} + type ( closeNotifierDelegator struct{ *responseWriterDelegator } flusherDelegator struct{ *responseWriterDelegator } diff --git a/prometheus/promhttp/delegator_test.go b/prometheus/promhttp/delegator_test.go new file mode 100644 index 0000000..4576ae7 --- /dev/null +++ b/prometheus/promhttp/delegator_test.go @@ -0,0 +1,78 @@ +// Copyright 2024 The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package promhttp + +import ( + "net/http" + "testing" + "time" +) + +type responseWriter struct { + flushErrorCalled bool + setWriteDeadlineCalled time.Time + setReadDeadlineCalled time.Time +} + +func (rw *responseWriter) Header() http.Header { + return nil +} + +func (rw *responseWriter) Write(p []byte) (int, error) { + return 0, nil +} + +func (rw *responseWriter) WriteHeader(statusCode int) { +} + +func (rw *responseWriter) FlushError() error { + rw.flushErrorCalled = true + + return nil +} + +func (rw *responseWriter) SetWriteDeadline(deadline time.Time) error { + rw.setWriteDeadlineCalled = deadline + + return nil +} + +func (rw *responseWriter) SetReadDeadline(deadline time.Time) error { + rw.setReadDeadlineCalled = deadline + + return nil +} + +func TestResponseWriterDelegatorUnwrap(t *testing.T) { + w := &responseWriter{} + rwd := &responseWriterDelegator{ResponseWriter: w} + + if rwd.Unwrap() != w { + t.Error("unwrapped responsewriter must equal to the original responsewriter") + } + + controller := http.NewResponseController(rwd) + if err := controller.Flush(); err != nil || !w.flushErrorCalled { + t.Error("FlushError must be propagated to the original responsewriter") + } + + timeNow := time.Now() + if err := controller.SetWriteDeadline(timeNow); err != nil || w.setWriteDeadlineCalled != timeNow { + t.Error("SetWriteDeadline must be propagated to the original responsewriter") + } + + if err := controller.SetReadDeadline(timeNow); err != nil || w.setReadDeadlineCalled != timeNow { + t.Error("SetReadDeadline must be propagated to the original responsewriter") + } +}