mirror of
https://github.com/coder/coder.git
synced 2026-10-30 18:12:29 +08:00
fix(coderd/x/chatd/chatprompt): keep NUL out of jsonb chat content (#28625)
PostgreSQL jsonb rejects \u0000, so persisting Gemini responses failed with "pq: unsupported Unicode escape sequence": binary thought signatures reach ProviderMetadata as strings containing NUL, and the sentinel codec only covered six fields. Real recorded Gemini signatures contain NUL in 22 of 34 cases. Classify every string-bearing part field in one table, partNulFields. Free-form and opaque fields encode NUL reversibly as PUA sentinel pairs; structured fields (identifiers, paths, URLs) reject it with an error naming the part, field, and offset, since their consumers cannot handle NUL anyway. Raw JSON is also normalized for lone surrogate escapes, which jsonb equally rejects. A reflection-driven test injects NUL into every field and fails when a future field is missing from the table. The Gemini regression asserts the sentinel-encoded row in the database and a byte-exact signature round-trip. Known limits: json.Marshal corrupts other invalid UTF-8 upstream of this codec (CODAGT-986), and chat debug steps have always dropped silently on NUL prompt text (CODAGT-987, pre-existing). Fixes coder/coder#25555 Fixes CODAGT-463
This commit is contained in:
@@ -1175,18 +1175,21 @@ func sanitizeToolCallID(id string) string {
|
||||
}
|
||||
|
||||
// MarshalParts encodes SDK chat message parts for persistence.
|
||||
// NUL characters in string fields are encoded as PUA sentinel
|
||||
// pairs (U+E000 U+E001) before marshaling so the resulting JSON
|
||||
// never contains \u0000 (rejected by PostgreSQL jsonb). The
|
||||
// encoding operates on Go string values, not JSON bytes, so it
|
||||
// survives jsonb text normalization.
|
||||
// NUL characters in free-form and opaque fields are encoded as PUA
|
||||
// sentinel pairs (U+E000 U+E001) so PostgreSQL jsonb can store them.
|
||||
// Structured fields such as identifiers, paths, and media types reject
|
||||
// NUL because their downstream representations cannot consume it.
|
||||
func MarshalParts(parts []codersdk.ChatMessagePart) (pqtype.NullRawMessage, error) {
|
||||
if len(parts) == 0 {
|
||||
return pqtype.NullRawMessage{}, nil
|
||||
}
|
||||
data, err := json.Marshal(encodeNulInParts(parts))
|
||||
encoded, err := encodeNulInParts(parts)
|
||||
if err != nil {
|
||||
return pqtype.NullRawMessage{}, xerrors.Errorf("encode chat message parts: %w", err)
|
||||
return pqtype.NullRawMessage{}, xerrors.Errorf("validate chat message parts: %w", err)
|
||||
}
|
||||
data, err := json.Marshal(encoded)
|
||||
if err != nil {
|
||||
return pqtype.NullRawMessage{}, xerrors.Errorf("marshal chat message parts: %w", err)
|
||||
}
|
||||
return pqtype.NullRawMessage{RawMessage: data, Valid: true}, nil
|
||||
}
|
||||
@@ -1797,16 +1800,12 @@ func decodeNulInValue(v any) any {
|
||||
}
|
||||
|
||||
// encodeNulInJSON walks all string values (and keys) inside a
|
||||
// json.RawMessage and applies encodeNulInString. Returns the
|
||||
// original unchanged when the raw message does not contain NUL
|
||||
// escapes or U+E000 bytes, or when parsing fails.
|
||||
// json.RawMessage and applies encodeNulInString. It also normalizes
|
||||
// unpaired UTF-16 surrogate escapes, which PostgreSQL jsonb rejects.
|
||||
// Returns the original unchanged when no transformation is needed or
|
||||
// when parsing fails.
|
||||
func encodeNulInJSON(raw json.RawMessage) json.RawMessage {
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
// Quick exit: no \u0000 escape and no U+E000 UTF-8 bytes.
|
||||
if !bytes.Contains(raw, []byte(`\u0000`)) &&
|
||||
!bytes.Contains(raw, []byte{0xEE, 0x80, 0x80}) {
|
||||
if len(raw) == 0 || !needsNulEncodingInJSON(raw) || !json.Valid(raw) {
|
||||
return raw
|
||||
}
|
||||
dec := json.NewDecoder(bytes.NewReader(raw))
|
||||
@@ -1822,6 +1821,80 @@ func encodeNulInJSON(raw json.RawMessage) json.RawMessage {
|
||||
return result
|
||||
}
|
||||
|
||||
// needsNulEncodingInJSON reports whether valid JSON contains a string
|
||||
// value or object key that encodeNulInJSON must transform. It only
|
||||
// detects candidates; encoding/json owns parsing and normalization.
|
||||
func needsNulEncodingInJSON(raw json.RawMessage) bool {
|
||||
// encoding/json emits U+E000 literally, while json.RawMessage may
|
||||
// contain either the literal UTF-8 bytes or a Unicode escape.
|
||||
if bytes.Contains(raw, []byte{0xEE, 0x80, 0x80}) {
|
||||
return true
|
||||
}
|
||||
|
||||
inString := false
|
||||
for i := 0; i < len(raw); {
|
||||
switch raw[i] {
|
||||
case '"':
|
||||
inString = !inString
|
||||
i++
|
||||
case '\\':
|
||||
if !inString || i+1 >= len(raw) {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if raw[i+1] != 'u' {
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
value, ok := parseJSONHexEscape(raw[i+2:])
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
switch {
|
||||
case value == 0 || value == 0xE000:
|
||||
return true
|
||||
case 0xD800 <= value && value <= 0xDBFF:
|
||||
if i+12 <= len(raw) && raw[i+6] == '\\' && raw[i+7] == 'u' {
|
||||
low, lowOK := parseJSONHexEscape(raw[i+8:])
|
||||
if lowOK && 0xDC00 <= low && low <= 0xDFFF {
|
||||
i += 12
|
||||
continue
|
||||
}
|
||||
}
|
||||
return true
|
||||
case 0xDC00 <= value && value <= 0xDFFF:
|
||||
return true
|
||||
default:
|
||||
i += 6
|
||||
}
|
||||
default:
|
||||
i++
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func parseJSONHexEscape(raw []byte) (uint16, bool) {
|
||||
if len(raw) < 4 {
|
||||
return 0, false
|
||||
}
|
||||
var value uint16
|
||||
for _, c := range raw[:4] {
|
||||
value <<= 4
|
||||
switch {
|
||||
case '0' <= c && c <= '9':
|
||||
value |= uint16(c - '0')
|
||||
case 'a' <= c && c <= 'f':
|
||||
value |= uint16(c-'a') + 10
|
||||
case 'A' <= c && c <= 'F':
|
||||
value |= uint16(c-'A') + 10
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
return value, true
|
||||
}
|
||||
|
||||
// decodeNulInJSON walks all string values (and keys) inside a
|
||||
// json.RawMessage and applies decodeNulInString.
|
||||
func decodeNulInJSON(raw json.RawMessage) json.RawMessage {
|
||||
@@ -1845,33 +1918,154 @@ func decodeNulInJSON(raw json.RawMessage) json.RawMessage {
|
||||
return result
|
||||
}
|
||||
|
||||
// encodeNulInParts returns a shallow copy of parts with all
|
||||
// string and json.RawMessage fields NUL-encoded. The caller's
|
||||
// slice is not modified.
|
||||
func encodeNulInParts(parts []codersdk.ChatMessagePart) []codersdk.ChatMessagePart {
|
||||
// nulPolicy classifies how MarshalParts treats NUL in one persisted
|
||||
// ChatMessagePart field.
|
||||
type nulPolicy int
|
||||
|
||||
const (
|
||||
// nulEncode reversibly encodes NUL as PUA sentinel pairs. Used for
|
||||
// free-form text and opaque provider data that must survive
|
||||
// persistence byte-for-byte, e.g. Gemini binary thought signatures
|
||||
// in ProviderMetadata.
|
||||
nulEncode nulPolicy = iota
|
||||
// nulReject fails MarshalParts with an actionable error. Used for
|
||||
// structured values (enums, identifiers, paths, filenames, URLs,
|
||||
// media types, OS labels, and parsed commands) whose downstream
|
||||
// consumers cannot handle NUL, so persisting an encoded form would
|
||||
// only hide invalid input.
|
||||
nulReject
|
||||
)
|
||||
|
||||
// partNulField binds one string-bearing ChatMessagePart field to its
|
||||
// NUL policy. Pointer accessors let encoding, decoding, and validation
|
||||
// share the table. Exactly one accessor is set per entry; grid is only
|
||||
// supported with nulReject and raw only with nulEncode. Unsupported
|
||||
// combinations and missing entries for new string-bearing fields are
|
||||
// caught by TestChatMessagePartNULCoverage.
|
||||
type partNulField struct {
|
||||
name string
|
||||
policy nulPolicy
|
||||
str func(*codersdk.ChatMessagePart) *string
|
||||
raw func(*codersdk.ChatMessagePart) *json.RawMessage
|
||||
grid func(*codersdk.ChatMessagePart) [][]string
|
||||
}
|
||||
|
||||
// partNulFields classifies every persisted string-bearing field of
|
||||
// codersdk.ChatMessagePart. Adding a string-bearing field to the struct
|
||||
// requires an entry here.
|
||||
var partNulFields = []partNulField{
|
||||
{name: "Text", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Text }},
|
||||
{name: "Args", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.Args }},
|
||||
{name: "ArgsDelta", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ArgsDelta }},
|
||||
{name: "Result", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.Result }},
|
||||
{name: "ResultDelta", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ResultDelta }},
|
||||
{name: "Title", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Title }},
|
||||
{name: "Content", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.Content }},
|
||||
{name: "ProviderMetadata", policy: nulEncode, raw: func(p *codersdk.ChatMessagePart) *json.RawMessage { return &p.ProviderMetadata }},
|
||||
{name: "ContextFileContent", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileContent }},
|
||||
{name: "SkillDescription", policy: nulEncode, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillDescription }},
|
||||
{name: "Type", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return (*string)(&p.Type) }},
|
||||
{name: "ToolCallID", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ToolCallID }},
|
||||
{name: "ToolName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ToolName }},
|
||||
{name: "ParsedCommands", policy: nulReject, grid: func(p *codersdk.ChatMessagePart) [][]string { return p.ParsedCommands }},
|
||||
{name: "SourceID", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SourceID }},
|
||||
{name: "URL", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.URL }},
|
||||
{name: "MediaType", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.MediaType }},
|
||||
{name: "Name", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.Name }},
|
||||
{name: "FileName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.FileName }},
|
||||
{name: "ContextFilePath", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFilePath }},
|
||||
{name: "ContextFileOS", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileOS }},
|
||||
{name: "ContextFileDirectory", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileDirectory }},
|
||||
{name: "ContextFileSkillMetaFile", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.ContextFileSkillMetaFile }},
|
||||
{name: "SkillName", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillName }},
|
||||
{name: "SkillDir", policy: nulReject, str: func(p *codersdk.ChatMessagePart) *string { return &p.SkillDir }},
|
||||
}
|
||||
|
||||
// encode applies the reversible NUL encoding to a nulEncode field.
|
||||
func (f partNulField) encode(p *codersdk.ChatMessagePart) {
|
||||
switch {
|
||||
case f.str != nil:
|
||||
*f.str(p) = encodeNulInString(*f.str(p))
|
||||
case f.raw != nil:
|
||||
*f.raw(p) = encodeNulInJSON(*f.raw(p))
|
||||
}
|
||||
}
|
||||
|
||||
// decode reverses encode.
|
||||
func (f partNulField) decode(p *codersdk.ChatMessagePart) {
|
||||
switch {
|
||||
case f.str != nil:
|
||||
*f.str(p) = decodeNulInString(*f.str(p))
|
||||
case f.raw != nil:
|
||||
*f.raw(p) = decodeNulInJSON(*f.raw(p))
|
||||
}
|
||||
}
|
||||
|
||||
// validateNulFree returns an actionable error when a nulReject field
|
||||
// contains NUL.
|
||||
func (f partNulField) validateNulFree(partIndex int, p *codersdk.ChatMessagePart) error {
|
||||
switch {
|
||||
case f.str != nil:
|
||||
if nul := strings.IndexByte(*f.str(p), 0); nul >= 0 {
|
||||
return xerrors.Errorf(
|
||||
"chat message part %d field %s contains NUL at byte %d",
|
||||
partIndex,
|
||||
f.name,
|
||||
nul,
|
||||
)
|
||||
}
|
||||
case f.grid != nil:
|
||||
for i, values := range f.grid(p) {
|
||||
for j, value := range values {
|
||||
if nul := strings.IndexByte(value, 0); nul >= 0 {
|
||||
return xerrors.Errorf(
|
||||
"chat message part %d field %s[%d][%d] contains NUL at byte %d",
|
||||
partIndex,
|
||||
f.name,
|
||||
i,
|
||||
j,
|
||||
nul,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// encodeNulInParts returns a copy of parts with every nulEncode field
|
||||
// encoded and every nulReject field validated per partNulFields. The
|
||||
// caller's values are not modified.
|
||||
//
|
||||
// Content version 1 cannot distinguish sentinel pairs written by this
|
||||
// codec from matching natural text stored before a field was covered, so
|
||||
// matching historical values in encoded fields decode as NUL.
|
||||
func encodeNulInParts(parts []codersdk.ChatMessagePart) ([]codersdk.ChatMessagePart, error) {
|
||||
encoded := make([]codersdk.ChatMessagePart, len(parts))
|
||||
copy(encoded, parts)
|
||||
for i := range encoded {
|
||||
p := &encoded[i]
|
||||
p.Text = encodeNulInString(p.Text)
|
||||
p.Content = encodeNulInString(p.Content)
|
||||
p.Args = encodeNulInJSON(p.Args)
|
||||
p.ArgsDelta = encodeNulInString(p.ArgsDelta)
|
||||
p.Result = encodeNulInJSON(p.Result)
|
||||
p.ResultDelta = encodeNulInString(p.ResultDelta)
|
||||
for _, f := range partNulFields {
|
||||
switch f.policy {
|
||||
case nulEncode:
|
||||
f.encode(p)
|
||||
case nulReject:
|
||||
if err := f.validateNulFree(i, p); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return encoded
|
||||
return encoded, nil
|
||||
}
|
||||
|
||||
// decodeNulInParts reverses encodeNulInParts in place.
|
||||
func decodeNulInParts(parts []codersdk.ChatMessagePart) {
|
||||
for i := range parts {
|
||||
p := &parts[i]
|
||||
p.Text = decodeNulInString(p.Text)
|
||||
p.Content = decodeNulInString(p.Content)
|
||||
p.Args = decodeNulInJSON(p.Args)
|
||||
p.ArgsDelta = decodeNulInString(p.ArgsDelta)
|
||||
p.Result = decodeNulInJSON(p.Result)
|
||||
p.ResultDelta = decodeNulInString(p.ResultDelta)
|
||||
for _, f := range partNulFields {
|
||||
if f.policy == nulEncode {
|
||||
f.decode(&parts[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"charm.land/fantasy"
|
||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||
fantasygoogle "charm.land/fantasy/providers/google"
|
||||
"github.com/google/uuid"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -2007,7 +2008,7 @@ func extractToolResultIDs(t *testing.T, msgs ...fantasy.Message) []string {
|
||||
func TestNulEscapeRoundTrip(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
|
||||
// Seed minimal dependencies for the DB round-trip path:
|
||||
// user, provider, model config, chat.
|
||||
@@ -2157,6 +2158,97 @@ func TestNulEscapeRoundTrip(t *testing.T) {
|
||||
assert.JSONEq(t, string(resultJSON), string(decoded[0].Result))
|
||||
})
|
||||
|
||||
// Gemini thought signatures are binary data converted directly to a
|
||||
// string by the Google provider. NUL bytes in the signature must survive
|
||||
// the full persistence and provider-metadata round-trip.
|
||||
t.Run("GeminiReasoningMetadataWithNul", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const signature = "binary\x00thought\x00signature"
|
||||
part := chatprompt.PartFromContent(fantasy.ReasoningContent{
|
||||
Text: "reasoning",
|
||||
ProviderMetadata: fantasy.ProviderMetadata{
|
||||
fantasygoogle.Name: &fantasygoogle.ReasoningMetadata{
|
||||
Signature: signature,
|
||||
ToolID: "tool-1",
|
||||
},
|
||||
},
|
||||
})
|
||||
encoded, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part})
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(encoded.RawMessage), `\u0000`)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
dbMsg := dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
Content: encoded,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
})
|
||||
readBack, err := db.GetChatMessageByID(ctx, dbMsg.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Physical evidence of the encoding: the stored jsonb row contains
|
||||
// one sentinel pair per NUL byte and no \u0000 escape.
|
||||
var stored string
|
||||
require.NoError(t, sqlDB.QueryRowContext(ctx,
|
||||
"SELECT content::text FROM chat_messages WHERE id = $1", dbMsg.ID,
|
||||
).Scan(&stored))
|
||||
require.NotContains(t, stored, `\u0000`)
|
||||
require.Equal(t, 2, strings.Count(stored, "\uE000\uE001"))
|
||||
|
||||
prompt, err := chatprompt.ConvertMessagesWithFiles(
|
||||
ctx,
|
||||
[]database.ChatMessage{readBack},
|
||||
nil,
|
||||
slogtest.Make(t, nil),
|
||||
nil,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, prompt, 1)
|
||||
require.Len(t, prompt[0].Content, 1)
|
||||
reasoning, ok := fantasy.AsMessagePart[fantasy.ReasoningPart](prompt[0].Content[0])
|
||||
require.True(t, ok)
|
||||
metadata := fantasygoogle.GetReasoningMetadata(reasoning.ProviderOptions)
|
||||
require.NotNil(t, metadata)
|
||||
require.Equal(t, signature, metadata.Signature)
|
||||
})
|
||||
|
||||
// PostgreSQL rejects lone UTF-16 surrogate escapes. The RawMessage
|
||||
// encoder normalizes them through encoding/json before persistence.
|
||||
t.Run("ToolArgsWithLoneSurrogate", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
parts := []codersdk.ChatMessagePart{{
|
||||
Type: codersdk.ChatMessagePartTypeToolCall,
|
||||
Args: json.RawMessage(`{"value":"before\uD800after","pair":"\uD83D\uDE00","literal":"\\uD800"}`),
|
||||
}}
|
||||
encoded, err := chatprompt.MarshalParts(parts)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(encoded.RawMessage), `\uD800after`)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
dbMsg := dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
Content: encoded,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
})
|
||||
readBack, err := db.GetChatMessageByID(ctx, dbMsg.ID)
|
||||
require.NoError(t, err)
|
||||
decoded, err := chatprompt.ParseContent(readBack)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, decoded, 1)
|
||||
assert.JSONEq(t,
|
||||
`{"value":"before�after","pair":"😀","literal":"\\uD800"}`,
|
||||
string(decoded[0].Args),
|
||||
)
|
||||
})
|
||||
|
||||
// Multiple parts in one message: one with NUL, one without.
|
||||
t.Run("MultiPartMixed", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -15,6 +15,12 @@ const SyntheticPasteTitleBudgetForTest = syntheticPasteTitleBudget
|
||||
// for external tests.
|
||||
var ToolResultPartToMessagePartForTest = toolResultPartToMessagePart
|
||||
|
||||
// EncodeNulInJSONForTest exposes encodeNulInJSON for external tests.
|
||||
var EncodeNulInJSONForTest = encodeNulInJSON
|
||||
|
||||
// NeedsNulEncodingInJSONForTest exposes needsNulEncodingInJSON for external tests.
|
||||
var NeedsNulEncodingInJSONForTest = needsNulEncodingInJSON
|
||||
|
||||
// ToolResultContentToPartForTest exposes toolResultContentToPart
|
||||
// for external tests.
|
||||
var ToolResultContentToPartForTest = func(logger slog.Logger, content fantasy.ToolResultContent) codersdk.ChatMessagePart {
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
package chatprompt_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
// Probe values injected into string-bearing fields. The NUL probe's
|
||||
// prefix "before" puts the NUL at byte 6, which rejection tests assert.
|
||||
const (
|
||||
nulProbe = "before\x00middle\uE000\uE001after"
|
||||
naturalSentinelProbe = "before\uE000\uE001after"
|
||||
)
|
||||
|
||||
func TestNeedsNulEncodingInJSON(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
raw json.RawMessage
|
||||
want bool
|
||||
}{
|
||||
{name: "clean", raw: json.RawMessage(`{"value":"clean"}`)},
|
||||
{name: "nul", raw: json.RawMessage(`{"value":"\u0000"}`), want: true},
|
||||
{name: "escaped nul text", raw: json.RawMessage(`{"value":"\\u0000"}`)},
|
||||
{name: "escaped backslash then nul", raw: json.RawMessage(`{"value":"\\\u0000"}`), want: true},
|
||||
{name: "ordinary unicode escape", raw: json.RawMessage(`{"value":"\u1234"}`)},
|
||||
{name: "truncated unicode escape", raw: json.RawMessage(`{"value":"\u12"}`), want: true},
|
||||
{name: "unicode escape without digits", raw: json.RawMessage(`{"value":"\u`), want: true},
|
||||
{name: "invalid unicode escape", raw: json.RawMessage(`{"value":"\uZZZZ"}`), want: true},
|
||||
{name: "backslash outside string", raw: json.RawMessage(`\`)},
|
||||
{name: "trailing backslash in string", raw: json.RawMessage(`{"value":"\`)},
|
||||
|
||||
{name: "two escaped backslashes then nul text", raw: json.RawMessage(`{"value":"\\\\u0000"}`)},
|
||||
{name: "escaped sentinel", raw: json.RawMessage(`{"value":"\uE000"}`), want: true},
|
||||
{name: "literal sentinel", raw: json.RawMessage(`{"value":""}`), want: true},
|
||||
{name: "valid surrogate pair", raw: json.RawMessage(`{"value":"\uD83D\uDE00"}`)},
|
||||
{name: "mixed case surrogate pair", raw: json.RawMessage(`{"value":"\ud83D\uDe00"}`)},
|
||||
{name: "lone high surrogate", raw: json.RawMessage(`{"value":"\uD800"}`), want: true},
|
||||
{name: "lone low surrogate", raw: json.RawMessage(`{"value":"\udc00"}`), want: true},
|
||||
{name: "reversed surrogate pair", raw: json.RawMessage(`{"value":"\uDC00\uD800"}`), want: true},
|
||||
{name: "surrogate literal text", raw: json.RawMessage(`{"value":"\\uD800"}`)},
|
||||
{name: "object key", raw: json.RawMessage(`{"key\u0000":"value"}`), want: true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := chatprompt.NeedsNulEncodingInJSONForTest(tt.raw)
|
||||
require.Equalf(t, tt.want, got, "needsNulEncodingInJSON(%s)", tt.raw)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeNulInJSONPreservesCleanInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
raw := json.RawMessage(`{"number":1.2300,"pair":"\uD83D\uDE00","literal":"\\u0000"}`)
|
||||
encoded := chatprompt.EncodeNulInJSONForTest(raw)
|
||||
require.Equal(t, raw, encoded, "clean JSON changed")
|
||||
require.Same(t, &raw[0], &encoded[0], "clean JSON did not preserve its backing array")
|
||||
}
|
||||
|
||||
func TestEncodeNulInJSONNormalizesSurrogates(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
raw := json.RawMessage(`{"high":"\uD800","low":"\uDC00","pair":"\uD83D\uDE00","literal":"\\uD800"}`)
|
||||
encoded := chatprompt.EncodeNulInJSONForTest(raw)
|
||||
require.Truef(t, json.Valid(encoded), "encoded invalid JSON: %s", encoded)
|
||||
var got map[string]string
|
||||
require.NoError(t, json.Unmarshal(encoded, &got))
|
||||
want := map[string]string{
|
||||
"high": "�",
|
||||
"low": "�",
|
||||
"pair": "😀",
|
||||
"literal": `\uD800`,
|
||||
}
|
||||
require.Equal(t, want, got, "normalized JSON mismatch")
|
||||
}
|
||||
|
||||
func TestEncodeNulInJSONPreservesInvalidInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
raw := json.RawMessage(`{"value":"\uD800"} trailing`)
|
||||
encoded := chatprompt.EncodeNulInJSONForTest(raw)
|
||||
require.Equal(t, raw, encoded, "invalid JSON changed")
|
||||
}
|
||||
|
||||
func TestChatMessagePartNaturalSentinelCoverage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
partType := reflect.TypeFor[codersdk.ChatMessagePart]()
|
||||
for i := range partType.NumField() {
|
||||
field := partType.Field(i)
|
||||
if _, transforms := chatMessagePartNULProbe(t, field); !transforms {
|
||||
continue
|
||||
}
|
||||
|
||||
t.Run(field.Name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
part := reflect.New(partType).Elem()
|
||||
part.Field(i).Set(chatMessagePartNaturalSentinelProbeValue(t, field.Type))
|
||||
parts := []codersdk.ChatMessagePart{part.Interface().(codersdk.ChatMessagePart)}
|
||||
|
||||
encoded, err := chatprompt.MarshalParts(parts)
|
||||
require.NoErrorf(t, err, "MarshalParts rejected natural sentinels in ChatMessagePart.%s (%s)", field.Name, field.Type)
|
||||
decoded, err := chatprompt.ParseContent(database.ChatMessage{
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
Content: encoded,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
})
|
||||
require.NoErrorf(t, err, "ParseContent with natural sentinels in ChatMessagePart.%s (%s)", field.Name, field.Type)
|
||||
require.Equalf(t, parts, decoded, "natural sentinels in ChatMessagePart.%s did not round-trip", field.Name)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatMessagePartNULCoverage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
partType := reflect.TypeFor[codersdk.ChatMessagePart]()
|
||||
for i := range partType.NumField() {
|
||||
field := partType.Field(i)
|
||||
probe, transforms := chatMessagePartNULProbe(t, field)
|
||||
|
||||
t.Run(field.Name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
part := reflect.New(partType).Elem()
|
||||
part.Field(i).Set(probe)
|
||||
parts := []codersdk.ChatMessagePart{part.Interface().(codersdk.ChatMessagePart)}
|
||||
|
||||
want := reflect.New(partType).Elem()
|
||||
want.Field(i).Set(chatMessagePartNULProbeValue(t, field.Type))
|
||||
wantParts := []codersdk.ChatMessagePart{want.Interface().(codersdk.ChatMessagePart)}
|
||||
|
||||
encoded, err := chatprompt.MarshalParts(parts)
|
||||
require.Equalf(t, wantParts, parts, "MarshalParts mutated caller-owned ChatMessagePart.%s data", field.Name)
|
||||
if err != nil {
|
||||
require.Truef(t, transforms, "MarshalParts rejected NUL-free ChatMessagePart.%s (%s): %v", field.Name, field.Type, err)
|
||||
require.ErrorContainsf(t, err, "chat message part 0", "MarshalParts rejection for ChatMessagePart.%s does not name the part", field.Name)
|
||||
require.ErrorContainsf(t, err, "field "+field.Name, "MarshalParts rejection for ChatMessagePart.%s does not name the field", field.Name)
|
||||
require.ErrorContainsf(t, err, "contains NUL at byte 6", "MarshalParts rejection for ChatMessagePart.%s does not locate the NUL", field.Name)
|
||||
return
|
||||
}
|
||||
if transforms {
|
||||
require.NotContainsf(t, string(encoded.RawMessage), `\u0000`, "ChatMessagePart.%s is neither rejected nor handled by encodeNulInParts", field.Name)
|
||||
}
|
||||
|
||||
decoded, err := chatprompt.ParseContent(database.ChatMessage{
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
Content: encoded,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
})
|
||||
require.NoErrorf(t, err, "ParseContent with ChatMessagePart.%s (%s)", field.Name, field.Type)
|
||||
require.Equalf(t, wantParts, decoded, "ChatMessagePart.%s is not handled by decodeNulInParts", field.Name)
|
||||
require.Equalf(t, wantParts, parts, "ParseContent mutated caller-owned ChatMessagePart.%s data through an encoded alias", field.Name)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func chatMessagePartNaturalSentinelProbeValue(t *testing.T, typ reflect.Type) reflect.Value {
|
||||
t.Helper()
|
||||
|
||||
if typ == reflect.TypeFor[json.RawMessage]() {
|
||||
return reflect.ValueOf(json.RawMessage(`{"key":"value"}`))
|
||||
}
|
||||
if typ.Kind() == reflect.String {
|
||||
value := reflect.New(typ).Elem()
|
||||
value.SetString(naturalSentinelProbe)
|
||||
return value
|
||||
}
|
||||
if chatMessagePartNULTransforms(typ) {
|
||||
return chatMessagePartStringContainerValue(t, typ, naturalSentinelProbe)
|
||||
}
|
||||
|
||||
t.Fatalf("no natural sentinel probe for ChatMessagePart field type %s", typ)
|
||||
return reflect.Value{}
|
||||
}
|
||||
|
||||
func chatMessagePartNULProbe(t *testing.T, field reflect.StructField) (reflect.Value, bool) {
|
||||
t.Helper()
|
||||
|
||||
transforms := chatMessagePartNULTransforms(field.Type)
|
||||
if transforms || chatMessagePartNULSafeType(field.Type) {
|
||||
return chatMessagePartNULProbeValue(t, field.Type), transforms
|
||||
}
|
||||
|
||||
t.Fatalf(
|
||||
"ChatMessagePart.%s has unknown NUL coverage shape %s; classify the type and add a partNulFields entry if it can contain strings",
|
||||
field.Name,
|
||||
field.Type,
|
||||
)
|
||||
return reflect.Value{}, false
|
||||
}
|
||||
|
||||
func chatMessagePartNULTransforms(typ reflect.Type) bool {
|
||||
if typ == reflect.TypeFor[json.RawMessage]() {
|
||||
return true
|
||||
}
|
||||
if typ.Kind() == reflect.String {
|
||||
return true
|
||||
}
|
||||
if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array {
|
||||
return false
|
||||
}
|
||||
return chatMessagePartNULStringContainer(typ.Elem())
|
||||
}
|
||||
|
||||
func chatMessagePartNULStringContainer(typ reflect.Type) bool {
|
||||
if typ.Kind() == reflect.String {
|
||||
return true
|
||||
}
|
||||
if typ.Kind() != reflect.Slice && typ.Kind() != reflect.Array {
|
||||
return false
|
||||
}
|
||||
return chatMessagePartNULStringContainer(typ.Elem())
|
||||
}
|
||||
|
||||
func chatMessagePartNULSafeType(typ reflect.Type) bool {
|
||||
switch typ {
|
||||
case reflect.TypeFor[bool](),
|
||||
reflect.TypeFor[int](),
|
||||
reflect.TypeFor[[]byte](),
|
||||
reflect.TypeFor[uuid.NullUUID](),
|
||||
reflect.TypeFor[*time.Time]():
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func chatMessagePartNULProbeValue(t *testing.T, typ reflect.Type) reflect.Value {
|
||||
t.Helper()
|
||||
|
||||
if typ == reflect.TypeFor[json.RawMessage]() {
|
||||
return reflect.ValueOf(json.RawMessage(`{"key\u0000":"value\u0000"}`))
|
||||
}
|
||||
if typ == reflect.TypeFor[[][]string]() {
|
||||
return reflect.ValueOf([][]string{{nulProbe}, nil})
|
||||
}
|
||||
if typ.Kind() == reflect.String {
|
||||
value := reflect.New(typ).Elem()
|
||||
value.SetString(nulProbe)
|
||||
return value
|
||||
}
|
||||
if chatMessagePartNULTransforms(typ) {
|
||||
return chatMessagePartStringContainerValue(t, typ, nulProbe)
|
||||
}
|
||||
|
||||
switch typ {
|
||||
case reflect.TypeFor[bool]():
|
||||
return reflect.ValueOf(true)
|
||||
case reflect.TypeFor[int]():
|
||||
return reflect.ValueOf(42)
|
||||
case reflect.TypeFor[[]byte]():
|
||||
return reflect.ValueOf([]byte{0, 1, 2})
|
||||
case reflect.TypeFor[uuid.NullUUID]():
|
||||
return reflect.ValueOf(uuid.NullUUID{UUID: uuid.MustParse("00000000-0000-0000-0000-000000000001"), Valid: true})
|
||||
case reflect.TypeFor[*time.Time]():
|
||||
return reflect.ValueOf(new(time.Date(2026, time.August, 25, 1, 2, 3, 0, time.UTC)))
|
||||
default:
|
||||
t.Fatalf("no probe value for recognized ChatMessagePart field type %s", typ)
|
||||
return reflect.Value{}
|
||||
}
|
||||
}
|
||||
|
||||
func chatMessagePartStringContainerValue(t *testing.T, typ reflect.Type, probe string) reflect.Value {
|
||||
t.Helper()
|
||||
|
||||
if typ.Kind() == reflect.String {
|
||||
value := reflect.New(typ).Elem()
|
||||
value.SetString(probe)
|
||||
return value
|
||||
}
|
||||
|
||||
switch typ.Kind() {
|
||||
case reflect.Slice:
|
||||
value := reflect.MakeSlice(typ, 1, 1)
|
||||
value.Index(0).Set(chatMessagePartStringContainerValue(t, typ.Elem(), probe))
|
||||
return value
|
||||
case reflect.Array:
|
||||
if typ.Len() == 0 {
|
||||
t.Fatalf("cannot inject a string probe into zero-length string container %s", typ)
|
||||
}
|
||||
value := reflect.New(typ).Elem()
|
||||
value.Index(0).Set(chatMessagePartStringContainerValue(t, typ.Elem(), probe))
|
||||
return value
|
||||
default:
|
||||
t.Fatalf("unsupported recursive string container %s", typ)
|
||||
return reflect.Value{}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user