package xfer
import (
"bytes"
"compress/zlib"
"context"
"encoding/binary"
"net/http"
"net/http/httptest"
"runtime"
"strings"
"testing"
)
// uvZeros is a container of n compressed zero bytes claiming the length
// claim: a few hundred KiB on the wire for every 100 MiB inflated.
func uvZeros(n int, claim uint32) []byte {
var buf bytes.Buffer
var head [4]byte
binary.BigEndian.PutUint32(head[:], claim)
buf.Write(head[:])
zw, _ := zlib.NewWriterLevel(&buf, zlib.BestCompression)
chunk := make([]byte, 1<<20)
for i := 0; i < n>>20; i++ {
zw.Write(chunk)
}
zw.Close()
return buf.Bytes()
}
// uvHeapGrowth runs f and reports how far the heap grew.
func uvHeapGrowth(f func()) uint64 {
var before, after runtime.MemStats
runtime.GC()
runtime.ReadMemStats(&before)
f()
runtime.ReadMemStats(&after)
return after.TotalAlloc - before.TotalAlloc
}
// Exchange refuses a reply that inflates past its MaxReply.
func TestUVExchangeRefusesABomb(t *testing.T) {
bomb := uvZeros(64<<20, 64<<20)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/x-fossil")
w.Write(bomb)
}))
defer srv.Close()
var err error
grew := uvHeapGrowth(func() {
_, err = UVUpstream{URL: srv.URL, MaxReply: 1 << 20}.Exchange(context.Background(), "uvgimme a\n")
})
if err == nil || grew > 16<<20 {
t.Fatalf("FAIL: a reply inflating to 64 MiB under a 1 MiB cap: %v, heap grew %d", err, grew)
}
t.Logf("ok: refused (%v), heap grew %d bytes", strings.SplitN(err.Error(), ";", 2)[0], grew)
}