coderd/httpapi/httpapi_test.go
Rama main · réplica de lectura del repositorio original.
package httpapi_test
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/mock/gomock"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/coderd/httpapi"
"github.com/coder/coder/v2/coderd/httpmw/loggermw"
"github.com/coder/coder/v2/coderd/httpmw/loggermw/loggermock"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
"github.com/coder/quartz"
)
func TestInternalServerError(t *testing.T) {
t.Parallel()
t.Run("NoError", func(t *testing.T) {
t.Parallel()
w := httptest.NewRecorder()
httpapi.InternalServerError(w, nil)
var resp codersdk.Response
err := json.NewDecoder(w.Body).Decode(&resp)
require.NoError(t, err)
require.Equal(t, http.StatusInternalServerError, w.Code)
require.NotEmpty(t, resp.Message)
require.Empty(t, resp.Detail)
})
t.Run("WithError", func(t *testing.T) {
t.Parallel()
var (
w = httptest.NewRecorder()
httpErr = xerrors.New("error!")
)
httpapi.InternalServerError(w, httpErr)
var resp codersdk.Response
err := json.NewDecoder(w.Body).Decode(&resp)
require.NoError(t, err)
require.Equal(t, http.StatusInternalServerError, w.Code)
require.NotEmpty(t, resp.Message)
require.Equal(t, httpErr.Error(), resp.Detail)
})
}
func TestWrite(t *testing.T) {
t.Parallel()
t.Run("NoErrors", func(t *testing.T) {
t.Parallel()
ctx := context.Background()
rw := httptest.NewRecorder()
httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
Message: "Wow.",
})
var m map[string]interface{}
err := json.NewDecoder(rw.Body).Decode(&m)
require.NoError(t, err)
_, ok := m["errors"]
require.False(t, ok)
})
}
func TestRead(t *testing.T) {
t.Parallel()
t.Run("EmptyStruct", func(t *testing.T) {
t.Parallel()
ctx := context.Background()
rw := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", bytes.NewBufferString("{}"))
v := struct{}{}
require.True(t, httpapi.Read(ctx, rw, r, &v))
})
t.Run("NoBody", func(t *testing.T) {
t.Parallel()
ctx := context.Background()
rw := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", nil)
var v json.RawMessage
require.False(t, httpapi.Read(ctx, rw, r, v))
})
t.Run("BodyTooLarge", func(t *testing.T) {
t.Parallel()
ctx := context.Background()
rw := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", strings.NewReader(`{"value":"too large"}`))
r.Body = http.MaxBytesReader(rw, r.Body, 4)
var v json.RawMessage
require.False(t, httpapi.Read(ctx, rw, r, &v))
require.Equal(t, http.StatusRequestEntityTooLarge, rw.Code)
})
t.Run("Validate", func(t *testing.T) {
t.Parallel()
type toValidate struct {
Value string `json:"value" validate:"required"`
}
ctx := context.Background()
rw := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", bytes.NewBufferString(`{"value":"hi"}`))
var validate toValidate
require.True(t, httpapi.Read(ctx, rw, r, &validate))
require.Equal(t, "hi", validate.Value)
})
t.Run("ValidateFailure", func(t *testing.T) {
t.Parallel()
type toValidate struct {
Value string `json:"value" validate:"required"`
}
ctx := context.Background()
rw := httptest.NewRecorder()
r := httptest.NewRequest("POST", "/", bytes.NewBufferString("{}"))
var validate toValidate
require.False(t, httpapi.Read(ctx, rw, r, &validate))
var v codersdk.Response
err := json.NewDecoder(rw.Body).Decode(&v)
require.NoError(t, err)
require.Len(t, v.Validations, 1)
require.Equal(t, "value", v.Validations[0].Field)
require.Equal(t, "Validation failed for tag \"required\" with value: \"\"", v.Validations[0].Detail)
})
}
// readBody is decoded by the request body limit tests. It carries no validate
// tags so that a decode failure is unambiguously a body-size failure.
type readBody struct {
Value string `json:"value"`
}
// jsonBodyOfSize returns a JSON object that decodes into readBody and is
// exactly size bytes long.
func jsonBodyOfSize(size int) string {
const (
prefix = `{"value":"`
suffix = `"}`
)
return prefix + strings.Repeat("a", size-len(prefix)-len(suffix)) + suffix
}
func TestReadDefaultLimit(t *testing.T) {
t.Parallel()
// requireTooLarge asserts the 413 response shape shared by every
// over-limit case.
requireTooLarge := func(t *testing.T, rw *httptest.ResponseRecorder, limit int) {
t.Helper()
require.Equal(t, http.StatusRequestEntityTooLarge, rw.Code)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(rw.Body).Decode(&resp))
require.Equal(t, "Request body too large.", resp.Message)
require.Contains(t, resp.Detail, strconv.Itoa(limit),
"the detail must name the limit so the error is actionable")
}
t.Run("AtDefaultLimit", func(t *testing.T) {
t.Parallel()
body := jsonBodyOfSize(httpapi.DefaultMaxRequestBodyBytes)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
var v readBody
require.True(t, httpapi.Read(context.Background(), rw, r, &v))
require.Len(t, v.Value, httpapi.DefaultMaxRequestBodyBytes-len(`{"value":""}`))
})
t.Run("OverDefaultLimitByOneByte", func(t *testing.T) {
t.Parallel()
body := jsonBodyOfSize(httpapi.DefaultMaxRequestBodyBytes + 1)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
var v readBody
require.False(t, httpapi.Read(context.Background(), rw, r, &v))
requireTooLarge(t, rw, httpapi.DefaultMaxRequestBodyBytes)
})
// The limit is enforced on bytes actually read, so neither an absent nor a
// dishonest Content-Length can raise it. Asserted in-process rather than
// over a connection because a client still streaming an oversized body may
// see the connection reset instead of the 413, which would make a
// network-driven assertion racy.
t.Run("NoContentLength", func(t *testing.T) {
t.Parallel()
body := jsonBodyOfSize(httpapi.DefaultMaxRequestBodyBytes + 1)
rw := httptest.NewRecorder()
// io.NopCloser hides the length, which is what net/http sees for a
// chunked request.
r := httptest.NewRequest(http.MethodPost, "/", io.NopCloser(strings.NewReader(body)))
r.TransferEncoding = []string{"chunked"}
require.EqualValues(t, -1, r.ContentLength)
var v readBody
require.False(t, httpapi.Read(context.Background(), rw, r, &v))
requireTooLarge(t, rw, httpapi.DefaultMaxRequestBodyBytes)
})
t.Run("UnderstatedContentLength", func(t *testing.T) {
t.Parallel()
body := jsonBodyOfSize(httpapi.DefaultMaxRequestBodyBytes + 1)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
r.ContentLength = 10
var v readBody
require.False(t, httpapi.Read(context.Background(), rw, r, &v))
requireTooLarge(t, rw, httpapi.DefaultMaxRequestBodyBytes)
})
t.Run("ChunkedUnderLimit", func(t *testing.T) {
t.Parallel()
body := jsonBodyOfSize(httpapi.DefaultMaxRequestBodyBytes)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", io.NopCloser(strings.NewReader(body)))
r.TransferEncoding = []string{"chunked"}
var v readBody
require.True(t, httpapi.Read(context.Background(), rw, r, &v))
})
}
func TestReadLimit(t *testing.T) {
t.Parallel()
// A limit above the default must not be tightened by the default that Read
// installs. This is the unit-level regression test for the endpoints that
// legitimately accept more than DefaultMaxRequestBodyBytes; without
// ReadLimit they would be silently capped at the default.
t.Run("AboveDefaultIsNotTightened", func(t *testing.T) {
t.Parallel()
const limit = 8 << 20
body := jsonBodyOfSize(6 << 20)
require.Greater(t, len(body), httpapi.DefaultMaxRequestBodyBytes)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
var v readBody
require.True(t, httpapi.ReadLimit(context.Background(), rw, r, limit, &v))
require.Len(t, v.Value, len(body)-len(`{"value":""}`))
})
t.Run("BelowDefaultIsEnforced", func(t *testing.T) {
t.Parallel()
const limit = 1024
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(jsonBodyOfSize(limit+1)))
var v readBody
require.False(t, httpapi.ReadLimit(context.Background(), rw, r, limit, &v))
require.Equal(t, http.StatusRequestEntityTooLarge, rw.Code)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(rw.Body).Decode(&resp))
require.Contains(t, resp.Detail, strconv.Itoa(limit))
})
t.Run("AtLimit", func(t *testing.T) {
t.Parallel()
const limit = 1024
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(jsonBodyOfSize(limit)))
var v readBody
require.True(t, httpapi.ReadLimit(context.Background(), rw, r, limit, &v))
})
// The limit lands on the request's existing log line rather than one of its
// own: a caller can produce 413s at will, so a dedicated line would let them
// drive log volume.
t.Run("RecordsLimitOnRequestLog", func(t *testing.T) {
t.Parallel()
const limit = 1024
ctrl := gomock.NewController(t)
requestLogger := loggermock.NewMockRequestLogger(ctrl)
requestLogger.EXPECT().
WithFields(slog.F("max_request_body_bytes", int64(limit))).
Times(1)
ctx := loggermw.WithRequestLogger(context.Background(), requestLogger)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(jsonBodyOfSize(limit+1))).WithContext(ctx)
var v readBody
require.False(t, httpapi.ReadLimit(ctx, rw, r, limit, &v))
require.Equal(t, http.StatusRequestEntityTooLarge, rw.Code)
})
// A caller that installs its own reader is doing what the docstring
// forbids, but the number reported must still be the one that rejected the
// request. Nested readers compose as tightest-wins, so reporting the limit
// this call installed would tell the client and the log a cap that is not
// the one it hit.
t.Run("ReportsLimitThatTripped", func(t *testing.T) {
t.Parallel()
const (
tight = 1024
loose = 1 << 20
)
for _, tc := range []struct {
name string
callerWrap int64
readLimit int64
}{
{name: "CallerWrapsTighter", callerWrap: tight, readLimit: loose},
{name: "CallerWrapsLooser", callerWrap: loose, readLimit: tight},
} {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctrl := gomock.NewController(t)
requestLogger := loggermock.NewMockRequestLogger(ctrl)
requestLogger.EXPECT().
WithFields(slog.F("max_request_body_bytes", int64(tight))).
Times(1)
ctx := loggermw.WithRequestLogger(context.Background(), requestLogger)
rw := httptest.NewRecorder()
r := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(jsonBodyOfSize(tight+1))).WithContext(ctx)
r.Body = http.MaxBytesReader(rw, r.Body, tc.callerWrap)
var v readBody
require.False(t, httpapi.ReadLimit(ctx, rw, r, tc.readLimit, &v))
require.Equal(t, http.StatusRequestEntityTooLarge, rw.Code)
var resp codersdk.Response
require.NoError(t, json.NewDecoder(rw.Body).Decode(&resp))
require.Contains(t, resp.Detail, strconv.Itoa(tight))
require.NotContains(t, resp.Detail, strconv.Itoa(loose))
})
}
})
}
func TestWebsocketCloseMsg(t *testing.T) {
t.Parallel()
t.Run("Sprintf", func(t *testing.T) {
t.Parallel()
var (
msg = "this is my message %q %q"
opts = []any{"colin", "kyle"}
)
expected := fmt.Sprintf(msg, opts...)
got := httpapi.WebsocketCloseSprintf(msg, opts...)
assert.Equal(t, expected, got)
})
t.Run("TruncateSingleByteCharacters", func(t *testing.T) {
t.Parallel()
msg := strings.Repeat("d", 255)
trunc := httpapi.WebsocketCloseSprintf("%s", msg)
assert.Equal(t, len(trunc), 123)
})
t.Run("TruncateMultiByteCharacters", func(t *testing.T) {
t.Parallel()
msg := strings.Repeat("こんにちは", 10)
trunc := httpapi.WebsocketCloseSprintf("%s", msg)
assert.Equal(t, len(trunc), 123)
})
}
// Our WebSocket library accepts any arbitrary ResponseWriter at the type level,
// but the writer must also implement http.Hijacker for long-lived connections.
type mockOneWaySocketWriter struct {
serverRecorder *httptest.ResponseRecorder
serverConn net.Conn
clientConn net.Conn
serverReadWriter *bufio.ReadWriter
testContext *testing.T
}
func (m mockOneWaySocketWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return m.serverConn, m.serverReadWriter, nil
}
func (m mockOneWaySocketWriter) Flush() {
err := m.serverReadWriter.Flush()
require.NoError(m.testContext, err)
}
func (m mockOneWaySocketWriter) Header() http.Header {
return m.serverRecorder.Header()
}
func (m mockOneWaySocketWriter) Write(b []byte) (int, error) {
return m.serverReadWriter.Write(b)
}
func (m mockOneWaySocketWriter) WriteHeader(code int) {
m.serverRecorder.WriteHeader(code)
}
func TestOneWayWebSocketEventSender(t *testing.T) {
t.Parallel()
newBaseRequest := func(ctx context.Context) *http.Request {
url := "ws://www.fake-website.com/logs"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
h := req.Header
h.Add("Connection", "Upgrade")
h.Add("Upgrade", "websocket")
h.Add("Sec-WebSocket-Version", "13")
h.Add("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") // Just need any string
return req
}
newOneWayWriter := func(t *testing.T) mockOneWaySocketWriter {
mockServer, mockClient := net.Pipe()
recorder := httptest.NewRecorder()
return mockOneWaySocketWriter{
testContext: t,
serverConn: mockServer,
clientConn: mockClient,
serverRecorder: recorder,
serverReadWriter: bufio.NewReadWriter(
bufio.NewReader(mockServer),
bufio.NewWriter(mockServer),
),
}
}
t.Run("Produces error if the socket connection could not be established", func(t *testing.T) {
t.Parallel()
incorrectProtocols := []struct {
major int
minor int
proto string
}{
{0, 9, "HTTP/0.9"},
{1, 0, "HTTP/1.0"},
}
for _, p := range incorrectProtocols {
ctx := testutil.Context(t, testutil.WaitShort)
req := newBaseRequest(ctx)
req.ProtoMajor = p.major
req.ProtoMinor = p.minor
req.Proto = p.proto
writer := newOneWayWriter(t)
_, _, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), nil)(writer, req)
require.ErrorContains(t, err, p.proto)
}
})
t.Run("Returned callback can publish new event to WebSocket connection", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
wsw := httpapi.NewWSWatcher(quartz.NewReal(), nil)
req := newBaseRequest(ctx)
writer := newOneWayWriter(t)
send, _, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), wsw)(writer, req)
require.NoError(t, err)
serverPayload := codersdk.ServerSentEvent{
Type: codersdk.ServerSentEventTypeData,
Data: "Blah",
}
err = send(serverPayload)
require.NoError(t, err)
// The client connection will receive a little bit of additional data on
// top of the main payload. Have to make sure check has tolerance for
// extra data being present. Read until the payload shows up; EOF only
// arrives after the close handshake times out against this raw pipe.
serverBytes, err := json.Marshal(serverPayload)
require.NoError(t, err)
deadline, ok := ctx.Deadline()
require.True(t, ok)
require.NoError(t, writer.clientConn.SetReadDeadline(deadline))
var clientBytes []byte
buf := make([]byte, 1024)
for !bytes.Contains(clientBytes, serverBytes) {
n, err := writer.clientConn.Read(buf)
require.NoError(t, err)
clientBytes = append(clientBytes, buf[:n]...)
}
})
t.Run("Signals to outside consumer when socket has been closed", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
wsw := httpapi.NewWSWatcher(quartz.NewReal(), nil)
req := newBaseRequest(ctx)
writer := newOneWayWriter(t)
_, done, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), wsw)(writer, req)
require.NoError(t, err)
successC := make(chan bool)
ticker := time.NewTicker(testutil.WaitShort)
go func() {
select {
case <-done:
successC <- true
case <-ticker.C:
successC <- false
}
}()
cancel()
require.True(t, <-successC)
})
t.Run("Socket will immediately close if client sends any message", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
wsw := httpapi.NewWSWatcher(quartz.NewReal(), nil)
req := newBaseRequest(ctx)
writer := newOneWayWriter(t)
_, done, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), wsw)(writer, req)
require.NoError(t, err)
successC := make(chan bool)
ticker := time.NewTicker(testutil.WaitShort)
go func() {
select {
case <-done:
successC <- true
case <-ticker.C:
successC <- false
}
}()
type JunkClientEvent struct {
Value string
}
b, err := json.Marshal(JunkClientEvent{"Hi :)"})
require.NoError(t, err)
_, err = writer.clientConn.Write(b)
require.NoError(t, err)
require.True(t, <-successC)
})
t.Run("Renders the socket inert if the request context cancels", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
wsw := httpapi.NewWSWatcher(quartz.NewReal(), nil)
req := newBaseRequest(ctx)
writer := newOneWayWriter(t)
send, done, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), wsw)(writer, req)
require.NoError(t, err)
successC := make(chan bool)
ticker := time.NewTicker(testutil.WaitShort)
go func() {
select {
case <-done:
successC <- true
case <-ticker.C:
successC <- false
}
}()
cancel()
require.True(t, <-successC)
err = send(codersdk.ServerSentEvent{
Type: codersdk.ServerSentEventTypeData,
Data: "Didn't realize you were closed - sorry! I'll try coming back tomorrow.",
})
require.Equal(t, err, ctx.Err())
_, open := <-done
require.False(t, open)
_, err = writer.serverConn.Write([]byte{})
require.Equal(t, err, io.ErrClosedPipe)
_, err = writer.clientConn.Read([]byte{})
require.Equal(t, err, io.EOF)
})
t.Run("Sends a heartbeat to the socket on a fixed internal of time to keep connections alive", func(t *testing.T) {
t.Parallel()
// One tick of the mock clock must produce a heartbeat; read a few
// bytes of it to prove it was written.
heartbeatCount := 3
timeout := testutil.WaitShort
ctx := testutil.Context(t, timeout)
mClock := quartz.NewMock(t)
trap := mClock.Trap().NewTicker("WSWatcher")
defer trap.Close()
wsw := httpapi.NewWSWatcher(mClock, nil)
req := newBaseRequest(ctx)
writer := newOneWayWriter(t)
_, _, err := httpapi.OneWayWebSocketEventSender(slogtest.Make(t, nil), wsw)(writer, req)
require.NoError(t, err)
trap.MustWait(ctx).MustRelease(ctx)
type Result struct {
Err error
Success bool
}
resultC := make(chan Result)
go func() {
err := writer.
clientConn.
SetReadDeadline(time.Now().Add(timeout))
if err != nil {
resultC <- Result{err, false}
return
}
for range heartbeatCount {
pingBuffer := make([]byte, 1)
pingSize, err := writer.clientConn.Read(pingBuffer)
if err != nil || pingSize != 1 {
resultC <- Result{err, false}
return
}
}
resultC <- Result{nil, true}
}()
mClock.Advance(httpapi.HeartbeatInterval).MustWait(ctx)
result := <-resultC
require.NoError(t, result.Err)
require.True(t, result.Success)
})
}
// ServerSentEventSender accepts any arbitrary ResponseWriter at the type level,
// but the writer must also implement http.Flusher for long-lived connections
type mockServerSentWriter struct {
serverRecorder *httptest.ResponseRecorder
serverConn net.Conn
clientConn net.Conn
buffer *bytes.Buffer
testContext *testing.T
}
func (m mockServerSentWriter) Flush() {
b := m.buffer.Bytes()
_, err := m.serverConn.Write(b)
require.NoError(m.testContext, err)
m.buffer.Reset()
// Must close server connection to indicate EOF for any reads from the
// client connection; otherwise reads block forever. This is a testing
// limitation compared to the one-way websockets, since we have no way to
// frame the data and auto-indicate EOF for each message
err = m.serverConn.Close()
require.NoError(m.testContext, err)
}
func (m mockServerSentWriter) Header() http.Header {
return m.serverRecorder.Header()
}
func (m mockServerSentWriter) Write(b []byte) (int, error) {
return m.buffer.Write(b)
}
func (m mockServerSentWriter) WriteHeader(code int) {
m.serverRecorder.WriteHeader(code)
}
func TestServerSentEventSender(t *testing.T) {
t.Parallel()
newBaseRequest := func(ctx context.Context) *http.Request {
url := "ws://www.fake-website.com/logs"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
require.NoError(t, err)
return req
}
newServerSentWriter := func(t *testing.T) mockServerSentWriter {
mockServer, mockClient := net.Pipe()
return mockServerSentWriter{
testContext: t,
serverRecorder: httptest.NewRecorder(),
clientConn: mockClient,
serverConn: mockServer,
buffer: &bytes.Buffer{},
}
}
t.Run("Mutates response headers to support SSE connections", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
req := newBaseRequest(ctx)
writer := newServerSentWriter(t)
_, _, err := httpapi.ServerSentEventSender(writer, req)
require.NoError(t, err)
h := writer.Header()
require.Equal(t, h.Get("Content-Type"), "text/event-stream")
require.Equal(t, h.Get("Cache-Control"), "no-cache")
require.Equal(t, h.Get("Connection"), "keep-alive")
require.Equal(t, h.Get("X-Accel-Buffering"), "no")
})
t.Run("Returned callback can publish new event to SSE connection", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
req := newBaseRequest(ctx)
writer := newServerSentWriter(t)
send, _, err := httpapi.ServerSentEventSender(writer, req)
require.NoError(t, err)
serverPayload := codersdk.ServerSentEvent{
Type: codersdk.ServerSentEventTypeData,
Data: "Blah",
}
err = send(serverPayload)
require.NoError(t, err)
clientBytes, err := io.ReadAll(writer.clientConn)
require.NoError(t, err)
require.Equal(
t,
string(clientBytes),
"event: data\ndata: \"Blah\"\n\n",
)
})
t.Run("Signals to outside consumer when connection has been closed", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitShort))
req := newBaseRequest(ctx)
writer := newServerSentWriter(t)
_, done, err := httpapi.ServerSentEventSender(writer, req)
require.NoError(t, err)
successC := make(chan bool)
ticker := time.NewTicker(testutil.WaitShort)
go func() {
select {
case <-done:
successC <- true
case <-ticker.C:
successC <- false
}
}()
cancel()
require.True(t, <-successC)
})
t.Run("Sends a heartbeat to the client on a fixed internal of time to keep connections alive", func(t *testing.T) {
t.Parallel()
// One tick of the mock clock must produce a heartbeat; read a few
// bytes of it to prove it was written.
heartbeatCount := 3
timeout := testutil.WaitShort
ctx := testutil.Context(t, timeout)
mClock := quartz.NewMock(t)
trap := mClock.Trap().NewTicker("ServerSentEventSender")
defer trap.Close()
req := newBaseRequest(ctx)
writer := newServerSentWriter(t)
_, _, err := httpapi.ServerSentEventSenderWithClock(mClock)(writer, req)
require.NoError(t, err)
trap.MustWait(ctx).MustRelease(ctx)
type Result struct {
Err error
Success bool
}
resultC := make(chan Result)
go func() {
err := writer.
clientConn.
SetReadDeadline(time.Now().Add(timeout))
if err != nil {
resultC <- Result{err, false}
return
}
for range heartbeatCount {
pingBuffer := make([]byte, 1)
pingSize, err := writer.clientConn.Read(pingBuffer)
if err != nil || pingSize != 1 {
resultC <- Result{err, false}
return
}
}
resultC <- Result{nil, true}
}()
mClock.Advance(httpapi.HeartbeatInterval).MustWait(ctx)
result := <-resultC
require.NoError(t, result.Err)
require.True(t, result.Success)
})
}
// TestRecordRequestBodyLimit pins both halves of the call every oversized-body
// 413 site shares: the log field naming the limit, and the metric tracker.
func TestRecordRequestBodyLimit(t *testing.T) {
t.Parallel()
t.Run("RecordsFieldAndMarksTracker", func(t *testing.T) {
t.Parallel()
const limit = int64(4096)
ctrl := gomock.NewController(t)
requestLogger := loggermock.NewMockRequestLogger(ctrl)
requestLogger.EXPECT().
WithFields(slog.F("max_request_body_bytes", limit)).
Times(1)
tracker := &httpapi.RequestBodyLimitTracker{}
ctx := httpapi.WithRequestBodyLimitTracker(
loggermw.WithRequestLogger(context.Background(), requestLogger), tracker)
require.False(t, tracker.Exceeded())
httpapi.RecordRequestBodyLimit(ctx, limit)
require.True(t, tracker.Exceeded())
})
// The middleware that installs the tracker is not mounted on every route, so
// a call without one must not panic.
t.Run("NoTrackerInContext", func(t *testing.T) {
t.Parallel()
httpapi.RecordRequestBodyLimit(context.Background(), 4096)
})
}