http2buffer_test.go

package http2buffer

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
)

func TestBufferedRead(t *testing.T) {
	BufferedRead(t, 0)
}

func TestBufferedReadWithPause(t *testing.T) {
	BufferedRead(t, 5*time.Millisecond)
}

func BufferedRead(t *testing.T, d time.Duration) {
	handler := h2c.NewHandler(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
		<-time.After(100 * time.Millisecond)
		_, _ = rw.Write([]byte{0xAA, 0xBB})
		rw.(http.Flusher).Flush()
		if d > 0 {
			<-time.After(d)
		}
		_, _ = rw.Write([]byte{0xCC, 0xDD})
		rw.(http.Flusher).Flush()
	}), &http2.Server{})

	// launch server on a random port
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		t.Fatal(err)
	}
	go func() {
		_ = http.Serve(listener, handler)
	}()
	addr := listener.Addr().String()

	// flacky test, launch several times to increase reliability
	for i := 0; i < 100; i++ {
		launchRequest(t, addr)
	}
}

func launchRequest(t *testing.T, addr string) {
	// create http2 client
	transport := &http2.Transport{
		AllowHTTP: true,
		DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
			return net.Dial(network, addr)
		},
	}
	client := http.Client{
		Transport: transport,
		Timeout:   1 * time.Second,
	}

	// create request
	url := fmt.Sprintf("http://%s", addr)
	req, err := http.NewRequest("POST", url, &delayedReader{})
	if err != nil {
		t.Fatal(err)
	}

	// launch request
	res, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer res.Body.Close()

	// read data
	recv := make([]byte, 1024)
	n, err := res.Body.Read(recv)
	if err != nil {
		t.Fatal(err)
	}
	want := []byte{0xAA, 0xBB}
	if !bytes.Equal(recv[:n], want) {
		t.Errorf("got %x, want %x", recv[:n], want)
	}
	want = []byte{0xCC, 0xDD}
	n, err = res.Body.Read(recv)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(recv[:n], want) {
		t.Errorf("got %x, want %x", recv[:n], want)
	}
}

type delayedReader struct{}

func (d *delayedReader) Read(b []byte) (int, error) {
	<-time.After(100 * time.Millisecond)
	return 2, nil
}