fix: cap websocket body log growth
This commit is contained in:
@@ -34,6 +34,8 @@ const (
|
|||||||
wsTurnStateHeader = "x-codex-turn-state"
|
wsTurnStateHeader = "x-codex-turn-state"
|
||||||
wsRequestBodyKey = "REQUEST_BODY_OVERRIDE"
|
wsRequestBodyKey = "REQUEST_BODY_OVERRIDE"
|
||||||
wsPayloadLogMaxSize = 2048
|
wsPayloadLogMaxSize = 2048
|
||||||
|
wsBodyLogMaxSize = 64 * 1024
|
||||||
|
wsBodyLogTruncated = "\n[websocket log truncated]\n"
|
||||||
)
|
)
|
||||||
|
|
||||||
var responsesWebsocketUpgrader = websocket.Upgrader{
|
var responsesWebsocketUpgrader = websocket.Upgrader{
|
||||||
@@ -825,18 +827,71 @@ func appendWebsocketEvent(builder *strings.Builder, eventType string, payload []
|
|||||||
if builder == nil {
|
if builder == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if builder.Len() >= wsBodyLogMaxSize {
|
||||||
|
return
|
||||||
|
}
|
||||||
trimmedPayload := bytes.TrimSpace(payload)
|
trimmedPayload := bytes.TrimSpace(payload)
|
||||||
if len(trimmedPayload) == 0 {
|
if len(trimmedPayload) == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if builder.Len() > 0 {
|
if builder.Len() > 0 {
|
||||||
builder.WriteString("\n")
|
if !appendWebsocketLogString(builder, "\n") {
|
||||||
|
return
|
||||||
|
}
|
||||||
}
|
}
|
||||||
builder.WriteString("websocket.")
|
if !appendWebsocketLogString(builder, "websocket.") {
|
||||||
builder.WriteString(eventType)
|
return
|
||||||
builder.WriteString("\n")
|
}
|
||||||
builder.Write(trimmedPayload)
|
if !appendWebsocketLogString(builder, eventType) {
|
||||||
builder.WriteString("\n")
|
return
|
||||||
|
}
|
||||||
|
if !appendWebsocketLogString(builder, "\n") {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !appendWebsocketLogBytes(builder, trimmedPayload, len(wsBodyLogTruncated)) {
|
||||||
|
appendWebsocketLogString(builder, wsBodyLogTruncated)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
appendWebsocketLogString(builder, "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendWebsocketLogString(builder *strings.Builder, value string) bool {
|
||||||
|
if builder == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
remaining := wsBodyLogMaxSize - builder.Len()
|
||||||
|
if remaining <= 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if len(value) <= remaining {
|
||||||
|
builder.WriteString(value)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
builder.WriteString(value[:remaining])
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func appendWebsocketLogBytes(builder *strings.Builder, value []byte, reserveForSuffix int) bool {
|
||||||
|
if builder == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
remaining := wsBodyLogMaxSize - builder.Len()
|
||||||
|
if remaining <= 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if len(value) <= remaining {
|
||||||
|
builder.Write(value)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
limit := remaining - reserveForSuffix
|
||||||
|
if limit < 0 {
|
||||||
|
limit = 0
|
||||||
|
}
|
||||||
|
if limit > len(value) {
|
||||||
|
limit = len(value)
|
||||||
|
}
|
||||||
|
builder.Write(value[:limit])
|
||||||
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
func websocketPayloadEventType(payload []byte) string {
|
func websocketPayloadEventType(payload []byte) string {
|
||||||
|
|||||||
@@ -266,6 +266,34 @@ func TestAppendWebsocketEvent(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
func TestAppendWebsocketEventTruncatesAtLimit(t *testing.T) {
|
||||||
|
var builder strings.Builder
|
||||||
|
payload := bytes.Repeat([]byte("x"), wsBodyLogMaxSize)
|
||||||
|
|
||||||
|
appendWebsocketEvent(&builder, "request", payload)
|
||||||
|
|
||||||
|
got := builder.String()
|
||||||
|
if len(got) > wsBodyLogMaxSize {
|
||||||
|
t.Fatalf("body log len = %d, want <= %d", len(got), wsBodyLogMaxSize)
|
||||||
|
}
|
||||||
|
if !strings.Contains(got, wsBodyLogTruncated) {
|
||||||
|
t.Fatalf("expected truncation marker in body log")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAppendWebsocketEventNoGrowthAfterLimit(t *testing.T) {
|
||||||
|
var builder strings.Builder
|
||||||
|
appendWebsocketEvent(&builder, "request", bytes.Repeat([]byte("x"), wsBodyLogMaxSize))
|
||||||
|
initial := builder.String()
|
||||||
|
|
||||||
|
appendWebsocketEvent(&builder, "response", []byte(`{"type":"response.completed"}`))
|
||||||
|
|
||||||
|
if builder.String() != initial {
|
||||||
|
t.Fatalf("builder grew after reaching limit")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestSetWebsocketRequestBody(t *testing.T) {
|
func TestSetWebsocketRequestBody(t *testing.T) {
|
||||||
gin.SetMode(gin.TestMode)
|
gin.SetMode(gin.TestMode)
|
||||||
recorder := httptest.NewRecorder()
|
recorder := httptest.NewRecorder()
|
||||||
|
|||||||
Reference in New Issue
Block a user