response_test.go

package middleware

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"testing"
)

// response must satisfy io.ReaderFrom so that io.Copy (used by http.FileServer)
// can reach the underlying sendfile(2) fast path through the wrapper.
var _ io.ReaderFrom = (*response)(nil)

// readerFromWriter is a ResponseWriter that also implements io.ReaderFrom, just
// like the net/http response that supports the sendfile fast path.
type readerFromWriter struct {
	buf          bytes.Buffer
	header       http.Header
	status       int
	usedReadFrom bool
}

func (w *readerFromWriter) Header() http.Header {
	if w.header == nil {
		w.header = http.Header{}
	}
	return w.header
}

func (w *readerFromWriter) Write(b []byte) (int, error) { return w.buf.Write(b) }

func (w *readerFromWriter) WriteHeader(status int) { w.status = status }

func (w *readerFromWriter) ReadFrom(src io.Reader) (int64, error) {
	w.usedReadFrom = true
	return w.buf.ReadFrom(src)
}

// plainWriter is a ResponseWriter that does not implement io.ReaderFrom, which
// forces the response wrapper to fall back to a regular copy.
type plainWriter struct {
	buf    bytes.Buffer
	header http.Header
	status int
}

func (w *plainWriter) Header() http.Header {
	if w.header == nil {
		w.header = http.Header{}
	}
	return w.header
}

func (w *plainWriter) Write(b []byte) (int, error) { return w.buf.Write(b) }

func (w *plainWriter) WriteHeader(status int) { w.status = status }

func TestResponseReadFromDelegates(t *testing.T) {
	underlying := &readerFromWriter{}
	w := &response{ResponseWriter: underlying}

	data := "hello sendfile world"
	n, err := w.ReadFrom(strings.NewReader(data))

	if err != nil {
		t.Fatalf("ReadFrom returned error: %s", err)
	}
	if !underlying.usedReadFrom {
		t.Fatal("expected ReadFrom to delegate to the underlying io.ReaderFrom")
	}
	if n != int64(len(data)) {
		t.Fatalf("expected %d bytes read, got %d", len(data), n)
	}
	if w.Length != len(data) {
		t.Fatalf("expected recorded length %d, got %d", len(data), w.Length)
	}
	if w.Status != http.StatusOK {
		t.Fatalf("expected status %d, got %d", http.StatusOK, w.Status)
	}
	if underlying.buf.String() != data {
		t.Fatalf("expected body %q, got %q", data, underlying.buf.String())
	}
}

func TestResponseReadFromFallback(t *testing.T) {
	underlying := &plainWriter{}
	w := &response{ResponseWriter: underlying}

	data := "fallback copy path"
	n, err := w.ReadFrom(strings.NewReader(data))

	if err != nil {
		t.Fatalf("ReadFrom returned error: %s", err)
	}
	if n != int64(len(data)) {
		t.Fatalf("expected %d bytes read, got %d", len(data), n)
	}
	if w.Length != len(data) {
		t.Fatalf("expected recorded length %d, got %d", len(data), w.Length)
	}
	if w.Status != http.StatusOK {
		t.Fatalf("expected status %d, got %d", http.StatusOK, w.Status)
	}
	if underlying.buf.String() != data {
		t.Fatalf("expected body %q, got %q", data, underlying.buf.String())
	}
}