| 1 | package extension |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "crypto/sha256" |
| 6 | "encoding/base64" |
| 7 | "encoding/hex" |
| 8 | "encoding/json" |
| 9 | "errors" |
| 10 | "fmt" |
| 11 | "strings" |
| 12 | "sync" |
| 13 | "testing" |
| 14 | "time" |
| 15 | ) |
| 16 | |
| 17 | // contentStore scripts the host's content store for host/content/read. |
| 18 | type contentStore struct { |
| 19 | mu sync.Mutex |
| 20 | objects map[string][]byte |
| 21 | requests []ContentReadParams |
| 22 | tamper bool // flip one byte in the first served chunk |
| 23 | } |
| 24 | |
| 25 | func newContentStore() *contentStore { |
| 26 | return &contentStore{objects: make(map[string][]byte)} |
| 27 | } |
| 28 | |
| 29 | func (s *contentStore) put(data string) (ref string, descriptor ExternalizedField) { |
| 30 | sum := sha256.Sum256([]byte(data)) |
| 31 | ref = "content_test_" + hex.EncodeToString(sum[:4]) |
| 32 | s.mu.Lock() |
| 33 | s.objects[ref] = []byte(data) |
| 34 | s.mu.Unlock() |
| 35 | return ref, ExternalizedField{ |
| 36 | JSONPointer: "/payload", ContentRef: ref, |
| 37 | TotalBytes: int64(len(data)), SHA256: hex.EncodeToString(sum[:]), |
| 38 | } |
| 39 | } |
| 40 | |
| 41 | // handler pages like the real host: at most ContentRefChunkBytes per answer, |
| 42 | // NextOffset null at the end, content_ref_expired for unknown refs. |
| 43 | func (s *contentStore) handler(params json.RawMessage) (any, *hostError) { |
| 44 | var p ContentReadParams |
| 45 | if err := json.Unmarshal(params, &p); err != nil { |
| 46 | return nil, &hostError{Code: CodeInvalidParams, Message: "bad params"} |
| 47 | } |
| 48 | s.mu.Lock() |
| 49 | s.requests = append(s.requests, p) |
| 50 | data, ok := s.objects[p.ContentRef] |
| 51 | tamper := s.tamper |
| 52 | s.mu.Unlock() |
| 53 | if !ok { |
| 54 | return nil, &hostError{ |
| 55 | Code: DomainErrorCode, |
| 56 | Message: "The referenced content has expired.", |
| 57 | Data: ProtocolErrorData{Reason: ErrContentRefExpired, Retryable: true}, |
| 58 | } |
| 59 | } |
| 60 | if p.Offset < 0 || p.Offset > int64(len(data)) { |
| 61 | return nil, &hostError{ |
| 62 | Code: DomainErrorCode, |
| 63 | Message: "The referenced content has expired.", |
| 64 | Data: ProtocolErrorData{Reason: ErrContentRefExpired, Retryable: true}, |
| 65 | } |
| 66 | } |
| 67 | end := p.Offset + ContentRefChunkBytes |
| 68 | if end > int64(len(data)) { |
| 69 | end = int64(len(data)) |
| 70 | } |
| 71 | chunk := append([]byte(nil), data[p.Offset:end]...) |
| 72 | if tamper && p.Offset == 0 && len(chunk) > 0 { |
| 73 | chunk[0] ^= 0xFF |
| 74 | } |
| 75 | var next *int64 |
| 76 | if end < int64(len(data)) { |
| 77 | value := end |
| 78 | next = &value |
| 79 | } |
| 80 | sum := sha256.Sum256(data) |
| 81 | return ContentReadResult{ |
| 82 | ContentRef: p.ContentRef, Offset: p.Offset, |
| 83 | DataBase64: base64.StdEncoding.EncodeToString(chunk), |
| 84 | NextOffset: next, TotalBytes: int64(len(data)), |
| 85 | SHA256: hex.EncodeToString(sum[:]), Encoding: ContentUTF8, |
| 86 | }, nil |
| 87 | } |
| 88 | |
| 89 | func (s *contentStore) requestedOffsets() []int64 { |
| 90 | s.mu.Lock() |
| 91 | defer s.mu.Unlock() |
| 92 | var out []int64 |
| 93 | for _, r := range s.requests { |
| 94 | out = append(out, r.Offset) |
| 95 | } |
| 96 | return out |
| 97 | } |
| 98 | |
| 99 | // TestReadContentRefMultiChunk reads a payload spanning several chunks and |
| 100 | // verifies paging offsets and SHA-256. |
| 101 | func TestReadContentRefMultiChunk(t *testing.T) { |
| 102 | store := newContentStore() |
| 103 | big := strings.Repeat("abcdefghij", ContentRefChunkBytes/4) // exactly 2.5 chunks → 3 pages |
| 104 | ref, _ := store.put(big) |
| 105 | var data []byte |
| 106 | var readErr error |
| 107 | hook := map[string]InterceptorFunc{ |
| 108 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 109 | data, readErr = ReadContentRef(ctx, ref) |
| 110 | return Continue(), nil |
| 111 | }, |
| 112 | } |
| 113 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: hook}) |
| 114 | host.onRequest(MethodHostContentRead, store.handler) |
| 115 | host.handshake(t) |
| 116 | host.request(MethodExtensionIntercept, InterceptParams{ |
| 117 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 118 | }) |
| 119 | if readErr != nil { |
| 120 | t.Fatalf("ReadContentRef: %v", readErr) |
| 121 | } |
| 122 | if string(data) != big { |
| 123 | t.Fatalf("reassembled %d bytes, want %d identical bytes", len(data), len(big)) |
| 124 | } |
| 125 | offsets := store.requestedOffsets() |
| 126 | if len(offsets) != 3 { |
| 127 | t.Fatalf("read offsets = %v, want 3 pages", offsets) |
| 128 | } |
| 129 | for i, offset := range offsets { |
| 130 | if offset != int64(i)*ContentRefChunkBytes { |
| 131 | t.Fatalf("offset %d = %d, want %d", i, offset, int64(i)*ContentRefChunkBytes) |
| 132 | } |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | // TestReadContentRefTamper detects a SHA-256 mismatch. |
| 137 | func TestReadContentRefTamper(t *testing.T) { |
| 138 | store := newContentStore() |
| 139 | store.tamper = true |
| 140 | ref, _ := store.put(strings.Repeat("x", ContentRefChunkBytes+10)) |
| 141 | var readErr error |
| 142 | interceptors := map[string]InterceptorFunc{ |
| 143 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 144 | _, readErr = ReadContentRef(ctx, ref) |
| 145 | return Continue(), nil |
| 146 | }, |
| 147 | } |
| 148 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 149 | host.onRequest(MethodHostContentRead, store.handler) |
| 150 | host.handshake(t) |
| 151 | host.request(MethodExtensionIntercept, InterceptParams{ |
| 152 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 153 | }) |
| 154 | var protocolErr *ProtocolError |
| 155 | if !errors.As(readErr, &protocolErr) || !strings.Contains(protocolErr.Message, "SHA-256") { |
| 156 | t.Fatalf("readErr = %v, want SHA-256 mismatch protocol error", readErr) |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | // TestReadContentRefExpired maps the wire reason to a *ProtocolError. |
| 161 | func TestReadContentRefExpired(t *testing.T) { |
| 162 | store := newContentStore() |
| 163 | var readErr error |
| 164 | interceptors := map[string]InterceptorFunc{ |
| 165 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 166 | _, readErr = ReadContentRef(ctx, "content_gone") |
| 167 | return Continue(), nil |
| 168 | }, |
| 169 | } |
| 170 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 171 | host.onRequest(MethodHostContentRead, store.handler) |
| 172 | host.handshake(t) |
| 173 | host.request(MethodExtensionIntercept, InterceptParams{ |
| 174 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 175 | }) |
| 176 | var protocolErr *ProtocolError |
| 177 | if !errors.As(readErr, &protocolErr) { |
| 178 | t.Fatalf("readErr = %v, want *ProtocolError", readErr) |
| 179 | } |
| 180 | if protocolErr.Reason != ErrContentRefExpired { |
| 181 | t.Fatalf("reason = %q, want content_ref_expired", protocolErr.Reason) |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | // TestInterceptExternalizedPayload runs the full transparent rehydration: |
| 186 | // the host sends payload:null plus the externalized envelope, and the |
| 187 | // interceptor receives the reassembled bytes. |
| 188 | func TestInterceptExternalizedPayload(t *testing.T) { |
| 189 | store := newContentStore() |
| 190 | big := `{"text":"` + strings.Repeat("lorem ", ContentRefChunkBytes/3) + `"}` |
| 191 | _, descriptor := store.put(big) |
| 192 | var got json.RawMessage |
| 193 | interceptors := map[string]InterceptorFunc{ |
| 194 | "input.receive": func(_ context.Context, _ string, payload json.RawMessage) (*InterceptResult, error) { |
| 195 | got = payload |
| 196 | return Continue(), nil |
| 197 | }, |
| 198 | } |
| 199 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 200 | host.onRequest(MethodHostContentRead, store.handler) |
| 201 | host.handshake(t) |
| 202 | resp := host.request(MethodExtensionIntercept, map[string]any{ |
| 203 | "event": "input.receive", "seq": 1, "payload": nil, "timeoutMillis": 0, |
| 204 | "externalized": []ExternalizedField{descriptor}, |
| 205 | }) |
| 206 | if resp.Err != nil { |
| 207 | t.Fatalf("intercept failed: %+v", resp.Err) |
| 208 | } |
| 209 | if string(got) != big { |
| 210 | t.Fatalf("payload = %d bytes, want rehydrated %d bytes", len(got), len(big)) |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | // TestInterceptExternalizedViolation rejects an inline payload alongside an |
| 215 | // envelope. |
| 216 | func TestInterceptExternalizedViolation(t *testing.T) { |
| 217 | interceptors := map[string]InterceptorFunc{ |
| 218 | "*": func(context.Context, string, json.RawMessage) (*InterceptResult, error) { return Continue(), nil }, |
| 219 | } |
| 220 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 221 | host.handshake(t) |
| 222 | resp := host.request(MethodExtensionIntercept, map[string]any{ |
| 223 | "event": "input.receive", "seq": 1, "payload": json.RawMessage(`{"x":1}`), "timeoutMillis": 0, |
| 224 | "externalized": []ExternalizedField{{ |
| 225 | JSONPointer: "/payload", ContentRef: "content_fake", |
| 226 | TotalBytes: 7, SHA256: strings.Repeat("0", 64), |
| 227 | }}, |
| 228 | }) |
| 229 | if resp.Err == nil { |
| 230 | t.Fatal("expected a protocol error for inline payload plus envelope") |
| 231 | } |
| 232 | data, _ := resp.Err.Data.(ProtocolErrorData) |
| 233 | if data.Reason != ErrProtocolError { |
| 234 | t.Fatalf("reason = %q, want protocol_error", data.Reason) |
| 235 | } |
| 236 | } |
| 237 | |
| 238 | // TestResolveExternalizedHelper covers the exported helper directly, |
| 239 | // including the pointer check. |
| 240 | func TestResolveExternalizedHelper(t *testing.T) { |
| 241 | store := newContentStore() |
| 242 | _, descriptor := store.put(`{"hello":"world"}`) |
| 243 | var resolved json.RawMessage |
| 244 | var resolveErr error |
| 245 | interceptors := map[string]InterceptorFunc{ |
| 246 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 247 | resolved, resolveErr = ResolveExternalized(ctx, nil, []ExternalizedField{descriptor}, "/payload") |
| 248 | return Continue(), nil |
| 249 | }, |
| 250 | } |
| 251 | host, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors}) |
| 252 | host.onRequest(MethodHostContentRead, store.handler) |
| 253 | host.handshake(t) |
| 254 | host.request(MethodExtensionIntercept, InterceptParams{ |
| 255 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 256 | }) |
| 257 | if resolveErr != nil { |
| 258 | t.Fatalf("ResolveExternalized: %v", resolveErr) |
| 259 | } |
| 260 | if string(resolved) != `{"hello":"world"}` { |
| 261 | t.Fatalf("resolved = %s", resolved) |
| 262 | } |
| 263 | |
| 264 | // A wrong pointer must fail without any content read. |
| 265 | before := len(store.requestedOffsets()) |
| 266 | var wrongErr error |
| 267 | interceptors2 := map[string]InterceptorFunc{ |
| 268 | "tool.before": func(ctx context.Context, _ string, _ json.RawMessage) (*InterceptResult, error) { |
| 269 | _, wrongErr = ResolveExternalized(ctx, nil, []ExternalizedField{descriptor}, "/replacement") |
| 270 | return Continue(), nil |
| 271 | }, |
| 272 | } |
| 273 | host2, _ := startFakeHost(t, basicHandler(), Options{Interceptors: interceptors2}) |
| 274 | host2.onRequest(MethodHostContentRead, store.handler) |
| 275 | host2.handshake(t) |
| 276 | host2.request(MethodExtensionIntercept, InterceptParams{ |
| 277 | Event: EventToolBefore, Seq: 1, Payload: json.RawMessage(`{}`), |
| 278 | }) |
| 279 | var protocolErr *ProtocolError |
| 280 | if !errors.As(wrongErr, &protocolErr) || protocolErr.Reason != ErrProtocolError { |
| 281 | t.Fatalf("wrongErr = %v, want protocol_error", wrongErr) |
| 282 | } |
| 283 | if got := len(store.requestedOffsets()); got != before { |
| 284 | t.Fatalf("content reads happened despite the pointer violation: %d → %d", before, got) |
| 285 | } |
| 286 | } |
| 287 | |
| 288 | // TestResolveExternalizedNoConnection requires an SDK callback context. |
| 289 | func TestResolveExternalizedNoConnection(t *testing.T) { |
| 290 | if _, err := ResolveExternalized(context.Background(), nil, nil, "/payload"); !errors.Is(err, ErrNoConnection) { |
| 291 | t.Fatalf("err = %v, want ErrNoConnection", err) |
| 292 | } |
| 293 | if _, err := ReadContentRef(context.Background(), "content_x"); !errors.Is(err, ErrNoConnection) { |
| 294 | t.Fatalf("err = %v, want ErrNoConnection", err) |
| 295 | } |
| 296 | } |
| 297 | |
| 298 | // TestEventExternalizedPayload rehydrates event payloads too. |
| 299 | func TestEventExternalizedPayload(t *testing.T) { |
| 300 | store := newContentStore() |
| 301 | big := fmt.Sprintf(`{"blob":"%s"}`, strings.Repeat("z", ContentRefChunkBytes+100)) |
| 302 | _, descriptor := store.put(big) |
| 303 | seen := make(chan json.RawMessage, 1) |
| 304 | opts := Options{Observer: func(_ context.Context, _ string, payload json.RawMessage) { seen <- payload }} |
| 305 | host, _ := startFakeHost(t, basicHandler(), opts) |
| 306 | host.onRequest(MethodHostContentRead, store.handler) |
| 307 | host.handshake(t) |
| 308 | host.notify(MethodExtensionEvent, map[string]any{ |
| 309 | "event": "session.end", "payload": nil, "externalized": []ExternalizedField{descriptor}, |
| 310 | }) |
| 311 | select { |
| 312 | case payload := <-seen: |
| 313 | if string(payload) != big { |
| 314 | t.Fatalf("payload = %d bytes, want %d", len(payload), len(big)) |
| 315 | } |
| 316 | case <-time.After(5 * time.Second): |
| 317 | t.Fatal("observer not called for the externalized event") |
| 318 | } |
| 319 | } |
| 320 |