Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

module/apmhttp: implement io.ReaderFrom in wrapper #830

Merged
merged 2 commits into from
Oct 7, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.asciidoc
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ https://github.com/elastic/apm-agent-go/compare/v1.8.0...master[View commits]
- Relax Kubernetes pod UID discovery rules {pull}819[#(819)]
- Add transaction and span outcome {pull}820[#(820)]
- Add cloud metadata, configurable with ELASTIC_APM_CLOUD_PROVIDER {pull}823[#(823)]
- module/apmhttp: implement io.ReaderFrom in wrapped http.ResponseWriter {pull}830[#(830)]

[[release-notes-1.x]]
=== Go Agent version 1.x
Expand Down
61 changes: 51 additions & 10 deletions module/apmhttp/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ package apmhttp

import (
"context"
"io"
"net/http"

"go.elastic.co/apm"
Expand Down Expand Up @@ -153,39 +154,59 @@ func SetContext(ctx *apm.Context, req *http.Request, resp *Response, body *apm.B
// ResponseWriter's Write or WriteHeader methods are called, then the
// response's StatusCode field will be zero.
//
// The returned http.ResponseWriter implements http.Pusher and http.Hijacker
// if and only if the provided http.ResponseWriter does.
// The returned http.ResponseWriter implements http.Pusher, http.Hijacker,
// and io.ReaderFrom if and only if the provided http.ResponseWriter does.
func WrapResponseWriter(w http.ResponseWriter) (http.ResponseWriter, *Response) {
rw := responseWriter{
ResponseWriter: w,
resp: Response{
Headers: w.Header(),
},
}

h, _ := w.(http.Hijacker)
p, _ := w.(http.Pusher)
rf, _ := w.(io.ReaderFrom)

switch {
case h != nil && p != nil:
rwhp := &responseWriterHijackerPusher{
rwhp := responseWriterHijackerPusher{
responseWriter: rw,
Hijacker: h,
Pusher: p,
}
return rwhp, &rwhp.resp
if rf != nil {
rwhprf := responseWriterHijackerPusherReaderFrom{rwhp, rf}
return &rwhprf, &rwhprf.resp
}
return &rwhp, &rwhp.resp
case h != nil:
rwh := &responseWriterHijacker{
rwh := responseWriterHijacker{
responseWriter: rw,
Hijacker: h,
}
return rwh, &rwh.resp
if rf != nil {
rwhrf := responseWriterHijackerReaderFrom{rwh, rf}
return &rwhrf, &rwhrf.resp
}
return &rwh, &rwh.resp
case p != nil:
rwp := &responseWriterPusher{
rwp := responseWriterPusher{
responseWriter: rw,
Pusher: p,
}
return rwp, &rwp.resp
if rf != nil {
rwprf := responseWriterPusherReaderFrom{rwp, rf}
return &rwprf, &rwprf.resp
}
return &rwp, &rwp.resp
default:
if rf != nil {
rwrf := responseWriterReaderFrom{rw, rf}
return &rwrf, &rwrf.resp
}
return &rw, &rw.resp
}
return &rw, &rw.resp
}

// Response records details of the HTTP response.
Expand Down Expand Up @@ -229,30 +250,50 @@ func (w *responseWriter) CloseNotify() <-chan bool {
return nil
}

// Flush calls w.flush() if w.flush is non-nil, otherwise
// Flush calls w.ResponseWriter's Flush method if implemented, otherwise
// it does nothing.
func (w *responseWriter) Flush() {
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
}

type responseWriterReaderFrom struct {
responseWriter
io.ReaderFrom
}

type responseWriterHijacker struct {
responseWriter
http.Hijacker
}

type responseWriterHijackerReaderFrom struct {
responseWriterHijacker
io.ReaderFrom
}

type responseWriterPusher struct {
responseWriter
http.Pusher
}

type responseWriterPusherReaderFrom struct {
responseWriterPusher
io.ReaderFrom
}

type responseWriterHijackerPusher struct {
responseWriter
http.Hijacker
http.Pusher
}

type responseWriterHijackerPusherReaderFrom struct {
responseWriterHijackerPusher
io.ReaderFrom
}

// ServerOption sets options for tracing server requests.
type ServerOption func(*handler)

Expand Down
21 changes: 21 additions & 0 deletions module/apmhttp/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -526,6 +526,27 @@ func TestHandlerTracestateHeader(t *testing.T) {
assert.Equal(t, "", w.Body.String())
}

func TestHandlerReaderFrom(t *testing.T) {
recorder := apmtest.NewRecordingTracer()
defer recorder.Close()

mux := http.NewServeMux()
mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
assert.Implements(t, new(io.ReaderFrom), w)
rf := w.(io.ReaderFrom)
rf.ReadFrom(strings.NewReader("hello"))
}))

srv := httptest.NewServer(apmhttp.Wrap(mux, apmhttp.WithTracer(recorder.Tracer)))
defer srv.Close()

resp, err := http.Get(srv.URL)
require.NoError(t, err)
content, _ := ioutil.ReadAll(resp.Body)
assert.NoError(t, resp.Body.Close())
assert.Equal(t, "hello", string(content))
}

func panicHandler(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusTeapot)
panic("foo")
Expand Down
1 change: 1 addition & 0 deletions scripts/before_install.sh
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ function pin() {
}

if (! go run scripts/mingoversion.go 1.11 &>/dev/null); then
pin github.com/astaxie/beego v1.11.1
pin github.com/gin-gonic/gin v1.3.0
pin github.com/stretchr/testify v1.4.0
pin github.com/cucumber/godog v0.8.0
Expand Down