diff --git a/render/reader.go b/render/reader.go index 5752d8d8..da3d577f 100644 --- a/render/reader.go +++ b/render/reader.go @@ -21,13 +21,8 @@ type Reader struct { // Render (Reader) writes data with custom ContentType and headers. func (r Reader) Render(w http.ResponseWriter) (err error) { r.WriteContentType(w) - if r.ContentLength >= 0 { - if r.Headers == nil { - r.Headers = map[string]string{} - } - r.Headers["Content-Length"] = strconv.FormatInt(r.ContentLength, 10) - } r.writeHeaders(w, r.Headers) + _, err = io.Copy(w, r.Reader) return } @@ -45,4 +40,8 @@ func (r Reader) writeHeaders(w http.ResponseWriter, headers map[string]string) { header.Set(k, v) } } + + if r.ContentLength > 0 { + header.Set("Content-Length", strconv.FormatInt(r.ContentLength, 10)) + } } diff --git a/render/reader_test.go b/render/reader_test.go index aaceb9ea..33d1e4a6 100644 --- a/render/reader_test.go +++ b/render/reader_test.go @@ -6,6 +6,7 @@ package render import ( "net/http/httptest" + "strconv" "strings" "testing" @@ -21,3 +22,21 @@ func TestReaderRenderNoHeaders(t *testing.T) { err := r.Render(httptest.NewRecorder()) require.NoError(t, err) } + +func TestReaderRenderWithHeaders(t *testing.T) { + content := "test" + r := Reader{ + ContentLength: int64(len(content)), + Reader: strings.NewReader(content), + Headers: map[string]string{ + "Test-Content": "test/content", + }, + } + recorder := httptest.NewRecorder() + err := r.Render(recorder) + require.NoError(t, err) + + require.Contains(t, recorder.Header()["Content-Length"], strconv.FormatInt(r.ContentLength, 10)) + + require.Contains(t, recorder.Header()["Test-Content"], "test/content") +}