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)

})

}