Skip to content

Commit

Permalink
Implement Unwrap() for responseWriterDelegator
Browse files Browse the repository at this point in the history
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 <ihardrozdov@gmail.com>
  • Loading branch information
igor-drozdov authored and Igor Drozdov committed Mar 28, 2024
1 parent 50ab457 commit 9f59f26
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 0 deletions.
6 changes: 6 additions & 0 deletions prometheus/promhttp/delegator.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
90 changes: 90 additions & 0 deletions prometheus/promhttp/delegator_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
// 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"
"time"

"testing"
)

type responseWriter struct {
flushErrorCalled bool
enableFullDuplexCalled 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 (rw *responseWriter) EnableFullDuplex() error {
rw.enableFullDuplexCalled = true

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")
}

if err := controller.EnableFullDuplex(); err != nil || !w.enableFullDuplexCalled {

Check failure on line 78 in prometheus/promhttp/delegator_test.go

View workflow job for this annotation

GitHub Actions / lint

controller.EnableFullDuplex undefined (type *http.ResponseController has no field or method EnableFullDuplex) (typecheck)

Check failure on line 78 in prometheus/promhttp/delegator_test.go

View workflow job for this annotation

GitHub Actions / Tests (1.20)

controller.EnableFullDuplex undefined (type *http.ResponseController has no field or method EnableFullDuplex)
t.Error("EnableFullDuplex 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")
}
}

0 comments on commit 9f59f26

Please sign in to comment.