From 312f8a338d69005ee7c12b787ebdd414ed941d4c Mon Sep 17 00:00:00 2001 From: salmonumbrella <182032677+salmonumbrella@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:37:15 +0000 Subject: [PATCH] feat(web): add document processing and rendition views --- .gitattributes | 8 +- .roborev.toml | 10 + Makefile | 5 +- cmd/docbank/backfill_retry.go | 56 + cmd/docbank/backfill_retry_test.go | 38 + cmd/docbank/cli_test.go | 40 +- cmd/docbank/daemon.go | 267 +++ cmd/docbank/daemon_embedding_worker_test.go | 211 ++ cmd/docbank/daemon_processing_runtime_test.go | 130 ++ cmd/docbank/daemon_test.go | 19 + cmd/docbank/daemon_vector_index_test.go | 31 + cmd/docbank/embedding_runtime.go | 328 +++ cmd/docbank/processing.go | 307 +++ cmd/docbank/processing_test.go | 147 ++ cmd/docbank/rendition.go | 63 + cmd/docbank/rendition_runtime.go | 31 + cmd/docbank/rendition_test.go | 74 + cmd/docbank/search.go | 216 +- cmd/docbank/search_modes_test.go | 79 + cmd/docbank/stat.go | 39 + cmd/docbank/stat_test.go | 47 + cmd/docbank/tokenizer_runtime.go | 38 + cmd/docbank/versions.go | 3 + cmd/docbank/versions_test.go | 29 + document/bridge/client.go | 749 +++++++ document/bridge/client_test.go | 1036 +++++++++ document/bridge/contract_test.go | 45 + document/bridge/openapi.yaml | 218 ++ .../bridge/source-evidence-v1.schema.json | 153 ++ document/bridge/testdata/unknown-major.json | 1 + document/bridge/types.go | 135 ++ document/chunking.go | 1679 +++++++++++++++ document/chunking_test.go | 984 +++++++++ document/cohereembed/client.go | 607 ++++++ document/cohereembed/client_test.go | 478 +++++ document/cohereembed/errors.go | 53 + document/cohereembed/errors_test.go | 19 + document/cohereembed/lifecycle_test.go | 237 +++ document/cohereembed/profile.go | 328 +++ document/cohereembed/profile_test.go | 191 ++ document/datalab/client.go | 915 ++++++++ document/datalab/client_test.go | 497 +++++ .../datalab/testdata/convert-complete.json | 28 + .../testdata/convert-schema-drift.json | 9 + document/doc.go | 12 +- document/docling/client.go | 850 ++++++++ document/docling/client_test.go | 918 ++++++++ document/docling/testdata/docling-pages.json | 21 + .../testdata/docling-schema-drift.json | 12 + document/embedding.go | 576 +++++ document/embedding_input_test.go | 22 + document/embedding_test.go | 350 ++++ document/embeddingbridge/client.go | 471 +++++ document/embeddingbridge/client_test.go | 1105 ++++++++++ document/embeddingbridge/contract_test.go | 443 ++++ .../embedding-request-v1.schema.json | 195 ++ .../embedding-response-v1.schema.json | 51 + document/embeddingbridge/errors.go | 51 + document/embeddingbridge/openapi.yaml | 75 + document/embeddingbridge/profile.go | 316 +++ document/embeddingbridge/source_gate_test.go | 40 + document/embeddingbridge/types.go | 103 + document/evidence_codec.go | 4 +- document/evidence_codec_test.go | 35 + document/execution.go | 440 ++++ document/execution_test.go | 148 ++ document/geminiembed/client.go | 1101 ++++++++++ document/geminiembed/client_test.go | 290 +++ document/geminiembed/errors.go | 79 + document/geminiembed/errors_test.go | 36 + document/geminiembed/lifecycle_test.go | 1408 +++++++++++++ document/geminiembed/profile.go | 400 ++++ document/geminiembed/profile_test.go | 130 ++ document/glmocr/client.go | 26 +- document/glmocr/client_internal_test.go | 44 + document/glmocr/client_test.go | 2 +- document/input_contract.go | 315 +++ document/input_contract_test.go | 130 ++ document/internal/formatdetect/detect.go | 1862 +++++++++++++++++ document/internal/formatdetect/formats.go | 78 + document/llamaparse/client.go | 1224 +++++++++++ document/llamaparse/client_test.go | 957 +++++++++ document/llamaparse/doc.go | 3 + document/marker/client.go | 755 +++++++ document/marker/client_test.go | 530 +++++ document/media/detect.go | 303 ++- document/media/detect_test.go | 69 + document/media/doc.go | 16 +- document/media/inspect.go | 938 +++++++++ document/media/inspect_test.go | 1166 +++++++++++ document/media/mediatest/codec.go | 77 + document/media/mediatest/mediatest.go | 49 + document/media/mp4_codec.go | 162 +- document/media/policy.go | 2 +- document/media/types.go | 4 + document/mistral/client.go | 80 +- document/mistral/embedding.go | 515 +++++ document/mistral/embedding_review_test.go | 47 + document/mistral/embedding_test.go | 335 +++ document/mistral/formats.go | 48 +- document/mistral/rendition.go | 482 +++++ document/mistral/rendition_test.go | 570 +++++ document/mistral/sniff.go | 1069 +--------- document/mistral/sniff_test.go | 14 + document/mistral/test_helpers_test.go | 35 +- document/normalize.go | 340 +-- document/openaiembed/client.go | 720 +++++++ document/openaiembed/client_test.go | 640 ++++++ document/openaiembed/profiles.go | 193 ++ document/openaiembed/profiles_test.go | 170 ++ .../openaiembed/testdata/schema-drift.json | 1 + .../openaiembed/testdata/success-indexed.json | 1 + document/openaihosted/client.go | 210 ++ document/openaihosted/client_test.go | 269 +++ document/openaihosted/errors.go | 72 + document/openaihosted/profile.go | 319 +++ document/openaihosted/profile_test.go | 317 +++ document/plaintext/provider.go | 220 ++ document/plaintext/provider_test.go | 198 ++ document/profile.go | 48 +- document/profile_test.go | 32 +- document/provider.go | 1014 +++++++++ document/provider_test.go | 687 ++++++ document/providerhttp/policy.go | 154 ++ document/providerhttp/transport.go | 252 +++ document/providerhttp/transport_test.go | 450 ++++ document/pymupdf/doc.go | 2 + document/pymupdf/provider.go | 494 +++++ document/pymupdf/provider_test.go | 387 ++++ document/pymupdf/testdata/helper/main.go | 136 ++ document/reducto/client.go | 1342 ++++++++++++ document/reducto/client_test.go | 1045 +++++++++ document/reducto/doc.go | 3 + document/rendition_frontmatter.go | 338 +++ document/rendition_frontmatter_test.go | 85 + document/rendition_normalize_test.go | 2 +- document/source_metadata.go | 93 + document/source_metadata_codec.go | 356 ++++ document/source_metadata_test.go | 125 ++ .../testdata/chunks-bge-m3-v1.golden.json | 66 + document/testdata/chunks-e5-v1.golden.json | 66 + document/testdata/chunks-gte-v1.golden.json | 66 + document/testdata/chunks-nomic-v1.golden.json | 66 + document/testdata/chunks-qwen3-v1.golden.json | 66 + .../embedding-input-generation-v1.golden.json | 42 + document/testdata/profile-v1.golden.json | 2 +- .../testdata/source-metadata-v1.golden.json | 1 + document/testdata/vector-set-v1.golden.bin | Bin 0 -> 293 bytes document/tika/profile.go | 317 +++ document/tika/profile_test.go | 192 ++ document/tokenizer.go | 72 + document/trafilatura/doc.go | 16 + document/trafilatura/native_launcher_linux.go | 174 ++ document/trafilatura/native_runner_linux.go | 286 +++ .../trafilatura/native_runner_linux_test.go | 389 ++++ .../trafilatura/native_runner_unsupported.go | 7 + document/trafilatura/provider.go | 939 +++++++++ document/trafilatura/provider_test.go | 390 ++++ document/trafilatura/runner_contract_test.go | 341 +++ .../testdata/isolatedhelper/main.go | 92 + document/unstructured/profile.go | 292 +++ document/unstructured/profile_test.go | 161 ++ document/upload/doc.go | 5 + document/upload/spool.go | 384 ++++ document/upload/spool_test.go | 445 ++++ document/upload/spool_unix.go | 116 + document/upload/spool_windows.go | 106 + document/upload/spool_windows_test.go | 47 + document/vector_set.go | 411 ++++ document/vector_set_test.go | 125 ++ document/voyage/client.go | 3 +- document/voyage/embedding.go | 740 +++++++ document/voyage/embedding_review_test.go | 151 ++ document/voyage/embedding_test.go | 350 ++++ document/zeroentropyembed/client.go | 308 +++ document/zeroentropyembed/client_test.go | 218 ++ document/zeroentropyembed/errors.go | 70 + document/zeroentropyembed/profile.go | 360 ++++ document/zeroentropyembed/profile_test.go | 142 ++ frontend/package-lock.json | 15 +- frontend/package.json | 3 +- frontend/screenshots/README.md | 9 +- frontend/screenshots/processing-profile.go | 107 + frontend/screenshots/web-trash.screenshot.ts | 174 +- frontend/src/App.processing.test.ts | 40 + frontend/src/App.svelte | 55 + frontend/src/ProcessingDrawer.svelte | 298 +++ frontend/src/ProcessingDrawer.test.ts | 190 ++ frontend/src/RenditionDrawer.svelte | 147 ++ frontend/src/RenditionDrawer.test.ts | 66 + frontend/src/api.processing.test.ts | 81 + frontend/src/api.ts | 355 ++++ go.mod | 5 +- go.sum | 2 + internal/api/gate.go | 25 +- internal/api/gate_test.go | 57 + internal/api/middleware.go | 1 + internal/api/openapi_test.go | 15 + internal/api/routes_processing.go | 521 +++++ internal/api/routes_processing_test.go | 628 ++++++ internal/api/routes_read.go | 28 +- internal/api/routes_read_test.go | 47 +- internal/api/server.go | 3 + internal/api/types.go | 314 ++- internal/api/web_session.go | 17 +- internal/api/web_session_processing_test.go | 38 + internal/backupapp/app.go | 75 +- internal/backupapp/app_internal_test.go | 43 + internal/backupapp/app_test.go | 790 +++---- internal/backupapp/metadata.go | 11 +- internal/backupapp/restore.go | 50 +- .../restore_placement_internal_test.go | 24 - internal/blob/blob.go | 32 +- internal/blob/blob_test.go | 24 + internal/client/client.go | 54 +- internal/client/client_test.go | 31 + internal/client/processing.go | 448 ++++ internal/client/processing_test.go | 170 ++ internal/cohereapi/protocol.go | 105 + internal/cohereapi/protocol_test.go | 84 + internal/config/config.go | 253 ++- internal/config/config_test.go | 92 + internal/ingest/ingest.go | 2 +- internal/processing/artifacts.go | 99 +- internal/processing/artifacts_test.go | 84 +- internal/processing/embedding_runtime.go | 166 ++ internal/processing/embedding_worker.go | 814 +++++++ internal/processing/embedding_worker_test.go | 1062 ++++++++++ internal/processing/index_worker.go | 14 + internal/processing/index_worker_test.go | 512 +++++ internal/processing/rendition_worker.go | 1000 +++++++++ internal/processing/rendition_worker_test.go | 1107 ++++++++++ internal/processing/service.go | 1707 +++++++++++++++ internal/processing/service_test.go | 92 + internal/processing/source_metadata.go | 855 ++++++++ internal/processing/source_metadata_test.go | 319 +++ internal/processing/test_gate_test.go | 28 + internal/qmdexport/export.go | 575 +++++ internal/qmdexport/export_test.go | 332 +++ internal/qmdexport/load.go | 105 + internal/qmdexport/replace_unix.go | 7 + internal/qmdexport/replace_windows.go | 17 + internal/retrieval/coherererank/client.go | 305 +++ .../retrieval/coherererank/client_test.go | 228 ++ internal/retrieval/coherererank/errors.go | 53 + .../retrieval/coherererank/errors_test.go | 19 + internal/retrieval/coherererank/profile.go | 280 +++ .../retrieval/coherererank/profile_test.go | 108 + internal/retrieval/fusion.go | 119 ++ internal/retrieval/providers.go | 452 ++++ internal/retrieval/providers_test.go | 68 + internal/retrieval/qmdbridge/client.go | 269 +++ internal/retrieval/qmdbridge/client_test.go | 298 +++ internal/retrieval/qmdbridge/profile.go | 186 ++ internal/retrieval/search.go | 617 ++++++ internal/retrieval/search_stages_test.go | 512 +++++ internal/retrieval/search_test.go | 556 +++++ internal/retrieval/types.go | 142 ++ .../retrieval/zeroentropyrerank/client.go | 241 +++ .../zeroentropyrerank/client_test.go | 216 ++ .../retrieval/zeroentropyrerank/errors.go | 70 + .../retrieval/zeroentropyrerank/profile.go | 289 +++ .../zeroentropyrerank/profile_test.go | 104 + internal/store/audit_content_replace.go | 2 +- internal/store/consent.go | 437 ++++ internal/store/consent_test.go | 215 ++ internal/store/content_checksums.go | 155 ++ internal/store/content_checksums_test.go | 75 + internal/store/embedding_catalog.go | 1840 ++++++++++++++++ .../store/embedding_catalog_authority_test.go | 404 ++++ .../store/embedding_catalog_integrity_test.go | 547 +++++ internal/store/embedding_catalog_test.go | 1351 ++++++++++++ internal/store/embedding_jobs.go | 520 +++++ internal/store/embedding_metadata.go | 635 ++++++ internal/store/embedding_reconcile.go | 225 ++ internal/store/extraction.go | 18 +- internal/store/gc.go | 401 +++- internal/store/gc_test.go | 197 +- internal/store/ingest.go | 2 +- internal/store/metadata.go | 442 +++- internal/store/metadata_test.go | 259 ++- internal/store/node.go | 4 +- internal/store/pack_catalog.go | 5 + internal/store/pack_catalog_test.go | 6 +- internal/store/physical_storage.go | 8 + internal/store/processing_catalog.go | 268 ++- internal/store/processing_catalog_test.go | 292 ++- internal/store/processing_metadata.go | 1389 ++++++++++-- internal/store/processing_migration.go | 289 ++- internal/store/processing_migration_test.go | 159 +- internal/store/qmd_export.go | 183 ++ internal/store/qmd_export_test.go | 99 + internal/store/rendition_delete.go | 269 +++ internal/store/rendition_jobs.go | 1406 +++++++++++++ internal/store/rendition_jobs_test.go | 1053 ++++++++++ internal/store/schema.sql | 577 ++++- internal/store/search.go | 1039 ++++++++- internal/store/search_test.go | 330 ++- internal/store/source_metadata.go | 207 ++ internal/store/source_metadata_test.go | 129 ++ internal/store/store.go | 8 +- internal/store/store_registry_test.go | 63 + internal/store/trash_test.go | 7 +- internal/store/upgrade.go | 30 +- internal/store/upgrade_test.go | 163 +- internal/store/vector_index.go | 530 +++++ internal/store/vector_index_test.go | 248 +++ internal/store/version.go | 22 +- internal/store/version_prune_test.go | 203 +- internal/store/walk.go | 2 +- internal/store/write.go | 30 +- internal/vectorindex/generation.go | 888 ++++++++ internal/vectorindex/generation_test.go | 134 ++ internal/vectorindex/public_test.go | 51 + .../vectorindex/selection_benchmark_test.go | 348 +++ .../testdata/vector-index-v1.golden.bin | Bin 0 -> 615 bytes internal/vectorindex/validation_test.go | 290 +++ internal/vectorworker/worker.go | 553 +++++ maintenance_test.go | 130 +- processing.go | 176 ++ types.go | 228 ++ vault.go | 69 +- vault_external_test.go | 350 ++++ 323 files changed, 87540 insertions(+), 2882 deletions(-) create mode 100644 cmd/docbank/backfill_retry.go create mode 100644 cmd/docbank/backfill_retry_test.go create mode 100644 cmd/docbank/daemon_embedding_worker_test.go create mode 100644 cmd/docbank/daemon_processing_runtime_test.go create mode 100644 cmd/docbank/daemon_vector_index_test.go create mode 100644 cmd/docbank/embedding_runtime.go create mode 100644 cmd/docbank/processing.go create mode 100644 cmd/docbank/processing_test.go create mode 100644 cmd/docbank/rendition.go create mode 100644 cmd/docbank/rendition_runtime.go create mode 100644 cmd/docbank/rendition_test.go create mode 100644 cmd/docbank/search_modes_test.go create mode 100644 cmd/docbank/tokenizer_runtime.go create mode 100644 cmd/docbank/versions_test.go create mode 100644 document/bridge/client.go create mode 100644 document/bridge/client_test.go create mode 100644 document/bridge/contract_test.go create mode 100644 document/bridge/openapi.yaml create mode 100644 document/bridge/source-evidence-v1.schema.json create mode 100644 document/bridge/testdata/unknown-major.json create mode 100644 document/bridge/types.go create mode 100644 document/chunking.go create mode 100644 document/chunking_test.go create mode 100644 document/cohereembed/client.go create mode 100644 document/cohereembed/client_test.go create mode 100644 document/cohereembed/errors.go create mode 100644 document/cohereembed/errors_test.go create mode 100644 document/cohereembed/lifecycle_test.go create mode 100644 document/cohereembed/profile.go create mode 100644 document/cohereembed/profile_test.go create mode 100644 document/datalab/client.go create mode 100644 document/datalab/client_test.go create mode 100644 document/datalab/testdata/convert-complete.json create mode 100644 document/datalab/testdata/convert-schema-drift.json create mode 100644 document/docling/client.go create mode 100644 document/docling/client_test.go create mode 100644 document/docling/testdata/docling-pages.json create mode 100644 document/docling/testdata/docling-schema-drift.json create mode 100644 document/embedding.go create mode 100644 document/embedding_input_test.go create mode 100644 document/embedding_test.go create mode 100644 document/embeddingbridge/client.go create mode 100644 document/embeddingbridge/client_test.go create mode 100644 document/embeddingbridge/contract_test.go create mode 100644 document/embeddingbridge/embedding-request-v1.schema.json create mode 100644 document/embeddingbridge/embedding-response-v1.schema.json create mode 100644 document/embeddingbridge/errors.go create mode 100644 document/embeddingbridge/openapi.yaml create mode 100644 document/embeddingbridge/profile.go create mode 100644 document/embeddingbridge/source_gate_test.go create mode 100644 document/embeddingbridge/types.go create mode 100644 document/execution.go create mode 100644 document/execution_test.go create mode 100644 document/geminiembed/client.go create mode 100644 document/geminiembed/client_test.go create mode 100644 document/geminiembed/errors.go create mode 100644 document/geminiembed/errors_test.go create mode 100644 document/geminiembed/lifecycle_test.go create mode 100644 document/geminiembed/profile.go create mode 100644 document/geminiembed/profile_test.go create mode 100644 document/glmocr/client_internal_test.go create mode 100644 document/input_contract.go create mode 100644 document/input_contract_test.go create mode 100644 document/internal/formatdetect/detect.go create mode 100644 document/internal/formatdetect/formats.go create mode 100644 document/llamaparse/client.go create mode 100644 document/llamaparse/client_test.go create mode 100644 document/llamaparse/doc.go create mode 100644 document/marker/client.go create mode 100644 document/marker/client_test.go create mode 100644 document/media/inspect.go create mode 100644 document/media/inspect_test.go create mode 100644 document/media/mediatest/codec.go create mode 100644 document/mistral/embedding.go create mode 100644 document/mistral/embedding_review_test.go create mode 100644 document/mistral/embedding_test.go create mode 100644 document/mistral/rendition.go create mode 100644 document/mistral/rendition_test.go create mode 100644 document/openaiembed/client.go create mode 100644 document/openaiembed/client_test.go create mode 100644 document/openaiembed/profiles.go create mode 100644 document/openaiembed/profiles_test.go create mode 100644 document/openaiembed/testdata/schema-drift.json create mode 100644 document/openaiembed/testdata/success-indexed.json create mode 100644 document/openaihosted/client.go create mode 100644 document/openaihosted/client_test.go create mode 100644 document/openaihosted/errors.go create mode 100644 document/openaihosted/profile.go create mode 100644 document/openaihosted/profile_test.go create mode 100644 document/plaintext/provider.go create mode 100644 document/plaintext/provider_test.go create mode 100644 document/provider.go create mode 100644 document/provider_test.go create mode 100644 document/providerhttp/policy.go create mode 100644 document/providerhttp/transport.go create mode 100644 document/providerhttp/transport_test.go create mode 100644 document/pymupdf/doc.go create mode 100644 document/pymupdf/provider.go create mode 100644 document/pymupdf/provider_test.go create mode 100644 document/pymupdf/testdata/helper/main.go create mode 100644 document/reducto/client.go create mode 100644 document/reducto/client_test.go create mode 100644 document/reducto/doc.go create mode 100644 document/rendition_frontmatter.go create mode 100644 document/rendition_frontmatter_test.go create mode 100644 document/source_metadata.go create mode 100644 document/source_metadata_codec.go create mode 100644 document/source_metadata_test.go create mode 100644 document/testdata/chunks-bge-m3-v1.golden.json create mode 100644 document/testdata/chunks-e5-v1.golden.json create mode 100644 document/testdata/chunks-gte-v1.golden.json create mode 100644 document/testdata/chunks-nomic-v1.golden.json create mode 100644 document/testdata/chunks-qwen3-v1.golden.json create mode 100644 document/testdata/embedding-input-generation-v1.golden.json create mode 100644 document/testdata/source-metadata-v1.golden.json create mode 100644 document/testdata/vector-set-v1.golden.bin create mode 100644 document/tika/profile.go create mode 100644 document/tika/profile_test.go create mode 100644 document/tokenizer.go create mode 100644 document/trafilatura/doc.go create mode 100644 document/trafilatura/native_launcher_linux.go create mode 100644 document/trafilatura/native_runner_linux.go create mode 100644 document/trafilatura/native_runner_linux_test.go create mode 100644 document/trafilatura/native_runner_unsupported.go create mode 100644 document/trafilatura/provider.go create mode 100644 document/trafilatura/provider_test.go create mode 100644 document/trafilatura/runner_contract_test.go create mode 100644 document/trafilatura/testdata/isolatedhelper/main.go create mode 100644 document/unstructured/profile.go create mode 100644 document/unstructured/profile_test.go create mode 100644 document/upload/doc.go create mode 100644 document/upload/spool.go create mode 100644 document/upload/spool_test.go create mode 100644 document/upload/spool_unix.go create mode 100644 document/upload/spool_windows.go create mode 100644 document/upload/spool_windows_test.go create mode 100644 document/vector_set.go create mode 100644 document/vector_set_test.go create mode 100644 document/voyage/embedding.go create mode 100644 document/voyage/embedding_review_test.go create mode 100644 document/voyage/embedding_test.go create mode 100644 document/zeroentropyembed/client.go create mode 100644 document/zeroentropyembed/client_test.go create mode 100644 document/zeroentropyembed/errors.go create mode 100644 document/zeroentropyembed/profile.go create mode 100644 document/zeroentropyembed/profile_test.go create mode 100644 frontend/screenshots/processing-profile.go create mode 100644 frontend/src/App.processing.test.ts create mode 100644 frontend/src/ProcessingDrawer.svelte create mode 100644 frontend/src/ProcessingDrawer.test.ts create mode 100644 frontend/src/RenditionDrawer.svelte create mode 100644 frontend/src/RenditionDrawer.test.ts create mode 100644 frontend/src/api.processing.test.ts create mode 100644 internal/api/routes_processing.go create mode 100644 internal/api/routes_processing_test.go create mode 100644 internal/api/web_session_processing_test.go create mode 100644 internal/backupapp/app_internal_test.go create mode 100644 internal/client/processing.go create mode 100644 internal/client/processing_test.go create mode 100644 internal/cohereapi/protocol.go create mode 100644 internal/cohereapi/protocol_test.go create mode 100644 internal/processing/embedding_runtime.go create mode 100644 internal/processing/embedding_worker.go create mode 100644 internal/processing/embedding_worker_test.go create mode 100644 internal/processing/index_worker.go create mode 100644 internal/processing/index_worker_test.go create mode 100644 internal/processing/rendition_worker.go create mode 100644 internal/processing/rendition_worker_test.go create mode 100644 internal/processing/service.go create mode 100644 internal/processing/service_test.go create mode 100644 internal/processing/source_metadata.go create mode 100644 internal/processing/source_metadata_test.go create mode 100644 internal/processing/test_gate_test.go create mode 100644 internal/qmdexport/export.go create mode 100644 internal/qmdexport/export_test.go create mode 100644 internal/qmdexport/load.go create mode 100644 internal/qmdexport/replace_unix.go create mode 100644 internal/qmdexport/replace_windows.go create mode 100644 internal/retrieval/coherererank/client.go create mode 100644 internal/retrieval/coherererank/client_test.go create mode 100644 internal/retrieval/coherererank/errors.go create mode 100644 internal/retrieval/coherererank/errors_test.go create mode 100644 internal/retrieval/coherererank/profile.go create mode 100644 internal/retrieval/coherererank/profile_test.go create mode 100644 internal/retrieval/fusion.go create mode 100644 internal/retrieval/providers.go create mode 100644 internal/retrieval/providers_test.go create mode 100644 internal/retrieval/qmdbridge/client.go create mode 100644 internal/retrieval/qmdbridge/client_test.go create mode 100644 internal/retrieval/qmdbridge/profile.go create mode 100644 internal/retrieval/search.go create mode 100644 internal/retrieval/search_stages_test.go create mode 100644 internal/retrieval/search_test.go create mode 100644 internal/retrieval/types.go create mode 100644 internal/retrieval/zeroentropyrerank/client.go create mode 100644 internal/retrieval/zeroentropyrerank/client_test.go create mode 100644 internal/retrieval/zeroentropyrerank/errors.go create mode 100644 internal/retrieval/zeroentropyrerank/profile.go create mode 100644 internal/retrieval/zeroentropyrerank/profile_test.go create mode 100644 internal/store/consent.go create mode 100644 internal/store/consent_test.go create mode 100644 internal/store/content_checksums.go create mode 100644 internal/store/content_checksums_test.go create mode 100644 internal/store/embedding_catalog.go create mode 100644 internal/store/embedding_catalog_authority_test.go create mode 100644 internal/store/embedding_catalog_integrity_test.go create mode 100644 internal/store/embedding_catalog_test.go create mode 100644 internal/store/embedding_jobs.go create mode 100644 internal/store/embedding_metadata.go create mode 100644 internal/store/embedding_reconcile.go create mode 100644 internal/store/qmd_export.go create mode 100644 internal/store/qmd_export_test.go create mode 100644 internal/store/rendition_jobs.go create mode 100644 internal/store/rendition_jobs_test.go create mode 100644 internal/store/source_metadata.go create mode 100644 internal/store/source_metadata_test.go create mode 100644 internal/store/vector_index.go create mode 100644 internal/store/vector_index_test.go create mode 100644 internal/vectorindex/generation.go create mode 100644 internal/vectorindex/generation_test.go create mode 100644 internal/vectorindex/public_test.go create mode 100644 internal/vectorindex/selection_benchmark_test.go create mode 100644 internal/vectorindex/testdata/vector-index-v1.golden.bin create mode 100644 internal/vectorindex/validation_test.go create mode 100644 internal/vectorworker/worker.go create mode 100644 processing.go diff --git a/.gitattributes b/.gitattributes index c42f4218..338ddd3e 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,8 +1,8 @@ # The compatibility bundle is a byte-for-byte oracle, not a text file. document/internal/compattest/testdata/document-compat-v1.json -text -# The evidence golden bundle is a byte-for-byte oracle, not a text file. +# Document goldens are byte-for-byte protocol, rendering, and vector oracles. document/testdata/normalized-evidence-v1.golden.json -text -# The rendition golden bundle is a byte-for-byte oracle, not a text file. -document/testdata/rendition-v1.golden.md -text -# The processing profile golden bundle is a byte-for-byte oracle, not a text file. document/testdata/profile-v1.golden.json -text +document/testdata/rendition-v1.golden.md -text +document/testdata/source-metadata-v1.golden.json -text +document/testdata/vector-set-v1.golden.bin -text diff --git a/.roborev.toml b/.roborev.toml index b258b074..9c3fbb53 100644 --- a/.roborev.toml +++ b/.roborev.toml @@ -1,3 +1,9 @@ +agent = "codex" +model = "gpt-5.6-terra" +review_agent = "codex" +review_model = "gpt-5.6-terra" +review_reasoning = "high" + review_guidelines = """ docbank is a single-user personal document archive: a local daemon owns the vault (SQLite + content-addressed blobs) and the CLI, agents, and a @@ -138,3 +144,7 @@ Do NOT flag issues that only apply to public-facing, multi-tenant, or internet-exposed services. Focus on bugs, logic errors, data corruption risks, and code quality issues. """ + +# Cadence: reviews are batched, not per-commit. No post-commit hook is +# installed; every 5 commits on a working branch, run a range review: +# roborev review ..HEAD diff --git a/Makefile b/Makefile index 3e15ddf1..2ed16392 100644 --- a/Makefile +++ b/Makefile @@ -93,7 +93,10 @@ install-hooks: docs-install: cd docs && uv sync --frozen -docs-build: +bridge-contract: + go test -tags fts5 ./document/bridge -run '^TestBridgeContractNormativeDocuments' + +docs-build: bridge-contract cd docs && ./zensical-docs.sh build docs-serve: diff --git a/cmd/docbank/backfill_retry.go b/cmd/docbank/backfill_retry.go new file mode 100644 index 00000000..b8903861 --- /dev/null +++ b/cmd/docbank/backfill_retry.go @@ -0,0 +1,56 @@ +package main + +import "time" + +const maxBackfillRetryDelay = 10 * time.Minute + +type backfillRetryState struct { + failures uint + notBefore time.Time +} + +type backfillRetrySet map[string]backfillRetryState + +func newBackfillRetrySet() backfillRetrySet { + return make(backfillRetrySet) +} + +func (retries backfillRetrySet) ready(key string, now time.Time) bool { + state, found := retries[key] + return !found || !now.Before(state.notBefore) +} + +func (retries backfillRetrySet) failed(key string, now time.Time) { + state := retries[key] + state.failures++ + delay := 5 * time.Second + for attempt := uint(1); attempt < state.failures && delay < maxBackfillRetryDelay; attempt++ { + delay = min(delay*2, maxBackfillRetryDelay) + } + state.notBefore = now.Add(delay) + retries[key] = state +} + +func (retries backfillRetrySet) succeeded(key string) { + delete(retries, key) +} + +func (retries backfillRetrySet) waitDelay(now time.Time, maximum time.Duration) time.Duration { + delay := maximum + for _, state := range retries { + if !now.Before(state.notBefore) { + return 0 + } + delay = min(delay, state.notBefore.Sub(now)) + } + return delay +} + +func backfillBatchWaitDelay( + attempted int, retries backfillRetrySet, now time.Time, +) time.Duration { + if attempted != 0 || len(retries) == 0 { + return 100 * time.Millisecond + } + return retries.waitDelay(now, 10*time.Second) +} diff --git a/cmd/docbank/backfill_retry_test.go b/cmd/docbank/backfill_retry_test.go new file mode 100644 index 00000000..413655c3 --- /dev/null +++ b/cmd/docbank/backfill_retry_test.go @@ -0,0 +1,38 @@ +package main + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestBackfillRetrySetQuarantinesFailuresWithoutBlockingOtherTargets(t *testing.T) { + retries := newBackfillRetrySet() + now := time.Date(2026, time.August, 24, 12, 0, 0, 0, time.UTC) + assert.True(t, retries.ready("bad", now)) + assert.True(t, retries.ready("later", now)) + + retries.failed("bad", now) + assert.False(t, retries.ready("bad", now)) + assert.True(t, retries.ready("later", now)) + assert.True(t, retries.ready("bad", now.Add(5*time.Second))) + + retries.failed("bad", now.Add(5*time.Second)) + assert.False(t, retries.ready("bad", now.Add(14*time.Second))) + assert.True(t, retries.ready("bad", now.Add(15*time.Second))) + + retries.succeeded("bad") + assert.True(t, retries.ready("bad", now)) +} + +func TestBackfillBatchWaitHonorsRetryWhenEveryTargetIsSkipped(t *testing.T) { + retries := newBackfillRetrySet() + now := time.Date(2026, time.August, 25, 12, 0, 0, 0, time.UTC) + retries.failed("deferred", now) + + assert.Equal(t, 5*time.Second, + backfillBatchWaitDelay(0, retries, now)) + assert.Equal(t, 100*time.Millisecond, + backfillBatchWaitDelay(1, retries, now)) +} diff --git a/cmd/docbank/cli_test.go b/cmd/docbank/cli_test.go index d5661a88..ef884d97 100644 --- a/cmd/docbank/cli_test.go +++ b/cmd/docbank/cli_test.go @@ -686,13 +686,25 @@ func TestJobsShowsDaemonStatus(t *testing.T) { assert.Contains(t, out, "extract:plain-text") assert.Contains(t, out, "running") - out, err = runCLI(t, "jobs", "--json") - require.NoError(t, err) var got api.JobList - require.NoError(t, json.Unmarshal([]byte(out), &got)) - require.Len(t, got.Items, 1) + require.Eventually(t, func() bool { + out, err = runCLI(t, "jobs", "--json") + if err != nil || json.Unmarshal([]byte(out), &got) != nil || len(got.Items) != 4 { + return false + } + return got.Items[1].Name == "extract:source-metadata" && + got.Items[2].Name == "maintenance:auxiliary-checksums" && got.Items[2].Status == "completed" && + got.Items[3].Name == "process:vector-indexes" && got.Items[3].Status == "running" + }, 5*time.Second, 25*time.Millisecond) + require.Len(t, got.Items, 4) assert.Equal(t, "extract:plain-text", got.Items[0].Name) assert.Equal(t, "running", got.Items[0].Status) + assert.Equal(t, "extract:source-metadata", got.Items[1].Name) + assert.Equal(t, "running", got.Items[1].Status) + assert.Equal(t, "maintenance:auxiliary-checksums", got.Items[2].Name) + assert.Equal(t, "completed", got.Items[2].Status) + assert.Equal(t, "process:vector-indexes", got.Items[3].Name) + assert.Equal(t, "running", got.Items[3].Status) } func TestConfiguredAutomaticPackingPacksAndKeepsDaemonAlive(t *testing.T) { @@ -734,10 +746,14 @@ func TestConfiguredAutomaticPackingPacksAndKeepsDaemonAlive(t *testing.T) { require.NoError(t, err) var got api.JobList require.NoError(t, json.Unmarshal([]byte(out), &got)) - require.Len(t, got.Items, 2) + require.Len(t, got.Items, 5) assert.Equal(t, "extract:plain-text", got.Items[0].Name) - assert.Equal(t, "storage:pack", got.Items[1].Name) - assert.Equal(t, "running", got.Items[1].Status) + assert.Equal(t, "extract:source-metadata", got.Items[1].Name) + assert.Equal(t, "maintenance:auxiliary-checksums", got.Items[2].Name) + assert.Equal(t, "process:vector-indexes", got.Items[3].Name) + assert.Equal(t, "running", got.Items[3].Status) + assert.Equal(t, "storage:pack", got.Items[4].Name) + assert.Equal(t, "running", got.Items[4].Status) time.Sleep(100 * time.Millisecond) _, _, found, err := client.Find(t.Context(), home) @@ -778,11 +794,15 @@ func TestConfiguredWatchIngestsStableFilesAndRemainsObservable(t *testing.T) { require.NoError(t, err) var got api.JobList require.NoError(t, json.Unmarshal([]byte(out), &got)) - require.Len(t, got.Items, 2) + require.Len(t, got.Items, 5) assert.Equal(t, "extract:plain-text", got.Items[0].Name) assert.Equal(t, "running", got.Items[0].Status) - assert.Equal(t, "watch:sessions", got.Items[1].Name) - assert.Equal(t, "running", got.Items[1].Status) + assert.Equal(t, "extract:source-metadata", got.Items[1].Name) + assert.Equal(t, "maintenance:auxiliary-checksums", got.Items[2].Name) + assert.Equal(t, "process:vector-indexes", got.Items[3].Name) + assert.Equal(t, "running", got.Items[3].Status) + assert.Equal(t, "watch:sessions", got.Items[4].Name) + assert.Equal(t, "running", got.Items[4].Status) out, err = runCLI(t, "watch", "list", "--json") require.NoError(t, err) diff --git a/cmd/docbank/daemon.go b/cmd/docbank/daemon.go index fcc3f612..33fe3150 100644 --- a/cmd/docbank/daemon.go +++ b/cmd/docbank/daemon.go @@ -32,6 +32,7 @@ import ( "go.kenn.io/docbank/internal/ingest" "go.kenn.io/docbank/internal/jobs" internalmaintenance "go.kenn.io/docbank/internal/maintenance" + "go.kenn.io/docbank/internal/processing" "go.kenn.io/docbank/internal/store" docweb "go.kenn.io/docbank/internal/web" ) @@ -51,6 +52,49 @@ var daemonRunCmd = &cobra.Command{ }, } +type embeddingRuntimeReadiness interface{ Ready() bool } + +type embeddingJobRunner interface { + Run(ctx context.Context) error +} + +type embeddingJobStarter interface { + Start(name string, run func(context.Context) error) error +} + +// startEmbeddingWorkerIfReady preserves the daemon's normal supervisor-owned +// cancellation and drain lifecycle. An empty runtime registry leaves durable +// work untouched. A ready runtime must provide a real worker; configuration +// failures are not converted into a no-op polling loop. +func startEmbeddingWorkerIfReady(starter embeddingJobStarter, readiness embeddingRuntimeReadiness, + build func() (embeddingJobRunner, error), +) error { + if readiness == nil || !readiness.Ready() { + return nil + } + worker, err := build() + if err != nil { + return err + } + if worker == nil { + return errors.New("embedding worker builder returned nil") + } + return starter.Start("process:embeddings", worker.Run) +} + +func startVectorIndexWorker(starter embeddingJobStarter, + build func() (embeddingJobRunner, error), +) error { + worker, err := build() + if err != nil { + return err + } + if worker == nil { + return errors.New("vector index worker builder returned nil") + } + return starter.Start("process:vector-indexes", worker.Run) +} + func runServe(ctx context.Context) (retErr error) { layout, err := home.Resolve() if err != nil { @@ -136,6 +180,9 @@ func runServe(ctx context.Context) (retErr error) { } defer func() { _ = blobs.Close() }() // Exclusive lock holder: any stale tmp file is provably abandoned. + if err := recoverEmbeddingRuntimeSpool(sigCtx, layout.BlobTmpDir()); err != nil { + return err + } if err := blobs.CleanTmp(); err != nil { return err } @@ -149,6 +196,196 @@ func runServe(ctx context.Context) (retErr error) { } }() operationGate := api.NewOperationGate() + runtimeRegistry := processing.NewRenditionRuntimeRegistry() + embeddingRuntimes, err := configureEmbeddingRuntimeBundle(cfg, blobs, layout.BlobTmpDir()) + if err != nil { + return fmt.Errorf("configuring embedding runtimes: %w", err) + } + embeddingRuntimeRegistry := embeddingRuntimes.registry + processingProfiles, err := executableProcessingProfiles(cfg, embeddingRuntimes) + if err != nil { + return fmt.Errorf("configuring executable processing profiles: %w", err) + } + processingService, err := processing.NewService(processing.ServiceConfig{ + Catalog: s, Blobs: blobs, Gate: operationGate, Profiles: processingProfiles, + RenditionRuntimes: runtimeRegistry, + Principal: "daemon:operator", Scope: "document-processing", SpoolDirectory: layout.BlobTmpDir(), + Lifecycle: sigCtx, + }) + if err != nil { + return fmt.Errorf("configuring processing service: %w", err) + } + // Provider adapters register before this admission point. Until the daemon + // has one, leave restored jobs untouched instead of repeatedly claiming and + // delaying work that this process cannot execute. + if runtimeRegistry.Ready() { + renditionWorker, workerErr := processing.NewRenditionWorker(processing.RenditionWorkerConfig{ + Catalog: s, Blobs: blobs, Runtime: runtimeRegistry, Gate: operationGate, + Owner: "daemon-rendition-worker", LeaseDuration: 5 * time.Minute, + IdleDelay: time.Second, + }) + if workerErr != nil { + return fmt.Errorf("configuring rendition worker: %w", workerErr) + } + if err := jobSupervisor.Start("process:renditions", renditionWorker.Run); err != nil { + return fmt.Errorf("starting rendition worker: %w", err) + } + } + if err := startEmbeddingWorkerIfReady(jobSupervisor, embeddingRuntimeRegistry, + func() (embeddingJobRunner, error) { + worker, workerErr := processing.NewEmbeddingWorker(processing.EmbeddingWorkerConfig{ + Catalog: s, Authority: s, Blobs: blobs, GenerationBlobs: blobs, Runtime: embeddingRuntimeRegistry, + Gate: operationGate, Owner: "daemon-embedding-worker", + LeaseDuration: 5 * time.Minute, IdleDelay: time.Second, + RetryLimit: 3, RetryBaseDelay: time.Second, MaxRetryDelay: 30 * time.Second, + AttemptLifetime: 30 * time.Minute, MaxRows: 100_000, + MaxDimensions: 1_048_576, MaxVectorBlobBytes: 64 << 20, + DescriptorFingerprints: embeddingRuntimeRegistry.Fingerprints(), + }) + if workerErr != nil { + return nil, fmt.Errorf("configuring embedding worker: %w", workerErr) + } + return worker, nil + }); err != nil { + return err + } + if err := startVectorIndexWorker(jobSupervisor, func() (embeddingJobRunner, error) { + worker, workerErr := processing.NewIndexWorker(processing.IndexWorkerConfig{ + Catalog: s, Blobs: blobs, Gate: operationGate, + Owner: "daemon-vector-index-worker", BuildLease: 30 * time.Minute, + ReaderLease: 5 * time.Minute, IdleDelay: time.Second, + }) + if workerErr != nil { + return nil, fmt.Errorf("configuring vector index worker: %w", workerErr) + } + return worker, nil + }); err != nil { + return err + } + if err := jobSupervisor.Start("maintenance:auxiliary-checksums", func(ctx context.Context) error { + cursor := "" + retries := newBackfillRetrySet() + for { + targets, listErr := retryDaemonList(ctx, 5*time.Second, func() ( + []store.BlobChecksumTarget, error, + ) { + return s.MissingBlobChecksumTargetsAfter(ctx, cursor, 100) + }, func(err error) { + logger.Warn("listing auxiliary checksum targets will retry", "error", err) + }) + if listErr != nil { + return listErr + } + if len(targets) == 0 { + cursor = "" + if len(retries) == 0 { + return nil + } + if err := waitDaemonJob(ctx, retries.waitDelay(time.Now().UTC(), 10*time.Second)); err != nil { + return err + } + continue + } + cursor = targets[len(targets)-1].BlobSHA256 + completed := 0 + attempted := 0 + err := operationGate.MutateContext(ctx, func() error { + var batchErr error + for _, target := range targets { + now := time.Now().UTC() + if !retries.ready(target.BlobSHA256, now) { + continue + } + attempted++ + done, targetErr := processing.BackfillAuxiliaryChecksumTargets( + ctx, s, blobs, []store.BlobChecksumTarget{target}) + completed += done + if targetErr != nil { + retries.failed(target.BlobSHA256, now) + batchErr = errors.Join(batchErr, targetErr) + } else { + retries.succeeded(target.BlobSHA256) + } + } + return batchErr + }) + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + logger.Warn("auxiliary checksum backfill will retry", "completed", completed, "error", err) + } + delay := backfillBatchWaitDelay(attempted, retries, time.Now().UTC()) + if err := waitDaemonJob(ctx, delay); err != nil { + return err + } + } + }); err != nil { + return fmt.Errorf("starting auxiliary checksum backfill: %w", err) + } + if err := jobSupervisor.Start("extract:source-metadata", func(ctx context.Context) error { + idleDelay := time.Second + cursor := "" + retries := newBackfillRetrySet() + for { + targets, listErr := retryDaemonList(ctx, 5*time.Second, func() ( + []store.SourceMetadataTarget, error, + ) { + return s.MissingSourceMetadataTargetsAfter( + ctx, processing.SourceMetadataExtractorFingerprint, cursor, 10) + }, func(err error) { + logger.Warn("listing source metadata targets will retry", "error", err) + }) + if listErr != nil { + return listErr + } + if len(targets) == 0 { + cursor = "" + delay := retries.waitDelay(time.Now().UTC(), idleDelay) + if err := waitDaemonJob(ctx, delay); err != nil { + return err + } + idleDelay = min(idleDelay*2, 10*time.Second) + continue + } + cursor = targets[len(targets)-1].SourceSHA256 + idleDelay = time.Second + completed := 0 + attempted := 0 + err := operationGate.MutateContext(ctx, func() error { + var batchErr error + for _, target := range targets { + now := time.Now().UTC() + if !retries.ready(target.SourceSHA256, now) { + continue + } + attempted++ + done, targetErr := processing.BackfillSourceMetadataTargets( + ctx, s, blobs, []store.SourceMetadataTarget{target}) + completed += done + if targetErr != nil { + retries.failed(target.SourceSHA256, now) + batchErr = errors.Join(batchErr, targetErr) + } else { + retries.succeeded(target.SourceSHA256) + } + } + return batchErr + }) + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + logger.Warn("source metadata backfill will retry", "completed", completed, "error", err) + } + delay := backfillBatchWaitDelay(attempted, retries, time.Now().UTC()) + if err := waitDaemonJob(ctx, delay); err != nil { + return err + } + } + }); err != nil { + return fmt.Errorf("starting source metadata backfill: %w", err) + } placementRunner := blob.PlacementRunner{ Metadata: s, Blobs: blobs, Commit: operationGate.PhysicalMutate, } @@ -250,6 +487,7 @@ func runServe(ctx context.Context) (retErr error) { Store: s, Blobs: blobs, VaultRoot: layout.Root, Cfg: cfg, Logger: logger, StartedAt: time.Now(), ShutdownToken: shutdownToken, Shutdown: stop, Tracker: tracker, Jobs: jobSupervisor, Gate: operationGate, WebURL: webURL, BlobRegistry: blobRegistry, + Processing: processingService, }) defer srv.Close() newHTTPServer := func() *http.Server { @@ -326,6 +564,35 @@ func runServe(ctx context.Context) (retErr error) { return serveErr } +func retryDaemonList[T any]( + ctx context.Context, retryDelay time.Duration, list func() ([]T, error), onRetry func(error), +) ([]T, error) { + for { + items, err := list() + if err == nil { + return items, nil + } + if ctx.Err() != nil { + return nil, ctx.Err() + } + onRetry(err) + if err := waitDaemonJob(ctx, retryDelay); err != nil { + return nil, err + } + } +} + +func waitDaemonJob(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + func validateConfiguredWatchStores( cfg config.Config, stores []store.BlobStore, ) error { diff --git a/cmd/docbank/daemon_embedding_worker_test.go b/cmd/docbank/daemon_embedding_worker_test.go new file mode 100644 index 00000000..b3ef9870 --- /dev/null +++ b/cmd/docbank/daemon_embedding_worker_test.go @@ -0,0 +1,211 @@ +package main + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/openaiembed" + "go.kenn.io/docbank/document/voyage" + "go.kenn.io/docbank/document/voyage/voyagetest" + "go.kenn.io/docbank/internal/config" + "go.kenn.io/docbank/internal/processing" +) + +func TestStartEmbeddingWorkerIfReadySkipsMissingBindings(t *testing.T) { + starter := &fakeEmbeddingJobStarter{} + built := false + err := startEmbeddingWorkerIfReady(starter, embeddingReady(false), func() (embeddingJobRunner, error) { + built = true + return embeddingRunnerFunc(func(context.Context) error { return nil }), nil + }) + require.NoError(t, err) + assert.False(t, built) + assert.Empty(t, starter.name) +} + +func TestStartEmbeddingWorkerIfReadyUsesSupervisorCancellation(t *testing.T) { + starter := &fakeEmbeddingJobStarter{} + err := startEmbeddingWorkerIfReady(starter, embeddingReady(true), func() (embeddingJobRunner, error) { + return embeddingRunnerFunc(func(ctx context.Context) error { + <-ctx.Done() + return ctx.Err() + }), nil + }) + require.NoError(t, err) + assert.Equal(t, "process:embeddings", starter.name) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + require.ErrorIs(t, starter.run(ctx), context.Canceled) + + want := errors.New("synthetic configuration failure") + err = startEmbeddingWorkerIfReady(&fakeEmbeddingJobStarter{}, embeddingReady(true), func() (embeddingJobRunner, error) { + return nil, want + }) + require.ErrorIs(t, err, want) +} + +func TestConfigureEmbeddingRuntimesFailsWhenNamedCredentialEnvironmentIsMissing(t *testing.T) { + t.Setenv("DOCBANK_TEST_MISSING_EMBEDDING_KEY", "") + cfg := config.Default() + cfg.CredentialBindings["semantic"] = config.CredentialBindingConfig{ + EnvironmentVariable: "DOCBANK_TEST_MISSING_EMBEDDING_KEY", + } + cfg.EmbeddingProfiles["semantic"] = config.EmbeddingProfileConfig{ + CredentialBinding: "credential:semantic", + Runtime: &config.EmbeddingRuntimeConfig{AdapterContract: openAIEmbeddingAdapter}, + } + _, err := configureEmbeddingRuntimes(cfg, unavailableEmbeddingBlobs{}, t.TempDir()) + require.ErrorContains(t, err, "environment variable is missing") +} + +func TestConfigureEmbeddingRuntimesRegistersSyntheticLoopbackOpenAI(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + t.Cleanup(server.Close) + t.Setenv("DOCBANK_TEST_EMBEDDING_KEY", "synthetic-secret") + cfg := config.Default() + cfg.CredentialBindings["semantic"] = config.CredentialBindingConfig{EnvironmentVariable: "DOCBANK_TEST_EMBEDDING_KEY"} + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}) + require.NoError(t, err) + profile := config.EmbeddingProfileConfig{ + Activation: string(document.EmbeddingRequired), AuthorizationFingerprint: strings.Repeat("1", 64), + CompatibilityID: contract.CompatibilityID, CredentialBinding: "credential:semantic", + DescriptorID: openaiembed.ProviderID, Dimensions: 2, DisclosureFingerprint: strings.Repeat("2", 64), + DocumentFormatter: openaiembed.DocumentFormatterV1, InputKind: string(document.EmbeddingInputRenditionChunk), + MaxBatchItems: 8, MaxInputBytes: 1 << 20, MaxResponseBytes: 1 << 20, + Metric: document.VectorMetricCosine, Model: "synthetic-model", Normalization: document.VectorNormalizationNone, + QueryFormatter: openaiembed.QueryFormatterV1, ScalarEncoding: openaiembed.ScalarEncodingFloat32, + TrustBoundary: string(document.EmbeddingTrustOperatorNetwork), + Chunk: config.EmbeddingChunkConfig{ContextFingerprint: strings.Repeat("3", 64), Formatter: "synthetic/v1", + MaxTokens: 128, OverlapTokens: 8, Tokenizer: "synthetic@v1", TruncationPolicy: string(document.TruncationPolicyReject)}, + ModelInput: config.EmbeddingModelInputConfig{Profile: string(document.ModelInputProfileNomic)}, + Runtime: &config.EmbeddingRuntimeConfig{AdapterContract: openAIEmbeddingAdapter, Endpoint: server.URL, + ModelRevision: "deployment-v1", DeploymentEpoch: "deployment-v1", RequestTimeout: config.Duration(time.Second), + MaxRequestBytes: 1 << 20, MaxRetries: 1, AllowedCIDRs: []string{"127.0.0.0/8"}, ProxyMode: "disabled", + ConnectTimeout: config.Duration(time.Second), KeepAlive: config.Duration(time.Second), + TLSHandshakeTimeout: config.Duration(time.Second)}, + } + descriptor := configuredEmbeddingDescriptor(profile, contract) + final, _, err := finalizeOpenAIEmbeddingDescriptor(openaiembed.Profile{Origin: server.URL, Descriptor: descriptor, + ModelInput: contract, SecretBinding: profile.CredentialBinding, DeploymentEpoch: profile.Runtime.DeploymentEpoch, + RequestTimeout: profile.Runtime.RequestTimeout.Std(), MaxBatchItems: profile.MaxBatchItems, + MaxInputBytes: profile.MaxInputBytes, MaxRequestBytes: profile.Runtime.MaxRequestBytes, + MaxResponseBytes: profile.MaxResponseBytes, EgressPolicy: configuredEmbeddingEgress(*profile.Runtime)}) + require.NoError(t, err) + profile.DescriptorFingerprint = final.Fingerprint + cfg.EmbeddingProfiles["semantic"] = profile + require.NoError(t, cfg.Validate()) + + registry, err := configureEmbeddingRuntimes(cfg, unavailableEmbeddingBlobs{}, t.TempDir()) + require.NoError(t, err) + assert.True(t, registry.Ready()) + assert.Equal(t, []string{final.Fingerprint}, registry.Fingerprints()) + classification, _ := classifyOpenAIEmbeddingError(fmt.Errorf("%w: local request envelope", openaiembed.ErrCapacityResponse)) + assert.Equal(t, processing.EmbeddingProviderCapacity, classification) +} + +func TestConfigureEmbeddingRuntimesRegistersCapabilityAttestedVoyageOriginal(t *testing.T) { + t.Setenv("DOCBANK_TEST_VOYAGE_KEY", "synthetic-secret") + policy, err := voyage.NewPolicy(voyage.PolicyConfig{Model: voyage.DefaultModel, Dimension: voyage.DefaultDimension, + Media: media.Policy{MaxBytes: 1 << 20, AllowStill: true, AllowVideo: true}, MaxBatchItems: 8, + MaxRequestBytes: 1 << 20, MaxResponseBytes: 1 << 20}) + require.NoError(t, err) + manifest, err := voyagetest.SyntheticManifest(policy, voyage.CapabilityImagePNG) + require.NoError(t, err) + var encoded bytes.Buffer + require.NoError(t, voyage.EncodeCapabilityManifest(&encoded, manifest)) + manifestPath := filepath.Join(t.TempDir(), "voyage-capabilities.json") + require.NoError(t, os.WriteFile(manifestPath, encoded.Bytes(), 0o600)) + revision, err := voyage.DirectFileModelRevision(policy, manifest) + require.NoError(t, err) + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "voyage/synthetic-direct/v1", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "document: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "query: {{content}}"}, + }) + require.NoError(t, err) + profile := config.EmbeddingProfileConfig{ + Activation: string(document.EmbeddingRequired), AuthorizationFingerprint: strings.Repeat("1", 64), + CompatibilityID: contract.CompatibilityID, CredentialBinding: "credential:voyage", + DescriptorID: voyage.EmbeddingProviderID, Dimensions: voyage.DefaultDimension, + DisclosureFingerprint: strings.Repeat("2", 64), DocumentFormatter: voyage.EmbeddingDocumentFormatterV1, + InputKind: string(document.EmbeddingInputOriginalFile), MaxBatchItems: 8, MaxInputBytes: 1 << 20, + MaxResponseBytes: 1 << 20, Metric: document.VectorMetricCosine, Model: voyage.DefaultModel, + Normalization: document.VectorNormalizationUnitLength, QueryFormatter: voyage.EmbeddingQueryFormatterV1, + ScalarEncoding: voyage.EmbeddingScalarFloat32, TrustBoundary: string(document.EmbeddingTrustHostedProvider), + ModelInput: config.EmbeddingModelInputConfig{Profile: string(document.ModelInputProfileCustom), + CompatibilityID: contract.CompatibilityID, + Document: config.EmbeddingModelInputEncoderConfig{Mode: string(document.ModelInputModeDocument), Template: "document: {{content}}"}, + Query: config.EmbeddingModelInputEncoderConfig{Mode: string(document.ModelInputModeQuery), Template: "query: {{content}}"}}, + Runtime: &config.EmbeddingRuntimeConfig{AdapterContract: voyageEmbeddingAdapter, Endpoint: voyage.DefaultEndpoint, + ModelRevision: revision, CapabilityManifest: manifestPath, RequestTimeout: config.Duration(time.Second), + MaxRequestBytes: 1 << 20, MaxRetries: 1, AllowedCIDRs: []string{"0.0.0.0/0", "::/0"}, ProxyMode: "disabled", + ConnectTimeout: config.Duration(time.Second), KeepAlive: config.Duration(time.Second), + TLSHandshakeTimeout: config.Duration(time.Second)}, + } + secrets := environmentEmbeddingSecrets{variables: map[string]string{"credential:voyage": "DOCBANK_TEST_VOYAGE_KEY"}} + _, final, err := configuredVoyageProvider(profile, contract, secrets) + require.NoError(t, err) + profile.DescriptorFingerprint = final.Fingerprint + cfg := config.Default() + cfg.CredentialBindings["voyage"] = config.CredentialBindingConfig{EnvironmentVariable: "DOCBANK_TEST_VOYAGE_KEY"} + cfg.EmbeddingProfiles["voyage"] = profile + require.NoError(t, cfg.Validate()) + + registry, err := configureEmbeddingRuntimes(cfg, unavailableEmbeddingBlobs{}, t.TempDir()) + require.NoError(t, err) + assert.Equal(t, []string{final.Fingerprint}, registry.Fingerprints()) +} + +func TestRecoverEmbeddingRuntimeSpoolRemovesOnlyAbandonedUploadState(t *testing.T) { + root := t.TempDir() + stale := filepath.Join(root, ".docbank-upload-synthetic") + require.NoError(t, os.Mkdir(stale, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(stale, "source"), []byte("partial"), 0o600)) + unrelated := filepath.Join(root, "blob-partial") + require.NoError(t, os.WriteFile(unrelated, []byte("keep"), 0o600)) + + require.NoError(t, recoverEmbeddingRuntimeSpool(t.Context(), root)) + _, err := os.Stat(stale) + require.ErrorIs(t, err, os.ErrNotExist) + _, err = os.Stat(unrelated) + require.NoError(t, err) +} + +type unavailableEmbeddingBlobs struct{} + +func (unavailableEmbeddingBlobs) OpenContext(context.Context, string) (io.ReadSeekCloser, error) { + return nil, errors.New("unexpected blob read") +} + +type embeddingReady bool + +func (ready embeddingReady) Ready() bool { return bool(ready) } + +type fakeEmbeddingJobStarter struct { + name string + run func(context.Context) error +} + +func (starter *fakeEmbeddingJobStarter) Start(name string, run func(context.Context) error) error { + starter.name, starter.run = name, run + return nil +} + +type embeddingRunnerFunc func(context.Context) error + +func (run embeddingRunnerFunc) Run(ctx context.Context) error { return run(ctx) } diff --git a/cmd/docbank/daemon_processing_runtime_test.go b/cmd/docbank/daemon_processing_runtime_test.go new file mode 100644 index 00000000..3a4393c0 --- /dev/null +++ b/cmd/docbank/daemon_processing_runtime_test.go @@ -0,0 +1,130 @@ +package main + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/plaintext" + "go.kenn.io/docbank/internal/config" + "go.kenn.io/docbank/internal/processing" +) + +func TestExecutableProcessingProfilesRegistersPlaintextRendition(t *testing.T) { + provider, err := plaintext.New(plaintext.Profile{MaxDocumentBytes: plaintext.MaxDocumentBytes}) + require.NoError(t, err) + descriptor := provider.Descriptor() + cfg := plaintextProcessingConfig(descriptor.Fingerprint) + require.NoError(t, cfg.Validate()) + + profiles, err := executableProcessingProfiles(cfg, embeddingRuntimeBundle{}) + require.NoError(t, err) + require.Contains(t, profiles, "private-text") + require.NotNil(t, profiles["private-text"].RenditionProvider) + assert.Equal(t, descriptor, profiles["private-text"].RenditionProvider.Descriptor()) + assert.Empty(t, profiles["private-text"].EmbeddingProviders) +} + +func TestExecutableProcessingProfilesRejectsDriftedPlaintextDescriptor(t *testing.T) { + cfg := plaintextProcessingConfig(strings.Repeat("0", 64)) + require.NoError(t, cfg.Validate()) + + _, err := executableProcessingProfiles(cfg, embeddingRuntimeBundle{}) + require.ErrorContains(t, err, "descriptor differs from portable binding") +} + +func TestExecutableProcessingProfilesRegistersPinnedRenditionChunkTokenizer(t *testing.T) { + provider, err := plaintext.New(plaintext.Profile{MaxDocumentBytes: plaintext.MaxDocumentBytes}) + require.NoError(t, err) + cfg := plaintextProcessingConfig(provider.Descriptor().Fingerprint) + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileNomic, + }) + require.NoError(t, err) + descriptor, err := document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ + ID: "synthetic.embedding-v1", ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("e", 64), TrustBoundary: document.EmbeddingTrustOperatorNetwork, + Model: "synthetic-model", ModelRevision: "v1", Dimension: 2, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationNone, ScalarEncoding: "float32", + DocumentFormatter: "synthetic/document-v1", QueryFormatter: "synthetic/query-v1", + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{contract.Document.Mode}, + }) + require.NoError(t, err) + cfg.EmbeddingProfiles["semantic"] = config.EmbeddingProfileConfig{ + Activation: string(document.EmbeddingOptional), AuthorizationFingerprint: strings.Repeat("1", 64), + CompatibilityID: contract.CompatibilityID, CredentialBinding: "credential:semantic", + DescriptorID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, Dimensions: descriptor.Dimension, + DisclosureFingerprint: strings.Repeat("2", 64), DocumentFormatter: descriptor.DocumentFormatter, + InputKind: string(document.EmbeddingInputRenditionChunk), MaxBatchItems: 8, MaxInputBytes: 1 << 20, + MaxResponseBytes: 1 << 20, Metric: descriptor.Metric, Model: descriptor.Model, + Normalization: descriptor.Normalization, QueryFormatter: descriptor.QueryFormatter, + ScalarEncoding: descriptor.ScalarEncoding, TrustBoundary: string(descriptor.TrustBoundary), + Chunk: config.EmbeddingChunkConfig{ContextFingerprint: strings.Repeat("3", 64), Formatter: "synthetic/v1", + MaxTokens: 128, OverlapTokens: 8, Tokenizer: unicodeRuneSpec, + TruncationPolicy: string(document.TruncationPolicyReject)}, + ModelInput: config.EmbeddingModelInputConfig{Profile: string(document.ModelInputProfileNomic)}, + } + cfg.ProcessingProfiles["private-text"] = config.ProcessingProfileConfig{ + Rendition: "plaintext", Embeddings: []string{"semantic"}, Retrieval: "lexical", + AttachmentPolicyFingerprint: strings.Repeat("5", 64), CompletenessFingerprint: strings.Repeat("6", 64), + ConsentFingerprint: strings.Repeat("7", 64), LexicalSegmenterFingerprint: strings.Repeat("8", 64), + MaxSegmentRunes: 2000, MaxUnitRunes: 100000, NormalizerFingerprint: strings.Repeat("9", 64), + RetainSanitizedMarkdown: true, SanitizerFingerprint: strings.Repeat("a", 64), TrustBoundary: "vault-primary", + } + require.NoError(t, cfg.Validate()) + bundle := embeddingRuntimeBundle{ + providers: map[string]document.EmbeddingProvider{"semantic": inertEmbeddingProvider{descriptor: descriptor}}, + classifiers: map[string]func(error) (processing.EmbeddingProviderFailure, time.Duration){ + "semantic": func(error) (processing.EmbeddingProviderFailure, time.Duration) { + return processing.EmbeddingProviderPermanent, 0 + }, + }, + } + + profiles, err := executableProcessingProfiles(cfg, bundle) + require.NoError(t, err) + require.Contains(t, profiles, "private-text") + assert.Equal(t, unicodeRuneSpec, + profiles["private-text"].Tokenizers["semantic"].Identity().Name+"@"+ + profiles["private-text"].Tokenizers["semantic"].Identity().Revision) +} + +type inertEmbeddingProvider struct{ descriptor document.EmbeddingDescriptor } + +func (provider inertEmbeddingProvider) Descriptor() document.EmbeddingDescriptor { + return provider.descriptor +} +func (inertEmbeddingProvider) Embed(context.Context, []document.EmbeddingInput, + document.EmbeddingAuthorization, +) (document.EmbeddingResult, error) { + return document.EmbeddingResult{}, nil +} + +func plaintextProcessingConfig(descriptorFingerprint string) config.Config { + cfg := config.Default() + cfg.RenditionProfiles["plaintext"] = config.RenditionProfileConfig{ + AdapterContract: plaintextRenditionAdapter, AuthorizationFingerprint: strings.Repeat("1", 64), + CredentialBinding: "credential:none", DeploymentFingerprint: strings.Repeat("2", 64), + DescriptorID: "plaintext.in-process-v1", DescriptorFingerprint: descriptorFingerprint, + DisclosureFingerprint: strings.Repeat("3", 64), MaxDocumentBytes: plaintext.MaxDocumentBytes, + MaxResponseBytes: plaintext.MaxDocumentBytes, MaxUnits: 1, + RequestedArtifacts: []string{string(document.EvidenceArtifactStructured)}, + TrustBoundary: string(document.RenditionTrustLocalProcess), UploadOptionsFingerprint: strings.Repeat("4", 64), + } + cfg.RetrievalProfiles["lexical"] = config.RetrievalProfileConfig{LexicalLimit: 20, VectorLimit: 20} + cfg.ProcessingProfiles["private-text"] = config.ProcessingProfileConfig{ + Rendition: "plaintext", Retrieval: "lexical", AttachmentPolicyFingerprint: strings.Repeat("5", 64), + CompletenessFingerprint: strings.Repeat("6", 64), ConsentFingerprint: strings.Repeat("7", 64), + LexicalSegmenterFingerprint: strings.Repeat("8", 64), MaxSegmentRunes: 2000, MaxUnitRunes: 100000, + NormalizerFingerprint: strings.Repeat("9", 64), RetainSanitizedMarkdown: true, + SanitizerFingerprint: strings.Repeat("a", 64), TrustBoundary: "vault-primary", + } + return cfg +} diff --git a/cmd/docbank/daemon_test.go b/cmd/docbank/daemon_test.go index bb4ebd9d..fa56c2e1 100644 --- a/cmd/docbank/daemon_test.go +++ b/cmd/docbank/daemon_test.go @@ -2,6 +2,7 @@ package main import ( "context" + "errors" "net" "net/http" "net/url" @@ -22,6 +23,24 @@ import ( "go.kenn.io/docbank/internal/store" ) +func TestRetryDaemonListSurvivesTransientFailure(t *testing.T) { + attempts := 0 + retries := 0 + items, err := retryDaemonList(t.Context(), 0, func() ([]string, error) { + attempts++ + if attempts == 1 { + return nil, errors.New("synthetic transient listing failure") + } + return []string{"ready"}, nil + }, func(error) { + retries++ + }) + require.NoError(t, err) + assert.Equal(t, []string{"ready"}, items) + assert.Equal(t, 2, attempts) + assert.Equal(t, 1, retries) +} + func TestWebOriginUsesDedicatedEphemeralLoopbackListeners(t *testing.T) { first, firstURL, err := listenWebOriginWithIdentity( t.Context(), "00000000000000000000000000000000") diff --git a/cmd/docbank/daemon_vector_index_test.go b/cmd/docbank/daemon_vector_index_test.go new file mode 100644 index 00000000..37978778 --- /dev/null +++ b/cmd/docbank/daemon_vector_index_test.go @@ -0,0 +1,31 @@ +package main + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStartVectorIndexWorkerUsesSupervisorLifecycle(t *testing.T) { + starter := &fakeEmbeddingJobStarter{} + err := startVectorIndexWorker(starter, func() (embeddingJobRunner, error) { + return embeddingRunnerFunc(func(ctx context.Context) error { + <-ctx.Done() + return ctx.Err() + }), nil + }) + require.NoError(t, err) + assert.Equal(t, "process:vector-indexes", starter.name) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + require.ErrorIs(t, starter.run(ctx), context.Canceled) + + want := errors.New("synthetic vector index configuration failure") + err = startVectorIndexWorker(&fakeEmbeddingJobStarter{}, func() (embeddingJobRunner, error) { + return nil, want + }) + require.ErrorIs(t, err, want) +} diff --git a/cmd/docbank/embedding_runtime.go b/cmd/docbank/embedding_runtime.go new file mode 100644 index 00000000..72a14647 --- /dev/null +++ b/cmd/docbank/embedding_runtime.go @@ -0,0 +1,328 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/netip" + "net/url" + "os" + "slices" + "strconv" + "strings" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/openaiembed" + "go.kenn.io/docbank/document/providerhttp" + "go.kenn.io/docbank/document/upload" + "go.kenn.io/docbank/document/voyage" + "go.kenn.io/docbank/internal/config" + "go.kenn.io/docbank/internal/processing" +) + +const ( + openAIEmbeddingAdapter = "docbank-openai-compatible-embeddings/v1" + voyageEmbeddingAdapter = "docbank-voyage-embeddings/v1" +) + +func recoverEmbeddingRuntimeSpool(ctx context.Context, spoolDirectory string) error { + _, err := upload.RecoverStale(ctx, spoolDirectory) + return err +} + +type environmentEmbeddingSecrets struct{ variables map[string]string } + +func (resolver environmentEmbeddingSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + variable, ok := resolver.variables[name] + if !ok { + return "", errors.New("embedding credential binding is unavailable") + } + value, ok := os.LookupEnv(variable) + if !ok || value == "" { + return "", errors.New("embedding credential environment variable is unavailable") + } + return value, nil +} + +func configureEmbeddingRuntimes(cfg config.Config, blobs embeddingRuntimeBlobStore, + spoolDirectory string, +) (*processing.EmbeddingRuntimeRegistry, error) { + bundle, err := configureEmbeddingRuntimeBundle(cfg, blobs, spoolDirectory) + if err != nil { + return nil, err + } + return bundle.registry, nil +} + +type embeddingRuntimeBundle struct { + registry *processing.EmbeddingRuntimeRegistry + providers map[string]document.EmbeddingProvider + classifiers map[string]func(error) (processing.EmbeddingProviderFailure, time.Duration) +} + +func configureEmbeddingRuntimeBundle(cfg config.Config, blobs embeddingRuntimeBlobStore, + spoolDirectory string, +) (embeddingRuntimeBundle, error) { + bundle := embeddingRuntimeBundle{registry: processing.NewEmbeddingRuntimeRegistry(), + providers: make(map[string]document.EmbeddingProvider), + classifiers: make(map[string]func(error) (processing.EmbeddingProviderFailure, time.Duration))} + secrets := environmentEmbeddingSecrets{variables: make(map[string]string)} + for name, binding := range cfg.CredentialBindings { + portable := "credential:" + name + secrets.variables[portable] = binding.EnvironmentVariable + } + for name, configured := range cfg.EmbeddingProfiles { + if configured.Runtime == nil { + continue + } + variable, ok := secrets.variables[configured.CredentialBinding] + if !ok { + return embeddingRuntimeBundle{}, fmt.Errorf("embedding credential %q is not configured", configured.CredentialBinding) + } + if value, ok := os.LookupEnv(variable); !ok || value == "" { + return embeddingRuntimeBundle{}, fmt.Errorf("embedding credential %q environment variable is missing", configured.CredentialBinding) + } + modelInput, err := cfg.EmbeddingModelInput(name) + if err != nil { + return embeddingRuntimeBundle{}, err + } + descriptor := configuredEmbeddingDescriptor(configured, modelInput) + var provider document.EmbeddingProvider + var classify func(error) (processing.EmbeddingProviderFailure, time.Duration) + switch configured.Runtime.AdapterContract { + case openAIEmbeddingAdapter: + profile := openaiembed.Profile{Origin: configured.Runtime.Endpoint, Descriptor: descriptor, + ModelInput: modelInput, SecretBinding: configured.CredentialBinding, + DeploymentEpoch: configured.Runtime.DeploymentEpoch, + ProviderRevisionHeader: configured.Runtime.ProviderRevisionHeader, + RequestTimeout: configured.Runtime.RequestTimeout.Std(), MaxBatchItems: configured.MaxBatchItems, + MaxInputBytes: configured.MaxInputBytes, MaxRequestBytes: configured.Runtime.MaxRequestBytes, + MaxResponseBytes: configured.MaxResponseBytes, + EgressPolicy: configuredEmbeddingEgress(*configured.Runtime)} + descriptor, profile, err = finalizeOpenAIEmbeddingDescriptor(profile) + if err == nil { + transport, transportErr := providerhttp.NewTransport(configuredEmbeddingEgress(*configured.Runtime), nil) + if transportErr != nil { + err = transportErr + } else { + profile.Descriptor = descriptor + provider, err = openaiembed.New(profile, secrets, &http.Client{Transport: transport}) + } + } + classify = classifyOpenAIEmbeddingError + case voyageEmbeddingAdapter: + provider, descriptor, err = configuredVoyageProvider(configured, modelInput, secrets) + classify = classifyVoyageEmbeddingError + default: + err = errors.New("unsupported embedding runtime adapter") + } + if err != nil { + return embeddingRuntimeBundle{}, fmt.Errorf("configuring embedding runtime %q: %w", name, err) + } + if descriptor.Fingerprint != configured.DescriptorFingerprint || + descriptor.ID != configured.DescriptorID || descriptor.ModelRevision != configured.Runtime.ModelRevision { + return embeddingRuntimeBundle{}, fmt.Errorf("configuring embedding runtime %q: descriptor differs from portable binding", name) + } + runtime, err := processing.NewProviderEmbeddingRuntime(provider, blobs, spoolDirectory, classify) + if err != nil { + return embeddingRuntimeBundle{}, err + } + if err := bundle.registry.Register(descriptor.Fingerprint, runtime); err != nil { + return embeddingRuntimeBundle{}, err + } + bundle.providers[name] = provider + bundle.classifiers[name] = classify + } + return bundle, nil +} + +func executableProcessingProfiles(cfg config.Config, + bundle embeddingRuntimeBundle, +) (map[string]processing.ProfileConfig, error) { + renditionProviders, err := configureRenditionProviders(cfg) + if err != nil { + return nil, err + } + profiles := make(map[string]processing.ProfileConfig) + names := make([]string, 0, len(cfg.ProcessingProfiles)) + for name := range cfg.ProcessingProfiles { + names = append(names, name) + } + slices.Sort(names) + for _, name := range names { + portable, err := cfg.PortableProcessingProfile(name) + if err != nil { + return nil, err + } + if portable.Rendition == nil && len(portable.Embeddings) == 0 { + continue + } + executable := true + configured := processing.ProfileConfig{Profile: portable, + EmbeddingProviders: make(map[string]document.EmbeddingProvider), + EmbeddingClassifiers: make(map[string]func(error) (processing.EmbeddingProviderFailure, time.Duration)), + Tokenizers: make(map[string]document.Tokenizer)} + if portable.Rendition != nil { + configured.RenditionProvider = renditionProviders[portable.Rendition.Name] + if configured.RenditionProvider == nil { + continue + } + } + for _, binding := range portable.Embeddings { + provider := bundle.providers[binding.Name] + classifier := bundle.classifiers[binding.Name] + if provider == nil || classifier == nil { + executable = false + break + } + configured.EmbeddingProviders[binding.Name] = provider + configured.EmbeddingClassifiers[binding.Name] = classifier + if binding.InputKind == document.EmbeddingInputRenditionChunk { + tokenizer := configuredEmbeddingTokenizer(binding.Chunk.Tokenizer) + if tokenizer == nil { + executable = false + break + } + configured.Tokenizers[binding.Name] = tokenizer + } + } + if executable { + profiles[name] = configured + } + } + return profiles, nil +} + +type embeddingRuntimeBlobStore interface { + OpenContext(ctx context.Context, hash string) (io.ReadSeekCloser, error) +} + +func configuredEmbeddingDescriptor(profile config.EmbeddingProfileConfig, modelInput document.ModelInputContract) document.EmbeddingDescriptor { + runtime := profile.Runtime + modes := []document.ModelInputMode{modelInput.Document.Mode} + supportsQuery := runtime.AdapterContract == openAIEmbeddingAdapter + if supportsQuery && modelInput.Query.Mode != modelInput.Document.Mode { + modes = append(modes, modelInput.Query.Mode) + } + return document.EmbeddingDescriptor{ID: profile.DescriptorID, + ContractVersion: document.EmbeddingProviderContractVersion, PolicyFingerprint: strings.Repeat("0", 64), + TrustBoundary: document.EmbeddingTrustBoundary(profile.TrustBoundary), Model: profile.Model, + ModelRevision: runtime.ModelRevision, Dimension: profile.Dimensions, Metric: profile.Metric, + Normalization: profile.Normalization, ScalarEncoding: profile.ScalarEncoding, + DocumentFormatter: profile.DocumentFormatter, QueryFormatter: profile.QueryFormatter, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputKind(profile.InputKind)}, + CompatibilityID: profile.CompatibilityID, SupportsTextQuery: supportsQuery, + ModelInput: modelInput, SupportedRequestModes: modes} +} + +func finalizeOpenAIEmbeddingDescriptor(profile openaiembed.Profile) (document.EmbeddingDescriptor, openaiembed.Profile, error) { + temporary, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil { + return document.EmbeddingDescriptor{}, profile, err + } + profile.Descriptor = temporary + fingerprint, err := openaiembed.PolicyFingerprint(profile) + if err != nil { + return document.EmbeddingDescriptor{}, profile, err + } + temporary.PolicyFingerprint, temporary.Fingerprint = fingerprint, "" + final, err := document.NewEmbeddingDescriptor(temporary) + return final, profile, err +} + +func configuredVoyageProvider(profile config.EmbeddingProfileConfig, modelInput document.ModelInputContract, + secrets environmentEmbeddingSecrets, +) (document.EmbeddingProvider, document.EmbeddingDescriptor, error) { + file, err := os.Open(profile.Runtime.CapabilityManifest) + if err != nil { + return nil, document.EmbeddingDescriptor{}, errors.New("voyage capability manifest is unavailable") + } + manifest, decodeErr := voyage.DecodeCapabilityManifest(file) + closeErr := file.Close() + if decodeErr != nil || closeErr != nil { + return nil, document.EmbeddingDescriptor{}, errors.Join(decodeErr, closeErr) + } + policy, err := voyage.NewPolicy(voyage.PolicyConfig{Model: profile.Model, Dimension: profile.Dimensions, + Media: media.Policy{MaxBytes: profile.MaxInputBytes, AllowStill: true, AllowVideo: true}, + MaxBatchItems: profile.MaxBatchItems, MaxRequestBytes: profile.Runtime.MaxRequestBytes, + MaxResponseBytes: profile.MaxResponseBytes}) + if err != nil { + return nil, document.EmbeddingDescriptor{}, err + } + descriptor := configuredEmbeddingDescriptor(profile, modelInput) + configured := voyage.EmbeddingProfile{Mode: voyage.EmbeddingModeDirectFile, + Endpoint: profile.Runtime.Endpoint, EgressPolicy: configuredEmbeddingEgress(*profile.Runtime), + Descriptor: descriptor, ModelInput: modelInput, SecretBinding: profile.CredentialBinding, + RequestTimeout: profile.Runtime.RequestTimeout.Std(), MaxRetries: 1, + MaxBatchItems: profile.MaxBatchItems, MaxInputBytes: profile.MaxInputBytes, + MaxRequestBytes: profile.Runtime.MaxRequestBytes, MaxResponseBytes: profile.MaxResponseBytes, + Policy: policy, CapabilityManifest: manifest} + temporary, err := document.NewEmbeddingDescriptor(descriptor) + if err != nil { + return nil, document.EmbeddingDescriptor{}, err + } + configured.Descriptor = temporary + fingerprint, err := voyage.EmbeddingPolicyFingerprint(configured) + if err != nil { + return nil, document.EmbeddingDescriptor{}, err + } + temporary.PolicyFingerprint, temporary.Fingerprint = fingerprint, "" + final, err := document.NewEmbeddingDescriptor(temporary) + if err != nil { + return nil, document.EmbeddingDescriptor{}, err + } + configured.Descriptor = final + provider, err := voyage.NewEmbeddingProvider(configured, secrets, nil) + return provider, final, err +} + +func configuredEmbeddingEgress(runtime config.EmbeddingRuntimeConfig) providerhttp.EgressPolicy { + parsed, _ := url.Parse(runtime.Endpoint) + port := uint16(443) + if parsed.Scheme == "http" { + port = 80 + } + if parsed.Port() != "" { + value, _ := strconv.ParseUint(parsed.Port(), 10, 16) + port = uint16(value) + } + prefixes := make([]netip.Prefix, 0, len(runtime.AllowedCIDRs)) + for _, value := range runtime.AllowedCIDRs { + prefix, _ := netip.ParsePrefix(value) + prefixes = append(prefixes, prefix) + } + return providerhttp.EgressPolicy{Scheme: parsed.Scheme, Host: parsed.Hostname(), Port: port, + AllowedCIDRs: prefixes, ProxyMode: providerhttp.ProxyDisabled, + ConnectTimeout: runtime.ConnectTimeout.Std(), KeepAlive: runtime.KeepAlive.Std(), + TLSHandshakeTimeout: runtime.TLSHandshakeTimeout.Std(), + TLS: providerhttp.TLSPolicy{SPKISHA256: runtime.SPKISHA256}} +} + +func classifyOpenAIEmbeddingError(err error) (processing.EmbeddingProviderFailure, time.Duration) { + if errors.Is(err, openaiembed.ErrCapacityResponse) { + return processing.EmbeddingProviderCapacity, 0 + } + if errors.Is(err, openaiembed.ErrTransientResponse) { + delay, _ := openaiembed.RetryAfter(err) + return processing.EmbeddingProviderTransient, delay + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return processing.EmbeddingProviderTransient, 0 + } + return processing.EmbeddingProviderPermanent, 0 +} + +func classifyVoyageEmbeddingError(err error) (processing.EmbeddingProviderFailure, time.Duration) { + if errors.Is(err, voyage.ErrBatchTooLarge) { + return processing.EmbeddingProviderCapacity, 0 + } + if voyage.IsRetryable(err) || errors.Is(err, voyage.ErrMalformedResponse) { + delay, _ := voyage.RetryAfter(err) + return processing.EmbeddingProviderTransient, delay + } + return processing.EmbeddingProviderPermanent, 0 +} diff --git a/cmd/docbank/processing.go b/cmd/docbank/processing.go new file mode 100644 index 00000000..61e20b82 --- /dev/null +++ b/cmd/docbank/processing.go @@ -0,0 +1,307 @@ +package main + +import ( + "encoding/hex" + "errors" + "fmt" + "strings" + "text/tabwriter" + + "github.com/spf13/cobra" + + "go.kenn.io/docbank/internal/api" + "go.kenn.io/docbank/internal/client" + "go.kenn.io/docbank/internal/store" +) + +var ( + processingProfilesJSON bool + processingPlanProfile string + processingPlanJSON bool + processingBuildProfile string + processingBuildFingerprint string + processingBuildConsent bool + processingBuildJSON bool + processingBuildNDJSON bool + processingStatusJSON bool +) + +var processingCmd = &cobra.Command{ + Use: "processing", + Short: "Preview and run document processing", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + return cmd.Help() + }, +} + +var processingProfilesCmd = &cobra.Command{ + Use: "profiles", + Short: "List processing profiles this daemon can execute", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runProcessingProfiles(cmd, c, processingProfilesJSON) + }, +} + +var processingPlanCmd = &cobra.Command{ + Use: "plan ", + Short: "Preview provider disclosure for one exact document version", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if processingPlanProfile == "" { + return usageError(errors.New("--profile is required")) + } + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runProcessingPlan(cmd, c, args[0], processingPlanProfile, processingPlanJSON) + }, +} + +var processingBuildCmd = &cobra.Command{ + Use: "build ", + Short: "Run one exact reviewed processing plan", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if processingBuildProfile == "" { + return usageError(errors.New("--profile is required")) + } + if err := validateProcessingBuild(processingBuildFingerprint, processingBuildConsent); err != nil { + return err + } + if processingBuildJSON && processingBuildNDJSON { + return usageError(errors.New("--json and --ndjson are mutually exclusive")) + } + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runProcessingBuild(cmd, c, args[0], processingBuildProfile, + processingBuildFingerprint, processingBuildConsent, processingBuildJSON, processingBuildNDJSON) + }, +} + +var processingStatusCmd = &cobra.Command{ + Use: "status ", + Short: "Show aggregate processing status", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if !canonicalSHA256(args[0]) { + return usageError(errors.New("job ID must be lowercase SHA-256")) + } + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runProcessingStatus(cmd, c, args[0], processingStatusJSON) + }, +} + +func runProcessingProfiles(cmd *cobra.Command, c *client.Client, jsonOutput bool) error { + profiles, err := c.ProcessingProfiles(cmd.Context()) + if err != nil { + return err + } + if jsonOutput { + return writeCLIJSON(cmd.OutOrStdout(), profiles) + } + if len(profiles) == 0 { + _, err := fmt.Fprintln(cmd.OutOrStdout(), "no executable processing profiles") + if err != nil { + return fmt.Errorf("writing empty processing profile list: %w", err) + } + return nil + } + w := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(w, "PROFILE\tRENDITION\tEMBEDDINGS\tFINGERPRINT") + for _, profile := range profiles { + rendition := "-" + if profile.Rendition { + rendition = "rendition" + } + embeddings := "-" + if len(profile.EmbeddingBindings) > 0 { + embeddings = strings.Join(profile.EmbeddingBindings, ",") + } + _, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", + profile.Name, rendition, embeddings, profile.Fingerprint) + } + if err := w.Flush(); err != nil { + return fmt.Errorf("writing processing profiles: %w", err) + } + return nil +} + +func runProcessingPlan(cmd *cobra.Command, c *client.Client, rawSelector, profile string, jsonOutput bool) error { + selector, err := resolveProcessingSelector(cmd, c, rawSelector, profile) + if err != nil { + return err + } + plan, err := c.PlanProcessing(cmd.Context(), api.ProcessingPlanRequest{Selector: selector}) + if err != nil { + return err + } + if jsonOutput { + return writeCLIJSON(cmd.OutOrStdout(), plan) + } + return writeProcessingPlan(cmd, plan) +} + +func writeProcessingPlan(cmd *cobra.Command, plan api.ProcessingPlan) error { + w := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + _, _ = fmt.Fprintf(w, "plan:\t%s\n", plan.Fingerprint) + _, _ = fmt.Fprintf(w, "vault:\t%s\n", plan.VaultUID) + _, _ = fmt.Fprintf(w, "document:\t%s version %s\n", + formatNodeSelector(plan.Selector.NodeID), plan.Selector.ContentVersionID) + _, _ = fmt.Fprintf(w, "profile:\t%s (%s)\n", plan.Selector.Profile, plan.ProfileFingerprint) + _, _ = fmt.Fprintf(w, "discloses:\t%s\n", displayList(plan.DisclosedClasses)) + _, _ = fmt.Fprintf(w, "retains:\t%s\n", displayList(plan.RetainedClasses)) + _, _ = fmt.Fprintf(w, "estimate:\t%d provider call(s), %d vector space(s), %d source byte(s)\n", + plan.Estimate.ProviderCalls, plan.Estimate.VectorSpaces, plan.Estimate.SourceBytes) + _, _ = fmt.Fprintf(w, "consent required:\t%t\n", plan.ConsentRequired) + _, _ = fmt.Fprintf(w, "backup:\t%s\n", plan.BackupConsequence) + if err := w.Flush(); err != nil { + return fmt.Errorf("writing processing plan: %w", err) + } + if len(plan.Flow) == 0 { + return nil + } + flow := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(flow, "CAPABILITY\tPROVIDER\tBOUNDARY\tINPUTS") + for _, hop := range plan.Flow { + _, _ = fmt.Fprintf(flow, "%s\t%s\t%s\t%s\n", hop.Capability, hop.ProviderID, + hop.TrustBoundary, displayList(hop.InputClasses)) + } + if err := flow.Flush(); err != nil { + return fmt.Errorf("writing processing flow: %w", err) + } + return nil +} + +func runProcessingBuild(cmd *cobra.Command, c *client.Client, rawSelector, profile, fingerprint string, + consent, jsonOutput, ndjsonOutput bool, +) error { + if err := validateProcessingBuild(fingerprint, consent); err != nil { + return err + } + selector, err := resolveProcessingSelector(cmd, c, rawSelector, profile) + if err != nil { + return err + } + job, err := c.StartProcessing(cmd.Context(), api.StartProcessingRequest{ + Selector: selector, PlanFingerprint: fingerprint, Consent: true, + }) + if err != nil { + return err + } + if jsonOutput { + return writeCLIJSON(cmd.OutOrStdout(), job) + } + status, err := c.ProcessingStatus(cmd.Context(), job.ID) + if err != nil { + return err + } + if ndjsonOutput { + if err := writeCLIJSON(cmd.OutOrStdout(), api.ProcessingJobEvent{ + Sequence: 1, Type: "job", Job: &job, + }); err != nil { + return err + } + return writeCLIJSON(cmd.OutOrStdout(), api.ProcessingJobEvent{ + Sequence: 2, Type: "status", Status: &status, Terminal: true, + }) + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "processing job: %s\n", job.ID) + return writeProcessingStatus(cmd, status) +} + +func validateProcessingBuild(fingerprint string, consent bool) error { + if !canonicalSHA256(fingerprint) { + return usageError(errors.New("--plan-fingerprint must be the exact lowercase SHA-256 plan fingerprint")) + } + if !consent { + return usageError(errors.New("--consent is required to run the reviewed plan")) + } + return nil +} + +func runProcessingStatus(cmd *cobra.Command, c *client.Client, jobID string, jsonOutput bool) error { + status, err := c.ProcessingStatus(cmd.Context(), jobID) + if err != nil { + return err + } + if jsonOutput { + return writeCLIJSON(cmd.OutOrStdout(), status) + } + return writeProcessingStatus(cmd, status) +} + +func writeProcessingStatus(cmd *cobra.Command, status api.ProcessingStatus) error { + w := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + _, _ = fmt.Fprintf(w, "job:\t%s\n", status.JobID) + _, _ = fmt.Fprintf(w, "state:\t%s\n", status.State) + _, _ = fmt.Fprintf(w, "phase:\t%s\n", status.Phase) + _, _ = fmt.Fprintf(w, "embeddings:\t%d/%d complete\n", + status.CompletedBindings, len(status.EmbeddingJobIDs)) + if status.FailureCode != "" { + _, _ = fmt.Fprintf(w, "failure:\t%s\n", status.FailureCode) + } + if err := w.Flush(); err != nil { + return fmt.Errorf("writing processing status: %w", err) + } + return nil +} + +func resolveProcessingSelector(cmd *cobra.Command, c *client.Client, rawSelector, profile string) (api.ProcessingSelector, error) { + selector, err := parseNodeSelector(rawSelector) + if err != nil { + return api.ProcessingSelector{}, err + } + node, err := selector.resolve(cmd.Context(), c) + if err != nil { + return api.ProcessingSelector{}, err + } + if node.Kind != "file" { + return api.ProcessingSelector{}, fmt.Errorf("processing %q: %w", rawSelector, store.ErrNotFile) + } + if node.CurrentVersionID == "" { + return api.ProcessingSelector{}, errors.New("document has no current content version") + } + return api.ProcessingSelector{NodeID: node.ID, ContentVersionID: node.CurrentVersionID, Profile: profile}, nil +} + +func displayList(values []string) string { + if len(values) == 0 { + return "-" + } + return strings.Join(values, ",") +} + +func canonicalSHA256(value string) bool { + if len(value) != 64 || value != strings.ToLower(value) { + return false + } + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == 32 +} + +func init() { + processingProfilesCmd.Flags().BoolVar(&processingProfilesJSON, "json", false, "emit machine-readable JSON") + processingPlanCmd.Flags().StringVar(&processingPlanProfile, "profile", "", "named executable processing profile") + processingPlanCmd.Flags().BoolVar(&processingPlanJSON, "json", false, "emit machine-readable JSON") + processingBuildCmd.Flags().StringVar(&processingBuildProfile, "profile", "", "named executable processing profile") + processingBuildCmd.Flags().StringVar(&processingBuildFingerprint, "plan-fingerprint", "", "exact reviewed plan fingerprint") + processingBuildCmd.Flags().BoolVar(&processingBuildConsent, "consent", false, "consent to the exact reviewed provider flow") + processingBuildCmd.Flags().BoolVar(&processingBuildJSON, "json", false, "emit machine-readable JSON") + processingBuildCmd.Flags().BoolVar(&processingBuildNDJSON, "ndjson", false, + "emit one job record followed by one terminal status record") + processingStatusCmd.Flags().BoolVar(&processingStatusJSON, "json", false, "emit machine-readable JSON") + processingCmd.AddCommand(processingProfilesCmd, processingPlanCmd, processingBuildCmd, processingStatusCmd) + rootCmd.AddCommand(processingCmd) +} diff --git a/cmd/docbank/processing_test.go b/cmd/docbank/processing_test.go new file mode 100644 index 00000000..489ff8be --- /dev/null +++ b/cmd/docbank/processing_test.go @@ -0,0 +1,147 @@ +package main + +import ( + "bytes" + "context" + "encoding/json/v2" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/internal/api" + "go.kenn.io/docbank/internal/client" +) + +const processingTestVersionID = "123e4567-e89b-42d3-a456-426614174000" + +func TestProcessingCLIProfilesPlanBuildAndStatus(t *testing.T) { + jobID := strings.Repeat("a", 64) + profileFingerprint := strings.Repeat("b", 64) + planFingerprint := strings.Repeat("c", 64) + embeddingJobID := strings.Repeat("d", 64) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + assert.Equal(t, "test-key", request.Header.Get("X-Api-Key")) + w.Header().Set("Content-Type", "application/json") + switch request.Method + " " + request.URL.Path { + case "GET /api/v1/processing/profiles": + assert.NoError(t, json.MarshalWrite(w, []api.ProcessingProfileSummary{{ + Name: "private", Fingerprint: profileFingerprint, Rendition: true, + EmbeddingBindings: []string{"semantic"}, + }})) + case "GET /api/v1/nodes/42": + assert.NoError(t, json.MarshalWrite(w, api.Node{ + ID: 42, Kind: "file", Path: "/docs/report.pdf", CurrentVersionID: processingTestVersionID, + })) + case "POST /api/v1/processing/plans": + assert.NoError(t, json.MarshalWrite(w, api.ProcessingPlan{ + Fingerprint: planFingerprint, VaultUID: "123e4567-e89b-42d3-a456-426614174001", + Selector: api.ProcessingSelector{NodeID: 42, ContentVersionID: processingTestVersionID, + Profile: "private"}, + ProfileFingerprint: profileFingerprint, + Flow: []api.ProcessingFlowHop{{Capability: "rendition", ProviderID: "docling-local", + TrustBoundary: "private-network", InputClasses: []string{"document_bytes"}}}, + DisclosedClasses: []string{"document_bytes"}, RetainedClasses: []string{"sanitized_markdown"}, + Estimate: api.ProcessingEstimate{SourceBytes: 1024, ProviderCalls: 1, VectorSpaces: 1}, + ConsentRequired: true, BackupConsequence: "retained derivatives enter future snapshots", + })) + case "POST /api/v1/processing/jobs": + w.Header().Set("Content-Type", "application/x-ndjson") + _, err := w.Write([]byte(`{"sequence":1,"type":"job","job":{"id":"` + jobID + + `","embedding_job_ids":["` + embeddingJobID + `"],"profile_fingerprint":"` + + profileFingerprint + `","content_version_id":"` + processingTestVersionID + `"}}` + "\n" + + `{"sequence":2,"type":"status","status":{"job_id":"` + jobID + + `","state":"partial","phase":"embedding","failure_code":"provider_unavailable",` + + `"embedding_job_ids":["` + embeddingJobID + `"],"completed_bindings":0},"terminal":true}` + "\n")) + assert.NoError(t, err) + case "GET /api/v1/processing/jobs/" + jobID: + assert.NoError(t, json.MarshalWrite(w, api.ProcessingStatus{ + JobID: jobID, State: "partial", Phase: "embedding", FailureCode: "provider_unavailable", + EmbeddingJobIDs: []string{embeddingJobID}, CompletedBindings: 0, + })) + default: + http.Error(w, request.Method+" "+request.URL.Path, http.StatusNotFound) + } + })) + t.Cleanup(server.Close) + c := client.New(server.URL, "test-key") + + profilesCommand, profilesOutput := processingTestCommand() + require.NoError(t, runProcessingProfiles(profilesCommand, c, false)) + assert.Contains(t, profilesOutput.String(), "private") + assert.Contains(t, profilesOutput.String(), "rendition") + assert.Contains(t, profilesOutput.String(), "semantic") + + planCommand, planOutput := processingTestCommand() + require.NoError(t, runProcessingPlan(planCommand, c, "id:42", "private", false)) + assert.Contains(t, planOutput.String(), "docling-local") + assert.Contains(t, planOutput.String(), "private-network") + assert.Contains(t, planOutput.String(), "sanitized_markdown") + assert.Contains(t, planOutput.String(), planFingerprint) + + buildCommand, buildOutput := processingTestCommand() + require.NoError(t, runProcessingBuild(buildCommand, c, "id:42", "private", planFingerprint, true, false, false)) + assert.Contains(t, buildOutput.String(), jobID) + assert.Contains(t, buildOutput.String(), "partial") + assert.Contains(t, buildOutput.String(), "provider_unavailable") + + ndjsonCommand, ndjsonOutput := processingTestCommand() + require.NoError(t, runProcessingBuild(ndjsonCommand, c, "id:42", "private", planFingerprint, true, false, true)) + lines := strings.Split(strings.TrimSpace(ndjsonOutput.String()), "\n") + require.Len(t, lines, 2) + var first, second api.ProcessingJobEvent + require.NoError(t, json.Unmarshal([]byte(lines[0]), &first)) + require.NoError(t, json.Unmarshal([]byte(lines[1]), &second)) + assert.Equal(t, "job", first.Type) + assert.Equal(t, jobID, first.Job.ID) + assert.Equal(t, "status", second.Type) + assert.True(t, second.Terminal) + assert.Equal(t, "partial", second.Status.State) + + statusCommand, statusOutput := processingTestCommand() + require.NoError(t, runProcessingStatus(statusCommand, c, jobID, true)) + var status api.ProcessingStatus + require.NoError(t, json.Unmarshal(statusOutput.Bytes(), &status)) + assert.Equal(t, "partial", status.State) +} + +func TestProcessingCLIBuildRequiresReviewedFingerprintAndConsent(t *testing.T) { + command, _ := processingTestCommand() + err := runProcessingBuild(command, client.New("http://127.0.0.1:1", "test-key"), + "id:42", "private", "", true, false, false) + require.ErrorContains(t, err, "plan fingerprint") + + err = runProcessingBuild(command, client.New("http://127.0.0.1:1", "test-key"), + "id:42", "private", strings.Repeat("a", 64), false, false, false) + require.ErrorContains(t, err, "--consent") +} + +func TestProcessingCLIRegistersCommandsAndValidatesBeforeDaemon(t *testing.T) { + t.Setenv("DOCBANK_HOME", t.TempDir()) + out, err := runCLI(t, "processing", "--help") + require.NoError(t, err) + assert.Contains(t, out, "profiles") + assert.Contains(t, out, "plan") + assert.Contains(t, out, "build") + assert.Contains(t, out, "status") + + _, err = runCLI(t, "processing", "build", "id:42", "--profile", "private") + require.ErrorContains(t, err, "plan fingerprint") + _, err = runCLI(t, "rendition", "get", "not-a-hash") + require.ErrorContains(t, err, "attachment ID") + _, err = runCLI(t, "search", "query", "--mode", "auto", "--profile", "private") + require.ErrorContains(t, err, "--source-version") +} + +func processingTestCommand() (*cobra.Command, *bytes.Buffer) { + var output bytes.Buffer + command := &cobra.Command{} + command.SetContext(context.Background()) + command.SetOut(&output) + command.SetErr(&output) + return command, &output +} diff --git a/cmd/docbank/rendition.go b/cmd/docbank/rendition.go new file mode 100644 index 00000000..919cac4c --- /dev/null +++ b/cmd/docbank/rendition.go @@ -0,0 +1,63 @@ +package main + +import ( + "errors" + "fmt" + + "github.com/spf13/cobra" + + "go.kenn.io/docbank/internal/client" +) + +var renditionMaxBytes int64 + +var renditionCmd = &cobra.Command{ + Use: "rendition", + Short: "Read verified retained document renditions", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + return cmd.Help() + }, +} + +var renditionGetCmd = &cobra.Command{ + Use: "get ", + Short: "Write one verified self-describing Markdown rendition to stdout", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if renditionMaxBytes < 1 || renditionMaxBytes > 64<<20 { + return usageError(errors.New("--max-bytes must be between 1 and 67108864")) + } + if !canonicalSHA256(args[0]) { + return usageError(errors.New("attachment ID must be lowercase SHA-256")) + } + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runRenditionGet(cmd, c, args[0], renditionMaxBytes) + }, +} + +func runRenditionGet(cmd *cobra.Command, c *client.Client, attachmentID string, maxBytes int64) error { + stream, err := c.Rendition(cmd.Context(), attachmentID, maxBytes) + if err != nil { + return err + } + _, copyErr := stream.CopyVerified(cmd.OutOrStdout()) + closeErr := stream.Close() + if copyErr != nil { + return fmt.Errorf("reading rendition %s: %w", attachmentID, copyErr) + } + if closeErr != nil { + return fmt.Errorf("closing rendition %s: %w", attachmentID, closeErr) + } + return nil +} + +func init() { + renditionGetCmd.Flags().Int64Var(&renditionMaxBytes, "max-bytes", 64<<20, + "maximum complete rendition bytes to accept (1-67108864)") + renditionCmd.AddCommand(renditionGetCmd) + rootCmd.AddCommand(renditionCmd) +} diff --git a/cmd/docbank/rendition_runtime.go b/cmd/docbank/rendition_runtime.go new file mode 100644 index 00000000..26959b2e --- /dev/null +++ b/cmd/docbank/rendition_runtime.go @@ -0,0 +1,31 @@ +package main + +import ( + "fmt" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/plaintext" + "go.kenn.io/docbank/internal/config" +) + +const plaintextRenditionAdapter = "docbank-plaintext-rendition/v1" + +func configureRenditionProviders(cfg config.Config) (map[string]document.RenditionProvider, error) { + providers := make(map[string]document.RenditionProvider) + for name, configured := range cfg.RenditionProfiles { + if configured.AdapterContract != plaintextRenditionAdapter { + continue + } + provider, err := plaintext.New(plaintext.Profile{MaxDocumentBytes: configured.MaxDocumentBytes}) + if err != nil { + return nil, fmt.Errorf("configuring rendition runtime %q: %w", name, err) + } + descriptor := provider.Descriptor() + if descriptor.ID != configured.DescriptorID || descriptor.Fingerprint != configured.DescriptorFingerprint || + string(descriptor.TrustBoundary) != configured.TrustBoundary { + return nil, fmt.Errorf("configuring rendition runtime %q: descriptor differs from portable binding", name) + } + providers[name] = provider + } + return providers, nil +} diff --git a/cmd/docbank/rendition_test.go b/cmd/docbank/rendition_test.go new file mode 100644 index 00000000..9c9885d6 --- /dev/null +++ b/cmd/docbank/rendition_test.go @@ -0,0 +1,74 @@ +package main + +import ( + "bytes" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/internal/api" + "go.kenn.io/docbank/internal/client" +) + +func TestRenditionCLIEmitsExactSelfDescribingMarkdown(t *testing.T) { + body := []byte("# Synthetic report\n\nprivate fixture\n") + evidenceHash := sha256Hex("evidence") + rendering := document.RenditionV1{ContractVersion: document.RenditionContractV1, + Completeness: document.EvidenceComplete, EvidenceChecksum: evidenceHash, + Markdown: body, MarkdownChecksum: sha256HexBytes(body), + Units: []document.NormalizedUnitV1{{EvidenceUnitID: "page:000000", Order: 0, Text: string(body), + Locator: document.EvidenceLocatorV1{Kind: document.EvidenceLocatorPage, + IndexOrigin: document.EvidenceIndexOriginZero}}}} + rendering.Checksum = sha256Hex("rendition") + buildID := sha256Hex("build") + enveloped, _, err := document.EnvelopeRenditionV1(rendering, document.RenditionEnvelopeV1{ + BuildID: buildID, SourceSHA256: sha256Hex("source"), SourceFormat: "pdf", + SourceMediaType: "application/pdf", RenditionRequestFingerprint: sha256Hex("request"), + EvidenceLexicalFingerprint: sha256Hex("lexical"), + NormalizedEvidenceContract: document.NormalizedEvidenceContractV1, UnitKind: document.EvidenceUnitPage, + }) + require.NoError(t, err) + artifact := enveloped.Markdown + attachmentID := sha256Hex("attachment") + artifactID := sha256Hex("artifact") + artifactHash := sha256.Sum256(artifact) + digest := "sha-256=:" + base64.StdEncoding.EncodeToString(artifactHash[:]) + ":" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + assert.Equal(t, "/api/v1/renditions/"+attachmentID, request.URL.Path) + w.Header().Set("Content-Type", "text/markdown; charset=utf-8") + w.Header().Set(api.RenditionAttachmentHeader, attachmentID) + w.Header().Set(api.RenditionBuildHeader, buildID) + w.Header().Set(api.RenditionArtifactHeader, artifactID) + w.Header().Set(api.RenditionProfileHeader, sha256Hex("profile")) + w.Header().Set(api.RenditionCompletenessHeader, string(document.EvidenceComplete)) + w.Header().Set(api.BlobHashHeader, hex.EncodeToString(artifactHash[:])) + w.Header().Set(api.BlobSizeHeader, strconv.Itoa(len(artifact))) + w.Header().Set(api.ContentVersionHeader, processingTestVersionID) + w.Header().Set("Trailer", "Content-Digest") + _, writeErr := w.Write(artifact) + assert.NoError(t, writeErr) + w.Header().Set("Content-Digest", digest) + })) + t.Cleanup(server.Close) + + command, output := processingTestCommand() + require.NoError(t, runRenditionGet(command, client.New(server.URL, "test-key"), attachmentID, 1<<20)) + assert.True(t, bytes.Equal(artifact, output.Bytes())) + assert.True(t, bytes.HasPrefix(output.Bytes(), []byte("---\ndocbank:\n"))) +} + +func sha256Hex(value string) string { return sha256HexBytes([]byte(value)) } + +func sha256HexBytes(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/cmd/docbank/search.go b/cmd/docbank/search.go index 4d600ee8..819b6600 100644 --- a/cmd/docbank/search.go +++ b/cmd/docbank/search.go @@ -1,12 +1,15 @@ package main import ( + "errors" "fmt" + "slices" "strings" "text/tabwriter" "github.com/spf13/cobra" + "go.kenn.io/docbank/internal/api" "go.kenn.io/docbank/internal/client" "go.kenn.io/docbank/internal/store" ) @@ -17,20 +20,53 @@ const ( ) var ( - searchLimit int - searchJSON bool - searchTag string - searchMIME string - searchUnder string - searchSince string - searchBefore string + searchLimit int + searchJSON bool + searchTag string + searchMIME string + searchUnder string + searchSince string + searchBefore string + searchMode string + searchProfile string + searchBinding string + searchSourceVersions []string + searchExplain bool ) +type documentSearchCLIOptions struct { + Mode string + Profile string + BindingID string + ContentVersionIDs []string + Limit int + Explain bool + JSON bool +} + var searchCmd = &cobra.Command{ Use: "search ...", Short: "Search document names and extracted text", Args: cobra.MinimumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { + if documentSearchFlagsChanged(cmd) { + if searchTag != "" || searchMIME != "" || searchUnder != "" || searchSince != "" || searchBefore != "" { + return usageError(errors.New("--mode search cannot be combined with tag, MIME, directory, or time filters")) + } + options := documentSearchCLIOptions{ + Mode: searchMode, Profile: searchProfile, BindingID: searchBinding, + ContentVersionIDs: searchSourceVersions, Limit: searchLimit, + Explain: searchExplain, JSON: searchJSON, + } + if err := validateDocumentSearchOptions(strings.Join(args, " "), options); err != nil { + return err + } + c, err := client.Ensure(cmd.Context()) + if err != nil { + return err + } + return runDocumentSearch(cmd, c, strings.Join(args, " "), options) + } if searchLimit < 1 || searchLimit > maxSearchLimit { return usageError(fmt.Errorf("--limit must be between 1 and %d", maxSearchLimit)) } @@ -135,6 +171,162 @@ var searchCmd = &cobra.Command{ }, } +func documentSearchFlagsChanged(cmd *cobra.Command) bool { + return cmd.Flags().Changed("mode") || cmd.Flags().Changed("profile") || + cmd.Flags().Changed("binding") || cmd.Flags().Changed("source-version") || + cmd.Flags().Changed("explain") +} + +func runDocumentSearch(cmd *cobra.Command, c *client.Client, query string, options documentSearchCLIOptions) error { + if err := validateDocumentSearchOptions(query, options); err != nil { + return err + } + + profiles, err := c.ProcessingProfiles(cmd.Context()) + if err != nil { + return err + } + profile, found := findProcessingProfile(profiles, options.Profile) + if !found { + return usageError(fmt.Errorf("processing profile %q is not executable on this daemon", options.Profile)) + } + bindingID, err := selectDocumentSearchBinding(options.Mode, options.BindingID, profile.EmbeddingBindings) + if err != nil { + return usageError(err) + } + info, err := c.Info(cmd.Context()) + if err != nil { + return err + } + report, err := c.SearchDocuments(cmd.Context(), api.DocumentSearchRequest{ + Query: query, Mode: options.Mode, Limit: options.Limit, Profile: options.Profile, + BindingID: bindingID, Explain: options.Explain, + Fence: api.DocumentSourceFence{VaultUID: info.VaultID, ContentVersionIDs: options.ContentVersionIDs}, + }) + if err != nil { + return err + } + if options.JSON { + return writeCLIJSON(cmd.OutOrStdout(), report) + } + return writeDocumentSearchReport(cmd, report, options.Explain) +} + +func validateDocumentSearchOptions(query string, options documentSearchCLIOptions) error { + if !validDocumentSearchMode(options.Mode) { + if options.Mode == "" { + return usageError(errors.New("--mode is required with processing search options")) + } + return usageError(errors.New("--mode must be lexical, semantic, hybrid, or auto")) + } + if options.Profile == "" { + return usageError(errors.New("--profile is required for processing search")) + } + if options.Limit < 1 || options.Limit > 100 { + return usageError(errors.New("--limit must be between 1 and 100 for processing search")) + } + if strings.TrimSpace(query) == "" { + return usageError(errors.New("search query must not be empty")) + } + if len(options.ContentVersionIDs) == 0 { + return usageError(errors.New("at least one --source-version is required")) + } + if len(options.ContentVersionIDs) > 4096 { + return usageError(errors.New("at most 4096 --source-version values are allowed")) + } + seen := make(map[string]struct{}, len(options.ContentVersionIDs)) + for _, versionID := range options.ContentVersionIDs { + if !client.IsCanonicalUUIDv4(versionID) { + return usageError(fmt.Errorf("source version %q must be a canonical UUIDv4", versionID)) + } + if _, exists := seen[versionID]; exists { + return usageError(fmt.Errorf("source version %q is duplicated", versionID)) + } + seen[versionID] = struct{}{} + } + return nil +} + +func validDocumentSearchMode(mode string) bool { + switch mode { + case "lexical", "semantic", "hybrid", "auto": + return true + default: + return false + } +} + +func findProcessingProfile(profiles []api.ProcessingProfileSummary, name string) (api.ProcessingProfileSummary, bool) { + for _, profile := range profiles { + if profile.Name == name { + return profile, true + } + } + return api.ProcessingProfileSummary{}, false +} + +func selectDocumentSearchBinding(mode, requested string, bindings []string) (string, error) { + if mode == "lexical" { + if requested != "" { + return "", errors.New("--binding is not used by lexical search") + } + return "", nil + } + if requested != "" { + if slices.Contains(bindings, requested) { + return requested, nil + } + return "", fmt.Errorf("--binding %q is not available in this profile", requested) + } + if len(bindings) == 1 { + return bindings[0], nil + } + if len(bindings) > 1 { + return "", errors.New("--binding is required because this profile has multiple embedding bindings") + } + if mode == "semantic" || mode == "hybrid" { + return "", errors.New("selected profile has no embedding binding for semantic search") + } + return "", nil +} + +func writeDocumentSearchReport(cmd *cobra.Command, report api.DocumentSearchReport, explain bool) error { + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "mode: %s", report.ActualMode) + if report.RequestedMode != report.ActualMode { + _, _ = fmt.Fprintf(cmd.OutOrStdout(), " (requested %s)", report.RequestedMode) + } + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "\ncoverage: %s (%d/%d source document(s) complete)\n", + report.Coverage.State, report.Coverage.CompleteDocuments, report.Coverage.ScopedDocuments) + for _, degradation := range report.Degradations { + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "degraded: %s\n", degradation) + } + if len(report.Results) == 0 { + _, err := fmt.Fprintln(cmd.OutOrStdout(), "no matches inside the source fence") + if err != nil { + return fmt.Errorf("writing empty document search result: %w", err) + } + return nil + } + w := tabwriter.NewWriter(cmd.OutOrStdout(), 2, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(w, "RANK\tSCORE\tSELECTOR\tVERSION\tPATH\tEXCERPT") + for _, result := range report.Results { + _, _ = fmt.Fprintf(w, "%d\t%.6g\t%s\t%s\t%s\t%s\n", result.Rank, result.Score, + formatNodeSelector(result.NodeID), result.ContentVersionID, result.Path, result.Excerpt) + } + if err := w.Flush(); err != nil { + return fmt.Errorf("writing document search results: %w", err) + } + if explain { + for _, event := range report.Trace { + _, _ = fmt.Fprintf(cmd.OutOrStdout(), "trace: %s=%d\n", event.Code, event.Count) + } + } + if report.Truncated { + _, _ = fmt.Fprintln(cmd.OutOrStdout(), "results truncated at the requested limit") + } + return nil +} + func init() { searchCmd.Flags().IntVar(&searchLimit, "limit", defaultSearchLimit, "maximum results to return (1-1000)") @@ -148,6 +340,16 @@ func init() { "require modification at or after an absolute RFC3339 timestamp") searchCmd.Flags().StringVar(&searchBefore, "modified-before", "", "require modification before an absolute RFC3339 timestamp") + searchCmd.Flags().StringVar(&searchMode, "mode", "", + "processing search mode: lexical, semantic, hybrid, or auto") + searchCmd.Flags().StringVar(&searchProfile, "profile", "", + "named executable processing profile for --mode search") + searchCmd.Flags().StringVar(&searchBinding, "binding", "", + "embedding binding for semantic or hybrid search") + searchCmd.Flags().StringSliceVar(&searchSourceVersions, "source-version", nil, + "allowed content-version UUID (repeat for a bounded source fence)") + searchCmd.Flags().BoolVar(&searchExplain, "explain", false, + "show bounded retrieval stages without raw similarities or vectors") searchCmd.Flags().BoolVar(&searchJSON, "json", false, "emit machine-readable JSON") rootCmd.AddCommand(searchCmd) } diff --git a/cmd/docbank/search_modes_test.go b/cmd/docbank/search_modes_test.go new file mode 100644 index 00000000..e9771c81 --- /dev/null +++ b/cmd/docbank/search_modes_test.go @@ -0,0 +1,79 @@ +package main + +import ( + "encoding/json/v2" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/internal/api" + "go.kenn.io/docbank/internal/client" +) + +func TestSearchModesRequireExplicitBindingWhenProfileIsAmbiguous(t *testing.T) { + profileFingerprint := strings.Repeat("a", 64) + vaultID := "123e4567-e89b-42d3-a456-426614174001" + var searchRequests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch request.Method + " " + request.URL.Path { + case "GET /api/v1/processing/profiles": + assert.NoError(t, json.MarshalWrite(w, []api.ProcessingProfileSummary{{ + Name: "private", Fingerprint: profileFingerprint, + EmbeddingBindings: []string{"general", "multilingual"}, + }})) + case "GET /api/v1/info": + assert.NoError(t, json.MarshalWrite(w, api.VaultInfo{VaultID: vaultID})) + case "POST /api/v1/search": + searchRequests++ + var body api.DocumentSearchRequest + assert.NoError(t, json.UnmarshalRead(request.Body, &body)) + assert.Equal(t, "hybrid", body.Mode) + assert.Equal(t, "multilingual", body.BindingID) + assert.Equal(t, vaultID, body.Fence.VaultUID) + assert.Equal(t, []string{processingTestVersionID}, body.Fence.ContentVersionIDs) + assert.NoError(t, json.MarshalWrite(w, api.DocumentSearchReport{ + RequestedMode: "hybrid", ActualMode: "hybrid", + Coverage: api.DocumentSearchCoverage{BindingRequired: true, ScopedDocuments: 1, + CompleteDocuments: 1, State: "complete"}, + Results: []api.DocumentSearchResult{{VaultUID: vaultID, NodeID: 42, + ContentVersionID: processingTestVersionID, Rank: 1, Score: 0.75, + Path: "/docs/report.pdf", Excerpt: "synthetic match"}}, + Trace: []api.DocumentSearchTrace{{Code: "source_fence", Count: 1}}, + })) + default: + http.Error(w, request.Method+" "+request.URL.Path, http.StatusNotFound) + } + })) + t.Cleanup(server.Close) + c := client.New(server.URL, "test-key") + + command, _ := processingTestCommand() + err := runDocumentSearch(command, c, "synthetic", documentSearchCLIOptions{ + Mode: "hybrid", Profile: "private", ContentVersionIDs: []string{processingTestVersionID}, Limit: 10, + }) + require.ErrorContains(t, err, "--binding") + assert.Zero(t, searchRequests) + + command, output := processingTestCommand() + require.NoError(t, runDocumentSearch(command, c, "synthetic", documentSearchCLIOptions{ + Mode: "hybrid", Profile: "private", BindingID: "multilingual", + ContentVersionIDs: []string{processingTestVersionID}, Limit: 10, Explain: true, JSON: true, + })) + var report api.DocumentSearchReport + require.NoError(t, json.Unmarshal(output.Bytes(), &report)) + assert.Equal(t, "hybrid", report.ActualMode) + assert.Equal(t, "source_fence", report.Trace[0].Code) + assert.Equal(t, 1, searchRequests) +} + +func TestSearchModesRequireSourceFence(t *testing.T) { + command, _ := processingTestCommand() + err := runDocumentSearch(command, client.New("http://127.0.0.1:1", "test-key"), "query", + documentSearchCLIOptions{Mode: "auto", Profile: "private", Limit: 10}) + require.ErrorContains(t, err, "--source-version") +} diff --git a/cmd/docbank/stat.go b/cmd/docbank/stat.go index 7b4e9009..831590d1 100644 --- a/cmd/docbank/stat.go +++ b/cmd/docbank/stat.go @@ -3,10 +3,12 @@ package main import ( "fmt" "strconv" + "strings" "text/tabwriter" "github.com/spf13/cobra" + "go.kenn.io/docbank/document" "go.kenn.io/docbank/internal/api" "go.kenn.io/docbank/internal/client" ) @@ -68,12 +70,23 @@ func writeNodeStat(cmd *cobra.Command, node api.Node) error { if node.Kind == "file" { _, _ = fmt.Fprintf(w, "version:\t%s\n", node.CurrentVersionID) _, _ = fmt.Fprintf(w, "sha256:\t%s\n", node.BlobHash) + if node.MD5 != "" { + _, _ = fmt.Fprintf(w, "md5:\t%s\n", node.MD5) + } _, _ = fmt.Fprintf(w, "size:\t%d\n", node.Size) mimeType := "not recorded" if node.MimeType != "" { mimeType = strconv.Quote(node.MimeType) } _, _ = fmt.Fprintf(w, "mime:\t%s\n", mimeType) + if node.SourceMetadata != nil { + _, _ = fmt.Fprintf(w, "metadata:\t%s (%d fields, %d warnings)\n", + node.SourceMetadata.ContractVersion, len(node.SourceMetadata.Fields), len(node.SourceMetadata.Warnings)) + for _, field := range node.SourceMetadata.Fields { + _, _ = fmt.Fprintf(w, " %s:\t%s%s\n", field.Key, sourceMetadataDisplayValue(field.Value), + map[bool]string{true: " [sensitive]"}[field.Sensitive]) + } + } } if err := w.Flush(); err != nil { return fmt.Errorf("writing node details: %w", err) @@ -81,6 +94,32 @@ func writeNodeStat(cmd *cobra.Command, node api.Node) error { return nil } +func sourceMetadataDisplayValue(value document.SourceMetadataValueV1) string { + switch value.Kind { + case document.SourceMetadataString: + return strconv.Quote(value.String) + case document.SourceMetadataStringList: + return strconv.Quote(strings.Join(value.Strings, "; ")) + case document.SourceMetadataInteger: + if value.Integer != nil { + return strconv.FormatInt(*value.Integer, 10) + } + case document.SourceMetadataNumber: + if value.Number != nil { + return strconv.FormatFloat(*value.Number, 'g', -1, 64) + } + case document.SourceMetadataBoolean: + if value.Boolean != nil { + return strconv.FormatBool(*value.Boolean) + } + case document.SourceMetadataTimestamp: + if value.Timestamp != nil { + return strconv.Quote(value.Timestamp.Raw) + } + } + return "not recorded" +} + func nodeState(node api.Node) string { if node.TrashedAt != "" { return "trashed" diff --git a/cmd/docbank/stat_test.go b/cmd/docbank/stat_test.go index 98c91141..3c350d9a 100644 --- a/cmd/docbank/stat_test.go +++ b/cmd/docbank/stat_test.go @@ -1,9 +1,13 @@ package main import ( + "bytes" "context" "encoding/json/v2" + "fmt" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -32,12 +36,15 @@ func TestStatCLIInspectsLiveAndTrashedNodes(t *testing.T) { assert.Contains(t, out, "kind: file") assert.Contains(t, out, "version: "+node.CurrentVersionID) assert.Contains(t, out, "sha256: "+node.BlobHash) + assert.Contains(t, out, "md5: "+node.MD5) assert.Contains(t, out, "mime: \"text/plain; charset=utf-8\"") out, err = runCLI(t, "stat", selector, "--json") require.NoError(t, err, out) var got api.Node require.NoError(t, json.Unmarshal([]byte(out), &got)) + // The metadata worker may publish between the first stat and this read. + node.SourceMetadata = got.SourceMetadata assert.Equal(t, node, got) _, err = runCLI(t, "rm", selector) @@ -56,6 +63,46 @@ func TestStatCLIInspectsLiveAndTrashedNodes(t *testing.T) { assert.Empty(t, got.Path) } +func TestStatCLIShowsEmbeddedSourceMetadata(t *testing.T) { + _ = setupVaultHome(t) + source := writeSourceFile(t, "report.pdf", statMetadataPDF()) + _, err := runCLI(t, "add", source, "--dest", "/archive") + require.NoError(t, err) + var out string + require.Eventually(t, func() bool { + out, err = runCLI(t, "stat", "/archive/report.pdf") + return err == nil && strings.Contains(out, `"Synthetic report"`) + }, 5*time.Second, 25*time.Millisecond) + assert.Contains(t, out, "metadata: source-metadata/v1") + assert.Contains(t, out, "page_count: 2") +} + +func statMetadataPDF() string { + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + "<< /Type /Pages /Kids [3 0 R 4 0 R] /Count 2 >>", + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>", + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>", + "<< /Title (Synthetic report) /Author (Ada; Grace) >>", + } + var output bytes.Buffer + _, _ = output.WriteString("%PDF-1.4\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, + "trailer\n<< /Size %d /Root 1 0 R /Info %d 0 R >>\nstartxref\n%d\n%%%%EOF\n", + len(objects)+1, len(objects), xref) + return output.String() +} + func TestStatCLIValidatesSelectorBeforeDaemonStartup(t *testing.T) { t.Setenv("DOCBANK_HOME", t.TempDir()) _, err := runCLI(t, "stat", "relative/path") diff --git a/cmd/docbank/tokenizer_runtime.go b/cmd/docbank/tokenizer_runtime.go new file mode 100644 index 00000000..a2eb110d --- /dev/null +++ b/cmd/docbank/tokenizer_runtime.go @@ -0,0 +1,38 @@ +package main + +import ( + "unicode/utf8" + + "go.kenn.io/docbank/document" +) + +const unicodeRuneSpec = "unicode-runes@v1" + +type unicodeRuneTokenizer struct{} + +func (unicodeRuneTokenizer) Identity() document.TokenizerIdentity { + return document.TokenizerIdentity{ + Name: "unicode-runes", Revision: "v1", PrefixTokenCountsMonotonic: true, + } +} + +func (unicodeRuneTokenizer) Tokenize(text string, limit int) ([]document.TokenBoundary, error) { + runeCount := utf8.RuneCountInString(text) + if runeCount > limit { + return nil, document.ErrTokenizerLimit + } + boundaries := make([]document.TokenBoundary, runeCount) + for index := range boundaries { + boundaries[index] = document.TokenBoundary{Start: index, End: index + 1} + } + return boundaries, nil +} + +func configuredEmbeddingTokenizer(identity string) document.Tokenizer { + if identity == unicodeRuneSpec { + return unicodeRuneTokenizer{} + } + return nil +} + +var _ document.Tokenizer = unicodeRuneTokenizer{} diff --git a/cmd/docbank/versions.go b/cmd/docbank/versions.go index 237eb256..8ee48a0a 100644 --- a/cmd/docbank/versions.go +++ b/cmd/docbank/versions.go @@ -117,6 +117,9 @@ var versionsShowCmd = &cobra.Command{ _, _ = fmt.Fprintf(w, "Recorded:\t%s\n", version.RecordedAt) _, _ = fmt.Fprintf(w, "Kind:\t%s\n", version.TransitionKind) _, _ = fmt.Fprintf(w, "Blob:\t%s\n", version.BlobHash) + if version.MD5 != "" { + _, _ = fmt.Fprintf(w, "MD5:\t%s\n", version.MD5) + } _, _ = fmt.Fprintf(w, "Size:\t%d\n", version.Size) if version.MimeType != "" { _, _ = fmt.Fprintf(w, "Media type:\t%s\n", version.MimeType) diff --git a/cmd/docbank/versions_test.go b/cmd/docbank/versions_test.go new file mode 100644 index 00000000..975d1ef0 --- /dev/null +++ b/cmd/docbank/versions_test.go @@ -0,0 +1,29 @@ +package main + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/internal/client" +) + +func TestVersionsShowIncludesAuxiliaryMD5(t *testing.T) { + _ = setupVaultHome(t) + source := writeSourceFile(t, "versioned.txt", "versioned content") + _, err := runCLI(t, "add", source, "--dest", "/archive") + require.NoError(t, err) + + c, err := client.Ensure(context.Background()) + require.NoError(t, err) + node, err := c.Stat(context.Background(), "/archive/versioned.txt") + require.NoError(t, err) + require.NotEmpty(t, node.MD5) + + out, err := runCLI(t, "versions", "show", node.CurrentVersionID) + require.NoError(t, err, out) + assert.Contains(t, out, "Blob: "+node.BlobHash) + assert.Contains(t, out, "MD5: "+node.MD5) +} diff --git a/document/bridge/client.go b/document/bridge/client.go new file mode 100644 index 00000000..eed045c2 --- /dev/null +++ b/document/bridge/client.go @@ -0,0 +1,749 @@ +package bridge + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "reflect" + "strconv" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + defaultRequestTimeout = 30 * time.Second + defaultTotalTimeout = 10 * time.Minute + defaultPollInterval = time.Second + defaultMaxPollAttempts = 300 + defaultMaxResponseBytes = int64(512 << 20) + maxBridgeTimeout = 24 * time.Hour + maxBridgePollAttempts = 10_000 + maxBridgeResponseBytes = int64(512 << 20) + maxBridgeSourceBytes = int64(1 << 40) + maxBridgeMarkdownBytes = 64 << 20 + maxBridgeArtifactBytes = 256 << 20 + maxBridgeArtifacts = 64 + maxBridgeTotalBytes = 512 << 20 + maxBridgeEvidenceUnits = 100_000 + maxBridgeIdentifier = 128 + maxBridgeSecret = 64 << 10 +) + +var _ document.RenditionProvider = (*Client)(nil) + +// New validates a fixed bridge profile and returns an isolated HTTP client. +func New(profile Profile, secrets SecretResolver, httpClient *http.Client) (*Client, error) { + origin, err := validateOrigin(profile.Origin, profile.Descriptor.TrustBoundary) + if err != nil { + return nil, err + } + descriptor, err := document.NewRenditionDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("bridge: invalid descriptor: %w", err) + } + if profile.SecretBinding == "" { + if secrets != nil { + return nil, errors.New("bridge: secret resolver requires a named binding") + } + } else { + if secrets == nil { + return nil, errors.New("bridge: named secret binding requires a resolver") + } + if err := validateIdentifier(profile.SecretBinding, "secret binding"); err != nil { + return nil, err + } + } + if httpClient == nil { + return nil, errors.New("bridge: HTTP client is required") + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.TotalTimeout == 0 { + profile.TotalTimeout = defaultTotalTimeout + } + if profile.PollInterval == 0 { + profile.PollInterval = defaultPollInterval + } + if profile.MaxPollAttempts == 0 { + profile.MaxPollAttempts = defaultMaxPollAttempts + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxBridgeTimeout || + profile.TotalTimeout <= 0 || profile.TotalTimeout > maxBridgeTimeout || + profile.PollInterval <= 0 || profile.PollInterval > profile.TotalTimeout || + profile.MaxPollAttempts < 1 || profile.MaxPollAttempts > maxBridgePollAttempts || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maxBridgeResponseBytes { + return nil, errors.New("bridge: execution bounds are invalid") + } + if profile.MaxSourceBytes < 0 || profile.MaxSourceBytes > maxBridgeSourceBytes || + profile.MaxProviderMarkdownBytes < 0 || profile.MaxProviderMarkdownBytes > maxBridgeMarkdownBytes || + profile.MaxArtifactBytes < 0 || profile.MaxArtifactBytes > maxBridgeArtifactBytes || + profile.MaxArtifacts < 0 || profile.MaxArtifacts > maxBridgeArtifacts || + profile.MaxTotalResultBytes < 0 || profile.MaxTotalResultBytes > maxBridgeTotalBytes || + profile.MaxEvidenceUnits < 0 || profile.MaxEvidenceUnits > maxBridgeEvidenceUnits { + return nil, errors.New("bridge: provider profile ceilings are invalid") + } + isolatedHTTP := *httpClient + isolatedHTTP.CheckRedirect = providerhttp.RefuseRedirects + isolatedHTTP.Jar = nil + return &Client{ + origin: origin, descriptor: cloneDescriptor(descriptor), + secretBinding: profile.SecretBinding, secrets: secrets, http: &isolatedHTTP, + requestTimeout: profile.RequestTimeout, totalTimeout: profile.TotalTimeout, + pollInterval: profile.PollInterval, maxPollAttempts: profile.MaxPollAttempts, + maxResponseBytes: profile.MaxResponseBytes, + maxSourceBytes: profile.MaxSourceBytes, + maxProviderMarkdownBytes: profile.MaxProviderMarkdownBytes, + maxArtifactBytes: profile.MaxArtifactBytes, maxArtifacts: profile.MaxArtifacts, + maxTotalResultBytes: profile.MaxTotalResultBytes, maxEvidenceUnits: profile.MaxEvidenceUnits, + }, nil +} + +// Descriptor returns the immutable provider identity fixed by the profile. +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Render submits one exact upload and drives the bounded bridge job state machine. +func (client *Client) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (_ document.RenditionResult, retErr error) { + if client == nil { + return document.RenditionResult{}, errors.New("bridge: client is required") + } + if _, err := document.ValidateRenditionProviderRequest(client, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + if err := client.validateProfileAuthorization(upload.Metadata(), authorization); err != nil { + return document.RenditionResult{}, err + } + ctx, cancel := context.WithTimeout(ctx, client.totalTimeout) + defer cancel() + manifest := AuthorizationManifest{ + ContractVersion: ContractVersion, Source: upload.Metadata(), Authorization: authorization, + } + manifestJSON, err := json.Marshal(manifest, json.Deterministic(true)) + if err != nil { + return document.RenditionResult{}, fmt.Errorf("bridge: encode authorization: %w", err) + } + idempotencyDigest := sha256.Sum256(manifestJSON) + idempotencyKey := hex.EncodeToString(idempotencyDigest[:]) + envelope, err := client.submit(ctx, upload, manifestJSON, idempotencyKey) + if err != nil { + return document.RenditionResult{}, err + } + jobID := envelope.JobID + completed := false + defer func() { + if jobID != "" && !completed { + cancelCtx, cancelRemote := context.WithTimeout(context.WithoutCancel(ctx), client.requestTimeout) + defer cancelRemote() + _ = client.cancelJob(cancelCtx, jobID) + } + }() + + for attempt := 0; ; attempt++ { + switch envelope.Status { + case JobCompleted: + result, err := client.decodeCompleted(ctx, envelope, authorization) + if err != nil { + return document.RenditionResult{}, err + } + completed = true + return result, nil + case JobFailed, JobCanceled: + return document.RenditionResult{}, providerErrorFromEnvelope(envelope) + case JobQueued, JobRunning: + default: + return document.RenditionResult{}, malformedError("bridge returned an invalid job status", nil) + } + if attempt >= client.maxPollAttempts { + return document.RenditionResult{}, classifiedError( + document.RenditionErrorCapacity, "bridge polling limit reached", 0, nil) + } + delay := client.pollInterval + if envelope.RetryAfterMillis > 0 { + delay = min(time.Duration(envelope.RetryAfterMillis)*time.Millisecond, client.requestTimeout) + } + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + timer.Stop() + return document.RenditionResult{}, ctx.Err() + case <-timer.C: + } + envelope, err = client.getJob(ctx, jobID, authorization.SourceSHA256) + if err != nil { + if document.IsRenditionProviderErrorRetryable(err) { + envelope = jobEnvelope{Status: JobRunning, JobID: jobID} + continue + } + return document.RenditionResult{}, err + } + } +} + +func (client *Client) validateProfileAuthorization( + metadata document.AuthorizedUploadMetadata, authorization document.RenditionAuthorization, +) error { + if client.maxSourceBytes > 0 && (metadata.ByteLength > client.maxSourceBytes || + authorization.SourceBytes > client.maxSourceBytes) || + client.maxProviderMarkdownBytes > 0 && + authorization.MaxProviderMarkdownBytes > client.maxProviderMarkdownBytes || + client.maxArtifactBytes > 0 && authorization.MaxArtifactBytes > client.maxArtifactBytes || + client.maxArtifacts > 0 && authorization.MaxArtifacts > client.maxArtifacts || + client.maxTotalResultBytes > 0 && authorization.MaxTotalResultBytes > client.maxTotalResultBytes { + return classifiedError(document.RenditionErrorPolicyRejected, + "bridge authorization exceeds provider profile ceilings", 0, nil) + } + return nil +} + +func (client *Client) submit( + ctx context.Context, upload document.AuthorizedUpload, manifest []byte, idempotencyKey string, +) (jobEnvelope, error) { + bodyReader, bodyWriter := io.Pipe() + multipartWriter := multipart.NewWriter(bodyWriter) + contentType := multipartWriter.FormDataContentType() + go func() { + writeErr := writeMultipartUpload(multipartWriter, manifest, upload) + if closeErr := multipartWriter.Close(); writeErr == nil { + writeErr = closeErr + } + _ = bodyWriter.CloseWithError(writeErr) + }() + request, err := http.NewRequestWithContext(ctx, http.MethodPost, client.origin+jobsPath, bodyReader) + if err != nil { + _ = bodyReader.Close() + return jobEnvelope{}, fmt.Errorf("bridge: create submission: %w", err) + } + request.Header.Set("Content-Type", contentType) + request.Header.Set("Accept", jobMediaType) + request.Header.Set("Idempotency-Key", idempotencyKey) + envelope, status, err := client.doJobRequest(request) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return jobEnvelope{}, contextErr + } + if _, ok := errors.AsType[*document.RenditionProviderError](err); ok { + return jobEnvelope{}, err + } + return jobEnvelope{}, classifiedError(document.RenditionErrorAmbiguousSubmission, + "bridge submission outcome is ambiguous", 0, err) + } + if status != http.StatusOK && status != http.StatusAccepted { + return jobEnvelope{}, statusError(status, envelope) + } + if err := client.validateEnvelope(envelope, upload.Metadata()); err != nil { + return jobEnvelope{}, err + } + if status == http.StatusOK && envelope.Status != JobCompleted { + return jobEnvelope{}, malformedError("bridge synchronous response is not completed", nil) + } + if status == http.StatusAccepted && envelope.Status != JobQueued && envelope.Status != JobRunning { + return jobEnvelope{}, malformedError("bridge accepted response has an invalid status", nil) + } + return envelope, nil +} + +func writeMultipartUpload( + writer *multipart.Writer, manifest []byte, upload document.AuthorizedUpload, +) error { + manifestHeader := make(textproto.MIMEHeader) + manifestHeader.Set("Content-Disposition", `form-data; name="`+authorizationPartName+`"`) + manifestHeader.Set("Content-Type", "application/vnd.docbank.rendition-authorization+json;version=1") + manifestPart, err := writer.CreatePart(manifestHeader) + if err != nil { + return fmt.Errorf("bridge: create authorization multipart part: %w", err) + } + if _, err := manifestPart.Write(manifest); err != nil { + return err + } + metadata := upload.Metadata() + sourceHeader := make(textproto.MIMEHeader) + sourceHeader.Set("Content-Disposition", multipart.FileContentDisposition(sourcePartName, metadata.Filename)) + sourceHeader.Set("Content-Type", metadata.MediaType) + sourcePart, err := writer.CreatePart(sourceHeader) + if err != nil { + return fmt.Errorf("bridge: create source multipart part: %w", err) + } + written, err := io.Copy(sourcePart, io.LimitReader(upload, metadata.ByteLength+1)) + if err != nil { + return err + } + if written != metadata.ByteLength { + return errors.New("bridge: upload length changed during submission") + } + return nil +} + +func (client *Client) getJob(ctx context.Context, jobID, sourceSHA256 string) (jobEnvelope, error) { + if err := validateIdentifier(jobID, "job ID"); err != nil { + return jobEnvelope{}, malformedError("bridge returned an invalid job ID", err) + } + request, err := http.NewRequestWithContext(ctx, http.MethodGet, client.origin+jobsPath+"/"+jobID, nil) + if err != nil { + return jobEnvelope{}, err + } + request.Header.Set("Accept", jobMediaType) + envelope, status, err := client.doJobRequest(request) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return jobEnvelope{}, contextErr + } + if _, ok := errors.AsType[*document.RenditionProviderError](err); ok { + return jobEnvelope{}, err + } + return jobEnvelope{}, classifiedError( + document.RenditionErrorTransient, "bridge polling request failed", 0, err) + } + if status != http.StatusOK && status != http.StatusAccepted { + return jobEnvelope{}, statusError(status, envelope) + } + if err := client.validateEnvelope(envelope, document.AuthorizedUploadMetadata{ + SHA256: sourceSHA256, + }); err != nil { + return jobEnvelope{}, err + } + if envelope.JobID != jobID { + return jobEnvelope{}, malformedError("bridge job identity changed while polling", nil) + } + return envelope, nil +} + +func (client *Client) doJobRequest(request *http.Request) (jobEnvelope, int, error) { + parentCtx := request.Context() + requestCtx, cancel := context.WithTimeout(parentCtx, client.requestTimeout) + defer cancel() + request = request.Clone(requestCtx) + if err := client.authorizeRequest(request); err != nil { + if request.Body != nil { + _ = request.Body.Close() + } + return jobEnvelope{}, 0, err + } + response, err := client.http.Do(request) + if err != nil { + if contextErr := parentCtx.Err(); contextErr != nil { + return jobEnvelope{}, 0, contextErr + } + return jobEnvelope{}, 0, err + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode == http.StatusNotFound || response.StatusCode == http.StatusGone { + return jobEnvelope{}, response.StatusCode, nil + } + if err := requireMediaType(response.Header.Get("Content-Type"), jobMediaType); err != nil { + return jobEnvelope{}, response.StatusCode, err + } + body, err := readBounded(response.Body, client.maxResponseBytes) + if err != nil { + return jobEnvelope{}, response.StatusCode, err + } + var envelope jobEnvelope + if len(body) != 0 { + if err := json.Unmarshal(body, &envelope); err != nil { + return jobEnvelope{}, response.StatusCode, malformedError("bridge response JSON is invalid", err) + } + } + return envelope, response.StatusCode, nil +} + +func (client *Client) validateEnvelope( + envelope jobEnvelope, metadata document.AuthorizedUploadMetadata, +) error { + if envelope.ContractVersion != ContractVersion { + return malformedError("bridge contract version is unsupported", nil) + } + if err := validateIdentifier(envelope.JobID, "job ID"); err != nil { + return malformedError("bridge job ID is invalid", err) + } + if metadata.SHA256 != "" && envelope.SourceSHA256 != metadata.SHA256 { + return malformedError("bridge source identity does not match upload", nil) + } + if envelope.AdapterID != client.descriptor.ID || + envelope.DescriptorFingerprint != client.descriptor.Fingerprint || + envelope.PolicyFingerprint != client.descriptor.PolicyFingerprint { + return malformedError("bridge provider identity does not match profile", nil) + } + if envelope.RetryAfterMillis < 0 || envelope.RetryAfterMillis > int64(client.totalTimeout/time.Millisecond) { + return malformedError("bridge retry delay is outside bounds", nil) + } + return nil +} + +func (client *Client) decodeCompleted( + ctx context.Context, envelope jobEnvelope, authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if len(envelope.Result) == 0 { + return document.RenditionResult{}, malformedError("bridge completed response lacks a result", nil) + } + var wire completedResult + if err := json.Unmarshal(envelope.Result, &wire, json.RejectUnknownMembers(true)); err != nil { + return document.RenditionResult{}, malformedError("bridge completed result has an unknown member or invalid value", err) + } + evidence, err := decodeEvidence(wire.Evidence) + if err != nil { + return document.RenditionResult{}, err + } + if client.maxEvidenceUnits > 0 && len(evidence.Units) > client.maxEvidenceUnits { + return document.RenditionResult{}, malformedError( + "bridge evidence exceeds provider profile unit ceiling", nil) + } + var markdown []byte + if wire.ProviderMarkdown != nil { + markdown, err = decodeInlinePayload(*wire.ProviderMarkdown, authorization.MaxProviderMarkdownBytes) + if err != nil { + return document.RenditionResult{}, fmt.Errorf("bridge provider Markdown: %w", err) + } + if injectsDocbankFrontmatter(markdown) { + return document.RenditionResult{}, malformedError("bridge provider Markdown attempts Docbank frontmatter injection", nil) + } + } + if client.maxArtifacts > 0 && len(wire.Artifacts) > client.maxArtifacts { + return document.RenditionResult{}, malformedError( + "bridge artifacts exceed provider profile count ceiling", nil) + } + encodedEvidence, err := json.Marshal(evidence, json.Deterministic(true)) + if err != nil { + return document.RenditionResult{}, malformedError("bridge evidence cannot be measured", err) + } + totalBytes := int64(len(encodedEvidence) + len(markdown)) + artifacts := make([]document.RenditionArtifact, 0, len(wire.Artifacts)) + seenRoles := make(map[document.EvidenceArtifactRole]bool, len(wire.Artifacts)) + for _, artifact := range wire.Artifacts { + if seenRoles[artifact.Role] { + return document.RenditionResult{}, malformedError("bridge returned a duplicate artifact role", nil) + } + seenRoles[artifact.Role] = true + payload, err := client.resolveArtifact(ctx, envelope.JobID, artifact, authorization.MaxArtifactBytes) + if err != nil { + return document.RenditionResult{}, err + } + totalBytes += int64(len(payload)) + if client.maxTotalResultBytes > 0 && totalBytes > int64(client.maxTotalResultBytes) { + return document.RenditionResult{}, malformedError( + "bridge result exceeds provider profile total byte ceiling", nil) + } + artifacts = append(artifacts, document.RenditionArtifact{ + Role: artifact.Role, MediaType: artifact.MediaType, Payload: payload, SHA256: artifact.SHA256, + }) + } + if client.maxTotalResultBytes > 0 && totalBytes > int64(client.maxTotalResultBytes) { + return document.RenditionResult{}, malformedError( + "bridge result exceeds provider profile total byte ceiling", nil) + } + return document.RenditionResult{ + Evidence: evidence, ProviderMarkdown: markdown, Artifacts: artifacts, Receipt: wire.Receipt, + }, nil +} + +func decodeEvidence(payload evidencePayload) (document.SourceEvidenceV1, error) { + if payload.MediaType != evidenceMediaType { + return document.SourceEvidenceV1{}, malformedError("bridge evidence content type is invalid", nil) + } + if payload.ByteLength != int64(len(payload.Inline)) { + return document.SourceEvidenceV1{}, malformedError("bridge evidence length does not match declaration", nil) + } + if payload.SHA256 != sha256Hex(payload.Inline) { + return document.SourceEvidenceV1{}, malformedError("bridge evidence checksum does not match declaration", nil) + } + var evidence document.SourceEvidenceV1 + if err := json.Unmarshal(payload.Inline, &evidence, json.RejectUnknownMembers(true)); err != nil { + return document.SourceEvidenceV1{}, malformedError("bridge source evidence is invalid", err) + } + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.SourceEvidenceV1{}, malformedError("bridge source evidence failed validation", err) + } + return evidence, nil +} + +func (client *Client) resolveArtifact( + ctx context.Context, artifactJobID string, artifact artifactPayload, maxBytes int, +) ([]byte, error) { + if artifact.ByteLength < 0 || artifact.ByteLength > int64(maxBytes) { + return nil, malformedError("bridge artifact length is outside authorization", nil) + } + if artifact.Location == "inline" { + if artifact.ArtifactID != "" { + return nil, malformedError("bridge inline artifact has a result identity", nil) + } + return decodeInlinePayload(binaryPayloadRecord{ + MediaType: artifact.MediaType, ByteLength: artifact.ByteLength, + SHA256: artifact.SHA256, InlineBase64: artifact.InlineBase64, + }, maxBytes) + } + if artifact.Location != "result" || artifact.InlineBase64 != "" { + return nil, malformedError("bridge artifact location is invalid", nil) + } + if err := validateIdentifier(artifact.ArtifactID, "artifact ID"); err != nil { + return nil, malformedError("bridge artifact ID is invalid", err) + } + return client.fetchArtifact(ctx, artifactJobID, artifact) +} + +func (client *Client) fetchArtifact( + ctx context.Context, jobID string, artifact artifactPayload, +) ([]byte, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, + client.origin+jobsPath+"/"+jobID+"/artifacts/"+artifact.ArtifactID, nil) + if err != nil { + return nil, err + } + request.Header.Set("Accept", artifact.MediaType) + parentCtx := request.Context() + requestCtx, cancel := context.WithTimeout(parentCtx, client.requestTimeout) + defer cancel() + request = request.Clone(requestCtx) + if err := client.authorizeRequest(request); err != nil { + return nil, err + } + response, err := client.http.Do(request) + if err != nil { + if contextErr := parentCtx.Err(); contextErr != nil { + return nil, contextErr + } + return nil, classifiedError( + document.RenditionErrorTransient, "bridge artifact request failed", 0, err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, statusError(response.StatusCode, jobEnvelope{}) + } + if err := requireMediaType(response.Header.Get("Content-Type"), artifact.MediaType); err != nil { + return nil, err + } + if response.ContentLength >= 0 && response.ContentLength != artifact.ByteLength { + return nil, malformedError("bridge artifact HTTP length does not match declaration", nil) + } + payload, err := readBounded(response.Body, artifact.ByteLength) + if err != nil { + return nil, err + } + if int64(len(payload)) != artifact.ByteLength { + return nil, malformedError("bridge artifact length does not match declaration", nil) + } + if sha256Hex(payload) != artifact.SHA256 { + return nil, malformedError("bridge artifact checksum does not match declaration", nil) + } + return payload, nil +} + +func (client *Client) cancelJob(ctx context.Context, jobID string) error { + if err := validateIdentifier(jobID, "job ID"); err != nil { + return err + } + request, err := http.NewRequestWithContext(ctx, http.MethodDelete, client.origin+jobsPath+"/"+jobID, nil) + if err != nil { + return err + } + if err := client.authorizeRequest(request); err != nil { + return err + } + response, err := client.http.Do(request) + if err != nil { + return err + } + defer func() { _ = response.Body.Close() }() + _, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096)) + if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusNoContent && + response.StatusCode != http.StatusNotFound && response.StatusCode != http.StatusGone { + return fmt.Errorf("bridge cancellation returned HTTP %d", response.StatusCode) + } + return nil +} + +func (client *Client) authorizeRequest(request *http.Request) error { + if client.secretBinding == "" { + return nil + } + secret, err := client.secrets.ResolveSecret(request.Context(), client.secretBinding) + if err != nil { + return classifiedError(document.RenditionErrorAuthentication, + "bridge credential is unavailable", 0, err) + } + if secret == "" || len(secret) > maxBridgeSecret || strings.ContainsAny(secret, "\r\n\x00") { + return classifiedError(document.RenditionErrorAuthentication, + "bridge credential is invalid", 0, nil) + } + request.Header.Set("Authorization", "Bearer "+secret) + return nil +} + +func validateOrigin(raw string, trust document.RenditionTrustBoundary) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || + parsed.Opaque != "" || parsed.ForceQuery || + parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return "", errors.New("bridge: origin must be one absolute origin without path, credentials, query, or fragment") + } + if parsed.Scheme != "https" && (parsed.Scheme != "http" || trust != document.RenditionTrustOperatorNetwork) { + return "", errors.New("bridge: hosted origins require HTTPS; HTTP is operator-network only") + } + if trust != document.RenditionTrustOperatorNetwork && trust != document.RenditionTrustHostedProvider { + return "", errors.New("bridge: network origin requires an operator-network or hosted trust boundary") + } + return parsed.Scheme + "://" + parsed.Host, nil +} + +func validateIdentifier(value, subject string) error { + if value == "" || len(value) > maxBridgeIdentifier || value != strings.TrimSpace(value) || + !utf8.ValidString(value) { + return fmt.Errorf("bridge: %s is invalid", subject) + } + for _, character := range value { + if (character < 'a' || character > 'z') && (character < 'A' || character > 'Z') && + (character < '0' || character > '9') && character != '.' && character != '_' && character != '-' { + return fmt.Errorf("bridge: %s is invalid", subject) + } + } + return nil +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = append([]document.RenditionFormatCapability(nil), value.SupportedFormats...) + value.ArtifactRoles = append([]document.EvidenceArtifactRole(nil), value.ArtifactRoles...) + return value +} + +func decodeInlinePayload(value binaryPayloadRecord, maxBytes int) ([]byte, error) { + if value.ByteLength < 0 || value.ByteLength > int64(maxBytes) { + return nil, malformedError("inline payload length is outside authorization", nil) + } + payload, err := base64.StdEncoding.Strict().DecodeString(value.InlineBase64) + if err != nil { + return nil, malformedError("inline payload base64 is invalid", err) + } + if int64(len(payload)) != value.ByteLength { + return nil, malformedError("inline payload length does not match declaration", nil) + } + if sha256Hex(payload) != value.SHA256 { + return nil, malformedError("inline payload checksum does not match declaration", nil) + } + if _, _, err := mime.ParseMediaType(value.MediaType); err != nil { + return nil, malformedError("inline payload media type is invalid", err) + } + return payload, nil +} + +func injectsDocbankFrontmatter(markdown []byte) bool { + prefix := markdown + if len(prefix) > 4096 { + prefix = prefix[:4096] + } + return bytes.HasPrefix(prefix, []byte("---\n")) && + bytes.Contains(prefix, []byte("docbank-sanitized-markdown/v1")) +} + +func requireMediaType(got, want string) error { + gotType, gotParams, err := mime.ParseMediaType(got) + if err != nil { + return malformedError("bridge response content type is invalid", err) + } + wantType, wantParams, err := mime.ParseMediaType(want) + if err != nil || gotType != wantType || !reflect.DeepEqual(gotParams, wantParams) { + return malformedError("bridge response content type does not match protocol", err) + } + return nil +} + +func readBounded(reader io.Reader, maximum int64) ([]byte, error) { + if maximum < 0 { + return nil, errors.New("bridge: negative response bound") + } + value, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + return nil, err + } + if int64(len(value)) > maximum { + return nil, malformedError("bridge response exceeds byte limit", nil) + } + return value, nil +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} + +func providerErrorFromEnvelope(envelope jobEnvelope) error { + if envelope.Status == JobCanceled && len(envelope.Error) == 0 { + return classifiedError(document.RenditionErrorCanceled, "bridge job was canceled", 0, nil) + } + if len(envelope.Error) == 0 { + return malformedError("bridge failed response lacks a stable error", nil) + } + var providerError bridgeError + if err := json.Unmarshal(envelope.Error, &providerError, json.RejectUnknownMembers(true)); err != nil { + return malformedError("bridge failed response has an unknown member or invalid value", err) + } + if providerError.RetryAfterMillis < 0 || + providerError.RetryAfterMillis > int64(maxBridgeTimeout/time.Millisecond) { + return malformedError("bridge error retry delay is outside bounds", nil) + } + retry := time.Duration(providerError.RetryAfterMillis) * time.Millisecond + return classifiedError(providerError.Code, providerError.Message, retry, nil) +} + +func statusError(status int, envelope jobEnvelope) error { + if len(envelope.Error) != 0 { + return providerErrorFromEnvelope(envelope) + } + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return classifiedError(document.RenditionErrorAuthentication, "bridge authentication failed", 0, nil) + case http.StatusNotFound, http.StatusGone: + return classifiedError(document.RenditionErrorUnknownJob, "bridge job is unknown or expired", 0, nil) + case http.StatusTooManyRequests: + return classifiedError(document.RenditionErrorRateLimited, "bridge rate limit", 0, nil) + case http.StatusRequestTimeout, http.StatusInternalServerError, http.StatusBadGateway, + http.StatusServiceUnavailable, http.StatusGatewayTimeout: + return classifiedError(document.RenditionErrorTransient, "bridge is temporarily unavailable", 0, nil) + default: + return malformedError("bridge returned unexpected HTTP status "+strconv.Itoa(status), nil) + } +} + +func malformedError(message string, cause error) error { + return classifiedError(document.RenditionErrorMalformedEvidence, message, 0, cause) +} + +func classifiedError( + code document.RenditionErrorCode, message string, retry time.Duration, cause error, +) error { + providerError, err := document.NewRenditionProviderError(code, message, retry, cause) + if err != nil { + fallback, fallbackErr := document.NewRenditionProviderError( + document.RenditionErrorMalformedEvidence, "bridge returned an invalid error", 0, err) + if fallbackErr == nil { + return fallback + } + return errors.Join(err, fallbackErr) + } + return providerError +} diff --git a/document/bridge/client_test.go b/document/bridge/client_test.go new file mode 100644 index 00000000..9174fbfa --- /dev/null +++ b/document/bridge/client_test.go @@ -0,0 +1,1036 @@ +package bridge + +import ( + "bytes" + "context" + "crypto/sha256" + _ "embed" + "encoding/base64" + "encoding/hex" + "encoding/json/jsontext" + "encoding/json/v2" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "net/url" + "runtime" + "slices" + "strings" + "sync/atomic" + "testing" + "testing/iotest" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +type testUpload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (upload *testUpload) Close() error { return nil } +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("missing test secret") + } + return value, nil +} + +type secretResolverFunc func(context.Context, string) (string, error) + +func (resolver secretResolverFunc) ResolveSecret(ctx context.Context, name string) (string, error) { + return resolver(ctx, name) +} + +//go:embed testdata/unknown-major.json +var unknownMajorResponse []byte + +func TestBridgeContractSynchronousCompletion(t *testing.T) { + fixture := newBridgeFixture(t) + var idempotency string + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if !assert.Equal(t, jobsPath, request.URL.Path) || + !assert.Equal(t, http.MethodPost, request.Method) || + !assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) { + response.WriteHeader(http.StatusBadRequest) + return + } + idempotency = request.Header.Get("Idempotency-Key") + assert.Len(t, idempotency, 64) + assertMultipartRequest(t, request, fixture.authorization, fixture.source) + writeBridgeJSON(t, response, http.StatusOK, + completedEnvelope(t, fixture, "job-sync", nil)) + })) + t.Cleanup(server.Close) + + client := newTestBridgeClient(t, server.URL, fixture.descriptor, testSecrets{"bridge-api": "synthetic-secret"}) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, "synthetic bridge evidence", result.Evidence.Units[0].Text) + assert.Equal(t, "synthetic provider markdown\n", string(result.ProviderMarkdown)) + assert.Equal(t, "job-sync", result.Receipt.OperationID) + assert.NotEmpty(t, idempotency) +} + +func TestBridgeProfileCeilingsRejectAuthorizationBeforeEgress(t *testing.T) { + tests := []struct { + name string + prepare func(*testing.T, *bridgeFixture) + ceiling func(*Profile) + }{ + { + name: "source bytes", + ceiling: func(profile *Profile) { + profile.MaxSourceBytes = int64(len([]byte("synthetic bridge source")) - 1) + }, + }, + { + name: "provider Markdown", + ceiling: func(profile *Profile) { + profile.MaxProviderMarkdownBytes = 4095 + }, + }, + { + name: "artifact bytes", + prepare: func(t *testing.T, fixture *bridgeFixture) { + t.Helper() + *fixture = fixture.withStructuredArtifact(t) + }, + ceiling: func(profile *Profile) { + profile.MaxArtifactBytes = 4095 + }, + }, + { + name: "artifact count", + prepare: func(t *testing.T, fixture *bridgeFixture) { + t.Helper() + *fixture = fixture.withArtifactAuthorization(t, + []document.EvidenceArtifactRole{ + document.EvidenceArtifactImage, document.EvidenceArtifactStructured, + }, 2) + }, + ceiling: func(profile *Profile) { + profile.MaxArtifacts = 1 + }, + }, + { + name: "total result bytes", + ceiling: func(profile *Profile) { + profile.MaxTotalResultBytes = 16383 + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + fixture := newBridgeFixture(t) + if test.prepare != nil { + test.prepare(t, &fixture) + } + var requests atomic.Int64 + profile := testBridgeProfile("https://bridge.invalid", fixture.descriptor, nil) + test.ceiling(&profile) + client, err := New(profile, nil, &http.Client{Transport: roundTripFunc( + func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("ceiling violation reached egress") + }, + )}) + require.NoError(t, err) + + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.Zero(t, requests.Load()) + }) + } +} + +func TestBridgeProfileCeilingsAllowResultWithinEveryLimit(t *testing.T) { + fixture := newBridgeFixture(t).withStructuredArtifact(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + writeBridgeJSON(t, response, http.StatusOK, completedEnvelope(t, fixture, + "job-within-ceilings", []map[string]any{ + inlineArtifact(document.EvidenceArtifactStructured, fixture.artifact), + })) + })) + t.Cleanup(server.Close) + profile := testBridgeProfile(server.URL, fixture.descriptor, nil) + profile.MaxSourceBytes = fixture.metadata.ByteLength + profile.MaxProviderMarkdownBytes = fixture.authorization.MaxProviderMarkdownBytes + profile.MaxArtifactBytes = fixture.authorization.MaxArtifactBytes + profile.MaxArtifacts = fixture.authorization.MaxArtifacts + profile.MaxTotalResultBytes = fixture.authorization.MaxTotalResultBytes + profile.MaxEvidenceUnits = 1 + client, err := New(profile, nil, http.DefaultClient) + require.NoError(t, err) + + result, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 1) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, fixture.artifact, result.Artifacts[0].Payload) +} + +func TestBridgeProfileRejectsInvalidCeilings(t *testing.T) { + fixture := newBridgeFixture(t) + for _, mutate := range []func(*Profile){ + func(profile *Profile) { profile.MaxSourceBytes = -1 }, + func(profile *Profile) { profile.MaxProviderMarkdownBytes = -1 }, + func(profile *Profile) { profile.MaxArtifactBytes = -1 }, + func(profile *Profile) { profile.MaxArtifacts = -1 }, + func(profile *Profile) { profile.MaxTotalResultBytes = -1 }, + func(profile *Profile) { profile.MaxEvidenceUnits = -1 }, + } { + profile := testBridgeProfile("https://bridge.invalid", fixture.descriptor, nil) + mutate(&profile) + _, err := New(profile, nil, http.DefaultClient) + require.ErrorContains(t, err, "profile ceilings") + } +} + +func TestBridgeProfileCeilingsRejectDecodedResult(t *testing.T) { + t.Run("evidence units", func(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + value := completedEnvelope(t, fixture, "job-evidence-ceiling", nil) + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceComplete, + Family: "pdf", UnitKind: document.EvidenceUnitPage, + Units: []document.SourceEvidenceUnitV1{ + {Order: 0, ProviderID: "page-0", Text: "first", Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginZero, Start: 0, End: 0}}, + {Order: 1, ProviderID: "page-1", Text: "second", Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginZero, Start: 1, End: 1}}, + }, + } + setCompletedEvidence(t, value, evidence) + writeBridgeJSON(t, response, http.StatusOK, value) + })) + t.Cleanup(server.Close) + profile := testBridgeProfile(server.URL, fixture.descriptor, nil) + profile.MaxEvidenceUnits = 1 + client, err := New(profile, nil, http.DefaultClient) + require.NoError(t, err) + + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + + t.Run("artifact count", func(t *testing.T) { + fixture := newBridgeFixture(t).withArtifactAuthorization(t, + []document.EvidenceArtifactRole{ + document.EvidenceArtifactImage, document.EvidenceArtifactStructured, + }, 1) + payload := []byte("{}") + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + writeBridgeJSON(t, response, http.StatusOK, completedEnvelope(t, fixture, + "job-artifact-count", []map[string]any{ + inlineArtifact(document.EvidenceArtifactImage, payload), + inlineArtifact(document.EvidenceArtifactStructured, payload), + })) + })) + t.Cleanup(server.Close) + profile := testBridgeProfile(server.URL, fixture.descriptor, nil) + profile.MaxArtifacts = 1 + client, err := New(profile, nil, http.DefaultClient) + require.NoError(t, err) + + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + + t.Run("cumulative bytes include fetched artifacts", func(t *testing.T) { + fixture := newBridgeFixture(t).withStructuredArtifact(t) + fixture.artifact = bytes.Repeat([]byte("x"), 1800) + fixture.authorization.MaxArtifactBytes = 2000 + fixture.authorization.MaxTotalResultBytes = 2048 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case jobsPath: + writeBridgeJSON(t, response, http.StatusOK, completedEnvelope(t, fixture, + "job-total-ceiling", []map[string]any{{ + "role": string(document.EvidenceArtifactStructured), "media_type": "application/json", + "byte_length": len(fixture.artifact), "sha256": sha256String(fixture.artifact), + "location": "result", "artifact_id": "structured-1", + }})) + case jobsPath + "/job-total-ceiling/artifacts/structured-1": + response.Header().Set("Content-Type", "application/json") + _, err := response.Write(fixture.artifact) + assert.NoError(t, err) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + profile := testBridgeProfile(server.URL, fixture.descriptor, nil) + profile.MaxArtifactBytes = 2000 + profile.MaxArtifacts = 1 + profile.MaxTotalResultBytes = 2048 + client, err := New(profile, nil, http.DefaultClient) + require.NoError(t, err) + + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) +} + +func TestBridgeContractIdempotencyReplayAndForwardCompatibleEnvelope(t *testing.T) { + fixture := newBridgeFixture(t) + var keys []string + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + keys = append(keys, request.Header.Get("Idempotency-Key")) + value := completedEnvelope(t, fixture, "job-replay", nil) + value["future_minor_field"] = map[string]any{"ignored": true} + writeBridgeJSON(t, response, http.StatusOK, value) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + + for range 2 { + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + } + require.Len(t, keys, 2) + assert.Equal(t, keys[0], keys[1], "an exact replay must retain its idempotency identity") +} + +func TestBridgeContractAcceptsSchemaValidEvidenceEncoding(t *testing.T) { + fixture := newBridgeFixture(t) + value := completedEnvelope(t, fixture, "job-json-encoding", nil) + evidence := completedEvidence(value) + raw, ok := evidence["inline"].(jsontext.Value) + require.True(t, ok) + spaced := append([]byte("{\n "), raw[1:]...) + evidence["byte_length"] = len(spaced) + evidence["sha256"] = sha256String(spaced) + encoded, err := json.Marshal(value, json.Deterministic(true)) + require.NoError(t, err) + encoded = bytes.Replace(encoded, raw, spaced, 1) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + response.Header().Set("Content-Type", jobMediaType) + _, writeErr := response.Write(encoded) + assert.NoError(t, writeErr) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.NoError(t, err) +} + +func TestBridgeContractPollsAndFetchesFixedRouteArtifact(t *testing.T) { + fixture := newBridgeFixture(t) + fixture = fixture.withStructuredArtifact(t) + var polls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch { + case request.Method == http.MethodPost && request.URL.Path == jobsPath: + writeBridgeJSON(t, response, http.StatusAccepted, + pendingEnvelope(fixture, "job-async", JobQueued)) + case request.Method == http.MethodGet && request.URL.Path == jobsPath+"/job-async": + if polls.Add(1) == 1 { + writeBridgeJSON(t, response, http.StatusOK, + pendingEnvelope(fixture, "job-async", JobRunning)) + return + } + writeBridgeJSON(t, response, http.StatusOK, + completedEnvelope(t, fixture, "job-async", []map[string]any{{ + "role": string(document.EvidenceArtifactStructured), "media_type": "application/json", + "byte_length": len(fixture.artifact), "sha256": sha256String(fixture.artifact), + "location": "result", "artifact_id": "structured-1", + }})) + case request.Method == http.MethodGet && request.URL.Path == jobsPath+"/job-async/artifacts/structured-1": + response.Header().Set("Content-Type", "application/json") + _, err := response.Write(fixture.artifact) + assert.NoError(t, err) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, fixture.artifact, result.Artifacts[0].Payload) + assert.Equal(t, int64(2), polls.Load()) +} + +func TestBridgeContractRejectsUnsafeOrCorruptResponses(t *testing.T) { + tests := map[string]struct { + mutate func(map[string]any) + ctype string + want string + }{ + "source identity drift": { + mutate: func(value map[string]any) { value["source_sha256"] = strings.Repeat("9", 64) }, + want: "source identity", + }, + "evidence checksum mismatch": { + mutate: func(value map[string]any) { + completedEvidence(value)["sha256"] = strings.Repeat("0", 64) + }, + want: "evidence checksum", + }, + "evidence length mismatch": { + mutate: func(value map[string]any) { + completedEvidence(value)["byte_length"] = 1 + }, + want: "evidence length", + }, + "wrong response content type": { + ctype: "application/json", want: "content type", + }, + "Docbank frontmatter injection": { + mutate: func(value map[string]any) { + markdown := []byte("---\ncontract: docbank-sanitized-markdown/v1\n---\nsecret\n") + completedResultMap(value)["provider_markdown"] = binaryPayload(markdown, "text/markdown") + }, + want: "frontmatter", + }, + "artifact URL escapes origin": { + mutate: func(value map[string]any) { + completedResultMap(value)["artifacts"] = []map[string]any{{ + "role": string(document.EvidenceArtifactStructured), "media_type": "application/json", + "byte_length": 2, "sha256": sha256String([]byte("{}")), + "location": "result", "artifact_id": "structured-1", + "url": "https://attacker.invalid/secret", + }} + }, + want: "unknown member", + }, + } + for name, test := range tests { + t.Run(name, func(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + value := completedEnvelope(t, fixture, "job-corrupt", nil) + if test.mutate != nil { + test.mutate(value) + } + contentType := test.ctype + if contentType == "" { + contentType = jobMediaType + } + response.Header().Set("Content-Type", contentType) + response.WriteHeader(http.StatusOK) + assert.NoError(t, json.MarshalWrite(response, value, json.Deterministic(true))) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestBridgeContractRejectsRecordedUnknownMajorResponse(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + response.Header().Set("Content-Type", jobMediaType) + _, err := response.Write(unknownMajorResponse) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.ErrorContains(t, err, "contract version") +} + +func TestBridgeContractRejectsArtifactContentTypeLengthAndChecksumMismatch(t *testing.T) { + tests := map[string]struct { + mediaType string + declaredBytes int + declaredHash string + want string + }{ + "content type": {mediaType: "text/plain", want: "content type"}, + "length": {declaredBytes: 1, want: "HTTP length"}, + "checksum": {declaredHash: strings.Repeat("0", 64), want: "checksum"}, + } + for name, test := range tests { + t.Run(name, func(t *testing.T) { + fixture := newBridgeFixture(t).withStructuredArtifact(t) + declaredBytes := len(fixture.artifact) + if test.declaredBytes != 0 { + declaredBytes = test.declaredBytes + } + declaredHash := sha256String(fixture.artifact) + if test.declaredHash != "" { + declaredHash = test.declaredHash + } + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.URL.Path == jobsPath { + writeBridgeJSON(t, response, http.StatusOK, + completedEnvelope(t, fixture, "job-artifact", []map[string]any{{ + "role": string(document.EvidenceArtifactStructured), "media_type": "application/json", + "byte_length": declaredBytes, "sha256": declaredHash, + "location": "result", "artifact_id": "structured-1", + }})) + return + } + contentType := test.mediaType + if contentType == "" { + contentType = "application/json" + } + response.Header().Set("Content-Type", contentType) + _, err := response.Write(fixture.artifact) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.ErrorContains(t, err, test.want) + }) + } +} + +func TestBridgeContractPreservesCredentialFailuresAndStripsAmbientCookies(t *testing.T) { + t.Run("credential failure", func(t *testing.T) { + fixture := newBridgeFixture(t) + resolver := secretResolverFunc(func(context.Context, string) (string, error) { + return "", errors.New("synthetic resolver failure") + }) + client := newTestBridgeClientWithHTTP(t, "https://bridge.invalid", fixture.descriptor, + resolver, &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("request reached transport after credential failure") + return nil, errors.New("unreachable test transport") + })}) + baseline := runtime.NumGoroutine() + for range 10 { + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorAuthentication, providerError.Code()) + } + require.Eventually(t, func() bool { + return runtime.NumGoroutine() <= baseline+2 + }, time.Second, 10*time.Millisecond) + }) + + t.Run("ambient cookie jar", func(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + assert.Empty(t, request.Header.Get("Cookie")) + writeBridgeJSON(t, response, http.StatusOK, completedEnvelope(t, fixture, "job-cookie", nil)) + })) + t.Cleanup(server.Close) + jar, err := cookiejar.New(nil) + require.NoError(t, err) + origin, err := url.Parse(server.URL) + require.NoError(t, err) + jar.SetCookies(origin, []*http.Cookie{{ + Name: "ambient", Value: "private", Secure: true, HttpOnly: true, SameSite: http.SameSiteStrictMode, + }}) + httpClient := server.Client() + httpClient.Jar = jar + client := newTestBridgeClientWithHTTP(t, server.URL, fixture.descriptor, nil, httpClient) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.NoError(t, err) + }) +} + +func TestBridgeContractClassifiesPerRequestTimeouts(t *testing.T) { + t.Run("submission is ambiguous", func(t *testing.T) { + fixture := newBridgeFixture(t) + client := newTestBridgeClientWithHTTP(t, "https://bridge.invalid", fixture.descriptor, nil, + &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + <-request.Context().Done() + return nil, request.Context().Err() + })}) + client.requestTimeout = 10 * time.Millisecond + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorAmbiguousSubmission, providerError.Code()) + }) + + t.Run("poll is retried", func(t *testing.T) { + fixture := newBridgeFixture(t) + var polls atomic.Int64 + client := newTestBridgeClientWithHTTP(t, "https://bridge.invalid", fixture.descriptor, nil, + &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.Method == http.MethodPost { + return bridgeHTTPResponse(t, request, http.StatusAccepted, + pendingEnvelope(fixture, "job-slow", JobQueued)), nil + } + if request.Method == http.MethodGet { + polls.Add(1) + <-request.Context().Done() + return nil, request.Context().Err() + } + return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Request: request}, nil + })}) + client.requestTimeout = 10 * time.Millisecond + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorCapacity, providerError.Code()) + assert.Equal(t, int64(4), polls.Load()) + }) +} + +func TestBridgeContractRejectsUnboundedOrExtendedStableErrors(t *testing.T) { + tests := map[string]map[string]any{ + "unknown member": { + "code": string(document.RenditionErrorRateLimited), "message": "slow down", "internal": "private", + }, + "unbounded retry": { + "code": string(document.RenditionErrorRateLimited), "message": "slow down", + "retry_after_millis": int64(maxBridgeTimeout/time.Millisecond) + 1, + }, + } + for name, stableError := range tests { + t.Run(name, func(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + value := pendingEnvelope(fixture, "job-error", JobFailed) + value["error"] = stableError + writeBridgeJSON(t, response, http.StatusTooManyRequests, value) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorMalformedEvidence, providerError.Code()) + }) + } +} + +func TestBridgeContractClassifiesUnknownExpiredAndAmbiguousJobs(t *testing.T) { + for name, status := range map[string]int{ + "unknown": http.StatusNotFound, "expired": http.StatusGone, + } { + t.Run(name, func(t *testing.T) { + fixture := newBridgeFixture(t) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.Method == http.MethodPost { + writeBridgeJSON(t, response, http.StatusAccepted, + pendingEnvelope(fixture, "job-missing", JobQueued)) + return + } + response.WriteHeader(status) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorUnknownJob, providerError.Code()) + }) + } + + t.Run("ambiguous submission", func(t *testing.T) { + fixture := newBridgeFixture(t) + client := newTestBridgeClientWithHTTP(t, "https://bridge.invalid", fixture.descriptor, nil, + &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, io.ErrUnexpectedEOF + })}) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorAmbiguousSubmission, providerError.Code()) + }) + + t.Run("truncated submission response", func(t *testing.T) { + fixture := newBridgeFixture(t) + client := newTestBridgeClientWithHTTP(t, "https://bridge.invalid", fixture.descriptor, nil, + &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusAccepted, + Header: http.Header{"Content-Type": []string{jobMediaType}}, + Body: io.NopCloser(iotest.ErrReader(io.ErrUnexpectedEOF)), + Request: request, + }, nil + })}) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorAmbiguousSubmission, providerError.Code()) + assert.ErrorIs(t, err, io.ErrUnexpectedEOF) + }) +} + +func TestBridgeContractBoundsPollingRetriesAndRefusesRedirects(t *testing.T) { + for name, retryStatus := range map[string]int{ + "service unavailable": http.StatusServiceUnavailable, + "internal error": http.StatusInternalServerError, + } { + t.Run("bounded transient polling "+name, func(t *testing.T) { + fixture := newBridgeFixture(t) + var polls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.Method == http.MethodPost { + writeBridgeJSON(t, response, http.StatusAccepted, + pendingEnvelope(fixture, "job-retry", JobQueued)) + return + } + if request.Method == http.MethodGet { + polls.Add(1) + writeBridgeJSON(t, response, retryStatus, map[string]any{}) + return + } + response.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, document.RenditionErrorCapacity, providerError.Code()) + assert.Equal(t, int64(4), polls.Load()) + }) + } + + t.Run("redirect refused", func(t *testing.T) { + fixture := newBridgeFixture(t) + var escaped atomic.Int64 + target := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + escaped.Add(1) + })) + t.Cleanup(target.Close) + origin := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + http.Redirect(response, request, target.URL+jobsPath, http.StatusTemporaryRedirect) + })) + t.Cleanup(origin.Close) + client := newTestBridgeClient(t, origin.URL, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + assert.Zero(t, escaped.Load()) + }) +} + +func TestBridgeContractCancelsRemoteJobWhenContextEnds(t *testing.T) { + fixture := newBridgeFixture(t) + var deletes atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.Method { + case http.MethodPost: + writeBridgeJSON(t, response, http.StatusAccepted, + pendingEnvelope(fixture, "job-cancel", JobQueued)) + case http.MethodGet: + <-request.Context().Done() + case http.MethodDelete: + deletes.Add(1) + response.WriteHeader(http.StatusNoContent) + } + })) + t.Cleanup(server.Close) + client := newTestBridgeClient(t, server.URL, fixture.descriptor, nil) + ctx, cancel := context.WithTimeout(t.Context(), 30*time.Millisecond) + defer cancel() + _, err := client.Render(ctx, fixture.upload(), fixture.authorization) + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Eventually(t, func() bool { return deletes.Load() == 1 }, time.Second, time.Millisecond) +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type bridgeFixture struct { + descriptor document.RenditionDescriptor + authorization document.RenditionAuthorization + metadata document.AuthorizedUploadMetadata + source []byte + artifact []byte +} + +func newBridgeFixture(t *testing.T) bridgeFixture { + t.Helper() + source := []byte("synthetic bridge source") + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: "bridge.synthetic", ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), + TrustBoundary: document.RenditionTrustOperatorNetwork, + SupportedFormats: []document.RenditionFormatCapability{{ + MediaFamily: "pdf", MediaType: "application/pdf", + InputKind: document.RenditionInputOriginalFile, + }}, + ReturnsMarkdown: true, + }) + require.NoError(t, err) + metadata := document.AuthorizedUploadMetadata{ + Filename: "document.pdf", MediaFamily: "pdf", MediaType: "application/pdf", + ByteLength: int64(len(source)), SHA256: sha256String(source), + CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + } + authorizedAt := time.Now().UTC().Add(-time.Minute) + authorization := document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, + InputKind: metadata.InputKind, MaxProviderMarkdownBytes: 4096, + MaxTotalResultBytes: 16384, + AuthorizedAt: authorizedAt.Format("2006-01-02T15:04:05.000000000Z"), + ExpiresAt: authorizedAt.Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + } + return bridgeFixture{descriptor: descriptor, authorization: authorization, metadata: metadata, source: source} +} + +func (fixture bridgeFixture) withStructuredArtifact(t *testing.T) bridgeFixture { + t.Helper() + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: fixture.descriptor.ID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: fixture.descriptor.PolicyFingerprint, + TrustBoundary: fixture.descriptor.TrustBoundary, + SupportedFormats: fixture.descriptor.SupportedFormats, + ReturnsMarkdown: true, ReturnsStructured: true, + ArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + }) + require.NoError(t, err) + fixture.descriptor = descriptor + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.AllowedArtifactRoles = []document.EvidenceArtifactRole{document.EvidenceArtifactStructured} + fixture.authorization.MaxArtifactBytes = 4096 + fixture.authorization.MaxArtifacts = 1 + fixture.artifact = []byte(`{"synthetic":"value"}`) + return fixture +} + +func (fixture bridgeFixture) withArtifactAuthorization( + t *testing.T, roles []document.EvidenceArtifactRole, maxArtifacts int, +) bridgeFixture { + t.Helper() + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: fixture.descriptor.ID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: fixture.descriptor.PolicyFingerprint, + TrustBoundary: fixture.descriptor.TrustBoundary, + SupportedFormats: fixture.descriptor.SupportedFormats, + ReturnsMarkdown: true, ReturnsStructured: true, + ArtifactRoles: slices.Clone(roles), + }) + require.NoError(t, err) + fixture.descriptor = descriptor + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.AllowedArtifactRoles = slices.Clone(roles) + fixture.authorization.MaxArtifactBytes = 4096 + fixture.authorization.MaxArtifacts = maxArtifacts + return fixture +} + +func (fixture bridgeFixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func newTestBridgeClient( + t *testing.T, origin string, descriptor document.RenditionDescriptor, secrets SecretResolver, +) *Client { + t.Helper() + return newTestBridgeClientWithHTTP(t, origin, descriptor, secrets, http.DefaultClient) +} + +func newTestBridgeClientWithHTTP( + t *testing.T, origin string, descriptor document.RenditionDescriptor, + secrets SecretResolver, httpClient *http.Client, +) *Client { + t.Helper() + client, err := New(testBridgeProfile(origin, descriptor, secrets), secrets, httpClient) + require.NoError(t, err) + return client +} + +func testBridgeProfile( + origin string, descriptor document.RenditionDescriptor, secrets SecretResolver, +) Profile { + secretBinding := "" + if secrets != nil { + secretBinding = "bridge-api" + } + return Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: secretBinding, + RequestTimeout: time.Second, TotalTimeout: 2 * time.Second, + PollInterval: time.Millisecond, MaxPollAttempts: 4, MaxResponseBytes: 1 << 20, + } +} + +func completedEnvelope( + t *testing.T, fixture bridgeFixture, jobID string, artifacts []map[string]any, +) map[string]any { + t.Helper() + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceDegradedProvenance, + Family: fixture.authorization.MediaFamily, UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", + Reason: "synthetic bridge evidence has generic provenance", + }}, + Units: []document.SourceEvidenceUnitV1{{ + Order: 0, Text: "synthetic bridge evidence", + Locator: document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone, + }, + }}, + } + if len(artifacts) != 0 { + evidence.Artifacts = []document.SourceEvidenceArtifactV1{{ + ProviderID: "structured-1", Pointer: "provider/structured.json", + Role: document.EvidenceArtifactStructured, SHA256: sha256String(fixture.artifact), + }} + } + evidenceJSON, err := json.Marshal(evidence, json.Deterministic(true)) + require.NoError(t, err) + markdown := []byte("synthetic provider markdown\n") + authorizedAt, err := time.Parse("2006-01-02T15:04:05.000000000Z", fixture.authorization.AuthorizedAt) + require.NoError(t, err) + return map[string]any{ + "contract_version": ContractVersion, "status": JobCompleted, "job_id": jobID, + "source_sha256": fixture.authorization.SourceSHA256, + "adapter_id": fixture.descriptor.ID, + "descriptor_fingerprint": fixture.descriptor.Fingerprint, + "policy_fingerprint": fixture.descriptor.PolicyFingerprint, + "result": map[string]any{ + "evidence": map[string]any{ + "media_type": evidenceMediaType, "byte_length": len(evidenceJSON), + "sha256": sha256String(evidenceJSON), "inline": jsontext.Value(evidenceJSON), + }, + "provider_markdown": binaryPayload(markdown, "text/markdown"), + "artifacts": artifacts, + "receipt": document.RenditionReceipt{ + ProviderID: fixture.descriptor.ID, + DescriptorFingerprint: fixture.descriptor.Fingerprint, + PolicyFingerprint: fixture.descriptor.PolicyFingerprint, + SourceSHA256: fixture.authorization.SourceSHA256, OperationID: jobID, + StartedAt: authorizedAt.Add(time.Second).Format("2006-01-02T15:04:05.000000000Z"), + CompletedAt: authorizedAt.Add(2 * time.Second).Format("2006-01-02T15:04:05.000000000Z"), + Usage: document.RenditionUsage{Requests: 1, InputBytes: fixture.authorization.SourceBytes}, + }, + }, + } +} + +func pendingEnvelope(fixture bridgeFixture, jobID string, status JobStatus) map[string]any { + return map[string]any{ + "contract_version": ContractVersion, "status": status, "job_id": jobID, + "source_sha256": fixture.authorization.SourceSHA256, + "adapter_id": fixture.descriptor.ID, + "descriptor_fingerprint": fixture.descriptor.Fingerprint, + "policy_fingerprint": fixture.descriptor.PolicyFingerprint, + } +} + +func binaryPayload(payload []byte, mediaType string) map[string]any { + return map[string]any{ + "media_type": mediaType, "byte_length": len(payload), "sha256": sha256String(payload), + "inline_base64": base64.StdEncoding.EncodeToString(payload), + } +} + +func completedResultMap(value map[string]any) map[string]any { + result, ok := value["result"].(map[string]any) + if !ok { + panic("test completed envelope lacks a result map") + } + return result +} + +func completedEvidence(value map[string]any) map[string]any { + evidence, ok := completedResultMap(value)["evidence"].(map[string]any) + if !ok { + panic("test completed envelope lacks an evidence map") + } + return evidence +} + +func setCompletedEvidence(t *testing.T, value map[string]any, evidence document.SourceEvidenceV1) { + t.Helper() + encoded, err := json.Marshal(evidence, json.Deterministic(true)) + require.NoError(t, err) + payload := completedEvidence(value) + payload["inline"] = jsontext.Value(encoded) + payload["byte_length"] = len(encoded) + payload["sha256"] = sha256String(encoded) +} + +func inlineArtifact(role document.EvidenceArtifactRole, payload []byte) map[string]any { + return map[string]any{ + "role": string(role), "media_type": "application/json", + "byte_length": len(payload), "sha256": sha256String(payload), + "location": "inline", "inline_base64": base64.StdEncoding.EncodeToString(payload), + } +} + +func assertProviderCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + var providerError *document.RenditionProviderError + require.ErrorAs(t, err, &providerError) + assert.Equal(t, want, providerError.Code()) +} + +func assertMultipartRequest( + t *testing.T, request *http.Request, authorization document.RenditionAuthorization, source []byte, +) { + t.Helper() + mediaType, params, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, params["boundary"]) + manifestPart, err := reader.NextPart() + require.NoError(t, err) + assert.Equal(t, authorizationPartName, manifestPart.FormName()) + manifestBytes, err := io.ReadAll(manifestPart) + require.NoError(t, err) + var manifest AuthorizationManifest + require.NoError(t, json.Unmarshal(manifestBytes, &manifest)) + assert.Equal(t, authorization.SourceSHA256, manifest.Authorization.SourceSHA256) + sourcePart, err := reader.NextPart() + require.NoError(t, err) + assert.Equal(t, sourcePartName, sourcePart.FormName()) + assert.Equal(t, "document.pdf", sourcePart.FileName()) + gotSource, err := io.ReadAll(sourcePart) + require.NoError(t, err) + assert.Equal(t, source, gotSource) +} + +func writeBridgeJSON(t *testing.T, response http.ResponseWriter, status int, value any) { + t.Helper() + response.Header().Set("Content-Type", jobMediaType) + response.WriteHeader(status) + require.NoError(t, json.MarshalWrite(response, value, json.Deterministic(true))) +} + +func bridgeHTTPResponse( + t *testing.T, request *http.Request, status int, value any, +) *http.Response { + t.Helper() + encoded, err := json.Marshal(value, json.Deterministic(true)) + require.NoError(t, err) + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{jobMediaType}}, + Body: io.NopCloser(bytes.NewReader(encoded)), + Request: request, + } +} + +func sha256String(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/document/bridge/contract_test.go b/document/bridge/contract_test.go new file mode 100644 index 00000000..be45cc16 --- /dev/null +++ b/document/bridge/contract_test.go @@ -0,0 +1,45 @@ +package bridge + +import ( + "bytes" + _ "embed" + "encoding/json/jsontext" + "encoding/json/v2" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +//go:embed openapi.yaml +var openAPIContract []byte + +//go:embed source-evidence-v1.schema.json +var sourceEvidenceSchema []byte + +func TestBridgeContractNormativeDocumentsAreStrictAndVersioned(t *testing.T) { + var openAPI map[string]any + require.NoError(t, yaml.Unmarshal(openAPIContract, &openAPI)) + assert.Equal(t, "3.1.0", openAPI["openapi"]) + paths, ok := openAPI["paths"].(map[string]any) + require.True(t, ok) + for _, route := range []string{ + jobsPath, jobsPath + "/{job_id}", jobsPath + "/{job_id}/artifacts/{artifact_id}", + } { + assert.Contains(t, paths, route) + } + assert.NotContains(t, string(openAPIContract), "{{", + "the bridge contract must not define a template language") + assert.NotContains(t, string(openAPIContract), "artifact_url", + "artifacts are reachable only through the fixed job route") + + var schema map[string]jsontext.Value + require.NoError(t, json.Unmarshal(sourceEvidenceSchema, &schema, + json.RejectUnknownMembers(true))) + assert.Equal(t, `"https://json-schema.org/draft/2020-12/schema"`, string(schema["$schema"])) + assert.Equal(t, `"object"`, string(schema["type"])) + assert.Equal(t, "false", string(schema["additionalProperties"])) + assert.True(t, bytes.Contains(sourceEvidenceSchema, []byte(`"const": "source-evidence/v1"`))) + assert.True(t, jsontext.Value(sourceEvidenceSchema).IsValid()) +} diff --git a/document/bridge/openapi.yaml b/document/bridge/openapi.yaml new file mode 100644 index 00000000..abf1a870 --- /dev/null +++ b/document/bridge/openapi.yaml @@ -0,0 +1,218 @@ +openapi: 3.1.0 +info: + title: Docbank Rendition Bridge + version: 1.0.0 + description: Fixed provider-neutral rendition protocol. Implementations must refuse redirects and arbitrary result URLs. +paths: + /docbank-rendition/v1/jobs: + post: + operationId: submitRenditionJob + parameters: + - name: Idempotency-Key + in: header + required: true + schema: {$ref: '#/components/schemas/SHA256'} + requestBody: + required: true + content: + multipart/form-data: + schema: + type: object + additionalProperties: false + required: [authorization, source] + properties: + authorization: + $ref: '#/components/schemas/AuthorizationManifest' + source: + type: string + format: binary + encoding: + authorization: + contentType: application/vnd.docbank.rendition-authorization+json;version=1 + responses: + '200': {$ref: '#/components/responses/Job'} + '202': {$ref: '#/components/responses/Job'} + '400': {$ref: '#/components/responses/Job'} + '401': {$ref: '#/components/responses/Job'} + '403': {$ref: '#/components/responses/Job'} + '429': {$ref: '#/components/responses/Job'} + '500': {$ref: '#/components/responses/Job'} + '503': {$ref: '#/components/responses/Job'} + /docbank-rendition/v1/jobs/{job_id}: + parameters: + - $ref: '#/components/parameters/JobID' + get: + operationId: getRenditionJob + responses: + '200': {$ref: '#/components/responses/Job'} + '202': {$ref: '#/components/responses/Job'} + '404': {$ref: '#/components/responses/JobMissing'} + '410': {$ref: '#/components/responses/JobMissing'} + '429': {$ref: '#/components/responses/Job'} + '503': {$ref: '#/components/responses/Job'} + delete: + operationId: cancelRenditionJob + responses: + '200': {$ref: '#/components/responses/Job'} + '204': {description: Job cancellation is complete or was already complete.} + '404': {description: The job is already absent.} + '410': {description: The job result already expired.} + /docbank-rendition/v1/jobs/{job_id}/artifacts/{artifact_id}: + parameters: + - $ref: '#/components/parameters/JobID' + - name: artifact_id + in: path + required: true + schema: {$ref: '#/components/schemas/StableID'} + get: + operationId: getRenditionArtifact + description: Returns bytes only from this fixed same-origin route. Job manifests never contain URLs. + responses: + '200': + description: Exact artifact bytes with the media type declared by the completed job. + content: + application/octet-stream: + schema: {type: string, format: binary} + '404': {description: Artifact or job is unknown.} + '410': {description: Artifact result expired.} +components: + parameters: + JobID: + name: job_id + in: path + required: true + schema: {$ref: '#/components/schemas/StableID'} + responses: + Job: + description: Bounded rendition job state. + content: + application/vnd.docbank.rendition-job+json;version=1: + schema: {$ref: '#/components/schemas/JobEnvelope'} + JobMissing: + description: Job is unknown or its result expired. + schemas: + SHA256: + type: string + pattern: '^[0-9a-f]{64}$' + StableID: + type: string + minLength: 1 + maxLength: 128 + pattern: '^[A-Za-z0-9._-]+$' + AuthorizationManifest: + type: object + additionalProperties: false + required: [contract_version, source, authorization] + properties: + contract_version: {const: docbank-rendition/v1} + source: + type: object + additionalProperties: false + required: [filename, media_family, media_type, byte_length, sha256, capability_record_checksum, provider_metadata_checksum, input_kind] + properties: + filename: {type: string, minLength: 1, maxLength: 255} + media_family: {type: string, minLength: 1, maxLength: 63} + media_type: {type: string, minLength: 3, maxLength: 255} + byte_length: {type: integer, minimum: 1} + sha256: {$ref: '#/components/schemas/SHA256'} + capability_record_checksum: {$ref: '#/components/schemas/SHA256'} + provider_metadata_checksum: {$ref: '#/components/schemas/SHA256'} + input_kind: {enum: [original_file, derived_upload]} + authorization: + type: object + description: Exact canonical RenditionAuthorization from Docbank core. + additionalProperties: false + required: [provider_id, descriptor_fingerprint, policy_fingerprint, rendition_request_fingerprint, source_sha256, source_bytes, capability_record_checksum, provider_metadata_checksum, media_family, media_type, input_kind, allowed_artifact_roles, max_provider_markdown_bytes, max_artifact_bytes, max_artifacts, max_total_result_bytes, authorized_at, expires_at] + properties: + provider_id: {$ref: '#/components/schemas/StableID'} + descriptor_fingerprint: {$ref: '#/components/schemas/SHA256'} + policy_fingerprint: {$ref: '#/components/schemas/SHA256'} + rendition_request_fingerprint: {$ref: '#/components/schemas/SHA256'} + source_sha256: {$ref: '#/components/schemas/SHA256'} + source_bytes: {type: integer, minimum: 1} + capability_record_checksum: {$ref: '#/components/schemas/SHA256'} + provider_metadata_checksum: {$ref: '#/components/schemas/SHA256'} + media_family: {type: string, minLength: 1, maxLength: 63} + media_type: {type: string, minLength: 3, maxLength: 255} + input_kind: {enum: [original_file, derived_upload]} + allowed_artifact_roles: + type: array + uniqueItems: true + maxItems: 16 + items: {enum: [provider_image, provider_markdown, structured_evidence, provider_transcript]} + max_provider_markdown_bytes: {type: integer, minimum: 0} + max_artifact_bytes: {type: integer, minimum: 0} + max_artifacts: {type: integer, minimum: 0} + max_total_result_bytes: {type: integer, minimum: 1} + authorized_at: {type: string, format: date-time} + expires_at: {type: string, format: date-time} + JobEnvelope: + type: object + description: Top-level unknown fields are reserved for minor-version forward compatibility. Nested result and artifact records stay closed. + required: [contract_version, status, job_id, source_sha256, adapter_id, descriptor_fingerprint, policy_fingerprint] + properties: + contract_version: {const: docbank-rendition/v1} + status: {enum: [queued, running, completed, failed, canceled]} + job_id: {$ref: '#/components/schemas/StableID'} + source_sha256: {$ref: '#/components/schemas/SHA256'} + adapter_id: {$ref: '#/components/schemas/StableID'} + descriptor_fingerprint: {$ref: '#/components/schemas/SHA256'} + policy_fingerprint: {$ref: '#/components/schemas/SHA256'} + retry_after_millis: {type: integer, minimum: 0} + result: {$ref: '#/components/schemas/CompletedResult'} + error: {$ref: '#/components/schemas/StableError'} + CompletedResult: + type: object + additionalProperties: false + required: [evidence, receipt] + properties: + evidence: {$ref: '#/components/schemas/EvidencePayload'} + provider_markdown: {$ref: '#/components/schemas/InlinePayload'} + artifacts: + type: array + maxItems: 64 + items: {$ref: '#/components/schemas/ArtifactPayload'} + receipt: {type: object} + EvidencePayload: + type: object + additionalProperties: false + required: [media_type, byte_length, sha256, inline] + properties: + media_type: {const: 'application/vnd.docbank.source-evidence+json;version=1'} + byte_length: {type: integer, minimum: 2} + sha256: {$ref: '#/components/schemas/SHA256'} + inline: {$ref: './source-evidence-v1.schema.json'} + InlinePayload: + type: object + additionalProperties: false + required: [media_type, byte_length, sha256, inline_base64] + properties: + media_type: {type: string, minLength: 3, maxLength: 255} + byte_length: {type: integer, minimum: 0} + sha256: {$ref: '#/components/schemas/SHA256'} + inline_base64: {type: string, contentEncoding: base64} + ArtifactPayload: + type: object + additionalProperties: false + required: [role, media_type, byte_length, sha256, location] + properties: + role: {enum: [provider_image, provider_markdown, structured_evidence, provider_transcript]} + media_type: {type: string, minLength: 3, maxLength: 255} + byte_length: {type: integer, minimum: 0} + sha256: {$ref: '#/components/schemas/SHA256'} + location: {enum: [inline, result]} + inline_base64: {type: string, contentEncoding: base64} + artifact_id: {$ref: '#/components/schemas/StableID'} + oneOf: + - required: [inline_base64] + not: {required: [artifact_id]} + - required: [artifact_id] + not: {required: [inline_base64]} + StableError: + type: object + additionalProperties: false + required: [code, message] + properties: + code: {enum: [unsupported_input, policy_rejected, authentication, capacity, rate_limited, transient, malformed_evidence, unknown_job, canceled, ambiguous_submission]} + message: {type: string, minLength: 1, maxLength: 1024} + retry_after_millis: {type: integer, minimum: 0} diff --git a/document/bridge/source-evidence-v1.schema.json b/document/bridge/source-evidence-v1.schema.json new file mode 100644 index 00000000..860b2dc5 --- /dev/null +++ b/document/bridge/source-evidence-v1.schema.json @@ -0,0 +1,153 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://docbank.dev/schemas/source-evidence-v1.schema.json", + "title": "Docbank source-evidence/v1", + "type": "object", + "additionalProperties": false, + "required": ["contract_version", "completeness", "family", "unit_kind", "units"], + "properties": { + "contract_version": {"const": "source-evidence/v1"}, + "completeness": {"enum": ["complete", "partial", "degraded_provenance"]}, + "family": {"type": "string", "minLength": 1, "maxLength": 1024}, + "unit_kind": {"enum": ["generic", "line", "message", "page", "record", "section", "sheet", "slide", "spine", "time_range"]}, + "artifacts": {"type": "array", "maxItems": 10000, "items": {"$ref": "#/$defs/artifact"}}, + "omissions": {"type": "array", "maxItems": 100000, "items": {"$ref": "#/$defs/omission"}}, + "units": {"type": "array", "minItems": 1, "maxItems": 100000, "items": {"$ref": "#/$defs/unit"}} + }, + "$defs": { + "sha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"}, + "text_range": { + "type": "object", "additionalProperties": false, + "required": ["end", "start"], + "properties": {"end": {"type": "integer", "minimum": 0}, "start": {"type": "integer", "minimum": 0}} + }, + "point": { + "type": "object", "additionalProperties": false, + "required": ["x", "y"], + "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}} + }, + "polygon": { + "type": "object", "additionalProperties": false, + "required": ["points"], + "properties": {"points": {"type": "array", "minItems": 3, "maxItems": 10000, "items": {"$ref": "#/$defs/point"}}} + }, + "box": { + "type": "object", "additionalProperties": false, + "required": ["bottom", "left", "right", "top"], + "properties": { + "bottom": {"type": "integer"}, "left": {"type": "integer"}, + "right": {"type": "integer"}, "top": {"type": "integer"} + } + }, + "geometry": { + "type": "object", "additionalProperties": false, + "required": ["coordinate_origin", "coordinate_space", "height", "orientation", "scale", "unit", "width"], + "properties": { + "boxes": {"type": "array", "maxItems": 1000000, "items": {"$ref": "#/$defs/box"}}, + "coordinate_origin": {"enum": ["bottom_left", "top_left"]}, + "coordinate_space": {"enum": ["image", "page", "unit"]}, + "height": {"type": "integer", "minimum": 0}, + "orientation": {"type": "integer"}, + "polygons": {"type": "array", "maxItems": 1000000, "items": {"$ref": "#/$defs/polygon"}}, + "scale": {"type": "integer", "minimum": 1}, + "unit": {"enum": ["normalized", "pixel", "point"]}, + "width": {"type": "integer", "minimum": 0} + } + }, + "confidence": { + "type": "object", "additionalProperties": false, + "required": ["interpretation", "maximum", "minimum", "value"], + "properties": { + "interpretation": {"enum": ["higher_is_better", "lower_is_better", "probability"]}, + "maximum": {"type": "number"}, "minimum": {"type": "number"}, "value": {"type": "number"} + } + }, + "artifact": { + "type": "object", "additionalProperties": false, + "required": ["pointer", "provider_id", "role", "sha256"], + "properties": { + "pointer": {"type": "string", "minLength": 1, "maxLength": 1024}, + "provider_id": {"type": "string", "minLength": 1, "maxLength": 1024}, + "role": {"enum": ["provider_image", "provider_markdown", "structured_evidence", "provider_transcript"]}, + "sha256": {"$ref": "#/$defs/sha256"} + } + }, + "omission": { + "type": "object", "additionalProperties": false, + "required": ["kind", "reason"], + "properties": { + "field": {"type": "string", "maxLength": 1024}, + "kind": {"enum": ["field", "range", "unit"]}, + "range": {"$ref": "#/$defs/text_range"}, + "reason": {"type": "string", "minLength": 1, "maxLength": 4096}, + "unit_order": {"type": "integer", "minimum": 0} + } + }, + "locator": { + "type": "object", "additionalProperties": false, + "required": ["end", "index_origin", "kind", "start"], + "properties": { + "end": {"type": "integer", "minimum": 0}, + "index_origin": {"enum": ["none", "one", "zero"]}, + "kind": {"enum": ["generic", "line", "message", "page", "record", "section", "sheet", "slide", "spine", "time_range"]}, + "name": {"type": "string", "maxLength": 1024}, + "start": {"type": "integer", "minimum": 0} + } + }, + "region": { + "type": "object", "additionalProperties": false, + "required": ["kind", "order", "provider_id", "text_range"], + "properties": { + "artifact_provider_id": {"type": "string", "maxLength": 1024}, + "confidence": {"$ref": "#/$defs/confidence"}, + "geometry": {"$ref": "#/$defs/geometry"}, + "kind": {"enum": ["code", "figure", "footer", "heading", "header", "image", "list", "paragraph", "table", "table_cell"]}, + "order": {"type": "integer", "minimum": 0}, + "parent_provider_id": {"type": "string", "maxLength": 1024}, + "provider_id": {"type": "string", "minLength": 1, "maxLength": 1024}, + "text_range": {"$ref": "#/$defs/text_range"} + } + }, + "table_cell": { + "type": "object", "additionalProperties": false, + "required": ["column", "column_span", "header", "order", "row", "row_span", "text_range"], + "properties": { + "column": {"type": "integer", "minimum": 0}, + "column_span": {"type": "integer", "minimum": 1}, + "header": {"type": "boolean"}, + "order": {"type": "integer", "minimum": 0}, + "region_provider_id": {"type": "string", "maxLength": 1024}, + "row": {"type": "integer", "minimum": 0}, + "row_span": {"type": "integer", "minimum": 1}, + "text_range": {"$ref": "#/$defs/text_range"} + } + }, + "table": { + "type": "object", "additionalProperties": false, + "required": ["cells", "columns", "order", "provider_id", "rows"], + "properties": { + "cells": {"type": "array", "maxItems": 1000000, "items": {"$ref": "#/$defs/table_cell"}}, + "columns": {"type": "integer", "minimum": 1}, + "order": {"type": "integer", "minimum": 0}, + "provider_id": {"type": "string", "minLength": 1, "maxLength": 1024}, + "region_provider_id": {"type": "string", "maxLength": 1024}, + "rows": {"type": "integer", "minimum": 1} + } + }, + "unit": { + "type": "object", "additionalProperties": false, + "required": ["locator", "order", "text"], + "properties": { + "confidence": {"$ref": "#/$defs/confidence"}, + "heading_path": {"type": "array", "maxItems": 256, "items": {"type": "string", "minLength": 1}}, + "locator": {"$ref": "#/$defs/locator"}, + "omissions": {"type": "array", "maxItems": 100000, "items": {"$ref": "#/$defs/omission"}}, + "order": {"type": "integer", "minimum": 0}, + "provider_id": {"type": "string", "maxLength": 1024}, + "regions": {"type": "array", "maxItems": 1000000, "items": {"$ref": "#/$defs/region"}}, + "tables": {"type": "array", "maxItems": 100000, "items": {"$ref": "#/$defs/table"}}, + "text": {"type": "string", "maxLength": 268435456} + } + } + } +} diff --git a/document/bridge/testdata/unknown-major.json b/document/bridge/testdata/unknown-major.json new file mode 100644 index 00000000..6637f7d8 --- /dev/null +++ b/document/bridge/testdata/unknown-major.json @@ -0,0 +1 @@ +{"adapter_id":"bridge.synthetic","contract_version":"docbank-rendition/v2","descriptor_fingerprint":"0000000000000000000000000000000000000000000000000000000000000000","job_id":"job-future","policy_fingerprint":"1111111111111111111111111111111111111111111111111111111111111111","source_sha256":"2222222222222222222222222222222222222222222222222222222222222222","status":"completed"} diff --git a/document/bridge/types.go b/document/bridge/types.go new file mode 100644 index 00000000..f6decd76 --- /dev/null +++ b/document/bridge/types.go @@ -0,0 +1,135 @@ +// Package bridge implements the fixed docbank-rendition/v1 HTTP bridge. +package bridge + +import ( + "context" + "encoding/json/jsontext" + "net/http" + "time" + + "go.kenn.io/docbank/document" +) + +const ( + ContractVersion = "docbank-rendition/v1" + + jobsPath = "/docbank-rendition/v1/jobs" + authorizationPartName = "authorization" + sourcePartName = "source" + jobMediaType = "application/vnd.docbank.rendition-job+json;version=1" + evidenceMediaType = "application/vnd.docbank.source-evidence+json;version=1" +) + +// JobStatus is one protocol state returned by a bridge. +type JobStatus string + +const ( + JobQueued JobStatus = "queued" + JobRunning JobStatus = "running" + JobCompleted JobStatus = "completed" + JobFailed JobStatus = "failed" + JobCanceled JobStatus = "canceled" +) + +// SecretResolver resolves only a configured named binding. Secret values are +// used for the fixed Authorization header and never enter manifests or receipts. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile fixes one bridge origin, provider identity, and execution bounds. +// Provider ceilings are optional; zero preserves the generic bridge contract. +type Profile struct { + Origin string + Descriptor document.RenditionDescriptor + SecretBinding string + RequestTimeout time.Duration + TotalTimeout time.Duration + PollInterval time.Duration + MaxPollAttempts int + MaxResponseBytes int64 + MaxSourceBytes int64 + MaxProviderMarkdownBytes int + MaxArtifactBytes int + MaxArtifacts int + MaxTotalResultBytes int + MaxEvidenceUnits int +} + +// Client implements document.RenditionProvider through docbank-rendition/v1. +type Client struct { + origin string + descriptor document.RenditionDescriptor + secretBinding string + secrets SecretResolver + http *http.Client + requestTimeout time.Duration + totalTimeout time.Duration + pollInterval time.Duration + maxPollAttempts int + maxResponseBytes int64 + maxSourceBytes int64 + maxProviderMarkdownBytes int + maxArtifactBytes int + maxArtifacts int + maxTotalResultBytes int + maxEvidenceUnits int +} + +// AuthorizationManifest is the canonical multipart policy part sent beside +// the exact source bytes. +type AuthorizationManifest struct { + ContractVersion string `json:"contract_version"` + Source document.AuthorizedUploadMetadata `json:"source"` + Authorization document.RenditionAuthorization `json:"authorization"` +} + +type jobEnvelope struct { + ContractVersion string `json:"contract_version"` + Status JobStatus `json:"status"` + JobID string `json:"job_id"` + SourceSHA256 string `json:"source_sha256"` + AdapterID string `json:"adapter_id"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + RetryAfterMillis int64 `json:"retry_after_millis,omitempty"` + Result jsontext.Value `json:"result,omitempty"` + Error jsontext.Value `json:"error,omitempty"` +} + +type bridgeError struct { + Code document.RenditionErrorCode `json:"code"` + Message string `json:"message"` + RetryAfterMillis int64 `json:"retry_after_millis,omitempty"` +} + +type completedResult struct { + Evidence evidencePayload `json:"evidence"` + ProviderMarkdown *binaryPayloadRecord `json:"provider_markdown,omitempty"` + Artifacts []artifactPayload `json:"artifacts,omitempty"` + Receipt document.RenditionReceipt `json:"receipt"` +} + +type evidencePayload struct { + MediaType string `json:"media_type"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + Inline jsontext.Value `json:"inline"` +} + +type binaryPayloadRecord struct { + MediaType string `json:"media_type"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + InlineBase64 string `json:"inline_base64"` +} + +type artifactPayload struct { + Role document.EvidenceArtifactRole `json:"role"` + MediaType string `json:"media_type"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + Location string `json:"location"` + InlineBase64 string `json:"inline_base64,omitempty"` + ArtifactID string `json:"artifact_id,omitempty"` +} diff --git a/document/chunking.go b/document/chunking.go new file mode 100644 index 00000000..572b5974 --- /dev/null +++ b/document/chunking.go @@ -0,0 +1,1679 @@ +package document + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "hash" + "io" + "slices" + "strconv" + "strings" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +const ( + EmbeddingInputGenerationVersionV1 = 1 + EmbeddingInputGenerationVersion = 2 + + maxAttachmentTitleBytes = 512 + maxAttachmentContextBytes = 4096 + maxGeneratedInputs = 100_000 + maxInputFormatterBytes = 128 + maxTruncationSearchRunes = 4096 + maxNonMonotonicFitChecks = 4096 + maxGenerationTotalTokens = int64(10_000_000) + maxGenerationTotalBytes = int64(1 << 30) + maxGenerationEncodedBytes = int64(16 << 30) + maxGenerationWorkTokens = int64(100_000_000) + maxGenerationWorkBytes = int64(16 << 30) + maxGenerationJSONIntegerBytes = 20 +) + +// TruncationPolicy declares what happens when one tokenizer token cannot fit +// the provider's complete rendered request boundary. +type TruncationPolicy string + +const ( + TruncationPolicyReject TruncationPolicy = "reject_indivisible" + TruncationPolicyTruncateIndivisible TruncationPolicy = "truncate_indivisible" + // TruncateIndivisibleAtom is retained as the provisional E2 spelling. + TruncateIndivisibleAtom = TruncationPolicyTruncateIndivisible +) + +// AttachmentContextSnapshotConfig is explicitly supplied human-authored +// context. Paths, collection names, and other mutable navigation metadata do +// not belong here. +type AttachmentContextSnapshotConfig struct { + Title string + Context string +} + +// AttachmentContextSnapshot is immutable after construction: its bounded +// strings are private and only value accessors are exposed. +type AttachmentContextSnapshot struct { + title string + context string +} + +// NewAttachmentContextSnapshot validates and seals human-authored attachment +// context. An all-empty snapshot is not a declared attachment context. +func NewAttachmentContextSnapshot(config AttachmentContextSnapshotConfig) (AttachmentContextSnapshot, error) { + if config.Title == "" && config.Context == "" { + return AttachmentContextSnapshot{}, errors.New("attachment context snapshot must contain a title or context") + } + for _, field := range []struct { + name string + value string + limit int + }{{"title", config.Title, maxAttachmentTitleBytes}, {"context", config.Context, maxAttachmentContextBytes}} { + if !utf8.ValidString(field.value) || len(field.value) > field.limit || strings.ContainsAny(field.value, "\x00\r") || !norm.NFC.IsNormalString(field.value) { + return AttachmentContextSnapshot{}, fmt.Errorf("attachment context %s must be bounded valid UTF-8 NFC text", field.name) + } + } + return AttachmentContextSnapshot{title: config.Title, context: config.Context}, nil +} + +func (snapshot AttachmentContextSnapshot) Title() string { return snapshot.title } +func (snapshot AttachmentContextSnapshot) Context() string { return snapshot.context } +func (snapshot AttachmentContextSnapshot) declared() bool { + return snapshot.title != "" || snapshot.context != "" +} + +func (snapshot AttachmentContextSnapshot) MarshalJSON() ([]byte, error) { + return json.Marshal(struct { + Title string `json:"title,omitempty"` + Context string `json:"context,omitempty"` + }{snapshot.title, snapshot.context}) +} + +func (snapshot *AttachmentContextSnapshot) UnmarshalJSON(data []byte) error { + var encoded struct { + Title string `json:"title,omitempty"` + Context string `json:"context,omitempty"` + } + if err := decodeStrictJSON(data, &encoded); err != nil { + return fmt.Errorf("decode attachment context: %w", err) + } + canonical, err := NewAttachmentContextSnapshot(AttachmentContextSnapshotConfig{Title: encoded.Title, Context: encoded.Context}) + if err != nil { + return err + } + *snapshot = canonical + return nil +} + +// InputPolicy seals one deterministic, document-role input generation. +type InputPolicy struct { + Tokenizer Tokenizer + ContentTokenBudget int + OverlapTokens int + MaxProviderTokens int + MaxProviderBytes int64 + MaxGeneratedInputs int + MaxTotalContentTokens int64 + MaxTotalRenderedTokens int64 + MaxTotalContentBytes int64 + MaxTotalRenderedBytes int64 + MaxFittingWorkTokens int64 + MaxFittingWorkBytes int64 + ModelInput ModelInputContract + Formatter string + LexicalEvidenceFingerprint string + ContextFingerprint string + AttachmentContext *AttachmentContextSnapshot + TruncationPolicy TruncationPolicy +} + +// GeneratedEmbeddingInput is one exact provider-ready document input. +type GeneratedEmbeddingInput struct { + Key string `json:"key"` + Content string `json:"content"` + Rendered string `json:"rendered"` + ContentTokens int `json:"content_tokens"` + RenderedTokens int `json:"rendered_tokens"` + Checksum string `json:"checksum"` + HeadingPaths [][]string `json:"heading_paths,omitempty"` + SourceSpans []ChunkSpan `json:"source_spans"` + Truncated bool `json:"truncated"` +} + +type generatedEmbeddingInputJSON struct { + Key string `json:"key"` + Content string `json:"content"` + Rendered string `json:"rendered"` + ContentTokens int `json:"content_tokens"` + RenderedTokens int `json:"rendered_tokens"` + Checksum string `json:"checksum"` + HeadingPaths [][]string `json:"heading_paths,omitempty"` + SourceSpans []generatedSourceSpanJSON `json:"source_spans"` + Truncated bool `json:"truncated"` +} + +type generatedSourceSpanJSON struct { + UnitIndex int `json:"unit_index"` + CharStart int `json:"char_start"` + CharEnd int `json:"char_end"` +} + +// MarshalJSON gives generated inputs a stable wire contract without changing +// the deliberately unserialized legacy ChunkSpan contract. +func (input GeneratedEmbeddingInput) MarshalJSON() ([]byte, error) { + encoded := generatedEmbeddingInputJSON{ + Key: input.Key, Content: input.Content, Rendered: input.Rendered, + ContentTokens: input.ContentTokens, RenderedTokens: input.RenderedTokens, + Checksum: input.Checksum, HeadingPaths: input.HeadingPaths, Truncated: input.Truncated, + SourceSpans: make([]generatedSourceSpanJSON, len(input.SourceSpans)), + } + for index, span := range input.SourceSpans { + encoded.SourceSpans[index] = generatedSourceSpanJSON(span) + } + return json.Marshal(encoded) +} + +func (input *GeneratedEmbeddingInput) UnmarshalJSON(data []byte) error { + var encoded generatedEmbeddingInputJSON + if err := decodeStrictJSON(data, &encoded); err != nil { + return fmt.Errorf("decode generated embedding input: %w", err) + } + input.Key = encoded.Key + input.Content = encoded.Content + input.Rendered = encoded.Rendered + input.ContentTokens = encoded.ContentTokens + input.RenderedTokens = encoded.RenderedTokens + input.Checksum = encoded.Checksum + input.HeadingPaths = cloneHeadingPaths(encoded.HeadingPaths) + input.Truncated = encoded.Truncated + input.SourceSpans = make([]ChunkSpan, len(encoded.SourceSpans)) + for index, span := range encoded.SourceSpans { + input.SourceSpans[index] = ChunkSpan(span) + } + return nil +} + +// ToEmbeddingInputs reconstructs the contextualized pre-envelope text for E1. +// The complete model-input fingerprint must match before any input is exposed, +// even when two contracts happen to share a document envelope. +func (generation EmbeddingInputGeneration) ToEmbeddingInputs(contract ModelInputContract) ([]EmbeddingInput, error) { + if err := validateModelInputContract(contract); err != nil { + return nil, err + } + if contract.Fingerprint != generation.ModelInputFingerprint { + return nil, errors.New("embedding generation model-input fingerprint does not match contract") + } + if err := validateEmbeddingInputGeneration(generation); err != nil { + return nil, err + } + attachment := AttachmentContextSnapshot{} + if generation.AttachmentContext != nil { + attachment = *generation.AttachmentContext + } + result := make([]EmbeddingInput, len(generation.Inputs)) + for index, input := range generation.Inputs { + contextualized := contextualizeDocumentInput(attachment, input.Content) + if contract.EncodeDocument(contextualized) != input.Rendered || sha256Hex([]byte(input.Rendered)) != input.Checksum { + return nil, errors.New("generated embedding input does not match model-input contract") + } + result[index] = EmbeddingInput{ + Key: input.Key, Role: EmbeddingRoleDocument, Kind: EmbeddingInputRenditionChunk, + Text: contextualized, HeadingPath: slices.Clone(input.HeadingPaths[0]), SourceSpans: slices.Clone(input.SourceSpans), + } + } + return result, nil +} + +// EmbeddingInputGeneration is one ordered projection of normalized evidence. +// Callers persist its exact JSON or clone it before exposing mutable slices. +type EmbeddingInputGeneration struct { + Version int `json:"version"` + Checksum string `json:"checksum"` + PolicyFingerprint string `json:"policy_fingerprint"` + EvidenceChecksum string `json:"evidence_checksum"` + TokenizerIdentity TokenizerIdentity `json:"tokenizer_identity"` + LexicalEvidenceFingerprint string `json:"lexical_evidence_fingerprint"` + Formatter string `json:"formatter"` + ModelInputFingerprint string `json:"model_input_fingerprint"` + ContentTokenBudget int `json:"content_token_budget"` + OverlapTokens int `json:"overlap_tokens"` + TruncationPolicy TruncationPolicy `json:"truncation_policy"` + ContextFingerprint string `json:"context_fingerprint"` + AttachmentContext *AttachmentContextSnapshot `json:"attachment_context,omitempty"` + TotalContentTokens int64 `json:"total_content_tokens"` + TotalRenderedTokens int64 `json:"total_rendered_tokens"` + TotalContentBytes int64 `json:"total_content_bytes"` + TotalRenderedBytes int64 `json:"total_rendered_bytes"` + Inputs []GeneratedEmbeddingInput `json:"inputs"` +} + +// embeddingInputGenerationV1 is the exact published E2 wire contract. It is +// retained as a decode and checksum boundary; the policy-complete fields added +// for catalog authority belong only to v2. +type embeddingInputGenerationV1 struct { + Version int `json:"version"` + Checksum string `json:"checksum"` + PolicyFingerprint string `json:"policy_fingerprint"` + EvidenceChecksum string `json:"evidence_checksum"` + TokenizerIdentity TokenizerIdentity `json:"tokenizer_identity"` + LexicalEvidenceFingerprint string `json:"lexical_evidence_fingerprint"` + Formatter string `json:"formatter"` + ModelInputFingerprint string `json:"model_input_fingerprint"` + AttachmentContext *AttachmentContextSnapshot `json:"attachment_context,omitempty"` + TotalContentTokens int64 `json:"total_content_tokens"` + TotalRenderedTokens int64 `json:"total_rendered_tokens"` + TotalContentBytes int64 `json:"total_content_bytes"` + TotalRenderedBytes int64 `json:"total_rendered_bytes"` + Inputs []GeneratedEmbeddingInput `json:"inputs"` +} + +type embeddingInputGenerationV2 EmbeddingInputGeneration + +func (generation EmbeddingInputGeneration) MarshalJSON() ([]byte, error) { + if generation.Version == EmbeddingInputGenerationVersionV1 { + return json.Marshal(embeddingInputGenerationV1{ + Version: generation.Version, Checksum: generation.Checksum, + PolicyFingerprint: generation.PolicyFingerprint, EvidenceChecksum: generation.EvidenceChecksum, + TokenizerIdentity: generation.TokenizerIdentity, LexicalEvidenceFingerprint: generation.LexicalEvidenceFingerprint, + Formatter: generation.Formatter, ModelInputFingerprint: generation.ModelInputFingerprint, + AttachmentContext: generation.AttachmentContext, TotalContentTokens: generation.TotalContentTokens, + TotalRenderedTokens: generation.TotalRenderedTokens, TotalContentBytes: generation.TotalContentBytes, + TotalRenderedBytes: generation.TotalRenderedBytes, Inputs: generation.Inputs, + }) + } + return json.Marshal(embeddingInputGenerationV2(generation)) +} + +func (generation *EmbeddingInputGeneration) UnmarshalJSON(data []byte) error { + var header struct { + Version int `json:"version"` + } + if err := json.Unmarshal(data, &header); err != nil { + return err + } + switch header.Version { + case EmbeddingInputGenerationVersionV1: + var encoded embeddingInputGenerationV1 + if err := decodeStrictJSON(data, &encoded); err != nil { + return err + } + *generation = EmbeddingInputGeneration{ + Version: encoded.Version, Checksum: encoded.Checksum, + PolicyFingerprint: encoded.PolicyFingerprint, EvidenceChecksum: encoded.EvidenceChecksum, + TokenizerIdentity: encoded.TokenizerIdentity, LexicalEvidenceFingerprint: encoded.LexicalEvidenceFingerprint, + Formatter: encoded.Formatter, ModelInputFingerprint: encoded.ModelInputFingerprint, + AttachmentContext: encoded.AttachmentContext, TotalContentTokens: encoded.TotalContentTokens, + TotalRenderedTokens: encoded.TotalRenderedTokens, TotalContentBytes: encoded.TotalContentBytes, + TotalRenderedBytes: encoded.TotalRenderedBytes, Inputs: encoded.Inputs, + } + return nil + case EmbeddingInputGenerationVersion: + var encoded embeddingInputGenerationV2 + if err := decodeStrictJSON(data, &encoded); err != nil { + return err + } + *generation = EmbeddingInputGeneration(encoded) + return nil + default: + return fmt.Errorf("embedding input generation version must be %d or %d", EmbeddingInputGenerationVersionV1, EmbeddingInputGenerationVersion) + } +} + +// EmbeddingInputGenerationDecodeBounds authorizes the encoded artifact and +// its allocations before canonical decoding begins. +type EmbeddingInputGenerationDecodeBounds struct { + MaxEncodedBytes int64 + MaxInputs int + MaxObjectFields int + MaxStringBytes int64 + MaxTotalStringBytes int64 +} + +// DecodeEmbeddingInputGeneration strictly decodes and validates the complete +// sealed generation. Unknown fields, lossy spans, forged totals, and stale +// checksums are rejected. +func DecodeEmbeddingInputGeneration(data []byte, bounds EmbeddingInputGenerationDecodeBounds) (EmbeddingInputGeneration, error) { + if err := preflightEmbeddingInputGenerationJSON(data, bounds); err != nil { + return EmbeddingInputGeneration{}, fmt.Errorf("preflight embedding input generation: %w", err) + } + var generation EmbeddingInputGeneration + if err := decodeStrictJSON(data, &generation); err != nil { + return EmbeddingInputGeneration{}, fmt.Errorf("decode embedding input generation: %w", err) + } + if err := validateEmbeddingInputGeneration(generation); err != nil { + return EmbeddingInputGeneration{}, err + } + return generation, nil +} + +func preflightEmbeddingInputGenerationJSON(data []byte, bounds EmbeddingInputGenerationDecodeBounds) error { + if bounds.MaxEncodedBytes < 1 || bounds.MaxEncodedBytes > maxGenerationEncodedBytes || int64(len(data)) > bounds.MaxEncodedBytes { + return errors.New("embedding generation encoded bytes exceed bounds") + } + if bounds.MaxInputs < 1 || bounds.MaxInputs > maxGeneratedInputs { + return errors.New("embedding generation input decode bound is invalid") + } + if bounds.MaxObjectFields < 1 || bounds.MaxObjectFields > 64 { + return errors.New("embedding generation object field decode bound is invalid") + } + if bounds.MaxStringBytes < 1 || bounds.MaxStringBytes > bounds.MaxEncodedBytes || bounds.MaxTotalStringBytes < 1 || bounds.MaxTotalStringBytes > bounds.MaxEncodedBytes { + return errors.New("embedding generation string decode bounds are invalid") + } + scanner := generationJSONPreflight{data: data, bounds: bounds} + if err := scanner.value(generationJSONPathRoot, 0); err != nil { + return err + } + scanner.skipSpace() + if scanner.position != len(data) { + return errors.New("JSON contains trailing value") + } + return nil +} + +type generationJSONPreflight struct { + data []byte + bounds EmbeddingInputGenerationDecodeBounds + position int + rawStringBytes int64 +} + +type generationJSONPath uint8 + +const ( + generationJSONPathOther generationJSONPath = iota + generationJSONPathRoot + generationJSONPathInputs + generationJSONPathInput + generationJSONPathHeadingPaths + generationJSONPathHeadingPath + generationJSONPathSourceSpans +) + +func (scanner *generationJSONPreflight) value(path generationJSONPath, depth int) error { + if depth > 16 { + return errors.New("embedding generation JSON nesting exceeds bounds") + } + scanner.skipSpace() + if scanner.position >= len(scanner.data) { + return io.ErrUnexpectedEOF + } + switch scanner.data[scanner.position] { + case '{': + return scanner.object(path, depth) + case '[': + return scanner.array(path, depth) + case '"': + _, _, _, err := scanner.stringValue() + return err + case 't': + return scanner.keyword("true") + case 'f': + return scanner.keyword("false") + case 'n': + return scanner.keyword("null") + default: + return scanner.number() + } +} + +func (scanner *generationJSONPreflight) object(path generationJSONPath, depth int) error { + scanner.position++ + scanner.skipSpace() + if scanner.consume('}') { + return nil + } + for fields := 1; ; fields++ { + if fields > scanner.bounds.MaxObjectFields { + return errors.New("embedding generation JSON object fields exceed bounds") + } + scanner.skipSpace() + keyStart, keyEnd, escaped, err := scanner.stringValue() + if err != nil { + return err + } + if escaped || !asciiJSONField(scanner.data[keyStart:keyEnd]) { + return errors.New("embedding generation JSON field names must be unescaped ASCII") + } + scanner.skipSpace() + if !scanner.consume(':') { + return errors.New("embedding generation JSON object field is missing a colon") + } + childPath := generationJSONChildPath(path, scanner.data[keyStart:keyEnd]) + if err := scanner.value(childPath, depth+1); err != nil { + return err + } + scanner.skipSpace() + if scanner.consume('}') { + return nil + } + if !scanner.consume(',') { + return errors.New("embedding generation JSON object is not closed") + } + } +} + +func (scanner *generationJSONPreflight) array(path generationJSONPath, depth int) error { + scanner.position++ + scanner.skipSpace() + if scanner.consume(']') { + return nil + } + limit := scanner.arrayLimit(path) + for count := 1; ; count++ { + if count > limit { + return fmt.Errorf("embedding generation JSON %s collection exceeds bounds", generationJSONPathName(path)) + } + if err := scanner.value(generationJSONArrayChildPath(path), depth+1); err != nil { + return err + } + scanner.skipSpace() + if scanner.consume(']') { + return nil + } + if !scanner.consume(',') { + return errors.New("embedding generation JSON array is not closed") + } + } +} + +func (scanner *generationJSONPreflight) stringValue() (int, int, bool, error) { + if !scanner.consume('"') { + return 0, 0, false, errors.New("embedding generation JSON object key is invalid") + } + start := scanner.position + escaped := false + var rawLength int64 + for scanner.position < len(scanner.data) { + value := scanner.data[scanner.position] + if value == '"' { + end := scanner.position + scanner.position++ + return start, end, escaped, nil + } + if value < 0x20 { + return 0, 0, false, errors.New("embedding generation JSON string contains a control byte") + } + if err := scanner.addRawStringByte(&rawLength); err != nil { + return 0, 0, false, err + } + scanner.position++ + if value != '\\' { + continue + } + escaped = true + if scanner.position >= len(scanner.data) { + return 0, 0, false, io.ErrUnexpectedEOF + } + escape := scanner.data[scanner.position] + if err := scanner.addRawStringByte(&rawLength); err != nil { + return 0, 0, false, err + } + scanner.position++ + if escape == 'u' { + for range 4 { + if scanner.position >= len(scanner.data) { + return 0, 0, false, io.ErrUnexpectedEOF + } + if !isJSONHex(scanner.data[scanner.position]) { + return 0, 0, false, errors.New("embedding generation JSON unicode escape is invalid") + } + if err := scanner.addRawStringByte(&rawLength); err != nil { + return 0, 0, false, err + } + scanner.position++ + } + continue + } + if !strings.ContainsRune(`"\\/bfnrt`, rune(escape)) { + return 0, 0, false, errors.New("embedding generation JSON escape is invalid") + } + } + return 0, 0, false, io.ErrUnexpectedEOF +} + +func (scanner *generationJSONPreflight) addRawStringByte(rawLength *int64) error { + if *rawLength >= scanner.bounds.MaxStringBytes || scanner.rawStringBytes >= scanner.bounds.MaxTotalStringBytes { + return errors.New("embedding generation JSON raw string bytes exceed bounds") + } + *rawLength++ + scanner.rawStringBytes++ + return nil +} + +func (scanner *generationJSONPreflight) keyword(value string) error { + if len(scanner.data)-scanner.position < len(value) || string(scanner.data[scanner.position:scanner.position+len(value)]) != value { + return errors.New("embedding generation JSON literal is invalid") + } + scanner.position += len(value) + return nil +} + +func (scanner *generationJSONPreflight) number() error { + start := scanner.position + if scanner.position < len(scanner.data) && scanner.data[scanner.position] == '-' { + scanner.position++ + } + if scanner.position >= len(scanner.data) || scanner.data[scanner.position] < '0' || scanner.data[scanner.position] > '9' { + return errors.New("embedding generation JSON integer is invalid") + } + if scanner.data[scanner.position] == '0' { + scanner.position++ + if scanner.position < len(scanner.data) && scanner.data[scanner.position] >= '0' && scanner.data[scanner.position] <= '9' { + return errors.New("embedding generation JSON integer has a leading zero") + } + } else { + for scanner.position < len(scanner.data) && scanner.data[scanner.position] >= '0' && scanner.data[scanner.position] <= '9' { + scanner.position++ + if scanner.position-start > maxGenerationJSONIntegerBytes { + return errors.New("embedding generation JSON integer token exceeds bounds") + } + } + } + if scanner.position-start > maxGenerationJSONIntegerBytes { + return errors.New("embedding generation JSON integer token exceeds bounds") + } + if scanner.position < len(scanner.data) && !isJSONValueTerminator(scanner.data[scanner.position]) { + return errors.New("embedding generation JSON integer grammar is invalid") + } + return nil +} + +func (scanner *generationJSONPreflight) skipSpace() { + for scanner.position < len(scanner.data) && strings.ContainsRune(" \t\r\n", rune(scanner.data[scanner.position])) { + scanner.position++ + } +} + +func (scanner *generationJSONPreflight) consume(value byte) bool { + if scanner.position >= len(scanner.data) || scanner.data[scanner.position] != value { + return false + } + scanner.position++ + return true +} + +func (scanner *generationJSONPreflight) arrayLimit(path generationJSONPath) int { + switch path { + case generationJSONPathInputs: + return scanner.bounds.MaxInputs + case generationJSONPathHeadingPaths, generationJSONPathSourceSpans: + return 1 + case generationJSONPathHeadingPath: + return maxEvidenceHeadingDepth + default: + return max(scanner.bounds.MaxInputs, maxEvidenceHeadingDepth) + } +} + +func generationJSONChildPath(parent generationJSONPath, key []byte) generationJSONPath { + if parent == generationJSONPathRoot && bytes.Equal(key, []byte("inputs")) { + return generationJSONPathInputs + } + if parent == generationJSONPathInput { + switch { + case bytes.Equal(key, []byte("heading_paths")): + return generationJSONPathHeadingPaths + case bytes.Equal(key, []byte("source_spans")): + return generationJSONPathSourceSpans + } + } + return generationJSONPathOther +} + +func generationJSONArrayChildPath(parent generationJSONPath) generationJSONPath { + switch parent { + case generationJSONPathInputs: + return generationJSONPathInput + case generationJSONPathHeadingPaths: + return generationJSONPathHeadingPath + default: + return generationJSONPathOther + } +} + +func generationJSONPathName(path generationJSONPath) string { + switch path { + case generationJSONPathInputs: + return "inputs" + case generationJSONPathHeadingPaths: + return "heading paths" + case generationJSONPathHeadingPath: + return "heading parts" + case generationJSONPathSourceSpans: + return "source spans" + default: + return "unknown" + } +} + +func asciiJSONField(value []byte) bool { + if len(value) == 0 { + return false + } + for _, character := range value { + if character < 0x20 || character > 0x7e { + return false + } + } + return true +} + +func isJSONHex(value byte) bool { + return value >= '0' && value <= '9' || value >= 'a' && value <= 'f' || value >= 'A' && value <= 'F' +} + +func isJSONValueTerminator(value byte) bool { + return value == ',' || value == ']' || value == '}' || strings.ContainsRune(" \t\r\n", rune(value)) +} + +func decodeStrictJSON(data []byte, target any) error { + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return err + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return errors.New("JSON contains trailing value") + } + return err + } + return nil +} + +func validateEmbeddingInputGeneration(generation EmbeddingInputGeneration) error { + if generation.Version != EmbeddingInputGenerationVersionV1 && generation.Version != EmbeddingInputGenerationVersion { + return fmt.Errorf("embedding input generation version must be %d or %d", EmbeddingInputGenerationVersionV1, EmbeddingInputGenerationVersion) + } + for _, fingerprint := range []struct{ value, name string }{ + {generation.Checksum, "generation checksum"}, + {generation.PolicyFingerprint, "generation policy fingerprint"}, + {generation.EvidenceChecksum, "generation evidence checksum"}, + {generation.LexicalEvidenceFingerprint, "generation lexical evidence fingerprint"}, + {generation.ModelInputFingerprint, "generation model-input fingerprint"}, + } { + if err := validateFingerprint(fingerprint.value, fingerprint.name); err != nil { + return err + } + } + if err := validateTokenizerIdentity(generation.TokenizerIdentity); err != nil { + return err + } + if err := validateCompatibilityID(generation.Formatter); err != nil { + return fmt.Errorf("embedding generation formatter: %w", err) + } + if generation.AttachmentContext != nil && !generation.AttachmentContext.declared() { + return errors.New("embedding generation attachment context is invalid") + } + if generation.Version == EmbeddingInputGenerationVersionV1 { + if generation.ContentTokenBudget != 0 || generation.OverlapTokens != 0 || generation.TruncationPolicy != "" || generation.ContextFingerprint != "" { + return errors.New("embedding generation v1 contains v2 policy authority") + } + } else { + if err := validateFingerprint(generation.ContextFingerprint, "generation context fingerprint"); err != nil { + return err + } + if generation.ContentTokenBudget < 1 || generation.ContentTokenBudget > maxEmbeddingChunkTokens || + generation.OverlapTokens < 0 || generation.OverlapTokens >= generation.ContentTokenBudget { + return errors.New("embedding generation chunk policy is invalid") + } + switch generation.TruncationPolicy { + case TruncationPolicyReject, TruncationPolicyTruncateIndivisible: + default: + return errors.New("embedding generation truncation policy is invalid") + } + } + if len(generation.Inputs) > maxGeneratedInputs { + return errors.New("embedding generation input count exceeds bounds") + } + var contentTokens, renderedTokens, contentBytes, renderedBytes int64 + for index, input := range generation.Inputs { + if err := validateStableToken(input.Key, "generated embedding input key", 128); err != nil { + return err + } + if input.Key != fmt.Sprintf("chunk-%06d-%s", index, input.Checksum[:min(12, len(input.Checksum))]) { + return errors.New("generated embedding input key is not canonical") + } + if input.Content == "" || input.Rendered == "" || !utf8.ValidString(input.Content) || !utf8.ValidString(input.Rendered) { + return errors.New("generated embedding input text is invalid") + } + if input.ContentTokens < 1 || input.ContentTokens > maxEmbeddingChunkTokens || input.RenderedTokens < 1 || input.RenderedTokens > maxEmbeddingChunkTokens { + return errors.New("generated embedding input token counts are invalid") + } + if sha256Hex([]byte(input.Rendered)) != input.Checksum { + return errors.New("generated embedding input checksum is invalid") + } + if err := validateEmbeddingInputAuxiliaries(EmbeddingInput{Role: EmbeddingRoleDocument, Kind: EmbeddingInputRenditionChunk, Text: input.Content, HeadingPath: firstHeadingPath(input.HeadingPaths), SourceSpans: input.SourceSpans}); err != nil { + return fmt.Errorf("generated embedding input auxiliaries: %w", err) + } + if len(input.HeadingPaths) != 1 || len(input.SourceSpans) != 1 { + return errors.New("generated embedding input collection canonical cardinality is invalid") + } + for _, span := range input.SourceSpans { + if span.UnitIndex < 0 || span.CharStart < 0 || span.CharEnd <= span.CharStart { + return errors.New("generated embedding input source span is invalid") + } + } + if !addWithinAggregate(contentTokens, int64(input.ContentTokens), maxGenerationTotalTokens) || + !addWithinAggregate(renderedTokens, int64(input.RenderedTokens), maxGenerationTotalTokens) || + !addWithinAggregate(contentBytes, int64(len(input.Content)), maxGenerationTotalBytes) || + !addWithinAggregate(renderedBytes, int64(len(input.Rendered)), maxGenerationTotalBytes) { + return errors.New("embedding generation aggregate values exceed bounds") + } + contentTokens += int64(input.ContentTokens) + renderedTokens += int64(input.RenderedTokens) + contentBytes += int64(len(input.Content)) + renderedBytes += int64(len(input.Rendered)) + } + if generation.TotalContentTokens != contentTokens || generation.TotalRenderedTokens != renderedTokens || generation.TotalContentBytes != contentBytes || generation.TotalRenderedBytes != renderedBytes { + return errors.New("embedding generation aggregate totals are invalid") + } + checksum := generation.Checksum + generation.Checksum = "" + if generationFingerprint(generation) != checksum { + return errors.New("embedding input generation checksum is invalid") + } + return nil +} + +func cloneHeadingPaths(source [][]string) [][]string { + result := make([][]string, len(source)) + for index, heading := range source { + result[index] = slices.Clone(heading) + } + return result +} + +func firstHeadingPath(paths [][]string) []string { + if len(paths) == 0 { + return nil + } + return paths[0] +} + +type embeddingToken struct { + unitIndex int + start int + end int + unitEnd int +} + +// BuildEmbeddingInputs derives provider inputs from canonical evidence only. +// It never reads a retained Markdown or YAML rendition. +func BuildEmbeddingInputs(evidence NormalizedEvidenceV1, policy InputPolicy) (EmbeddingInputGeneration, error) { + if nilInterface(policy.Tokenizer) { + return EmbeddingInputGeneration{}, errors.New("embedding input policy requires a tokenizer") + } + tokenizerIdentity := policy.Tokenizer.Identity() + if _, checksum, err := MarshalNormalizedEvidenceV1(evidence); err != nil || checksum != evidence.Checksum { + if err != nil { + return EmbeddingInputGeneration{}, fmt.Errorf("validate normalized evidence: %w", err) + } + return EmbeddingInputGeneration{}, errors.New("normalized evidence checksum is invalid") + } + if err := validateInputPolicy(policy, tokenizerIdentity); err != nil { + return EmbeddingInputGeneration{}, err + } + + tokens, naturalEnds, err := tokenizeEvidence(evidence, policy.Tokenizer) + if err != nil { + return EmbeddingInputGeneration{}, err + } + attachment := AttachmentContextSnapshot{} + if policy.AttachmentContext != nil { + attachment = *policy.AttachmentContext + } + result := EmbeddingInputGeneration{ + Version: EmbeddingInputGenerationVersion, EvidenceChecksum: evidence.Checksum, + TokenizerIdentity: tokenizerIdentity, LexicalEvidenceFingerprint: policy.LexicalEvidenceFingerprint, + Formatter: policy.Formatter, ModelInputFingerprint: policy.ModelInput.Fingerprint, + ContentTokenBudget: policy.ContentTokenBudget, OverlapTokens: policy.OverlapTokens, + TruncationPolicy: policy.TruncationPolicy, ContextFingerprint: policy.ContextFingerprint, + Inputs: make([]GeneratedEmbeddingInput, 0, generatedInputCapacity(policy, len(tokens))), + } + if attachment.declared() { + result.AttachmentContext = &attachment + } + result.PolicyFingerprint = inputPolicyFingerprint(policy, attachment, tokenizerIdentity) + work := fittingWorkBudget{remainingTokens: policy.MaxFittingWorkTokens, remainingBytes: policy.MaxFittingWorkBytes} + + for start := 0; start < len(tokens); { + if len(result.Inputs) == policy.MaxGeneratedInputs { + return EmbeddingInputGeneration{}, errors.New("embedding generated input limit exceeded") + } + limits, err := remainingInputLimits(result, policy, &work) + if err != nil { + return EmbeddingInputGeneration{}, err + } + unitEnd := tokens[start].unitEnd + end := chooseChunkEnd(start, unitEnd, policy.ContentTokenBudget, policy.OverlapTokens, naturalEnds) + input, contentBoundaries, fittedEnd, truncated, err := fitGeneratedInput(evidence, tokens, naturalEnds, start, end, policy, tokenizerIdentity, attachment, limits, len(result.Inputs)) + if err != nil { + return EmbeddingInputGeneration{}, err + } + input.Truncated = truncated + if err := addGeneratedInputTotals(&result, input, policy); err != nil { + return EmbeddingInputGeneration{}, err + } + result.Inputs = append(result.Inputs, input) + if truncated || fittedEnd == unitEnd { + start = fittedEnd + continue + } + if policy.OverlapTokens == 0 { + start = fittedEnd + continue + } + if len(contentBoundaries) <= policy.OverlapTokens || len(input.SourceSpans) != 1 { + if limits.contentTokens <= policy.OverlapTokens { + return EmbeddingInputGeneration{}, errors.New("embedding generation exceeds aggregate limits while preserving configured token overlap") + } + return EmbeddingInputGeneration{}, errors.New("provider limits cannot preserve configured token overlap") + } + desiredRuneStart := input.SourceSpans[0].CharStart + contentBoundaries[len(contentBoundaries)-policy.OverlapTokens].Start + next := start + for next < fittedEnd && tokens[next].end <= desiredRuneStart { + next++ + } + if next == fittedEnd { + return EmbeddingInputGeneration{}, errors.New("exact token overlap does not advance within the source unit") + } + tokens[next].start = desiredRuneStart + start = next + } + if policy.Tokenizer.Identity() != tokenizerIdentity { + return EmbeddingInputGeneration{}, errors.New("embedding tokenizer identity changed during generation") + } + result.Checksum = generationFingerprint(result) + return result, nil +} + +type inputConstructionLimits struct { + contentTokens int + renderedTokens int + contentBytes int64 + renderedBytes int64 + work *fittingWorkBudget +} + +type fittingWorkBudget struct { + remainingTokens int64 + remainingBytes int64 +} + +func remainingInputLimits(generation EmbeddingInputGeneration, policy InputPolicy, work *fittingWorkBudget) (inputConstructionLimits, error) { + limits := inputConstructionLimits{ + contentTokens: int(policy.MaxTotalContentTokens - generation.TotalContentTokens), + renderedTokens: int(policy.MaxTotalRenderedTokens - generation.TotalRenderedTokens), + contentBytes: policy.MaxTotalContentBytes - generation.TotalContentBytes, + renderedBytes: policy.MaxTotalRenderedBytes - generation.TotalRenderedBytes, + work: work, + } + if limits.contentTokens < 1 || limits.renderedTokens < 1 || limits.contentBytes < 1 || limits.renderedBytes < 1 { + return inputConstructionLimits{}, errors.New("embedding generation exceeds aggregate limits") + } + limits.contentTokens = min(limits.contentTokens, policy.ContentTokenBudget) + limits.renderedTokens = min(limits.renderedTokens, policy.MaxProviderTokens) + limits.contentBytes = min(limits.contentBytes, policy.MaxProviderBytes) + limits.renderedBytes = min(limits.renderedBytes, policy.MaxProviderBytes) + return limits, nil +} + +func generatedInputCapacity(policy InputPolicy, tokenCount int) int { + return min( + policy.MaxGeneratedInputs, + tokenCount, + int(policy.MaxTotalContentTokens), + int(policy.MaxTotalRenderedTokens), + int(policy.MaxTotalContentBytes), + int(policy.MaxTotalRenderedBytes), + int(policy.MaxFittingWorkTokens), + int(policy.MaxFittingWorkBytes), + ) +} + +func validateInputPolicy(policy InputPolicy, tokenizerIdentity TokenizerIdentity) error { + if err := validateTokenizerIdentity(tokenizerIdentity); err != nil { + return err + } + if policy.ContentTokenBudget < 1 || policy.ContentTokenBudget > maxEmbeddingChunkTokens { + return errors.New("embedding content token budget is invalid") + } + if policy.OverlapTokens < 0 || policy.OverlapTokens >= policy.ContentTokenBudget { + return errors.New("embedding overlap token count is invalid") + } + if policy.MaxProviderTokens < 1 || policy.MaxProviderTokens > maxEmbeddingChunkTokens || policy.MaxProviderBytes < 1 || policy.MaxProviderBytes > maxEmbeddingInputBytes { + return errors.New("embedding provider limits are invalid") + } + if policy.MaxGeneratedInputs < 1 || policy.MaxGeneratedInputs > maxGeneratedInputs { + return errors.New("embedding generated input limit is invalid") + } + for _, total := range []struct { + name string + value int64 + limit int64 + }{ + {"content token", policy.MaxTotalContentTokens, maxGenerationTotalTokens}, + {"rendered token", policy.MaxTotalRenderedTokens, maxGenerationTotalTokens}, + {"content byte", policy.MaxTotalContentBytes, maxGenerationTotalBytes}, + {"rendered byte", policy.MaxTotalRenderedBytes, maxGenerationTotalBytes}, + {"fitting work token", policy.MaxFittingWorkTokens, maxGenerationWorkTokens}, + {"fitting work byte", policy.MaxFittingWorkBytes, maxGenerationWorkBytes}, + } { + if total.value < 1 || total.value > total.limit { + return fmt.Errorf("embedding aggregate %s limit is invalid", total.name) + } + } + if err := validateModelInputContract(policy.ModelInput); err != nil || policy.ModelInput.Profile == "" { + return errors.New("embedding input policy has invalid model-input contract") + } + if len(policy.Formatter) > maxInputFormatterBytes { + return errors.New("embedding input formatter is too long") + } + if err := validateCompatibilityID(policy.Formatter); err != nil { + return fmt.Errorf("embedding input formatter: %w", err) + } + if err := validateFingerprint(policy.LexicalEvidenceFingerprint, "lexical evidence fingerprint"); err != nil { + return err + } + if err := validateFingerprint(policy.ContextFingerprint, "context fingerprint"); err != nil { + return err + } + if policy.AttachmentContext != nil && !policy.AttachmentContext.declared() { + return errors.New("attachment context snapshot is not canonical") + } + switch policy.TruncationPolicy { + case TruncationPolicyReject, TruncationPolicyTruncateIndivisible: + default: + return errors.New("embedding truncation policy is invalid") + } + return nil +} + +func addGeneratedInputTotals(generation *EmbeddingInputGeneration, input GeneratedEmbeddingInput, policy InputPolicy) error { + contentBytes := int64(len(input.Content)) + renderedBytes := int64(len(input.Rendered)) + if !addWithinAggregate(generation.TotalContentTokens, int64(input.ContentTokens), policy.MaxTotalContentTokens) || + !addWithinAggregate(generation.TotalRenderedTokens, int64(input.RenderedTokens), policy.MaxTotalRenderedTokens) || + !addWithinAggregate(generation.TotalContentBytes, contentBytes, policy.MaxTotalContentBytes) || + !addWithinAggregate(generation.TotalRenderedBytes, renderedBytes, policy.MaxTotalRenderedBytes) { + return errors.New("embedding generation exceeds aggregate limits") + } + generation.TotalContentTokens += int64(input.ContentTokens) + generation.TotalRenderedTokens += int64(input.RenderedTokens) + generation.TotalContentBytes += contentBytes + generation.TotalRenderedBytes += renderedBytes + return nil +} + +func addWithinAggregate(current, addition, limit int64) bool { + return current >= 0 && addition >= 0 && limit >= 0 && current <= limit && addition <= limit-current +} + +func (budget *fittingWorkBudget) consume(byteParts, tokenParts []int64) error { + if budget == nil { + return errors.New("embedding fitting work budget is missing") + } + bytes, ok := boundedWorkTotal(byteParts, budget.remainingBytes) + if !ok { + return errors.New("embedding fitting work byte budget exceeded") + } + tokens, ok := boundedWorkTotal(tokenParts, budget.remainingTokens) + if !ok { + return errors.New("embedding fitting work token budget exceeded") + } + budget.remainingBytes -= bytes + budget.remainingTokens -= tokens + return nil +} + +func boundedWorkTotal(parts []int64, limit int64) (int64, bool) { + var total int64 + for _, part := range parts { + if !addWithinAggregate(total, part, limit) { + return 0, false + } + total += part + } + return total, true +} + +func tokenizeEvidence(evidence NormalizedEvidenceV1, tokenizer Tokenizer) ([]embeddingToken, map[int]struct{}, error) { + tokens := make([]embeddingToken, 0, min(maxEmbeddingTokensPerGeneration, len(evidence.Units))) + naturalEnds := make(map[int]struct{}, min(maxEmbeddingTokensPerGeneration, len(evidence.Units))) + for unitIndex, unit := range evidence.Units { + runeCount := utf8.RuneCountInString(unit.Text) + if runeCount == 0 { + continue + } + remaining := maxEmbeddingTokensPerGeneration - len(tokens) + if remaining < 1 { + return nil, nil, ErrTokenizerLimit + } + unitTokens, err := tokenizer.Tokenize(unit.Text, remaining) + if err != nil { + if errors.Is(err, ErrTokenizerLimit) { + return nil, nil, ErrTokenizerLimit + } + return nil, nil, fmt.Errorf("tokenize evidence unit %d: %w", unitIndex, err) + } + if err := validateTokenBoundaries(unitTokens, runeCount, remaining); err != nil { + return nil, nil, fmt.Errorf("tokenize evidence unit %d: %w", unitIndex, err) + } + base := len(tokens) + for _, token := range unitTokens { + tokens = append(tokens, embeddingToken{unitIndex: unitIndex, start: token.Start, end: token.End}) + } + for tokenIndex := base; tokenIndex < len(tokens); tokenIndex++ { + tokens[tokenIndex].unitEnd = len(tokens) + } + cuts := naturalRuneEnds(unit, runeCount) + for localIndex, token := range unitTokens { + if _, natural := cuts[token.End]; natural { + naturalEnds[base+localIndex+1] = struct{}{} + } + } + naturalEnds[len(tokens)] = struct{}{} + } + return tokens, naturalEnds, nil +} + +func naturalRuneEnds(unit NormalizedEvidenceUnitV1, runeCount int) map[int]struct{} { + result := map[int]struct{}{runeCount: {}} + for _, region := range unit.Regions { + result[region.TextRange.Start] = struct{}{} + result[region.TextRange.End] = struct{}{} + } + for _, table := range unit.Tables { + if len(table.Cells) == 0 { + continue + } + start, end := table.Cells[0].TextRange.Start, table.Cells[0].TextRange.End + for _, cell := range table.Cells[1:] { + start = min(start, cell.TextRange.Start) + end = max(end, cell.TextRange.End) + } + result[start] = struct{}{} + result[end] = struct{}{} + } + return result +} + +func chooseChunkEnd(start, total, budget, overlap int, naturalEnds map[int]struct{}) int { + hardEnd := min(total, start+budget) + for candidate := hardEnd; candidate > start; candidate-- { + if _, natural := naturalEnds[candidate]; natural && (candidate == total || candidate-start > overlap) { + return candidate + } + } + return hardEnd +} + +func fitGeneratedInput(evidence NormalizedEvidenceV1, tokens []embeddingToken, naturalEnds map[int]struct{}, start, end int, policy InputPolicy, tokenizerIdentity TokenizerIdentity, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (GeneratedEmbeddingInput, []TokenBoundary, int, bool, error) { + byteEnd, err := maximalByteFit(evidence, tokens, start, end, policy, attachment, limits) + if err != nil { + return GeneratedEmbeddingInput{}, nil, 0, false, err + } + if byteEnd == 0 { + return truncateOrRejectGeneratedInput(evidence, tokens[start], start+1, policy, tokenizerIdentity, attachment, limits, ordinal) + } + preferredEnd := preferNaturalFitEnd(start, byteEnd, policy.OverlapTokens, naturalEnds) + input, contentBoundaries, err := makeGeneratedInput(evidence, tokens[start:preferredEnd], policy, attachment, limits, ordinal) + if err == nil { + return input, contentBoundaries, preferredEnd, false, nil + } + if !fittingLimitError(err) { + return GeneratedEmbeddingInput{}, nil, 0, false, err + } + end = byteEnd + if preferredEnd != byteEnd { + input, contentBoundaries, err = makeGeneratedInput(evidence, tokens[start:byteEnd], policy, attachment, limits, ordinal) + if err == nil { + return input, contentBoundaries, byteEnd, false, nil + } + if !fittingLimitError(err) { + return GeneratedEmbeddingInput{}, nil, 0, false, err + } + } + + var bestEnd int + if tokenizerIdentity.PrefixTokenCountsMonotonic && strings.HasSuffix(policy.ModelInput.Document.Template, modelInputContentSlot) { + bestEnd, err = monotonicTokenFit(evidence, tokens, start, end, policy, attachment, limits, ordinal) + } else { + bestEnd, err = boundedNonMonotonicTokenFit(evidence, tokens, start, end, policy, attachment, limits, ordinal) + } + if err != nil { + return GeneratedEmbeddingInput{}, nil, 0, false, err + } + if bestEnd != 0 { + validated, validatedBoundaries, validateErr := makeGeneratedInput(evidence, tokens[start:bestEnd], policy, attachment, limits, ordinal) + if validateErr != nil { + return GeneratedEmbeddingInput{}, nil, 0, false, fmt.Errorf("validate fitted embedding input: %w", validateErr) + } + return validated, validatedBoundaries, bestEnd, false, nil + } + return truncateOrRejectGeneratedInput(evidence, tokens[start], start+1, policy, tokenizerIdentity, attachment, limits, ordinal) +} + +func maximalByteFit(evidence NormalizedEvidenceV1, tokens []embeddingToken, start, end int, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits) (int, error) { + low, high := start+1, end + bestEnd := 0 + for low <= high { + candidate := low + (high-low)/2 + fits, err := candidateFitsByteLimits(evidence, tokens[start:candidate], policy, attachment, limits) + if err != nil { + return 0, err + } + if fits { + bestEnd = candidate + low = candidate + 1 + } else { + high = candidate - 1 + } + } + return bestEnd, nil +} + +func candidateFitsByteLimits(evidence NormalizedEvidenceV1, tokens []embeddingToken, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits) (bool, error) { + spans, _, err := tokenMetadata(evidence, tokens) + if err != nil { + return false, err + } + contentBytes, err := contentByteLength(evidence, spans, limits.contentBytes) + if errors.Is(err, errProviderInputLimit) { + return false, nil + } + if err != nil { + return false, err + } + return renderedDocumentByteLength(policy.ModelInput.Document, attachment, contentBytes) <= limits.renderedBytes, nil +} + +func monotonicTokenFit(evidence NormalizedEvidenceV1, tokens []embeddingToken, start, end int, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (int, error) { + low, high := start+1, end-1 + bestEnd := 0 + for low <= high { + candidate := low + (high-low)/2 + _, _, err := makeGeneratedInput(evidence, tokens[start:candidate], policy, attachment, limits, ordinal) + if err == nil { + bestEnd = candidate + low = candidate + 1 + continue + } + if !fittingLimitError(err) { + return 0, err + } + high = candidate - 1 + } + return bestEnd, nil +} + +func boundedNonMonotonicTokenFit(evidence NormalizedEvidenceV1, tokens []embeddingToken, start, end int, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (int, error) { + lower := max(start+1, end-maxNonMonotonicFitChecks) + for candidate := end - 1; candidate >= lower; candidate-- { + _, _, err := makeGeneratedInput(evidence, tokens[start:candidate], policy, attachment, limits, ordinal) + if err == nil { + return candidate, nil + } + if !fittingLimitError(err) { + return 0, err + } + } + if lower > start+1 { + return 0, errors.New("non-monotonic tokenizer fit search exceeds bounded checks") + } + return 0, nil +} + +func truncateOrRejectGeneratedInput(evidence NormalizedEvidenceV1, token embeddingToken, fittedEnd int, policy InputPolicy, tokenizerIdentity TokenizerIdentity, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (GeneratedEmbeddingInput, []TokenBoundary, int, bool, error) { + if policy.TruncationPolicy == TruncationPolicyReject { + return GeneratedEmbeddingInput{}, nil, 0, false, errors.New("indivisible token exceeds content budget, provider limits, or aggregate limits") + } + input, contentBoundaries, err := truncateGeneratedInput(evidence, token, policy, tokenizerIdentity, attachment, limits, ordinal) + if err != nil { + return GeneratedEmbeddingInput{}, nil, 0, false, err + } + return input, contentBoundaries, fittedEnd, true, nil +} + +func preferNaturalFitEnd(start, fittedEnd, overlap int, naturalEnds map[int]struct{}) int { + for candidate := fittedEnd; candidate > start; candidate-- { + if _, natural := naturalEnds[candidate]; natural && candidate-start > overlap { + return candidate + } + } + return fittedEnd +} + +var errProviderInputLimit = errors.New("provider input limit") +var errContentTokenLimit = errors.New("content token limit") + +func fittingLimitError(err error) bool { + return errors.Is(err, errProviderInputLimit) || errors.Is(err, errContentTokenLimit) +} + +func makeGeneratedInput(evidence NormalizedEvidenceV1, tokens []embeddingToken, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (GeneratedEmbeddingInput, []TokenBoundary, error) { + spans, headings, err := tokenMetadata(evidence, tokens) + if err != nil { + return GeneratedEmbeddingInput{}, nil, err + } + contentBytes, err := contentByteLength(evidence, spans, limits.contentBytes) + if err != nil || renderedDocumentByteLength(policy.ModelInput.Document, attachment, contentBytes) > limits.renderedBytes { + return GeneratedEmbeddingInput{}, nil, errProviderInputLimit + } + contentRunes := int64(tokens[len(tokens)-1].end - tokens[0].start) + contextualizedBytes := contextualizedDocumentByteLength(attachment, contentBytes) + renderedBytes := renderedDocumentByteLength(policy.ModelInput.Document, attachment, contentBytes) + contextualizedRunes := contextualizedDocumentRuneLength(attachment, contentRunes) + renderedRunes := renderedDocumentRuneLength(policy.ModelInput.Document, attachment, contentRunes) + if err := limits.work.consume( + []int64{contentBytes, contextualizedBytes, renderedBytes}, + []int64{contentRunes, contextualizedRunes, renderedRunes}, + ); err != nil { + return GeneratedEmbeddingInput{}, nil, err + } + content := contentForSpans(evidence, spans, contentBytes) + contentBoundaries, err := policy.Tokenizer.Tokenize(content, limits.contentTokens) + if errors.Is(err, ErrTokenizerLimit) { + return GeneratedEmbeddingInput{}, nil, errContentTokenLimit + } + if err != nil { + return GeneratedEmbeddingInput{}, nil, fmt.Errorf("tokenize exact embedding content: %w", err) + } + if err := validateTokenBoundaries(contentBoundaries, utf8.RuneCountInString(content), limits.contentTokens); err != nil { + if errors.Is(err, ErrTokenizerLimit) { + return GeneratedEmbeddingInput{}, nil, errContentTokenLimit + } + return GeneratedEmbeddingInput{}, nil, fmt.Errorf("validate exact embedding content tokens: %w", err) + } + contextualized := contextualizeDocumentInput(attachment, content) + rendered := policy.ModelInput.EncodeDocument(contextualized) + renderedTokens, err := countRenderedTokens(policy.Tokenizer, rendered, limits.renderedTokens) + if errors.Is(err, ErrTokenizerLimit) || int64(len(rendered)) > limits.renderedBytes { + return GeneratedEmbeddingInput{}, nil, errProviderInputLimit + } + if err != nil { + return GeneratedEmbeddingInput{}, nil, fmt.Errorf("tokenize rendered embedding input: %w", err) + } + checksum := sha256Hex([]byte(rendered)) + return GeneratedEmbeddingInput{ + Key: fmt.Sprintf("chunk-%06d-%s", ordinal, checksum[:12]), Content: content, + Rendered: rendered, + ContentTokens: len(contentBoundaries), RenderedTokens: renderedTokens, Checksum: checksum, + HeadingPaths: headings, SourceSpans: spans, + }, contentBoundaries, nil +} + +func truncateGeneratedInput(evidence NormalizedEvidenceV1, token embeddingToken, policy InputPolicy, tokenizerIdentity TokenizerIdentity, attachment AttachmentContextSnapshot, limits inputConstructionLimits, ordinal int) (GeneratedEmbeddingInput, []TokenBoundary, error) { + if token.end-token.start > maxTruncationSearchRunes { + return GeneratedEmbeddingInput{}, nil, errors.New("indivisible token exceeds bounded truncation search") + } + byteEnd, err := maximalTruncatedByteFit(evidence, token, policy, attachment, limits) + if err != nil { + return GeneratedEmbeddingInput{}, nil, err + } + if byteEnd == 0 { + return GeneratedEmbeddingInput{}, nil, errors.New("provider or aggregate limits cannot hold attachment context and one source rune") + } + candidate := token + candidate.end = byteEnd + input, contentBoundaries, err := makeGeneratedInput(evidence, []embeddingToken{candidate}, policy, attachment, limits, ordinal) + if err == nil { + return input, contentBoundaries, nil + } + if !fittingLimitError(err) { + return GeneratedEmbeddingInput{}, nil, err + } + + low, high := token.start+1, byteEnd-1 + bestEnd := 0 + if tokenizerIdentity.PrefixTokenCountsMonotonic && strings.HasSuffix(policy.ModelInput.Document.Template, modelInputContentSlot) { + for low <= high { + candidateEnd := low + (high-low)/2 + candidate := token + candidate.end = candidateEnd + _, _, err := makeGeneratedInput(evidence, []embeddingToken{candidate}, policy, attachment, limits, ordinal) + if err == nil { + bestEnd = candidateEnd + low = candidateEnd + 1 + continue + } + if !fittingLimitError(err) { + return GeneratedEmbeddingInput{}, nil, err + } + high = candidateEnd - 1 + } + } else { + for candidateEnd := high; candidateEnd >= low; candidateEnd-- { + candidate := token + candidate.end = candidateEnd + _, _, err := makeGeneratedInput(evidence, []embeddingToken{candidate}, policy, attachment, limits, ordinal) + if err == nil { + bestEnd = candidateEnd + break + } + if !fittingLimitError(err) { + return GeneratedEmbeddingInput{}, nil, err + } + } + } + if bestEnd != 0 { + validatedToken := token + validatedToken.end = bestEnd + validated, validatedBoundaries, err := makeGeneratedInput(evidence, []embeddingToken{validatedToken}, policy, attachment, limits, ordinal) + if err != nil { + return GeneratedEmbeddingInput{}, nil, fmt.Errorf("validate truncated embedding input: %w", err) + } + return validated, validatedBoundaries, nil + } + return GeneratedEmbeddingInput{}, nil, errors.New("provider or aggregate limits cannot hold attachment context and one source rune") +} + +func maximalTruncatedByteFit(evidence NormalizedEvidenceV1, token embeddingToken, policy InputPolicy, attachment AttachmentContextSnapshot, limits inputConstructionLimits) (int, error) { + low, high := token.start+1, token.end-1 + bestEnd := 0 + for low <= high { + candidateEnd := low + (high-low)/2 + candidate := token + candidate.end = candidateEnd + fits, err := candidateFitsByteLimits(evidence, []embeddingToken{candidate}, policy, attachment, limits) + if err != nil { + return 0, err + } + if fits { + bestEnd = candidateEnd + low = candidateEnd + 1 + } else { + high = candidateEnd - 1 + } + } + return bestEnd, nil +} + +func tokenMetadata(evidence NormalizedEvidenceV1, tokens []embeddingToken) ([]ChunkSpan, [][]string, error) { + if len(tokens) == 0 { + return nil, nil, errors.New("embedding input requires at least one token") + } + first, last := tokens[0], tokens[len(tokens)-1] + if first.unitIndex != last.unitIndex || first.unitIndex < 0 || first.unitIndex >= len(evidence.Units) || first.start < 0 || last.end <= first.start { + return nil, nil, errors.New("embedding input tokens cross natural units") + } + span := ChunkSpan{UnitIndex: first.unitIndex, CharStart: first.start, CharEnd: last.end} + heading := slices.Clone(evidence.Units[first.unitIndex].HeadingPath) + return []ChunkSpan{span}, [][]string{heading}, nil +} + +func contentByteLength(evidence NormalizedEvidenceV1, spans []ChunkSpan, limit int64) (int64, error) { + var total int64 + for _, span := range spans { + part, ok := runeRange(evidence.Units[span.UnitIndex].Text, span.CharStart, span.CharEnd) + if !ok || int64(len(part)) > limit-total { + return 0, errProviderInputLimit + } + total += int64(len(part)) + } + return total, nil +} + +func contentForSpans(evidence NormalizedEvidenceV1, spans []ChunkSpan, byteLength int64) string { + var content strings.Builder + content.Grow(int(byteLength)) + for _, span := range spans { + part, _ := runeRange(evidence.Units[span.UnitIndex].Text, span.CharStart, span.CharEnd) + content.WriteString(part) + } + return content.String() +} + +func runeRange(value string, runeStart, runeEnd int) (string, bool) { + byteStart, byteEnd := -1, -1 + runeIndex := 0 + for byteIndex := range value { + if runeIndex == runeStart { + byteStart = byteIndex + } + if runeIndex == runeEnd { + byteEnd = byteIndex + break + } + runeIndex++ + } + if runeStart == runeIndex && byteStart < 0 { + byteStart = len(value) + } + if runeEnd == runeIndex && byteEnd < 0 { + byteEnd = len(value) + } + if byteStart < 0 || byteEnd < byteStart { + return "", false + } + return value[byteStart:byteEnd], true +} + +func renderedDocumentByteLength(encoder ModelInputEncoder, attachment AttachmentContextSnapshot, contentBytes int64) int64 { + slot := strings.Index(encoder.Template, modelInputContentSlot) + if slot < 0 { + return int64(^uint64(0) >> 1) + } + return int64(len(encoder.Template)-len(modelInputContentSlot)) + contextualizedDocumentByteLength(attachment, contentBytes) +} + +func contextualizedDocumentByteLength(attachment AttachmentContextSnapshot, contentBytes int64) int64 { + length := contentBytes + if attachment.title != "" { + length += int64(len("Title: ") + len(attachment.title) + 1) + } + if attachment.context != "" { + length += int64(len("Context: ") + len(attachment.context) + 1) + } + if attachment.declared() { + length++ + } + return length +} + +func contextualizedDocumentRuneLength(attachment AttachmentContextSnapshot, contentRunes int64) int64 { + length := contentRunes + if attachment.title != "" { + length += int64(utf8.RuneCountInString("Title: ") + utf8.RuneCountInString(attachment.title) + 1) + } + if attachment.context != "" { + length += int64(utf8.RuneCountInString("Context: ") + utf8.RuneCountInString(attachment.context) + 1) + } + if attachment.declared() { + length++ + } + return length +} + +func renderedDocumentRuneLength(encoder ModelInputEncoder, attachment AttachmentContextSnapshot, contentRunes int64) int64 { + return int64(utf8.RuneCountInString(encoder.Template)-utf8.RuneCountInString(modelInputContentSlot)) + contextualizedDocumentRuneLength(attachment, contentRunes) +} + +func contextualizeDocumentInput(attachment AttachmentContextSnapshot, content string) string { + if !attachment.declared() { + return content + } + var contextualized strings.Builder + if attachment.title != "" { + contextualized.WriteString("Title: ") + contextualized.WriteString(attachment.title) + contextualized.WriteByte('\n') + } + if attachment.context != "" { + contextualized.WriteString("Context: ") + contextualized.WriteString(attachment.context) + contextualized.WriteByte('\n') + } + contextualized.WriteByte('\n') + contextualized.WriteString(content) + return contextualized.String() +} + +func countRenderedTokens(tokenizer Tokenizer, rendered string, limit int) (int, error) { + tokens, err := tokenizer.Tokenize(rendered, limit) + if err != nil { + return 0, err + } + if err := validateTokenBoundaries(tokens, utf8.RuneCountInString(rendered), limit); err != nil { + return 0, err + } + return len(tokens), nil +} + +func inputPolicyFingerprint(policy InputPolicy, attachment AttachmentContextSnapshot, identity TokenizerIdentity) string { + frames := newFingerprintFrames("docbank/embedding-input-policy", EmbeddingInputGenerationVersion) + frames.text("tokenizer.name", identity.Name) + frames.text("tokenizer.revision", identity.Revision) + frames.boolean("tokenizer.prefix_token_counts_monotonic", identity.PrefixTokenCountsMonotonic) + frames.integer("content_token_budget", policy.ContentTokenBudget) + frames.integer("overlap_tokens", policy.OverlapTokens) + frames.integer("max_provider_tokens", policy.MaxProviderTokens) + frames.integer64("max_provider_bytes", policy.MaxProviderBytes) + frames.integer("max_generated_inputs", policy.MaxGeneratedInputs) + frames.integer64("max_total_content_tokens", policy.MaxTotalContentTokens) + frames.integer64("max_total_rendered_tokens", policy.MaxTotalRenderedTokens) + frames.integer64("max_total_content_bytes", policy.MaxTotalContentBytes) + frames.integer64("max_total_rendered_bytes", policy.MaxTotalRenderedBytes) + frames.integer64("max_fitting_work_tokens", policy.MaxFittingWorkTokens) + frames.integer64("max_fitting_work_bytes", policy.MaxFittingWorkBytes) + frames.text("model_input_fingerprint", policy.ModelInput.Fingerprint) + frames.text("formatter", policy.Formatter) + frames.text("lexical_evidence_fingerprint", policy.LexicalEvidenceFingerprint) + frames.text("context_fingerprint", policy.ContextFingerprint) + frames.boolean("attachment_context.declared", attachment.declared()) + frames.text("attachment_context.title", attachment.title) + frames.text("attachment_context.context", attachment.context) + frames.text("truncation_policy", string(policy.TruncationPolicy)) + return frames.sum() +} + +func generationFingerprint(generation EmbeddingInputGeneration) string { + frames := newFingerprintFrames("docbank/embedding-input-generation", generation.Version) + frames.text("evidence_checksum", generation.EvidenceChecksum) + frames.text("policy_fingerprint", generation.PolicyFingerprint) + frames.text("tokenizer.name", generation.TokenizerIdentity.Name) + frames.text("tokenizer.revision", generation.TokenizerIdentity.Revision) + frames.boolean("tokenizer.prefix_token_counts_monotonic", generation.TokenizerIdentity.PrefixTokenCountsMonotonic) + frames.text("lexical_evidence_fingerprint", generation.LexicalEvidenceFingerprint) + frames.text("formatter", generation.Formatter) + frames.text("model_input_fingerprint", generation.ModelInputFingerprint) + if generation.Version >= EmbeddingInputGenerationVersion { + frames.integer("content_token_budget", generation.ContentTokenBudget) + frames.integer("overlap_tokens", generation.OverlapTokens) + frames.text("truncation_policy", string(generation.TruncationPolicy)) + frames.text("context_fingerprint", generation.ContextFingerprint) + } + frames.boolean("attachment_context.declared", generation.AttachmentContext != nil) + if generation.AttachmentContext != nil { + frames.text("attachment_context.title", generation.AttachmentContext.title) + frames.text("attachment_context.context", generation.AttachmentContext.context) + } + frames.integer64("total_content_tokens", generation.TotalContentTokens) + frames.integer64("total_rendered_tokens", generation.TotalRenderedTokens) + frames.integer64("total_content_bytes", generation.TotalContentBytes) + frames.integer64("total_rendered_bytes", generation.TotalRenderedBytes) + frames.count("inputs", len(generation.Inputs)) + for inputIndex, input := range generation.Inputs { + frames.object("input") + frames.integer("input.index", inputIndex) + frames.text("input.key", input.Key) + frames.text("input.content", input.Content) + frames.text("input.rendered", input.Rendered) + frames.integer("input.content_tokens", input.ContentTokens) + frames.integer("input.rendered_tokens", input.RenderedTokens) + frames.text("input.checksum", input.Checksum) + frames.boolean("input.truncated", input.Truncated) + frames.count("input.heading_paths", len(input.HeadingPaths)) + for headingIndex, heading := range input.HeadingPaths { + frames.object("input.heading_path") + frames.integer("input.heading_path.index", headingIndex) + frames.count("input.heading_path.parts", len(heading)) + for partIndex, part := range heading { + frames.integer("input.heading_path.part.index", partIndex) + frames.text("input.heading_path.part", part) + } + } + frames.count("input.source_spans", len(input.SourceSpans)) + for spanIndex, span := range input.SourceSpans { + frames.object("input.source_span") + frames.integer("input.source_span.index", spanIndex) + frames.integer("input.source_span.unit_index", span.UnitIndex) + frames.integer("input.source_span.char_start", span.CharStart) + frames.integer("input.source_span.char_end", span.CharEnd) + } + } + return frames.sum() +} + +type fingerprintFrameType byte + +const ( + fingerprintFrameDomain fingerprintFrameType = iota + 1 + fingerprintFrameText + fingerprintFrameInteger + fingerprintFrameBoolean + fingerprintFrameObject + fingerprintFrameCount +) + +type fingerprintFrames struct { + digest hash.Hash +} + +func newFingerprintFrames(domain string, version int) *fingerprintFrames { + frames := &fingerprintFrames{digest: sha256.New()} + frames.writeStringFrame(fingerprintFrameDomain, "domain", domain) + frames.integer("version", version) + return frames +} +func (frames *fingerprintFrames) writeHeader(frameType fingerprintFrameType, field string, valueLength int) { + _, _ = frames.digest.Write([]byte{byte(frameType)}) + frames.writeLength(len(field)) + _, _ = io.WriteString(frames.digest, field) + frames.writeLength(valueLength) +} +func (frames *fingerprintFrames) writeLength(value int) { + _, _ = io.WriteString(frames.digest, strconv.Itoa(value)) + _, _ = frames.digest.Write([]byte{':'}) +} +func (frames *fingerprintFrames) writeStringFrame(frameType fingerprintFrameType, field, value string) { + frames.writeHeader(frameType, field, len(value)) + _, _ = io.WriteString(frames.digest, value) +} +func (frames *fingerprintFrames) text(field, value string) { + frames.writeStringFrame(fingerprintFrameText, field, value) +} +func (frames *fingerprintFrames) integer(field string, value int) { + encoded := strconv.Itoa(value) + frames.writeHeader(fingerprintFrameInteger, field, len(encoded)) + _, _ = io.WriteString(frames.digest, encoded) +} +func (frames *fingerprintFrames) integer64(field string, value int64) { + encoded := strconv.FormatInt(value, 10) + frames.writeHeader(fingerprintFrameInteger, field, len(encoded)) + _, _ = io.WriteString(frames.digest, encoded) +} +func (frames *fingerprintFrames) boolean(field string, value bool) { + frames.writeHeader(fingerprintFrameBoolean, field, 1) + encoded := byte(0) + if value { + encoded = 1 + } + _, _ = frames.digest.Write([]byte{encoded}) +} +func (frames *fingerprintFrames) object(field string) { + frames.writeHeader(fingerprintFrameObject, field, 0) +} +func (frames *fingerprintFrames) count(field string, value int) { + encoded := strconv.Itoa(value) + frames.writeHeader(fingerprintFrameCount, field, len(encoded)) + _, _ = io.WriteString(frames.digest, encoded) +} +func (frames *fingerprintFrames) sum() string { + return hex.EncodeToString(frames.digest.Sum(nil)) +} diff --git a/document/chunking_test.go b/document/chunking_test.go new file mode 100644 index 00000000..3178267d --- /dev/null +++ b/document/chunking_test.go @@ -0,0 +1,984 @@ +package document + +import ( + "context" + "encoding/json" + "fmt" + "os" + "slices" + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type pairTokenizer struct{} + +func (pairTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "synthetic-pairs", Revision: "v1", PrefixTokenCountsMonotonic: true} +} + +func (pairTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + runes := []rune(text) + result := make([]TokenBoundary, 0, min((len(runes)+1)/2, limit)) + for start := 0; start < len(runes); start += 2 { + if len(result) == limit { + return nil, ErrTokenizerLimit + } + result = append(result, TokenBoundary{Start: start, End: min(start+2, len(runes))}) + } + return result, nil +} + +func TestBuildEmbeddingInputsUsesNaturalBoundariesAndExactOverlap(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{ + {text: "AABBCCDDEEFF", heading: []string{"First"}, regions: []SourceEvidenceRegionV1{{ProviderID: "p1", Kind: EvidenceRegionParagraph, Order: 0, TextRange: EvidenceTextRangeV1{Start: 0, End: 12}}}}, + {text: "GGHH", heading: []string{"Second"}, regions: []SourceEvidenceRegionV1{{ProviderID: "p2", Kind: EvidenceRegionParagraph, Order: 0, TextRange: EvidenceTextRangeV1{Start: 0, End: 4}}}}, + }) + generation, err := BuildEmbeddingInputs(evidence, testInputPolicy(t, 4, 2)) + require.NoError(t, err) + require.Len(t, generation.Inputs, 3) + + assert.Equal(t, []string{"AABBCCDD", "CCDDEEFF", "GGHH"}, []string{ + generation.Inputs[0].Content, generation.Inputs[1].Content, generation.Inputs[2].Content, + }) + assert.Equal(t, []int{4, 4, 2}, []int{ + generation.Inputs[0].ContentTokens, generation.Inputs[1].ContentTokens, generation.Inputs[2].ContentTokens, + }) + assert.Equal(t, [][]string{{"First"}}, generation.Inputs[0].HeadingPaths) + assert.Equal(t, [][]string{{"First"}}, generation.Inputs[1].HeadingPaths) + assert.Equal(t, []ChunkSpan{{UnitIndex: 0, CharStart: 4, CharEnd: 12}}, generation.Inputs[1].SourceSpans) + assert.Equal(t, generation.Inputs[0].Content[len(generation.Inputs[0].Content)-4:], generation.Inputs[1].Content[:4]) + assertGenerationSpans(t, evidence, generation) +} + +func TestBuildEmbeddingInputsNeverCombinesNaturalUnitsBeforeExactTokenization(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "A"}, {text: "B"}}) + policy := testInputPolicy(t, 1, 0) + policy.Tokenizer = concatAdversarialTokenizer{} + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 2) + assert.Equal(t, "A", generation.Inputs[0].Content) + assert.Equal(t, "B", generation.Inputs[1].Content) + assert.Equal(t, 1, generation.Inputs[0].ContentTokens) + assert.Equal(t, 1, generation.Inputs[1].ContentTokens) + assert.Equal(t, []ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 1}}, generation.Inputs[0].SourceSpans) + assert.Equal(t, []ChunkSpan{{UnitIndex: 1, CharStart: 0, CharEnd: 1}}, generation.Inputs[1].SourceSpans) +} + +func TestBuildEmbeddingInputsDerivesOverlapFromExactEmittedTokenization(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "abcdef"}}) + policy := testInputPolicy(t, 2, 1) + policy.Tokenizer = overlapAdversarialTokenizer{} + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 3) + assert.Equal(t, []string{"abcd", "bcd", "cdef"}, []string{ + generation.Inputs[0].Content, generation.Inputs[1].Content, generation.Inputs[2].Content, + }) + for _, input := range generation.Inputs { + assert.Equal(t, 2, input.ContentTokens) + } + assert.Equal(t, "bcd", generation.Inputs[1].Content[:3], "the second input starts with the exact final token of the first") + assert.Equal(t, "cd", generation.Inputs[2].Content[:2], "the third input starts with the exact final token of the second") + assertGenerationSpans(t, evidence, generation) +} + +func TestBuildEmbeddingInputsPrefersRegionAndTableAtomsBeforeTokenSplits(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{ + text: "AAAABBBBCCCC", heading: []string{"Structured"}, + regions: []SourceEvidenceRegionV1{ + {ProviderID: "heading", Kind: EvidenceRegionHeading, Order: 0, TextRange: EvidenceTextRangeV1{Start: 0, End: 4}}, + {ProviderID: "table", Kind: EvidenceRegionTable, Order: 1, TextRange: EvidenceTextRangeV1{Start: 4, End: 8}}, + {ProviderID: "paragraph", Kind: EvidenceRegionParagraph, Order: 2, TextRange: EvidenceTextRangeV1{Start: 8, End: 12}}, + }, + tables: []SourceEvidenceTableV1{{ + ProviderID: "t1", RegionProviderID: "table", Order: 0, Rows: 1, Columns: 1, + Cells: []SourceEvidenceTableCellV1{{Order: 0, Row: 0, Column: 0, RowSpan: 1, ColumnSpan: 1, TextRange: EvidenceTextRangeV1{Start: 4, End: 8}}}, + }}, + }}) + policy := testInputPolicy(t, 3, 0) + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + assert.Equal(t, []string{"AAAA", "BBBB", "CCCC"}, []string{generation.Inputs[0].Content, generation.Inputs[1].Content, generation.Inputs[2].Content}) + assertGenerationSpans(t, evidence, generation) +} + +func TestBuildEmbeddingInputsKeepsContentBudgetSeparateFromAttachmentAndEnvelope(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABBCC", heading: []string{"Evidence"}}}) + attachment, err := NewAttachmentContextSnapshot(AttachmentContextSnapshotConfig{Title: "Human title", Context: "Human context"}) + require.NoError(t, err) + policy := testInputPolicy(t, 3, 0) + policy.AttachmentContext = &attachment + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 1) + + input := generation.Inputs[0] + assert.Equal(t, "AABBCC", input.Content) + assert.Equal(t, 3, input.ContentTokens) + assert.Greater(t, input.RenderedTokens, input.ContentTokens) + assert.Contains(t, input.Rendered, "Human title") + assert.Contains(t, input.Rendered, "Human context") + assert.True(t, strings.HasPrefix(input.Rendered, "document: Title: Human title")) + assert.True(t, strings.HasSuffix(input.Rendered, "AABBCC")) + require.NotNil(t, generation.AttachmentContext) + assert.Equal(t, "Human title", generation.AttachmentContext.Title()) + assert.Equal(t, "Human context", generation.AttachmentContext.Context()) +} + +func TestBuildEmbeddingInputsTruncatesOnlyAnIndivisibleTokenWhenDeclared(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "abcdefghij"}}) + policy := testInputPolicy(t, 1, 0) + policy.MaxProviderBytes = 15 + policy.TruncationPolicy = TruncationPolicyTruncateIndivisible + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 5) + assert.Equal(t, "ab", generation.Inputs[0].Content) + assert.False(t, generation.Inputs[0].Truncated) + + policy.MaxProviderBytes = 11 + generation, err = BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + assert.Equal(t, "a", generation.Inputs[0].Content) + assert.Equal(t, []ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 1}}, generation.Inputs[0].SourceSpans) + assert.True(t, generation.Inputs[0].Truncated) + + policy.TruncationPolicy = TruncationPolicyReject + _, err = BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "indivisible token") +} + +func TestBuildEmbeddingInputsRetokenizesExactTruncatedContent(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "abcd"}}) + policy := testInputPolicy(t, 1, 0) + policy.Tokenizer = truncationAdversarialTokenizer{} + policy.MaxProviderBytes = 13 + policy.TruncationPolicy = TruncationPolicyTruncateIndivisible + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 1) + assert.Equal(t, "a", generation.Inputs[0].Content) + assert.Equal(t, 1, generation.Inputs[0].ContentTokens) + assert.Equal(t, []ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 1}}, generation.Inputs[0].SourceSpans) + assert.True(t, generation.Inputs[0].Truncated) +} + +func TestBuildEmbeddingInputsSkipsEmptyNaturalUnits(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AA"}, {text: ""}, {text: "BB"}}) + generation, err := BuildEmbeddingInputs(evidence, testInputPolicy(t, 2, 0)) + require.NoError(t, err) + require.Len(t, generation.Inputs, 2) + assert.Equal(t, []ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 2}}, generation.Inputs[0].SourceSpans) + assert.Equal(t, []ChunkSpan{{UnitIndex: 2, CharStart: 0, CharEnd: 2}}, generation.Inputs[1].SourceSpans) + + allEmpty := testChunkEvidence(t, []sourceChunkUnit{{text: ""}, {text: ""}}) + emptyGeneration, err := BuildEmbeddingInputs(allEmpty, testInputPolicy(t, 2, 0)) + require.NoError(t, err) + assert.Empty(t, emptyGeneration.Inputs) + assert.NotEmpty(t, emptyGeneration.Checksum) +} + +func TestBuildEmbeddingInputsSealsEveryIdentityInput(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB", heading: []string{"Heading"}}}) + basePolicy := testInputPolicy(t, 2, 0) + base, err := BuildEmbeddingInputs(evidence, basePolicy) + require.NoError(t, err) + repeat, err := BuildEmbeddingInputs(evidence, basePolicy) + require.NoError(t, err) + assert.Equal(t, base, repeat) + + mutations := []func(*InputPolicy){ + func(policy *InputPolicy) { policy.Formatter = "evidence-text/v2" }, + func(policy *InputPolicy) { policy.Tokenizer = namedPairTokenizer{"other", "v1"} }, + func(policy *InputPolicy) { policy.Tokenizer = nonMonotonicNamedPairTokenizer{} }, + func(policy *InputPolicy) { policy.ContentTokenBudget++ }, + func(policy *InputPolicy) { policy.OverlapTokens = 1 }, + func(policy *InputPolicy) { policy.MaxProviderTokens++ }, + func(policy *InputPolicy) { policy.MaxProviderBytes++ }, + func(policy *InputPolicy) { policy.MaxGeneratedInputs++ }, + func(policy *InputPolicy) { policy.MaxTotalContentTokens++ }, + func(policy *InputPolicy) { policy.MaxTotalRenderedTokens++ }, + func(policy *InputPolicy) { policy.MaxTotalContentBytes++ }, + func(policy *InputPolicy) { policy.MaxTotalRenderedBytes++ }, + func(policy *InputPolicy) { policy.MaxFittingWorkTokens++ }, + func(policy *InputPolicy) { policy.MaxFittingWorkBytes++ }, + } + for index, mutate := range mutations { + policy := basePolicy + mutate(&policy) + changed, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err, index) + assert.NotEqual(t, base.Checksum, changed.Checksum, index) + } + + changedHeading := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB", heading: []string{"Other heading"}}}) + changed, err := BuildEmbeddingInputs(changedHeading, basePolicy) + require.NoError(t, err) + assert.NotEqual(t, base.Checksum, changed.Checksum) +} + +func TestBuildEmbeddingInputsDoesNotLeakFrontmatterOrProvenanceMetadata(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "Actual evidence", heading: []string{"Visible heading"}}}) + policy := testInputPolicy(t, 100, 0) + policy.LexicalEvidenceFingerprint = strings.Repeat("d", 64) + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 1) + + rendered := generation.Inputs[0].Rendered + assert.Equal(t, "document: Actual evidence", rendered) + for _, forbidden := range []string{"---", "checksum", evidence.Checksum, evidence.Units[0].ID, policy.LexicalEvidenceFingerprint, "Visible heading", "page:", "line:", "byte:"} { + assert.NotContains(t, rendered, forbidden) + } +} + +func TestBuildEmbeddingInputsRejectsUnboundedOrNoncanonicalTokenizerOutput(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}) + for _, testCase := range []struct { + name string + tokenizer Tokenizer + want string + }{ + {"gap", brokenTokenizer{tokens: []TokenBoundary{{Start: 0, End: 1}, {Start: 2, End: 4}}}, "contiguous"}, + {"outside", brokenTokenizer{tokens: []TokenBoundary{{Start: 0, End: 5}}}, "bounds"}, + {"empty", brokenTokenizer{}, "at least one"}, + {"too many", brokenTokenizer{err: ErrTokenizerLimit}, "token limit"}, + } { + t.Run(testCase.name, func(t *testing.T) { + policy := testInputPolicy(t, 2, 0) + policy.Tokenizer = testCase.tokenizer + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, testCase.want) + }) + } + policy := testInputPolicy(t, 2, 0) + policy.MaxGeneratedInputs = 1 + _, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABBCCDD"}}), policy) + require.ErrorContains(t, err, "generated input limit") +} + +func TestBuildEmbeddingInputsBoundsAggregateAmplificationBeforeAppend(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABBCCDDEEFF"}}) + policy := testInputPolicy(t, 4, 3) + policy.MaxTotalContentBytes = 12 + policy.MaxTotalRenderedBytes = 44 + policy.MaxTotalContentTokens = 6 + policy.MaxTotalRenderedTokens = 40 + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "aggregate") + + policy.MaxTotalContentTokens = 100 + policy.MaxTotalRenderedTokens = 100 + policy.MaxTotalContentBytes = 100 + policy.MaxTotalRenderedBytes = 200 + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + var contentTokens, renderedTokens, contentBytes, renderedBytes int64 + for _, input := range generation.Inputs { + contentTokens += int64(input.ContentTokens) + renderedTokens += int64(input.RenderedTokens) + contentBytes += int64(len(input.Content)) + renderedBytes += int64(len(input.Rendered)) + } + assert.Equal(t, contentTokens, generation.TotalContentTokens) + assert.Equal(t, renderedTokens, generation.TotalRenderedTokens) + assert.Equal(t, contentBytes, generation.TotalContentBytes) + assert.Equal(t, renderedBytes, generation.TotalRenderedBytes) + + assert.False(t, addWithinAggregate(10, 1, 10)) + assert.False(t, addWithinAggregate(int64(^uint64(0)>>1), 1, int64(^uint64(0)>>1))) +} + +func TestBuildEmbeddingInputsAppliesRemainingAggregateBytesBeforeConstruction(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AA"}, {text: strings.Repeat("b", 1024)}}) + tokenizer := &recordingRuneTokenizer{contentLimit: 10_000} + policy := testInputPolicy(t, tokenizer.contentLimit, 0) + policy.Tokenizer = tokenizer + policy.MaxProviderBytes = 20_000 + policy.MaxProviderTokens = 20_000 + policy.MaxGeneratedInputs = 16 + policy.MaxTotalContentBytes = 6 + policy.MaxTotalRenderedBytes = 26 + policy.MaxTotalContentTokens = 20_000 + policy.MaxTotalRenderedTokens = 40_000 + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "aggregate") + assert.LessOrEqual(t, tokenizer.maxExactContentRunes, 4, "remaining aggregate bytes must shrink a candidate before content construction") +} + +func TestBuildEmbeddingInputsPassesRemainingAggregateTokenLimitToTokenizer(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AA"}, {text: "BBBB"}}) + tokenizer := &recordingRuneTokenizer{contentLimit: 4} + policy := testInputPolicy(t, tokenizer.contentLimit, 0) + policy.Tokenizer = tokenizer + policy.MaxTotalContentTokens = 3 + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "aggregate") + assert.True(t, tokenizer.sawOneTokenContentLimit, "the exact tokenizer must receive the remaining aggregate content-token limit") +} + +func TestBuildEmbeddingInputsRejectsTypedNilAndTokenizerIdentityDrift(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AA"}}) + policy := testInputPolicy(t, 2, 0) + var typedNil *nilTokenizer + policy.Tokenizer = typedNil + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "requires a tokenizer") + + policy = testInputPolicy(t, 2, 0) + policy.Tokenizer = &identityDriftTokenizer{} + _, err = BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "identity changed") +} + +func TestBuildEmbeddingInputsUsesBoundedFittingSearch(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: strings.Repeat("a", 1024)}}) + tokenizer := &countingRuneTokenizer{} + policy := testInputPolicy(t, 1024, 0) + policy.Tokenizer = tokenizer + policy.MaxProviderBytes = 20 + policy.MaxProviderTokens = 2048 + policy.MaxGeneratedInputs = 128 + policy.MaxTotalContentBytes = 4096 + policy.MaxTotalRenderedBytes = 8192 + policy.MaxTotalContentTokens = 2048 + policy.MaxTotalRenderedTokens = 4096 + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.NotEmpty(t, generation.Inputs) + assert.LessOrEqual(t, tokenizer.calls, 5000, "fitting must not retry one token at a time") +} + +func TestBuildEmbeddingInputsCapsCumulativeNonMonotonicFittingWork(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: strings.Repeat("a", 2048)}}) + tokenizer := &fittingWorkAdversarialTokenizer{} + policy := testInputPolicy(t, 2048, 0) + policy.Tokenizer = tokenizer + policy.MaxProviderBytes = 8192 + policy.MaxProviderTokens = 4096 + policy.MaxFittingWorkBytes = 20_000 + policy.MaxFittingWorkTokens = 100_000 + _, err := BuildEmbeddingInputs(evidence, policy) + require.ErrorContains(t, err, "fitting work") + assert.LessOrEqual(t, tokenizer.exactCalls, 3, "the work budget must be consumed before reconstruction and tokenization") +} + +func TestBuildEmbeddingInputsDoesNotAssumePrefixTokenCountsAreMonotonic(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "abcd"}}) + policy := testInputPolicy(t, 4, 0) + policy.Tokenizer = nonMonotonicPrefixTokenizer{} + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.NotEmpty(t, generation.Inputs) + assert.Equal(t, "abc", generation.Inputs[0].Content, "the longer fitting prefix must not be discarded after the middle prefix exceeds the token limit") +} + +func TestBuildEmbeddingInputsDoesNotBinarySearchTruncationWithTemplateSuffix(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "abcdef"}}) + contract, err := NewModelInputContract(ModelInputContractConfig{ + Profile: ModelInputProfileCustom, CompatibilityID: "suffix-space", + Document: ModelInputEncoder{Mode: ModelInputModeText, Template: "{{content}} suffix"}, + Query: ModelInputEncoder{Mode: ModelInputModeText, Template: "{{content}} suffix"}, + }) + require.NoError(t, err) + policy := testInputPolicy(t, 1, 0) + policy.Tokenizer = suffixTruncationAdversarialTokenizer{} + policy.ModelInput = contract + policy.MaxProviderTokens = 1 + policy.TruncationPolicy = TruncationPolicyTruncateIndivisible + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.Len(t, generation.Inputs, 1) + assert.Equal(t, "abcd", generation.Inputs[0].Content, "a longer fitting truncation must survive a failing middle prefix") + assert.True(t, generation.Inputs[0].Truncated) +} + +func TestBuildEmbeddingInputsPreservesNaturalCutWhenProviderLimitShrinksChunk(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{ + {text: "AAAABBBBCCCC", regions: []SourceEvidenceRegionV1{ + {ProviderID: "first", Kind: EvidenceRegionParagraph, Order: 0, TextRange: EvidenceTextRangeV1{Start: 0, End: 4}}, + {ProviderID: "second", Kind: EvidenceRegionParagraph, Order: 1, TextRange: EvidenceTextRangeV1{Start: 4, End: 8}}, + {ProviderID: "third", Kind: EvidenceRegionParagraph, Order: 2, TextRange: EvidenceTextRangeV1{Start: 8, End: 12}}, + }}, + }) + policy := testInputPolicy(t, 6, 0) + policy.MaxProviderBytes = int64(len("document: ") + 7) + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + require.NotEmpty(t, generation.Inputs) + assert.Equal(t, "AAAA", generation.Inputs[0].Content) +} + +func TestEmbeddingInputGenerationRoundTripsAndRejectsMalformedJSON(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB", heading: []string{"Evidence"}}}) + attachment, err := NewAttachmentContextSnapshot(AttachmentContextSnapshotConfig{Title: "Human title", Context: "Human context"}) + require.NoError(t, err) + policy := testInputPolicy(t, 2, 0) + policy.AttachmentContext = &attachment + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + encoded, err := json.Marshal(generation) + require.NoError(t, err) + decoded, err := DecodeEmbeddingInputGeneration(encoded, testGenerationDecodeBounds()) + require.NoError(t, err) + assert.Equal(t, generation, decoded) + reencoded, err := json.Marshal(decoded) + require.NoError(t, err) + assert.JSONEq(t, string(encoded), string(reencoded)) + + for _, testCase := range []struct{ name, value, want string }{ + {"unknown generation field", strings.Replace(string(encoded), fmt.Sprintf(`{"version":%d`, EmbeddingInputGenerationVersion), fmt.Sprintf(`{"unknown":1,"version":%d`, EmbeddingInputGenerationVersion), 1), "unknown field"}, + {"forged checksum", strings.Replace(string(encoded), generation.Checksum, strings.Repeat("0", 64), 1), "checksum"}, + {"empty attachment", strings.Replace(string(encoded), `"title":"Human title","context":"Human context"`, `"title":"","context":""`, 1), "attachment"}, + {"negative span", strings.Replace(string(encoded), `"unit_index":0`, `"unit_index":-1`, 1), "source span"}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := DecodeEmbeddingInputGeneration([]byte(testCase.value), testGenerationDecodeBounds()) + require.ErrorContains(t, err, testCase.want) + }) + } +} + +func TestEmbeddingInputGenerationDecodesFrozenPublishedV1WithoutReinterpretation(t *testing.T) { + encoded, err := os.ReadFile("testdata/embedding-input-generation-v1.golden.json") + require.NoError(t, err) + generation, err := DecodeEmbeddingInputGeneration(encoded, testGenerationDecodeBounds()) + require.NoError(t, err) + assert.Equal(t, 1, generation.Version) + assert.Zero(t, generation.ContentTokenBudget) + assert.Zero(t, generation.OverlapTokens) + assert.Empty(t, generation.TruncationPolicy) + assert.Empty(t, generation.ContextFingerprint) + reencoded, err := json.MarshalIndent(generation, "", " ") + require.NoError(t, err) + assert.JSONEq(t, string(encoded), string(reencoded)) + assert.NotContains(t, string(reencoded), "content_token_budget") +} + +func TestEmbeddingInputGenerationBuildsPolicyCompleteCurrentVersion(t *testing.T) { + generation, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}), testInputPolicy(t, 2, 0)) + require.NoError(t, err) + assert.Equal(t, 2, generation.Version) + assert.Positive(t, generation.ContentTokenBudget) + assert.NotEmpty(t, generation.TruncationPolicy) + assert.NotEmpty(t, generation.ContextFingerprint) +} + +func TestEmbeddingInputGenerationChecksumFramesEmptyHeadingCardinality(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}) + generation, err := BuildEmbeddingInputs(evidence, testInputPolicy(t, 2, 0)) + require.NoError(t, err) + encoded, err := json.Marshal(generation) + require.NoError(t, err) + tampered := strings.Replace(string(encoded), `"heading_paths":[[]]`, `"heading_paths":[]`, 1) + require.NotEqual(t, string(encoded), tampered) + _, err = DecodeEmbeddingInputGeneration([]byte(tampered), testGenerationDecodeBounds()) + require.ErrorContains(t, err, "canonical cardinality") +} + +func TestEmbeddingInputGenerationDecodePreflightsCallerBounds(t *testing.T) { + generation, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}), testInputPolicy(t, 2, 0)) + require.NoError(t, err) + encoded, err := json.Marshal(generation) + require.NoError(t, err) + + tooManyInputs := []byte(`{"inputs":[` + strings.TrimSuffix(strings.Repeat(`{},`, 64), ",") + `]}`) + bounds := testGenerationDecodeBounds() + bounds.MaxInputs = 2 + _, err = DecodeEmbeddingInputGeneration(tooManyInputs, bounds) + require.ErrorContains(t, err, "inputs") + + hugeString := []byte(`{"formatter":"` + strings.Repeat("x", 4096) + `"}`) + bounds = testGenerationDecodeBounds() + bounds.MaxStringBytes = 32 + _, err = DecodeEmbeddingInputGeneration(hugeString, bounds) + require.ErrorContains(t, err, "raw string") + + escapeHeavyString := []byte(`{"formatter":"` + strings.Repeat(`\u0061`, 64) + `"}`) + _, err = DecodeEmbeddingInputGeneration(escapeHeavyString, bounds) + require.ErrorContains(t, err, "raw string") + + invalidEscapeAfterBound := []byte(`{"formatter":"` + strings.Repeat("x", 64) + `\q"}`) + _, err = DecodeEmbeddingInputGeneration(invalidEscapeAfterBound, bounds) + require.ErrorContains(t, err, "raw string") + assert.NotContains(t, err.Error(), "escape", "the lexical bound must reject before JSON unmarshal reaches the invalid escape") + + bounds = testGenerationDecodeBounds() + bounds.MaxObjectFields = 2 + _, err = DecodeEmbeddingInputGeneration([]byte(`{"a":null,"b":null,"c":null}`), bounds) + require.ErrorContains(t, err, "object fields") + + bounds = testGenerationDecodeBounds() + bounds.MaxEncodedBytes = int64(len(encoded) - 1) + _, err = DecodeEmbeddingInputGeneration(encoded, bounds) + require.ErrorContains(t, err, "encoded") +} + +func TestEmbeddingInputGenerationDecodePreflightsIntegerTokens(t *testing.T) { + bounds := testGenerationDecodeBounds() + for _, value := range []string{"0", "-0", "9223372036854775807", "-9223372036854775808"} { + t.Run("valid_"+value, func(t *testing.T) { + err := preflightEmbeddingInputGenerationJSON([]byte(`{"value":`+value+`}`), bounds) + require.NoError(t, err) + }) + } + + for _, testCase := range []struct { + name string + value string + }{ + {"overlong", "123456789012345678901"}, + {"decimal", "1.0"}, + {"exponent", "1e3"}, + {"leading_zero", "01"}, + {"negative_leading_zero", "-01"}, + {"sign_only", "-"}, + {"plus_sign", "+1"}, + {"double_sign", "--1"}, + } { + t.Run(testCase.name, func(t *testing.T) { + err := preflightEmbeddingInputGenerationJSON([]byte(`{"value":`+testCase.value+`}`), bounds) + require.ErrorContains(t, err, "integer") + }) + } +} + +func TestEmbeddingInputGenerationChecksumFramesTypedCollectionsWithoutCollisions(t *testing.T) { + generation, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}), testInputPolicy(t, 2, 0)) + require.NoError(t, err) + span := generation.Inputs[0].SourceSpans[0] + + headingValues := generation + headingValues.Inputs = slices.Clone(generation.Inputs) + headingValues.Inputs[0].HeadingPaths = [][]string{{"span", "0", "1", "2"}} + headingValues.Inputs[0].SourceSpans = []ChunkSpan{span} + + spanValues := generation + spanValues.Inputs = slices.Clone(generation.Inputs) + spanValues.Inputs[0].HeadingPaths = [][]string{{}} + spanValues.Inputs[0].SourceSpans = []ChunkSpan{{UnitIndex: 0, CharStart: 1, CharEnd: 2}, span} + + assert.NotEqual(t, generationFingerprint(headingValues), generationFingerprint(spanValues), "heading text must not collide with span markers and integer frames") +} + +func TestEmbeddingInputGenerationRequiresCanonicalCollectionCardinalities(t *testing.T) { + generation, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}), testInputPolicy(t, 2, 0)) + require.NoError(t, err) + for _, testCase := range []struct { + name string + mutate func(*GeneratedEmbeddingInput) + }{ + {"heading paths", func(input *GeneratedEmbeddingInput) { input.HeadingPaths = nil }}, + {"source spans", func(input *GeneratedEmbeddingInput) { + input.SourceSpans = append(input.SourceSpans, input.SourceSpans[0]) + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + tampered := generation + tampered.Inputs = slices.Clone(generation.Inputs) + testCase.mutate(&tampered.Inputs[0]) + tampered.Checksum = generationFingerprint(tampered) + require.ErrorContains(t, validateEmbeddingInputGeneration(tampered), "canonical cardinality") + }) + } +} + +func TestEmbeddingInputGenerationMapsContextIntoE1ExactlyOnce(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: "AABB", heading: []string{"Evidence"}}}) + attachment, err := NewAttachmentContextSnapshot(AttachmentContextSnapshotConfig{Title: "Human title", Context: "Human context"}) + require.NoError(t, err) + policy := testInputPolicy(t, 2, 0) + policy.AttachmentContext = &attachment + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + generated := generation.Inputs[0] + inputs, err := generation.ToEmbeddingInputs(policy.ModelInput) + require.NoError(t, err) + require.Len(t, inputs, 1) + input := inputs[0] + assert.Contains(t, input.Text, "Human title") + assert.Contains(t, input.Text, "Human context") + assert.NotEqual(t, generated.Rendered, input.Text) + assert.Equal(t, generated.Rendered, policy.ModelInput.EncodeDocument(input.Text)) + + descriptor, err := NewEmbeddingDescriptor(EmbeddingDescriptor{ + ID: "synthetic-embedder", ContractVersion: EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("b", 64), TrustBoundary: EmbeddingTrustLocalProcess, + Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: VectorMetricCosine, + InputKinds: []EmbeddingInputKind{EmbeddingInputRenditionChunk}, CompatibilityID: policy.ModelInput.CompatibilityID, + ModelInput: policy.ModelInput, SupportedRequestModes: []ModelInputMode{ModelInputModeText}, + DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: VectorNormalizationUnitLength, + ScalarEncoding: "float32", + }) + require.NoError(t, err) + authorization := EmbeddingAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 1, + MaxInputBytes: int64(len(generated.Rendered)), MaxResponseBytes: 1024, + } + require.NoError(t, ValidateEmbeddingProviderRequest(chunkingTestProvider{descriptor: descriptor}, []EmbeddingInput{input}, authorization)) + authorization.MaxInputBytes-- + require.ErrorContains(t, ValidateEmbeddingProviderRequest(chunkingTestProvider{descriptor: descriptor}, []EmbeddingInput{input}, authorization), "input bytes") +} + +func TestEmbeddingInputGenerationRejectsEqualEnvelopeWithDifferentModelInputFingerprint(t *testing.T) { + bge, err := NewModelInputContract(ModelInputContractConfig{Profile: ModelInputProfileBGEM3}) + require.NoError(t, err) + gte, err := NewModelInputContract(ModelInputContractConfig{Profile: ModelInputProfileGTE}) + require.NoError(t, err) + assert.Equal(t, bge.Document, gte.Document, "the regression requires equal document envelopes") + assert.NotEqual(t, bge.Fingerprint, gte.Fingerprint) + + policy := testInputPolicy(t, 2, 0) + policy.ModelInput = bge + generation, err := BuildEmbeddingInputs(testChunkEvidence(t, []sourceChunkUnit{{text: "AABB"}}), policy) + require.NoError(t, err) + _, err = generation.ToEmbeddingInputs(gte) + require.ErrorContains(t, err, "fingerprint") +} + +type chunkingTestProvider struct{ descriptor EmbeddingDescriptor } + +func (provider chunkingTestProvider) Descriptor() EmbeddingDescriptor { return provider.descriptor } +func (chunkingTestProvider) Embed(context.Context, []EmbeddingInput, EmbeddingAuthorization) (EmbeddingResult, error) { + return EmbeddingResult{}, nil +} + +func TestBuildEmbeddingInputsGoldenProfiles(t *testing.T) { + evidence := testChunkEvidence(t, []sourceChunkUnit{ + {text: "Alpha beta.", heading: []string{"Intro"}}, + {text: "Gamma delta.", heading: []string{"Details"}}, + }) + profiles := []ModelInputContractConfig{ + {Profile: ModelInputProfileNomic}, + {Profile: ModelInputProfileE5}, + {Profile: ModelInputProfileBGEM3}, + {Profile: ModelInputProfileGTE}, + {Profile: ModelInputProfileQwen3, QueryInstruction: "Retrieve supporting passages"}, + } + for _, config := range profiles { + t.Run(string(config.Profile), func(t *testing.T) { + contract, err := NewModelInputContract(config) + require.NoError(t, err) + policy := testInputPolicy(t, 6, 1) + policy.ModelInput = contract + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + encoded, err := json.MarshalIndent(generation, "", " ") + require.NoError(t, err) + encoded = append(encoded, '\n') + goldenPath := "testdata/chunks-" + strings.ReplaceAll(string(config.Profile), "/", "-") + ".golden.json" + if os.Getenv("UPDATE_GOLDEN") == "1" { + require.NoError(t, os.WriteFile(goldenPath, encoded, 0o644)) + } + golden, err := os.ReadFile(goldenPath) + require.NoError(t, err) + assert.JSONEq(t, string(golden), string(encoded)) + }) + } +} + +func FuzzEmbeddingSpans(f *testing.F) { + f.Add("alpha beta", uint8(4), uint8(1)) + f.Add("éclair 世界", uint8(3), uint8(0)) + f.Fuzz(func(t *testing.T, text string, budgetByte, overlapByte uint8) { + if text == "" || !utf8.ValidString(text) || strings.ContainsAny(text, "\x00\r") || len([]rune(text)) > 256 { + t.Skip() + } + evidence := testChunkEvidence(t, []sourceChunkUnit{{text: text, heading: []string{"Synthetic"}}}) + budget := max(1, int(budgetByte%32)) + overlap := int(overlapByte) % budget + policy := testInputPolicy(t, budget, overlap) + policy.MaxGeneratedInputs = 512 + policy.MaxTotalContentTokens = 512 * 32 + policy.MaxTotalRenderedTokens = 512 * int64(policy.MaxProviderTokens) + policy.MaxTotalContentBytes = 512 * policy.MaxProviderBytes + policy.MaxTotalRenderedBytes = 512 * policy.MaxProviderBytes + generation, err := BuildEmbeddingInputs(evidence, policy) + require.NoError(t, err) + assertGenerationSpans(t, evidence, generation) + for index, input := range generation.Inputs { + assert.LessOrEqual(t, input.ContentTokens, budget) + assert.LessOrEqual(t, input.RenderedTokens, policy.MaxProviderTokens) + assert.LessOrEqual(t, int64(len(input.Rendered)), policy.MaxProviderBytes) + exactTokens, err := policy.Tokenizer.Tokenize(input.Content, budget) + require.NoError(t, err) + require.NoError(t, validateTokenBoundaries(exactTokens, utf8.RuneCountInString(input.Content), budget)) + assert.Equal(t, len(exactTokens), input.ContentTokens) + if index == 0 || overlap == 0 || generation.Inputs[index-1].Truncated || generation.Inputs[index-1].SourceSpans[0].UnitIndex != input.SourceSpans[0].UnitIndex { + continue + } + previous := generation.Inputs[index-1] + previousTokens, err := policy.Tokenizer.Tokenize(previous.Content, budget) + require.NoError(t, err) + overlapRunes := []rune(previous.Content)[previousTokens[len(previousTokens)-overlap].Start:] + assert.True(t, strings.HasPrefix(input.Content, string(overlapRunes))) + } + }) +} + +type namedPairTokenizer struct{ name, revision string } + +func (tokenizer namedPairTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: tokenizer.name, Revision: tokenizer.revision, PrefixTokenCountsMonotonic: true} +} + +type nonMonotonicNamedPairTokenizer struct{} + +func (nonMonotonicNamedPairTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "synthetic-pairs", Revision: "v1"} +} +func (nonMonotonicNamedPairTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + return pairTokenizer{}.Tokenize(text, limit) +} + +type concatAdversarialTokenizer struct{} + +func (concatAdversarialTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "concat-adversarial", Revision: "v1"} +} + +type overlapAdversarialTokenizer struct{} + +func (overlapAdversarialTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "overlap-adversarial", Revision: "v1"} +} +func (overlapAdversarialTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + var result []TokenBoundary + switch text { + case "abcdef": + result = []TokenBoundary{{Start: 0, End: 2}, {Start: 2, End: 4}, {Start: 4, End: 6}} + case "abcd": + result = []TokenBoundary{{Start: 0, End: 1}, {Start: 1, End: 4}} + case "bcd": + result = []TokenBoundary{{Start: 0, End: 1}, {Start: 1, End: 3}} + case "cdef": + result = []TokenBoundary{{Start: 0, End: 2}, {Start: 2, End: 4}} + default: + return runeBoundaries(text, limit) + } + if len(result) > limit { + return nil, ErrTokenizerLimit + } + return result, nil +} +func (concatAdversarialTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if text == "AB" { + if limit < 1 { + return nil, ErrTokenizerLimit + } + return []TokenBoundary{{Start: 0, End: 2}}, nil + } + return runeBoundaries(text, limit) +} + +type truncationAdversarialTokenizer struct{} + +func (truncationAdversarialTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "truncation-adversarial", Revision: "v1"} +} +func (truncationAdversarialTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if text == "abcd" { + if limit < 1 { + return nil, ErrTokenizerLimit + } + return []TokenBoundary{{Start: 0, End: 4}}, nil + } + return runeBoundaries(text, limit) +} + +type nonMonotonicPrefixTokenizer struct{} + +func (nonMonotonicPrefixTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "nonmonotonic-prefix", Revision: "v1"} +} + +type suffixTruncationAdversarialTokenizer struct{} + +func (suffixTruncationAdversarialTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "suffix-truncation-adversarial", Revision: "v1", PrefixTokenCountsMonotonic: true} +} + +type fittingWorkAdversarialTokenizer struct{ exactCalls int } + +func (*fittingWorkAdversarialTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "fitting-work-adversarial", Revision: "v1"} +} +func (tokenizer *fittingWorkAdversarialTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if limit == maxEmbeddingTokensPerGeneration { + return runeBoundaries(text, limit) + } + tokenizer.exactCalls++ + return nil, ErrTokenizerLimit +} +func (suffixTruncationAdversarialTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if limit > 1 || !strings.HasSuffix(text, " suffix") { + return []TokenBoundary{{Start: 0, End: utf8.RuneCountInString(text)}}, nil + } + switch text { + case "a suffix", "abcd suffix": + return []TokenBoundary{{Start: 0, End: utf8.RuneCountInString(text)}}, nil + default: + return nil, ErrTokenizerLimit + } +} +func (nonMonotonicPrefixTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if limit > 4 || strings.HasPrefix(text, "document: ") { + return runeBoundaries(text, limit) + } + switch text { + case "abcd", "ab": + return nil, ErrTokenizerLimit + default: + return runeBoundaries(text, limit) + } +} + +func runeBoundaries(text string, limit int) ([]TokenBoundary, error) { + runes := []rune(text) + if len(runes) > limit { + return nil, ErrTokenizerLimit + } + result := make([]TokenBoundary, len(runes)) + for index := range runes { + result[index] = TokenBoundary{Start: index, End: index + 1} + } + return result, nil +} +func (namedPairTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + return pairTokenizer{}.Tokenize(text, limit) +} + +type brokenTokenizer struct { + tokens []TokenBoundary + err error +} + +type nilTokenizer struct{} + +func (*nilTokenizer) Identity() TokenizerIdentity { panic("typed nil tokenizer must not be called") } +func (*nilTokenizer) Tokenize(string, int) ([]TokenBoundary, error) { + panic("typed nil tokenizer must not be called") +} + +type identityDriftTokenizer struct{ calls int } + +func (tokenizer *identityDriftTokenizer) Identity() TokenizerIdentity { + tokenizer.calls++ + return TokenizerIdentity{Name: "identity-drift", Revision: fmt.Sprintf("v%d", tokenizer.calls), PrefixTokenCountsMonotonic: true} +} +func (*identityDriftTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + return pairTokenizer{}.Tokenize(text, limit) +} + +type countingRuneTokenizer struct{ calls int } + +func (*countingRuneTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "counting-runes", Revision: "v1", PrefixTokenCountsMonotonic: true} +} +func (tokenizer *countingRuneTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + tokenizer.calls++ + return runeBoundaries(text, limit) +} + +type recordingRuneTokenizer struct { + contentLimit int + maxExactContentRunes int + sawOneTokenContentLimit bool +} + +func (*recordingRuneTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "recording-runes", Revision: "v1", PrefixTokenCountsMonotonic: true} +} +func (tokenizer *recordingRuneTokenizer) Tokenize(text string, limit int) ([]TokenBoundary, error) { + if limit == tokenizer.contentLimit { + tokenizer.maxExactContentRunes = max(tokenizer.maxExactContentRunes, utf8.RuneCountInString(text)) + } + if limit == 1 && !strings.HasPrefix(text, "document: ") { + tokenizer.sawOneTokenContentLimit = true + } + return runeBoundaries(text, limit) +} + +func (brokenTokenizer) Identity() TokenizerIdentity { + return TokenizerIdentity{Name: "broken", Revision: "v1"} +} +func (tokenizer brokenTokenizer) Tokenize(string, int) ([]TokenBoundary, error) { + return tokenizer.tokens, tokenizer.err +} + +func testInputPolicy(t *testing.T, budget, overlap int) InputPolicy { + t.Helper() + contract, err := NewModelInputContract(ModelInputContractConfig{ + Profile: ModelInputProfileCustom, CompatibilityID: "synthetic-space", + Document: ModelInputEncoder{Mode: ModelInputModeText, Template: "document: {{content}}"}, + Query: ModelInputEncoder{Mode: ModelInputModeText, Template: "query: {{content}}"}, + }) + require.NoError(t, err) + return InputPolicy{ + Tokenizer: pairTokenizer{}, ContentTokenBudget: budget, OverlapTokens: overlap, + MaxProviderTokens: 256, MaxProviderBytes: 4096, MaxGeneratedInputs: 128, + MaxTotalContentTokens: 4096, MaxTotalRenderedTokens: 8192, + MaxTotalContentBytes: 1 << 20, MaxTotalRenderedBytes: 2 << 20, + MaxFittingWorkBytes: 8 << 20, MaxFittingWorkTokens: 1 << 20, + ModelInput: contract, Formatter: "evidence-text/v1", LexicalEvidenceFingerprint: strings.Repeat("a", 64), + ContextFingerprint: strings.Repeat("b", 64), + TruncationPolicy: TruncationPolicyReject, + } +} + +func testGenerationDecodeBounds() EmbeddingInputGenerationDecodeBounds { + return EmbeddingInputGenerationDecodeBounds{ + MaxEncodedBytes: 1 << 20, + MaxInputs: 128, + MaxObjectFields: 32, + MaxStringBytes: 1 << 16, + MaxTotalStringBytes: 1 << 20, + } +} + +type sourceChunkUnit struct { + text string + heading []string + regions []SourceEvidenceRegionV1 + tables []SourceEvidenceTableV1 +} + +func testChunkEvidence(t *testing.T, units []sourceChunkUnit) NormalizedEvidenceV1 { + t.Helper() + sourceUnits := make([]SourceEvidenceUnitV1, len(units)) + for index, unit := range units { + sourceUnits[index] = SourceEvidenceUnitV1{ + Order: index, HeadingPath: unit.heading, Text: unit.text, Regions: unit.regions, Tables: unit.tables, + Locator: SourceEvidenceLocatorV1{Kind: EvidenceLocatorPage, IndexOrigin: EvidenceIndexOriginOne, Start: int64(index + 1), End: int64(index + 1)}, + } + } + source := SourceEvidenceV1{ContractVersion: SourceEvidenceContractV1, Completeness: EvidenceComplete, Family: "pdf", UnitKind: EvidenceUnitPage, Units: sourceUnits} + policy, err := NewEvidencePolicy(4096) + require.NoError(t, err) + evidence, err := NormalizeEvidenceV1(source, policy) + require.NoError(t, err) + return evidence +} + +func assertGenerationSpans(t *testing.T, evidence NormalizedEvidenceV1, generation EmbeddingInputGeneration) { + t.Helper() + for _, input := range generation.Inputs { + var reconstructed strings.Builder + for _, span := range input.SourceSpans { + require.GreaterOrEqual(t, span.UnitIndex, 0) + require.Less(t, span.UnitIndex, len(evidence.Units)) + runes := []rune(evidence.Units[span.UnitIndex].Text) + require.GreaterOrEqual(t, span.CharStart, 0) + require.Greater(t, span.CharEnd, span.CharStart) + require.LessOrEqual(t, span.CharEnd, len(runes)) + reconstructed.WriteString(string(runes[span.CharStart:span.CharEnd])) + } + assert.Equal(t, input.Content, reconstructed.String()) + } +} diff --git a/document/cohereembed/client.go b/document/cohereembed/client.go new file mode 100644 index 00000000..dc47ad48 --- /dev/null +++ b/document/cohereembed/client.go @@ -0,0 +1,607 @@ +package cohereembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "net/http" + "slices" + "sync" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/internal/cohereapi" +) + +const ( + maxSecretBytes = 64 << 10 + maxUsageValue = float64(1 << 50) +) + +var _ document.EmbeddingProvider = (*Client)(nil) + +type wireRequest struct { + Model string `json:"model"` + Texts []string `json:"texts,omitempty"` + Inputs []wireInput `json:"inputs,omitempty"` + InputType string `json:"input_type"` + EmbeddingTypes []string `json:"embedding_types"` + Truncate string `json:"truncate"` + OutputDimension int `json:"output_dimension"` +} + +type wireInput struct { + Content []wireContent `json:"content"` +} + +type wireContent struct { + Type string `json:"type"` + ImageURL wireImageURL `json:"image_url"` +} + +type wireImageURL struct { + URL string `json:"url"` +} + +type wireResponse struct { + ID string `json:"id"` + ResponseType string `json:"response_type,omitempty"` + Embeddings struct { + Float [][]float32 `json:"float"` + Int8 [][]int8 `json:"int8,omitempty"` + Uint8 [][]uint8 `json:"uint8,omitempty"` + Binary [][]int8 `json:"binary,omitempty"` + UBinary [][]uint8 `json:"ubinary,omitempty"` + Base64 []string `json:"base64,omitempty"` + } `json:"embeddings"` + Texts []string `json:"texts,omitempty"` + Images []wireImage `json:"images,omitempty"` + Meta *wireResponseMeta `json:"meta,omitempty"` +} + +type wireImage struct { + Width int64 `json:"width"` + Height int64 `json:"height"` + Format string `json:"format"` + BitDepth int `json:"bit_depth"` +} + +type wireResponseMeta struct { + APIVersion *struct { + Version string `json:"version"` + IsDeprecated *bool `json:"is_deprecated,omitempty"` + IsExperimental *bool `json:"is_experimental,omitempty"` + } `json:"api_version,omitempty"` + BilledUnits *struct { + Images *float64 `json:"images,omitempty"` + InputTokens *float64 `json:"input_tokens,omitempty"` + ImageTokens *float64 `json:"image_tokens,omitempty"` + OutputTokens *float64 `json:"output_tokens,omitempty"` + SearchUnits *float64 `json:"search_units,omitempty"` + Classifications *float64 `json:"classifications,omitempty"` + Pages *float64 `json:"pages,omitempty"` + } `json:"billed_units,omitempty"` + Tokens *struct { + InputTokens *float64 `json:"input_tokens,omitempty"` + OutputTokens *float64 `json:"output_tokens,omitempty"` + } `json:"tokens,omitempty"` + CachedTokens *float64 `json:"cached_tokens,omitempty"` + Warnings []string `json:"warnings,omitempty"` +} + +type preparedRequest struct { + positions []int + payload []byte + texts []string + images []string + imageMeta []wireImage +} + +// Receipt is bounded provider provenance and numeric usage without request, +// response, vector, source, or credential material. +type Receipt struct { + ProviderID string + DescriptorFingerprint string + PolicyFingerprint string + Model string + ModelRevision string + RequestCount int + ImageInputs int + BilledImages float64 + InputTokens float64 + ImageTokens float64 + OutputTokens float64 + SearchUnits float64 + Classifications float64 + Pages float64 + CachedTokens float64 + ProviderResponseIDs []string +} + +// Execution contains the E1 result plus its sanitized provider receipt. +type Execution struct { + Result document.EmbeddingResult + Receipt Receipt +} + +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + return client.embed(ctx, inputs, authorization, nil) +} + +// EmbedWithReceipt executes the same E1 boundary while returning bounded, +// sanitized provider provenance for callers that persist execution receipts. +func (client *Client) EmbedWithReceipt(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (Execution, error) { + receipt := Receipt{ProviderID: ProviderID, DescriptorFingerprint: client.Descriptor().Fingerprint, + PolicyFingerprint: client.Descriptor().PolicyFingerprint, Model: Model, + ModelRevision: client.Descriptor().ModelRevision} + result, err := client.embed(ctx, inputs, authorization, &receipt) + if err != nil { + return Execution{}, err + } + return Execution{Result: result, Receipt: receipt}, nil +} + +func (client *Client) embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization, receipt *Receipt) (document.EmbeddingResult, error) { + if client == nil || ctx == nil { + return document.EmbeddingResult{}, errors.New("cohere embed: client and context are required") + } + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + frozenInputs, enrolled := enrollOriginalUploads(inputs) + defer closeEnrolledUploads(enrolled) + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, contextErr + } + if err := document.ValidateEmbeddingProviderRequest(client, frozenInputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || authorization.MaxInputBytes > client.profile.MaxInputBytes || + authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrCapacityResponse} + } + sourceGate := newActiveSourceGate() + closeFinished := make(chan struct{}) + stopClose := context.AfterFunc(requestCtx, func() { + sourceGate.Cancel() + close(closeFinished) + }) + defer func() { + if !stopClose() { + <-closeFinished + } + }() + prepared, err := client.prepareRequests(requestCtx, frozenInputs, sourceGate) + if err != nil { + return document.EmbeddingResult{}, err + } + defer func() { + for index := range prepared { + clear(prepared[index].payload) + for image := range prepared[index].images { + prepared[index].images[image] = "" + } + } + }() + secret, err := client.secrets.ResolveSecret(requestCtx, client.profile.SecretBinding) + if err != nil || !cohereapi.ValidToken(secret, maxSecretBytes) { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("cohere embed: credential resolution canceled: %w", contextErr) + } + return document.EmbeddingResult{}, errors.New("cohere embed: API-key resolution failed") + } + result := document.EmbeddingResult{Vectors: make([]document.EmbeddingVector, len(frozenInputs))} + for index := range prepared { + vectors, err := client.execute(requestCtx, prepared[index], secret, receipt) + if err != nil { + return document.EmbeddingResult{}, err + } + for local, global := range prepared[index].positions { + result.Vectors[global] = document.EmbeddingVector{Key: frozenInputs[global].Key, Values: vectors[local]} + } + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, frozenInputs, authorization, result); err != nil { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + return result, nil +} + +func (client *Client) prepareRequests(ctx context.Context, inputs []document.EmbeddingInput, sourceGate *activeSourceGate) ([]preparedRequest, error) { + documentPositions, queryPositions, imagePositions := []int{}, []int{}, []int{} + documents, queries := []string{}, []string{} + sources := make([]*enrolledUpload, 0, len(inputs)) + for _, input := range inputs { + if input.Kind == document.EmbeddingInputOriginalFile { + frozen, ok := input.Source.(*frozenUpload) + if !ok { + return nil, errors.New("cohere embed: original upload was not frozen") + } + sources = append(sources, frozen.enrolled) + } + } + var imageBytes int64 + for index, input := range inputs { + switch { + case input.Kind == document.EmbeddingInputOriginalFile: + metadata := input.Source.Metadata() + if metadata.ByteLength > client.profile.MaxInputItemBytes || metadata.ByteLength > client.profile.MaxImageBytes-imageBytes { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + imageBytes += metadata.ByteLength + imagePositions = append(imagePositions, index) + case input.Role == document.EmbeddingRoleDocument: + rendered := client.descriptor.ModelInput.EncodeDocument(input.Text) + if int64(len(rendered)) > client.profile.MaxInputItemBytes { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + documentPositions, documents = append(documentPositions, index), append(documents, rendered) + case input.Role == document.EmbeddingRoleQuery: + rendered := client.descriptor.ModelInput.EncodeQuery(input.Text) + if int64(len(rendered)) > client.profile.MaxInputItemBytes { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + queryPositions, queries = append(queryPositions, index), append(queries, rendered) + } + } + images := make([]string, len(imagePositions)) + imageMeta := make([]wireImage, len(imagePositions)) + for local := range sources { + source := sources[local] + if source.liveMetadata() != source.metadata { + clearStrings(images) + return nil, errors.New("cohere embed: image source changed or could not be read exactly") + } + token, ok := sourceGate.Begin(source) + if !ok { + clearStrings(images) + if contextErr := ctx.Err(); contextErr != nil { + return nil, contextErr + } + return nil, errors.New("cohere embed: image source transfer stopped") + } + dataURL, detected, readErr := client.readImage(source.source, source.metadata) + sourceGate.End(token) + metadataChanged := source.liveMetadata() != source.metadata + closeErr := source.Close() + if contextErr := ctx.Err(); contextErr != nil { + clearStrings(images) + return nil, contextErr + } + if readErr != nil || closeErr != nil || metadataChanged { + clearStrings(images) + return nil, errors.New("cohere embed: image source changed or could not be read exactly") + } + images[local] = dataURL + imageMeta[local] = detected + } + requests := make([]preparedRequest, 0, 3) + for _, candidate := range []struct { + positions []int + texts []string + images []string + inputType string + }{{documentPositions, documents, nil, "search_document"}, {imagePositions, nil, images, "search_document"}, {queryPositions, queries, nil, "search_query"}} { + if len(candidate.positions) == 0 { + continue + } + payload, err := json.Marshal(wireRequest{Model: Model, Texts: candidate.texts, + Inputs: wireInputs(candidate.images), + InputType: candidate.inputType, EmbeddingTypes: []string{"float"}, Truncate: "NONE", + OutputDimension: client.descriptor.Dimension}) + if err != nil { + return nil, errors.New("cohere embed: request encoding failed") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + clear(payload) + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + requests = append(requests, preparedRequest{positions: slices.Clone(candidate.positions), payload: payload, + texts: slices.Clone(candidate.texts), images: slices.Clone(candidate.images)}) + if len(candidate.images) != 0 { + requests[len(requests)-1].imageMeta = slices.Clone(imageMeta) + } + } + return requests, nil +} + +func wireInputs(images []string) []wireInput { + if len(images) == 0 { + return nil + } + inputs := make([]wireInput, len(images)) + for index, image := range images { + inputs[index] = wireInput{Content: []wireContent{{Type: "image_url", ImageURL: wireImageURL{URL: image}}}} + } + return inputs +} + +type enrolledUpload struct { + source document.AuthorizedUpload + metadata document.AuthorizedUploadMetadata + + closeOnce sync.Once + closeErr error +} + +func (upload *enrolledUpload) Close() error { + upload.closeOnce.Do(func() { upload.closeErr = upload.source.Close() }) + return upload.closeErr +} + +func (upload *enrolledUpload) liveMetadata() document.AuthorizedUploadMetadata { + return upload.source.Metadata() +} + +type frozenUpload struct { + enrolled *enrolledUpload +} + +func (upload *frozenUpload) Read(value []byte) (int, error) { + return upload.enrolled.source.Read(value) +} +func (upload *frozenUpload) Close() error { return upload.enrolled.Close() } +func (upload *frozenUpload) Metadata() document.AuthorizedUploadMetadata { + return upload.enrolled.metadata +} + +func enrollOriginalUploads(inputs []document.EmbeddingInput) ([]document.EmbeddingInput, []*enrolledUpload) { + frozen := slices.Clone(inputs) + enrolled := make([]*enrolledUpload, 0, len(inputs)) + for index := range frozen { + frozen[index].HeadingPath = slices.Clone(frozen[index].HeadingPath) + frozen[index].SourceSpans = slices.Clone(frozen[index].SourceSpans) + if frozen[index].Kind != document.EmbeddingInputOriginalFile || nilInterface(frozen[index].Source) { + continue + } + upload := &enrolledUpload{source: frozen[index].Source, metadata: frozen[index].Source.Metadata()} + enrolled = append(enrolled, upload) + frozen[index].Source = &frozenUpload{enrolled: upload} + } + return frozen, enrolled +} + +func closeEnrolledUploads(uploads []*enrolledUpload) { + for _, upload := range uploads { + _ = upload.Close() + } +} + +func clearStrings(values []string) { + for index := range values { + values[index] = "" + } +} + +type activeSourceGate struct { + mu sync.Mutex + canceled bool + nextToken uint64 + activeToken uint64 + active *enrolledUpload +} + +func newActiveSourceGate() *activeSourceGate { return new(activeSourceGate) } + +func (gate *activeSourceGate) Begin(source *enrolledUpload) (uint64, bool) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.canceled { + return 0, false + } + gate.nextToken++ + gate.activeToken = gate.nextToken + gate.active = source + return gate.activeToken, true +} + +func (gate *activeSourceGate) End(token uint64) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.activeToken != token { + return + } + gate.active = nil + gate.activeToken = 0 +} + +func (gate *activeSourceGate) Cancel() { + gate.mu.Lock() + if gate.canceled { + gate.mu.Unlock() + return + } + gate.canceled = true + active := gate.active + gate.active = nil + gate.activeToken = 0 + gate.mu.Unlock() + if active != nil { + _ = active.Close() + } +} + +func (client *Client) readImage(source document.AuthorizedUpload, metadata document.AuthorizedUploadMetadata) (string, wireImage, error) { + data, readErr := io.ReadAll(io.LimitReader(source, metadata.ByteLength+1)) + defer clear(data) + if readErr != nil || int64(len(data)) != metadata.ByteLength { + return "", wireImage{}, errors.New("cohere embed: image source changed or could not be read exactly") + } + digest := sha256.Sum256(data) + if hex.EncodeToString(digest[:]) != metadata.SHA256 { + return "", wireImage{}, errors.New("cohere embed: image source checksum changed") + } + detected, reason := media.InspectBytes(data, metadata.MediaType, client.profile.MediaPolicy) + if reason != media.ReasonEligible || detected.Kind != media.KindImage || detected.MediaType != metadata.MediaType || + detected.Size != metadata.ByteLength || !slices.Contains(acceptedImageFormats, detected.MediaType) { + return "", wireImage{}, errors.New("cohere embed: image source media identity is invalid") + } + encoded := make([]byte, 0, len("data:")+len(detected.MediaType)+len(";base64,")+base64.StdEncoding.EncodedLen(len(data))) + encoded = append(encoded, "data:"...) + encoded = append(encoded, detected.MediaType...) + encoded = append(encoded, ";base64,"...) + encoded = base64.StdEncoding.AppendEncode(encoded, data) + result := string(encoded) + clear(encoded) + return result, wireImage{Width: detected.Width, Height: detected.Height, Format: string(detected.Format)}, nil +} + +func (client *Client) execute(ctx context.Context, prepared preparedRequest, secret string, receipt *Receipt) ([][]float32, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodPost, origin+embedPath, bytes.NewReader(prepared.payload)) + if err != nil { + return nil, errors.New("cohere embed: request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+secret) + response, err := client.http.Do(request) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("cohere embed: request canceled: %w", contextErr) + } + return nil, &ProviderError{Kind: ErrTransientResponse} + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if !cohereapi.IsJSONContentType(response.Header.Get("Content-Type")) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + body, outcome, readErr := cohereapi.ReadBounded(ctx, response.Body, client.profile.MaxResponseBytes) + switch outcome { + case cohereapi.ReadOK: + case cohereapi.ReadCanceled: + return nil, fmt.Errorf("cohere embed: response read canceled: %w", readErr) + case cohereapi.ReadTransient: + return nil, &ProviderError{Kind: ErrTransientResponse} + case cohereapi.ReadCapacity: + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + defer clear(body) + var decoded wireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if !cohereapi.ValidToken(decoded.ID, 128) || decoded.ResponseType != "" && decoded.ResponseType != "embeddings_by_type" || + len(decoded.Embeddings.Float) != len(prepared.positions) || + len(decoded.Embeddings.Int8) != 0 || len(decoded.Embeddings.Uint8) != 0 || + len(decoded.Embeddings.Binary) != 0 || len(decoded.Embeddings.UBinary) != 0 || len(decoded.Embeddings.Base64) != 0 || + decoded.Texts != nil && !slices.Equal(decoded.Texts, prepared.texts) || + !validImageMetadata(decoded.Images, prepared.imageMeta) || !validMeta(decoded.Meta, len(prepared.imageMeta)) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + for _, vector := range decoded.Embeddings.Float { + if len(vector) != client.descriptor.Dimension { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + } + } + if receipt != nil && !addReceipt(receipt, decoded, len(prepared.imageMeta)) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + return decoded.Embeddings.Float, nil +} + +func addReceipt(receipt *Receipt, response wireResponse, imageInputs int) bool { + values := make([]*float64, 8) + if response.Meta != nil { + if response.Meta.BilledUnits != nil { + values = []*float64{response.Meta.BilledUnits.Images, response.Meta.BilledUnits.InputTokens, + response.Meta.BilledUnits.ImageTokens, response.Meta.BilledUnits.OutputTokens, + response.Meta.BilledUnits.SearchUnits, response.Meta.BilledUnits.Classifications, + response.Meta.BilledUnits.Pages, response.Meta.CachedTokens} + } else { + values[7] = response.Meta.CachedTokens + } + if response.Meta.Tokens != nil { + if values[1] == nil { + values[1] = response.Meta.Tokens.InputTokens + } + if values[3] == nil { + values[3] = response.Meta.Tokens.OutputTokens + } + } + } + totals := []*float64{&receipt.BilledImages, &receipt.InputTokens, &receipt.ImageTokens, + &receipt.OutputTokens, &receipt.SearchUnits, &receipt.Classifications, &receipt.Pages, &receipt.CachedTokens} + for index, value := range values { + if value != nil { + if *value > maxUsageValue-*totals[index] { + return false + } + *totals[index] += *value + } + } + receipt.RequestCount++ + receipt.ImageInputs += imageInputs + receipt.ProviderResponseIDs = append(receipt.ProviderResponseIDs, response.ID) + return true +} + +func validImageMetadata(actual, expected []wireImage) bool { + if actual == nil { + return len(expected) == 0 + } + if len(actual) != len(expected) { + return false + } + for index, image := range actual { + if image.Width != expected[index].Width || image.Height != expected[index].Height || + image.Format != expected[index].Format || image.BitDepth < 1 || image.BitDepth > 64 { + return false + } + } + return true +} + +func validMeta(metadata *wireResponseMeta, expectedImages int) bool { + if metadata == nil { + return true + } + if metadata.APIVersion != nil && metadata.APIVersion.Version == "" { + return false + } + if metadata.BilledUnits != nil && metadata.Tokens != nil && + (!matchingUsage(metadata.BilledUnits.InputTokens, metadata.Tokens.InputTokens) || + !matchingUsage(metadata.BilledUnits.OutputTokens, metadata.Tokens.OutputTokens)) { + return false + } + values := []*float64{} + if metadata.BilledUnits != nil { + values = append(values, metadata.BilledUnits.Images, metadata.BilledUnits.InputTokens, + metadata.BilledUnits.ImageTokens, metadata.BilledUnits.OutputTokens, + metadata.BilledUnits.SearchUnits, metadata.BilledUnits.Classifications, metadata.BilledUnits.Pages) + } + if metadata.BilledUnits != nil && metadata.BilledUnits.Images != nil && *metadata.BilledUnits.Images != float64(expectedImages) { + return false + } + if expectedImages == 0 && metadata.BilledUnits != nil && metadata.BilledUnits.ImageTokens != nil && + *metadata.BilledUnits.ImageTokens != 0 { + return false + } + if metadata.Tokens != nil { + values = append(values, metadata.Tokens.InputTokens, metadata.Tokens.OutputTokens) + } + values = append(values, metadata.CachedTokens) + for _, value := range values { + if value != nil && (math.IsNaN(*value) || math.IsInf(*value, 0) || *value < 0 || *value > maxUsageValue) { + return false + } + } + return true +} + +func matchingUsage(left, right *float64) bool { + return left == nil || right == nil || *left == *right +} diff --git a/document/cohereembed/client_test.go b/document/cohereembed/client_test.go new file mode 100644 index 00000000..aaf1f2ee --- /dev/null +++ b/document/cohereembed/client_test.go @@ -0,0 +1,478 @@ +package cohereembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/v2" + "errors" + "image/color" + "io" + "math" + "net/http" + "net/netip" + "slices" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media/mediatest" +) + +func TestEmbedSendsExactRoleAndImageRequestsAndRestoresCallerOrder(t *testing.T) { + profile := testProfile(t, 256) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, profile, secrets, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + assert.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "https://api.cohere.com/v2/embed", request.URL.String()) + assert.Equal(t, "Bearer synthetic-key", request.Header.Get("Authorization")) + assert.Equal(t, "application/json", request.Header.Get("Accept")) + assert.Equal(t, "application/json", request.Header.Get("Content-Type")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + var payload struct { + Model string `json:"model"` + Texts []string `json:"texts"` + Inputs []struct { + Content []struct { + Type string `json:"type"` + ImageURL struct { + URL string `json:"url"` + } `json:"image_url"` + } `json:"content"` + } `json:"inputs"` + InputType string `json:"input_type"` + EmbeddingTypes []string `json:"embedding_types"` + Truncate string `json:"truncate"` + OutputDimension int `json:"output_dimension"` + } + require.NoError(t, json.Unmarshal(body, &payload, json.RejectUnknownMembers(true))) + assert.Equal(t, Model, payload.Model) + assert.Equal(t, []string{"float"}, payload.EmbeddingTypes) + assert.Equal(t, "NONE", payload.Truncate) + assert.Equal(t, 256, payload.OutputDimension) + vector := make([]float32, 256) + switch { + case len(payload.Inputs) == 1: + assert.Empty(t, payload.Texts) + assert.Equal(t, "search_document", payload.InputType) + require.Len(t, payload.Inputs[0].Content, 1) + assert.Equal(t, "image_url", payload.Inputs[0].Content[0].Type) + assert.Equal(t, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(tinyPNG(t)), payload.Inputs[0].Content[0].ImageURL.URL) + vector[2] = 1 + case payload.InputType == "search_document": + assert.Equal(t, []string{"document text"}, payload.Texts) + assert.Empty(t, payload.Inputs) + vector[0] = 1 + case payload.InputType == "search_query": + assert.Equal(t, []string{"query text"}, payload.Texts) + assert.Empty(t, payload.Inputs) + vector[1] = 1 + default: + t.Fatalf("unexpected request: %#v", payload) + } + responseValue := map[string]any{"id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}} + if len(payload.Inputs) != 0 { + responseValue["images"] = []map[string]any{{"width": 1, "height": 1, "format": "png", "bit_depth": 8}} + responseValue["response_type"] = "embeddings_by_type" + responseValue["meta"] = map[string]any{"billed_units": map[string]any{"images": 1}} + } + response, err := json.Marshal(responseValue) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, response), nil + })) + inputs := []document.EmbeddingInput{ + {Key: "document", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "document text"}, + {Key: "image", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: imageUpload(t, tinyPNG(t), "image/png")}, + {Key: "query", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "query text"}, + } + original := slices.Clone(inputs) + execution, err := client.EmbedWithReceipt(context.Background(), inputs, authorization(client.Descriptor(), len(inputs))) + require.NoError(t, err) + result := execution.Result + assert.Equal(t, int32(3), requests.Load()) + assert.Equal(t, int32(1), secrets.calls.Load()) + assert.Equal(t, original, inputs) + require.Len(t, result.Vectors, 3) + assert.Equal(t, "document", result.Vectors[0].Key) + assert.InDelta(t, 1, result.Vectors[0].Values[0], 0) + assert.Equal(t, "image", result.Vectors[1].Key) + assert.InDelta(t, 1, result.Vectors[1].Values[2], 0) + assert.Equal(t, "query", result.Vectors[2].Key) + assert.InDelta(t, 1, result.Vectors[2].Values[1], 0) + assert.Equal(t, Receipt{ProviderID: ProviderID, DescriptorFingerprint: client.Descriptor().Fingerprint, + PolicyFingerprint: client.Descriptor().PolicyFingerprint, Model: Model, + ModelRevision: "deployment-2026-08", RequestCount: 3, ImageInputs: 1, BilledImages: 1, + ProviderResponseIDs: []string{"synthetic-id", "synthetic-id", "synthetic-id"}}, execution.Receipt) +} + +func TestEmbedRejectsImageOverPerItemBoundBeforeSecretOrRequest(t *testing.T) { + data := tinyPNG(t) + profile := testProfile(t, 256) + profile.MaxInputItemBytes = int64(len(data) - 1) + profile.Descriptor = descriptorFor(t, profile) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + first := imageUpload(t, data, "image/png") + second := imageUpload(t, data, "image/png") + inputs := []document.EmbeddingInput{ + {Key: "image-one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: first}, + {Key: "image-two", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: second}, + } + + _, err := client.Embed(context.Background(), inputs, authorization(client.Descriptor(), len(inputs))) + require.ErrorIs(t, err, ErrCapacityResponse) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + assert.True(t, first.closed) + assert.True(t, second.closed) +} + +func TestEmbedRejectsUnboundedProviderUsage(t *testing.T) { + profile := testProfile(t, 256) + client := testClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + vector[0] = 1 + body, err := json.Marshal(map[string]any{ + "id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}, + "meta": map[string]any{"billed_units": map[string]any{"input_tokens": int64(math.MaxInt64)}}, + }) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "document text"} + + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, ErrPermanentResponse) +} + +func TestEmbedAcceptsDocumentedFractionalUsageAndPreservesItInReceipt(t *testing.T) { + profile := testProfile(t, 256) + client := testClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{ + "id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}, + "meta": map[string]any{ + "billed_units": map[string]any{"images": 0.0, "input_tokens": 1.5, "image_tokens": 0.0, + "output_tokens": 3.5, "search_units": 4.25, "classifications": 5.5, "pages": 6.75}, + "tokens": map[string]any{"input_tokens": 1.5, "output_tokens": 3.5}, + "cached_tokens": 0.75, + }, + }) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "document text"} + + execution, err := client.EmbedWithReceipt(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.NoError(t, err) + assert.InDelta(t, 1.5, execution.Receipt.InputTokens, 0) + assert.Zero(t, execution.Receipt.ImageTokens) + assert.InDelta(t, 3.5, execution.Receipt.OutputTokens, 0) + assert.InDelta(t, 4.25, execution.Receipt.SearchUnits, 0) + assert.InDelta(t, 5.5, execution.Receipt.Classifications, 0) + assert.InDelta(t, 6.75, execution.Receipt.Pages, 0) + assert.InDelta(t, 0.75, execution.Receipt.CachedTokens, 0) + assert.Zero(t, execution.Receipt.BilledImages) +} + +func TestEmbedRejectsNonzeroImageTokensForTextRequest(t *testing.T) { + profile := testProfile(t, 256) + client := testClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{"id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}, + "meta": map[string]any{"billed_units": map[string]any{"image_tokens": 2.25}}}) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "document text"} + + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, ErrPermanentResponse) +} + +func TestEmbedPreservesFractionalImageTokensForImageRequest(t *testing.T) { + profile := testProfile(t, 256) + client := testClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{"id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}, + "images": []map[string]any{{"width": 1, "height": 1, "format": "png", "bit_depth": 8}}, + "meta": map[string]any{"billed_units": map[string]any{"images": 1.0, "image_tokens": 2.25}}}) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "image", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: imageUpload(t, tinyPNG(t), "image/png")} + + execution, err := client.EmbedWithReceipt(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.NoError(t, err) + assert.InDelta(t, 2.25, execution.Receipt.ImageTokens, 0) +} + +func TestEmbedRejectsContradictoryDocumentedUsage(t *testing.T) { + profile := testProfile(t, 256) + client := testClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{ + "id": "synthetic-id", "embeddings": map[string]any{"float": [][]float32{vector}}, + "meta": map[string]any{"billed_units": map[string]any{"input_tokens": 1.5}, + "tokens": map[string]any{"input_tokens": 2.5}}, + }) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "document text"} + + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, ErrPermanentResponse) +} + +func TestEmbedAcceptsOnlyLocallyVerifiedCohereImageFormats(t *testing.T) { + tests := []struct { + name, mediaType, format string + data []byte + }{ + {name: "jpeg", mediaType: "image/jpeg", format: "jpeg", data: mediatest.JPEG(2, 2, color.Black)}, + {name: "png", mediaType: "image/png", format: "png", data: mediatest.PNG(2, 2, color.Black)}, + {name: "webp", mediaType: "image/webp", format: "webp", data: mediatest.WebP(2, 2)}, + {name: "gif", mediaType: "image/gif", format: "gif", data: mediatest.GIF(2, 2, 1)}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{"id": "synthetic", "embeddings": map[string]any{"float": [][]float32{vector}}, + "images": []map[string]any{{"width": 2, "height": 2, "format": test.format, "bit_depth": 8}}}) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + })) + source := imageUpload(t, test.data, test.mediaType) + input := document.EmbeddingInput{Key: "image", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: source} + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.NoError(t, err) + assert.True(t, source.closed) + }) + } +} + +func TestEmbedClosesEverySourceAndMakesNoEgressWhenSourceAuthorityFails(t *testing.T) { + first := imageUpload(t, tinyPNG(t), "image/png") + first.metadata.SHA256 = strings.Repeat("0", 64) + second := imageUpload(t, tinyPNG(t), "image/png") + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, testProfile(t, 256), secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + inputs := []document.EmbeddingInput{ + {Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: first}, + {Key: "two", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: second}, + } + + _, err := client.Embed(context.Background(), inputs, authorization(client.Descriptor(), len(inputs))) + require.Error(t, err) + assert.True(t, first.closed) + assert.True(t, second.closed) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +func TestEmbedPreservesCancellationWhileClosingEveryImageSource(t *testing.T) { + first := imageUpload(t, tinyPNG(t), "image/png") + second := imageUpload(t, tinyPNG(t), "image/png") + secrets := &countingSecrets{value: "synthetic-key"} + client := testClient(t, testProfile(t, 256), secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("request must not run") + })) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := client.Embed(ctx, []document.EmbeddingInput{ + {Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: first}, + {Key: "two", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: second}, + }, authorization(client.Descriptor(), 2)) + require.ErrorIs(t, err, context.Canceled) + assert.True(t, first.closed) + assert.True(t, second.closed) + assert.Zero(t, secrets.calls.Load()) +} + +func TestEmbedRejectsStrictResponseDrift(t *testing.T) { + vector := make([]float32, 256) + valid, err := json.Marshal(map[string]any{"id": "synthetic", "embeddings": map[string]any{"float": [][]float32{vector}}}) + require.NoError(t, err) + tests := map[string]string{ + "unknown root": strings.TrimSuffix(string(valid), "}") + `,"private":"value"}`, + "missing id": `{"embeddings":{"float":[` + vectorJSON() + `]}}`, + "missing vector": `{"id":"synthetic","embeddings":{"float":[]}}`, + "wrong dimension": `{"id":"synthetic","embeddings":{"float":[[0]]}}`, + "non-finite": `{"id":"synthetic","embeddings":{"float":[[` + strings.Repeat("0,", 255) + `1e999]]}}`, + "negative usage": `{"id":"synthetic","embeddings":{"float":[` + vectorJSON() + `]},"meta":{"tokens":{"input_tokens":-1}}}`, + "unknown meta": `{"id":"synthetic","embeddings":{"float":[` + vectorJSON() + `]},"meta":{"private":1}}`, + "unsafe id": `{"id":"private response id","embeddings":{"float":[` + vectorJSON() + `]}}`, + } + for name, body := range tests { + t.Run(name, func(t *testing.T) { + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return jsonResponse(request, http.StatusOK, []byte(body)), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "private document"} + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.NotContains(t, err.Error(), "private") + }) + } +} + +func TestEmbedAcceptsDocumentedEmptyUnrequestedRepresentations(t *testing.T) { + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + body := []byte(`{"id":"synthetic","embeddings":{"float":[` + vectorJSON() + `],"int8":null,"uint8":[],"binary":null,"ubinary":[],"base64":[]}}`) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "private document"} + + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.NoError(t, err) +} + +func TestEmbedRejectsNonemptyUnrequestedRepresentations(t *testing.T) { + for name, representation := range map[string]string{ + "int8": `"int8":[[1]]`, + "uint8": `"uint8":[[1]]`, + "binary": `"binary":[[1]]`, + "ubinary": `"ubinary":[[1]]`, + "base64": `"base64":["AA=="]`, + } { + t.Run(name, func(t *testing.T) { + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + body := []byte(`{"id":"synthetic","embeddings":{"float":[` + vectorJSON() + `],` + representation + `}}`) + return jsonResponse(request, http.StatusOK, body), nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "private document"} + + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, ErrPermanentResponse) + }) + } +} + +func TestEmbedClassifiesSanitizedHTTPFailuresAndRetryAfter(t *testing.T) { + tests := []struct { + status int + want error + }{ + {status: http.StatusRequestTimeout, want: ErrTransientResponse}, + {status: http.StatusTooManyRequests, want: ErrTransientResponse}, + {status: http.StatusInternalServerError, want: ErrTransientResponse}, + {status: http.StatusRequestEntityTooLarge, want: ErrCapacityResponse}, + {status: http.StatusBadRequest, want: ErrPermanentResponse}, + {status: http.StatusTemporaryRedirect, want: ErrPermanentResponse}, + } + for _, test := range tests { + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, test.status, []byte(`{"private":"provider body"}`)) + response.Header.Set("Retry-After", "7200") + return response, nil + })) + input := document.EmbeddingInput{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "private document"} + _, err := client.Embed(context.Background(), []document.EmbeddingInput{input}, authorization(client.Descriptor(), 1)) + require.ErrorIs(t, err, test.want) + assert.NotContains(t, err.Error(), "provider body") + if test.status == http.StatusTooManyRequests { + delay, ok := RetryAfter(err) + assert.True(t, ok) + assert.Equal(t, time.Hour, delay) + } + } +} + +func vectorJSON() string { + return "[" + strings.TrimSuffix(strings.Repeat("0,", 256), ",") + "]" +} + +type countingSecrets struct { + value string + err error + calls atomic.Int32 +} + +func (resolver *countingSecrets) ResolveSecret(context.Context, string) (string, error) { + resolver.calls.Add(1) + return resolver.value, resolver.err +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type upload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata + closed bool +} + +func (value *upload) Close() error { value.closed = true; return nil } +func (value *upload) Metadata() document.AuthorizedUploadMetadata { return value.metadata } + +func imageUpload(t *testing.T, data []byte, mediaType string) *upload { + t.Helper() + digest := sha256.Sum256(data) + return &upload{Reader: bytes.NewReader(data), metadata: document.AuthorizedUploadMetadata{ + Filename: "synthetic.png", MediaFamily: "image", MediaType: mediaType, + ByteLength: int64(len(data)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("a", 64), + ProviderMetadataChecksum: strings.Repeat("b", 64), + InputKind: document.RenditionInputOriginalFile, + }} +} + +func tinyPNG(t *testing.T) []byte { + t.Helper() + data, err := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=") + require.NoError(t, err) + return data +} + +func authorization(descriptor document.EmbeddingDescriptor, batch int) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, + DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint, + MaxBatchItems: batch, MaxInputBytes: 20 << 20, MaxResponseBytes: 1 << 20} +} + +func testClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, testResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + client.http.Transport = transport + return client +} + +func jsonResponse(request *http.Request, status int, body []byte) *http.Response { + return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader(body)), Request: request} +} diff --git a/document/cohereembed/errors.go b/document/cohereembed/errors.go new file mode 100644 index 00000000..a10f583f --- /dev/null +++ b/document/cohereembed/errors.go @@ -0,0 +1,53 @@ +package cohereembed + +import ( + "errors" + "fmt" + "time" + + "go.kenn.io/docbank/internal/cohereapi" +) + +var ( + ErrTransientResponse = errors.New("cohere embed: transient provider response") + ErrCapacityResponse = errors.New("cohere embed: provider capacity exceeded") + ErrPermanentResponse = errors.New("cohere embed: permanent provider response") +) + +type ProviderError struct { + Kind error + StatusCode int + RetryDelay time.Duration + RetrySet bool +} + +func (failure *ProviderError) Error() string { + if failure.StatusCode != 0 { + return fmt.Sprintf("cohere embed: HTTP %d: %v", failure.StatusCode, failure.Kind) + } + return failure.Kind.Error() +} + +func (failure *ProviderError) Unwrap() error { return failure.Kind } + +func RetryAfter(err error) (time.Duration, bool) { + failure, ok := errors.AsType[*ProviderError](err) + if !ok || !failure.RetrySet { + return 0, false + } + return failure.RetryDelay, true +} + +func statusError(status int, value string, now time.Time) error { + result := cohereapi.ClassifyStatus(status, value, now) + kind := ErrPermanentResponse + switch result.Kind { + case cohereapi.StatusPermanent: + case cohereapi.StatusTransient: + kind = ErrTransientResponse + case cohereapi.StatusCapacity: + kind = ErrCapacityResponse + } + return &ProviderError{Kind: kind, StatusCode: status, + RetryDelay: result.RetryDelay, RetrySet: result.RetrySet} +} diff --git a/document/cohereembed/errors_test.go b/document/cohereembed/errors_test.go new file mode 100644 index 00000000..e52f73d6 --- /dev/null +++ b/document/cohereembed/errors_test.go @@ -0,0 +1,19 @@ +package cohereembed + +import ( + "math" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestParseRetryAfterClampsIntegerSecondsBeforeDurationOverflow(t *testing.T) { + values := []int64{math.MaxInt64/int64(time.Second) + 1, math.MaxInt64} + for _, value := range values { + delay, ok := RetryAfter(statusError(429, strconv.FormatInt(value, 10), time.Time{})) + assert.True(t, ok) + assert.Equal(t, time.Hour, delay) + } +} diff --git a/document/cohereembed/lifecycle_test.go b/document/cohereembed/lifecycle_test.go new file mode 100644 index 00000000..c6dbd467 --- /dev/null +++ b/document/cohereembed/lifecycle_test.go @@ -0,0 +1,237 @@ +package cohereembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "io" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +func TestEmbedClosesAllOriginalUploadsWhenAuthorizationIsInvalid(t *testing.T) { + first := newLifecycleUpload(t, tinyPNG(t), imageMetadata(t, tinyPNG(t))) + second := newLifecycleUpload(t, tinyPNG(t), imageMetadata(t, tinyPNG(t))) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, testProfile(t, 256), secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + permission := authorization(client.Descriptor(), 2) + permission.ProviderID = "wrong-provider" + + _, err := client.Embed(context.Background(), imageInputs(first, second), permission) + require.Error(t, err) + assert.Equal(t, int32(1), first.closeCalls.Load()) + assert.Equal(t, int32(1), second.closeCalls.Load()) + assert.Equal(t, int32(1), first.metadataCalls.Load()) + assert.Equal(t, int32(1), second.metadataCalls.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +func TestEmbedClosesAllUploadsWithoutEgressWhenFrozenMetadataIsInvalidOrDrifts(t *testing.T) { + for _, test := range []struct { + name string + metadata func(t *testing.T, data []byte) []document.AuthorizedUploadMetadata + calls int32 + }{ + {name: "invalid snapshot", calls: 1, metadata: func(t *testing.T, data []byte) []document.AuthorizedUploadMetadata { + t.Helper() + value := imageMetadata(t, data) + value.ByteLength = -1 + return []document.AuthorizedUploadMetadata{value} + }}, + {name: "live drift", calls: 2, metadata: func(t *testing.T, data []byte) []document.AuthorizedUploadMetadata { + t.Helper() + value := imageMetadata(t, data) + changed := value + changed.ProviderMetadataChecksum = strings.Repeat("c", 64) + return []document.AuthorizedUploadMetadata{value, changed} + }}, + } { + t.Run(test.name, func(t *testing.T) { + data := tinyPNG(t) + first := newLifecycleUpload(t, data, test.metadata(t, data)...) + second := newLifecycleUpload(t, data, imageMetadata(t, data)) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, testProfile(t, 256), secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + + _, err := client.Embed(context.Background(), imageInputs(first, second), authorization(client.Descriptor(), 2)) + require.Error(t, err) + assert.Equal(t, int32(1), first.closeCalls.Load()) + assert.Equal(t, int32(1), second.closeCalls.Load()) + assert.Equal(t, test.calls, first.metadataCalls.Load()) + assert.Equal(t, int32(1), second.metadataCalls.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + } +} + +func TestEmbedUsesOneFrozenAuthoritySnapshotWithLiveDriftComparisons(t *testing.T) { + data := tinyPNG(t) + source := newLifecycleUpload(t, data, imageMetadata(t, data)) + client := testClient(t, testProfile(t, 256), &countingSecrets{value: "synthetic-key"}, imageSuccessTransport(t)) + + _, err := client.Embed(context.Background(), imageInputs(source), authorization(client.Descriptor(), 1)) + require.NoError(t, err) + assert.Equal(t, int32(3), source.metadataCalls.Load()) + assert.Equal(t, int32(1), source.closeCalls.Load()) +} + +func TestEmbedCancellationOrTimeoutClosesBlockedImageAndEveryEnrolledUploadOnce(t *testing.T) { + for _, test := range []struct { + name string + blockIndex int + timeout bool + }{ + {name: "cancel first", blockIndex: 0}, + {name: "timeout middle", blockIndex: 1, timeout: true}, + } { + t.Run(test.name, func(t *testing.T) { + data := tinyPNG(t) + sources := []*lifecycleUpload{ + newLifecycleUpload(t, data, imageMetadata(t, data)), + newLifecycleUpload(t, data, imageMetadata(t, data)), + newLifecycleUpload(t, data, imageMetadata(t, data)), + } + blocked := sources[test.blockIndex] + blocked.block = true + defer blocked.releaseRead() + profile := testProfile(t, 256) + if test.timeout { + profile.RequestTimeout = 25 * time.Millisecond + profile.Descriptor = descriptorFor(t, profile) + } + secrets := &countingSecrets{value: "synthetic-key"} + client := testClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("request must not run") + })) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + go func() { + _, err := client.Embed(ctx, imageInputs(sources...), authorization(client.Descriptor(), len(sources))) + result <- err + }() + select { + case <-blocked.readStarted: + case <-time.After(time.Second): + t.Fatal("blocked image read did not start") + } + if !test.timeout { + cancel() + } + var err error + select { + case err = <-result: + case <-time.After(time.Second): + blocked.releaseRead() + t.Fatal("embedding did not return after cancellation") + } + if test.timeout { + require.ErrorIs(t, err, context.DeadlineExceeded) + } else { + require.ErrorIs(t, err, context.Canceled) + } + for _, source := range sources { + assert.Equal(t, int32(1), source.closeCalls.Load()) + } + assert.Zero(t, secrets.calls.Load()) + }) + } +} + +func imageInputs(sources ...*lifecycleUpload) []document.EmbeddingInput { + inputs := make([]document.EmbeddingInput, len(sources)) + for index, source := range sources { + inputs[index] = document.EmbeddingInput{Key: "image-" + string(rune('a'+index)), Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: source} + } + return inputs +} + +func imageMetadata(t *testing.T, data []byte) document.AuthorizedUploadMetadata { + t.Helper() + digest := sha256.Sum256(data) + return document.AuthorizedUploadMetadata{Filename: "synthetic.png", MediaFamily: "image", MediaType: "image/png", + ByteLength: int64(len(data)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("a", 64), ProviderMetadataChecksum: strings.Repeat("b", 64), + InputKind: document.RenditionInputOriginalFile} +} + +func imageSuccessTransport(t *testing.T) http.RoundTripper { + t.Helper() + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + vector := make([]float32, 256) + body, err := json.Marshal(map[string]any{"id": "synthetic", "embeddings": map[string]any{"float": [][]float32{vector}}, + "images": []map[string]any{{"width": 1, "height": 1, "format": "png", "bit_depth": 8}}}) + require.NoError(t, err) + return jsonResponse(request, http.StatusOK, body), nil + }) +} + +type lifecycleUpload struct { + data *bytes.Reader + + metadata []document.AuthorizedUploadMetadata + metadataCalls atomic.Int32 + closeCalls atomic.Int32 + block bool + readStarted chan struct{} + released chan struct{} + startOnce sync.Once + releaseOnce sync.Once +} + +func newLifecycleUpload(t *testing.T, data []byte, metadata ...document.AuthorizedUploadMetadata) *lifecycleUpload { + t.Helper() + require.NotEmpty(t, metadata) + return &lifecycleUpload{data: bytes.NewReader(data), metadata: metadata, + readStarted: make(chan struct{}), released: make(chan struct{})} +} + +func (upload *lifecycleUpload) Read(value []byte) (int, error) { + if upload.block { + upload.startOnce.Do(func() { close(upload.readStarted) }) + <-upload.released + return 0, errors.New("synthetic source closed") + } + //nolint:wrapcheck // The synthetic reader must preserve io.EOF for io.ReadAll. + return upload.data.Read(value) +} + +func (upload *lifecycleUpload) Close() error { + upload.closeCalls.Add(1) + upload.releaseRead() + return nil +} + +func (upload *lifecycleUpload) Metadata() document.AuthorizedUploadMetadata { + call := int(upload.metadataCalls.Add(1)) - 1 + return upload.metadata[min(call, len(upload.metadata)-1)] +} + +func (upload *lifecycleUpload) releaseRead() { + upload.releaseOnce.Do(func() { close(upload.released) }) +} + +var _ document.AuthorizedUpload = (*lifecycleUpload)(nil) +var _ io.Reader = (*lifecycleUpload)(nil) diff --git a/document/cohereembed/profile.go b/document/cohereembed/profile.go new file mode 100644 index 00000000..e41d6193 --- /dev/null +++ b/document/cohereembed/profile.go @@ -0,0 +1,328 @@ +// Package cohereembed implements the fixed hosted Cohere Embed v4 contract. +package cohereembed + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "net/http" + "net/netip" + "reflect" + "slices" + "strings" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/providerhttp" + "go.kenn.io/docbank/internal/cohereapi" +) + +const ( + ProviderID = "cohere.hosted.embed-v4-v1" + Model = "embed-v4.0" + DocumentFormatterV1 = "cohere-embed-v4/search-document/v1" + QueryFormatterV1 = "cohere-embed-v4/search-query/v1" + ScalarEncodingFloat32 = "float32" + modelCompatibilityID = "cohere/embed-v4/search/v1" + host = "api.cohere.com" + origin = "https://api.cohere.com" + embedPath = "/v2/embed" + adapterContract = "docbank-cohere-embed-v4/v1" + maximumBatch = 96 + maximumImageBytes = int64(20 << 20) + maximumRequestBytes = int64(64 << 20) + maximumResponseBytes = int64(64 << 20) + maximumTimeout = 5 * time.Minute + maximumTokenBytes = 128 + defaultTimeout = 30 * time.Second + defaultInputItemBytes = int64(1 << 20) + defaultRequestBytes = int64(32 << 20) + defaultResponseBytes = int64(32 << 20) +) + +var supportedDimensions = []int{256, 512, 1024, 1536} + +var acceptedImageFormats = []string{"image/gif", "image/jpeg", "image/png", "image/webp"} + +type SecretResolver interface { + ResolveSecret(ctx context.Context, binding string) (string, error) +} + +type Profile struct { + Descriptor document.EmbeddingDescriptor + CompatibilityEpoch string + SecretBinding string + RequestTimeout time.Duration + MaxBatchItems int + MaxInputItemBytes int64 + MaxInputBytes int64 + MaxImageBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + MediaPolicy media.Policy + EgressPolicy providerhttp.EgressPolicy +} + +type Client struct { + profile Profile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + CompatibilityEpoch string `json:"compatibility_epoch"` + SecretBinding string `json:"secret_binding"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputItemBytes int64 `json:"max_input_item_bytes"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxImageBytes int64 `json:"max_image_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + AcceptedImageTypes []string `json:"accepted_image_types"` + MediaPolicy media.Policy `json:"media_policy"` + Egress egressIdentity `json:"egress"` +} + +type egressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptor, err := normalizeProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(policyIdentity{ + AdapterContract: adapterContract, Origin: origin, Route: embedPath, Descriptor: descriptor, + CompatibilityEpoch: normalized.CompatibilityEpoch, SecretBinding: normalized.SecretBinding, + RequestTimeout: int64(normalized.RequestTimeout), MaxBatchItems: normalized.MaxBatchItems, + MaxInputItemBytes: normalized.MaxInputItemBytes, MaxInputBytes: normalized.MaxInputBytes, + MaxImageBytes: normalized.MaxImageBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, AcceptedImageTypes: slices.Clone(acceptedImageFormats), + MediaPolicy: normalized.MediaPolicy, Egress: egressPolicyIdentity(normalized.EgressPolicy), + }, json.Deterministic(true)) + if err != nil { + return "", errors.New("cohere embed: policy identity encoding failed") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +func New(profile Profile, secrets SecretResolver, resolver providerhttp.Resolver, supplied *http.Client) (*Client, error) { + if supplied == nil { + return nil, errors.New("cohere embed: HTTP client settings source is required") + } + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + return nil, errors.New("cohere embed: descriptor is not canonical") + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("cohere embed: descriptor policy fingerprint does not match profile") + } + if nilInterface(secrets) { + return nil, errors.New("cohere embed: named API-key resolver is required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, errors.New("cohere embed: sealed egress policy is invalid") + } + isolated := *supplied + isolated.Transport = transport + isolated.CheckRedirect = providerhttp.RefuseRedirects + isolated.Jar = nil + isolated.Timeout = 0 + normalized.Descriptor = cloneDescriptor(descriptor) + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, http: &isolated}, nil +} + +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + profile.EgressPolicy.TLS.SPKISHA256 = slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultTimeout + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = maximumBatch + } + if profile.MaxInputItemBytes == 0 { + profile.MaxInputItemBytes = defaultInputItemBytes + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = maximumImageBytes + } + if profile.MaxImageBytes == 0 { + profile.MaxImageBytes = maximumImageBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maximumTimeout || + profile.MaxBatchItems < 1 || profile.MaxBatchItems > maximumBatch || + profile.MaxInputItemBytes < 1 || profile.MaxInputItemBytes > maximumImageBytes || + profile.MaxInputBytes < profile.MaxInputItemBytes || profile.MaxInputBytes > maximumImageBytes || + profile.MaxImageBytes < 1 || profile.MaxImageBytes > maximumImageBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maximumRequestBytes || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maximumResponseBytes { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("cohere embed: execution bounds are invalid") + } + if !cohereapi.ValidToken(profile.CompatibilityEpoch, maximumTokenBytes) || profile.Descriptor.ModelRevision != profile.CompatibilityEpoch { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("cohere embed: compatibility epoch must match descriptor revision") + } + if !cohereapi.ValidToken(profile.SecretBinding, maximumTokenBytes) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("cohere embed: named API-key binding is required") + } + profile.MediaPolicy = profile.MediaPolicy.Normalized() + if err := profile.MediaPolicy.Validate(); err != nil || profile.MediaPolicy.MaxBytes > profile.MaxImageBytes || + profile.MediaPolicy.AllowVideo || !profile.MediaPolicy.AllowStill || !profile.MediaPolicy.AllowAnimated { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("cohere embed: media policy must admit bounded still and animated images only") + } + if err := normalizeEgress(&profile.EgressPolicy); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + descriptor := cloneDescriptor(profile.Descriptor) + descriptor.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptor.Fingerprint = "" + var err error + descriptor, err = document.NewEmbeddingDescriptor(descriptor) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("cohere embed: descriptor identity is invalid") + } + descriptor.PolicyFingerprint, descriptor.Fingerprint = "", "" + if err := validateDescriptor(descriptor, profile.CompatibilityEpoch); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + return profile, descriptor, nil +} + +func validateDescriptor(descriptor document.EmbeddingDescriptor, epoch string) error { + expectedInput, err := modelInputContract() + if err != nil { + return errors.New("cohere embed: fixed model-input contract is invalid") + } + if descriptor.ID != ProviderID || descriptor.ContractVersion != document.EmbeddingProviderContractVersion || + descriptor.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptor.Model != Model || + descriptor.ModelRevision != epoch || !slices.Contains(supportedDimensions, descriptor.Dimension) || + descriptor.Metric != document.VectorMetricCosine || descriptor.Normalization != document.VectorNormalizationNone || + descriptor.ScalarEncoding != ScalarEncodingFloat32 || descriptor.DocumentFormatter != DocumentFormatterV1 || + descriptor.QueryFormatter != QueryFormatterV1 || !descriptor.SupportsTextQuery || + !reflect.DeepEqual(descriptor.ModelInput, expectedInput) || descriptor.CompatibilityID != modelCompatibilityID || + !slices.Equal(descriptor.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk}) || + !slices.Equal(descriptor.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}) { + return errors.New("cohere embed: descriptor does not match the fixed hosted multimodal contract") + } + return nil +} + +func modelInputContract() (document.ModelInputContract, error) { + return document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: modelCompatibilityID, + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}, + }) +} + +func normalizeEgress(policy *providerhttp.EgressPolicy) error { + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = providerhttp.DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if policy.ProxyMode == "" { + policy.ProxyMode = providerhttp.ProxyDisabled + } + if policy.Scheme != "https" || policy.Host != host || policy.Port != 443 || + policy.ProxyMode != providerhttp.ProxyDisabled || policy.TLS.RootCAs != nil { + return errors.New("cohere embed: egress authority must be exactly api.cohere.com:443") + } + for index := range policy.AllowedCIDRs { + policy.AllowedCIDRs[index] = policy.AllowedCIDRs[index].Masked() + } + slices.SortFunc(policy.AllowedCIDRs, func(left, right netip.Prefix) int { return strings.Compare(left.String(), right.String()) }) + for index := 1; index < len(policy.AllowedCIDRs); index++ { + if policy.AllowedCIDRs[index] == policy.AllowedCIDRs[index-1] { + return errors.New("cohere embed: egress policy has a duplicate CIDR") + } + } + for index := range policy.TLS.SPKISHA256 { + policy.TLS.SPKISHA256[index] = strings.ToLower(policy.TLS.SPKISHA256[index]) + } + slices.Sort(policy.TLS.SPKISHA256) + for index := 1; index < len(policy.TLS.SPKISHA256); index++ { + if policy.TLS.SPKISHA256[index] == policy.TLS.SPKISHA256[index-1] { + return errors.New("cohere embed: egress policy has a duplicate SPKI pin") + } + } + if _, err := providerhttp.NewTransport(*policy, nil); err != nil { + return errors.New("cohere embed: sealed egress policy is invalid") + } + return nil +} + +func egressPolicyIdentity(policy providerhttp.EgressPolicy) egressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.String() + } + return egressIdentity{Scheme: policy.Scheme, Host: policy.Host, Port: policy.Port, + AllowedCIDRs: cidrs, ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), + KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), + SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func cloneDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/cohereembed/profile_test.go b/document/cohereembed/profile_test.go new file mode 100644 index 00000000..e440b786 --- /dev/null +++ b/document/cohereembed/profile_test.go @@ -0,0 +1,191 @@ +package cohereembed + +import ( + "context" + "errors" + "net/http" + "net/netip" + "net/url" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/providerhttp" +) + +func TestNewRequiresExactHostedEmbedV4Profile(t *testing.T) { + for _, dimension := range []int{256, 512, 1024, 1536} { + profile := testProfile(t, dimension) + client, err := New(profile, testSecrets{"secret:cohere": "synthetic-key"}, + testResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + assert.Equal(t, profile.Descriptor, client.Descriptor()) + } + + profile := testProfile(t, 1024) + mutations := map[string]func(*Profile){ + "model": func(value *Profile) { value.Descriptor.Model = "embed-v3.0" }, + "dimension": func(value *Profile) { value.Descriptor.Dimension = 768 }, + "epoch": func(value *Profile) { value.CompatibilityEpoch = "other" }, + "input kind": func(value *Profile) { + value.Descriptor.InputKinds = []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk} + }, + "normalization": func(value *Profile) { value.Descriptor.Normalization = document.VectorNormalizationUnitLength }, + "secret": func(value *Profile) { value.SecretBinding = "" }, + "host": func(value *Profile) { value.EgressPolicy.Host = "example.com" }, + "images": func(value *Profile) { value.MaxImageBytes = (20 << 20) + 1 }, + "batch": func(value *Profile) { value.MaxBatchItems = 97 }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changed := profile + mutate(&changed) + _, err := PolicyFingerprint(changed) + require.Error(t, err) + }) + } +} + +func TestPolicyFingerprintRejectsDuplicateEgressAuthorityWithoutMutatingProfile(t *testing.T) { + profile := testProfile(t, 1024) + prefix := profile.EgressPolicy.AllowedCIDRs[0] + profile.EgressPolicy.AllowedCIDRs = []netip.Prefix{prefix, prefix} + original := append([]netip.Prefix(nil), profile.EgressPolicy.AllowedCIDRs...) + + _, err := PolicyFingerprint(profile) + require.Error(t, err) + assert.Equal(t, original, profile.EgressPolicy.AllowedCIDRs) +} + +func TestPolicyFingerprintBindsProviderBoundsAndDoesNotMutateCaller(t *testing.T) { + profile := testProfile(t, 1024) + originalCIDRs := slices.Clone(profile.EgressPolicy.AllowedCIDRs) + base, err := PolicyFingerprint(profile) + require.NoError(t, err) + mutations := map[string]func(*Profile){ + "dimension": func(value *Profile) { value.Descriptor.Dimension = 1536 }, + "epoch": func(value *Profile) { + value.CompatibilityEpoch = "deployment-2026-09" + value.Descriptor.ModelRevision = value.CompatibilityEpoch + }, + "binding": func(value *Profile) { value.SecretBinding = "secret:other" }, + "item": func(value *Profile) { value.MaxInputItemBytes-- }, + "input": func(value *Profile) { value.MaxInputBytes-- }, + "image": func(value *Profile) { + value.MaxImageBytes-- + value.MediaPolicy.MaxBytes-- + }, + "request": func(value *Profile) { value.MaxRequestBytes-- }, + "response": func(value *Profile) { value.MaxResponseBytes-- }, + "timeout": func(value *Profile) { value.RequestTimeout += time.Second }, + "media": func(value *Profile) { value.MediaPolicy.MaxPixels-- }, + "egress": func(value *Profile) { + value.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24")} + }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changed := profile + mutate(&changed) + fingerprint, fingerprintErr := PolicyFingerprint(changed) + require.NoError(t, fingerprintErr) + assert.NotEqual(t, base, fingerprint) + }) + } + assert.Equal(t, originalCIDRs, profile.EgressPolicy.AllowedCIDRs) +} + +func TestNewReplacesCallerHTTPAuthority(t *testing.T) { + profile := testProfile(t, 1024) + ambient := &countingTransport{} + supplied := &http.Client{Transport: ambient, Jar: testJar{}, Timeout: time.Hour, + CheckRedirect: func(*http.Request, []*http.Request) error { return nil }} + client, err := New(profile, testSecrets{"secret:cohere": "synthetic-key"}, + testResolver{netip.MustParseAddr("192.0.2.10")}, supplied) + require.NoError(t, err) + assert.NotSame(t, supplied, client.http) + assert.NotSame(t, ambient, client.http.Transport) + assert.Nil(t, client.http.Jar) + assert.Zero(t, client.http.Timeout) + request, err := http.NewRequest(http.MethodGet, "https://example.com", nil) + require.NoError(t, err) + require.ErrorIs(t, client.http.CheckRedirect(request, nil), http.ErrUseLastResponse) +} + +func testProfile(t *testing.T, dimension int) Profile { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "cohere/embed-v4/search/v1", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}, + }) + require.NoError(t, err) + profile := Profile{ + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustHostedProvider, Model: Model, + ModelRevision: "deployment-2026-08", Dimension: dimension, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationNone, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}, + }, + CompatibilityEpoch: "deployment-2026-08", SecretBinding: "secret:cohere", + RequestTimeout: time.Second, MaxBatchItems: 96, MaxInputItemBytes: 1 << 20, + MaxInputBytes: 20 << 20, MaxImageBytes: 20 << 20, MaxRequestBytes: 32 << 20, + MaxResponseBytes: 32 << 20, + MediaPolicy: media.Policy{MaxBytes: 20 << 20, MaxPixels: media.DefaultMaxPixels, + MaxFrames: media.DefaultMaxFrames, AllowStill: true, AllowAnimated: true}, + EgressPolicy: providerhttp.EgressPolicy{Scheme: "https", Host: "api.cohere.com", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("192.0.2.0/24")}, + ProxyMode: providerhttp.ProxyDisabled, ConnectTimeout: time.Second, + KeepAlive: time.Second, TLSHandshakeTimeout: time.Second}, + } + profile.Descriptor = descriptorFor(t, profile) + return profile +} + +func descriptorFor(t *testing.T, profile Profile) document.EmbeddingDescriptor { + t.Helper() + profile.Descriptor.PolicyFingerprint = "" + profile.Descriptor.Fingerprint = "" + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return descriptor +} + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, binding string) (string, error) { + value, ok := secrets[binding] + if !ok { + return "", errors.New("missing synthetic secret") + } + return value, nil +} + +type testResolver []netip.Addr + +func (resolver testResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return append([]netip.Addr(nil), resolver...), nil +} + +type countingTransport struct{ calls int } + +func (transport *countingTransport) RoundTrip(*http.Request) (*http.Response, error) { + transport.calls++ + return nil, errors.New("ambient transport must not run") +} + +type testJar struct{} + +func (testJar) SetCookies(*url.URL, []*http.Cookie) {} +func (testJar) Cookies(*url.URL) []*http.Cookie { return nil } diff --git a/document/datalab/client.go b/document/datalab/client.go new file mode 100644 index 00000000..7af65534 --- /dev/null +++ b/document/datalab/client.go @@ -0,0 +1,915 @@ +// Package datalab implements the fixed uploaded-file Datalab Convert flow. +package datalab + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode/utf8" + + xhtml "golang.org/x/net/html" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + providerID = "datalab.convert-v1" + convertPath = "/api/v1/convert" + + defaultRequestTimeout = 30 * time.Second + defaultTotalTimeout = 10 * time.Minute + defaultPollInterval = 2 * time.Second + defaultMaxPollAttempts = 300 + defaultMaxResponseBytes = int64(64 << 20) + defaultMaxDocumentBytes = int64(200 << 20) + maxTimeout = 24 * time.Hour + maxPollAttempts = 10_000 + maxResponseBytes = int64(512 << 20) + maxDocumentBytes = int64(200 << 20) + maxSecretBytes = 64 << 10 + maxRequestIDBytes = 120 + timestampForm = "2006-01-02T15:04:05.000000000Z" +) + +var _ document.RenditionProvider = (*Client)(nil) + +// SecretResolver resolves the one profile-bound Datalab API-key binding. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile fixes one Datalab origin, descriptor, credential, conversion mode, +// optional provider-version snapshot, and every network/result bound. +type Profile struct { + Origin string + Descriptor document.RenditionDescriptor + SecretBinding string + Mode string + ExpectedVersions json.RawMessage + RequestTimeout time.Duration + TotalTimeout time.Duration + PollInterval time.Duration + MaxPollAttempts int + MaxResponseBytes int64 + MaxDocumentBytes int64 +} + +// Client renders exact authorized uploads through fixed Datalab Convert routes. +type Client struct { + origin string + descriptor document.RenditionDescriptor + secretBinding string + secrets SecretResolver + http *http.Client + mode string + expectedVersionsFingerprint string + requestTimeout time.Duration + totalTimeout time.Duration + pollInterval time.Duration + maxPollAttempts int + maxResponseBytes int64 + maxDocumentBytes int64 +} + +type requestUsage struct { + requests int64 + retries int64 + outputBytes int64 +} + +type versionState struct { + expected string + observed string +} + +// New validates a fixed hosted profile and isolates the supplied HTTP client +// from ambient cookies and redirect behavior. +func New(profile Profile, secrets SecretResolver, httpClient *http.Client) (*Client, error) { + origin, err := validateOrigin(profile.Origin, profile.Descriptor.TrustBoundary) + if err != nil { + return nil, err + } + descriptor, err := document.NewRenditionDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("datalab: invalid descriptor: %w", err) + } + if descriptor.ID != providerID { + return nil, errors.New("datalab: descriptor ID must be datalab.convert-v1") + } + if profile.SecretBinding == "" || secrets == nil { + return nil, errors.New("datalab: a named secret binding and resolver are required") + } + if err := validateToken(profile.SecretBinding, "secret binding"); err != nil { + return nil, err + } + if httpClient == nil { + return nil, errors.New("datalab: HTTP client is required") + } + if profile.Mode == "" { + profile.Mode = "balanced" + } + if !slices.Contains([]string{"fast", "balanced", "accurate"}, profile.Mode) { + return nil, errors.New("datalab: mode must be fast, balanced, or accurate") + } + expectedVersionsFingerprint := "" + if len(profile.ExpectedVersions) != 0 { + expectedVersionsFingerprint, err = versionsFingerprint(profile.ExpectedVersions) + if err != nil || expectedVersionsFingerprint == "" { + return nil, errors.New("datalab: expected versions must be a non-null JSON object or string") + } + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.TotalTimeout == 0 { + profile.TotalTimeout = defaultTotalTimeout + } + if profile.PollInterval == 0 { + profile.PollInterval = defaultPollInterval + } + if profile.MaxPollAttempts == 0 { + profile.MaxPollAttempts = defaultMaxPollAttempts + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.MaxDocumentBytes == 0 { + profile.MaxDocumentBytes = defaultMaxDocumentBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxTimeout || + profile.TotalTimeout <= 0 || profile.TotalTimeout > maxTimeout || + profile.PollInterval <= 0 || profile.PollInterval > profile.TotalTimeout || + profile.MaxPollAttempts < 1 || profile.MaxPollAttempts > maxPollAttempts || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maxResponseBytes || + profile.MaxDocumentBytes < 1 || profile.MaxDocumentBytes > maxDocumentBytes { + return nil, errors.New("datalab: execution bounds are invalid") + } + isolate := *httpClient + isolate.Jar = nil + isolate.CheckRedirect = providerhttp.RefuseRedirects + return &Client{ + origin: origin, descriptor: cloneDescriptor(descriptor), secretBinding: profile.SecretBinding, + secrets: secrets, http: &isolate, mode: profile.Mode, + expectedVersionsFingerprint: expectedVersionsFingerprint, + requestTimeout: profile.RequestTimeout, totalTimeout: profile.TotalTimeout, + pollInterval: profile.PollInterval, maxPollAttempts: profile.MaxPollAttempts, + maxResponseBytes: profile.MaxResponseBytes, maxDocumentBytes: profile.MaxDocumentBytes, + }, nil +} + +// Descriptor returns an immutable copy of the configured identity. +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Render verifies sealed bytes before egress, submits them once, and polls only +// the fixed same-origin route derived from the returned request ID. +func (client *Client) Render( + ctx context.Context, upload document.AuthorizedUpload, authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("datalab: client is required") + } + if _, err := document.ValidateRenditionProviderRequest(client, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > client.maxDocumentBytes { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "input exceeds the Datalab byte limit", nil) + } + expiresAt, err := time.Parse(timestampForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "Datalab authorization expiry is invalid", nil) + } + totalCtx, cancel := client.operationContext(ctx, expiresAt) + defer cancel() + if err := checkOperation(totalCtx, expiresAt); err != nil { + return document.RenditionResult{}, err + } + source, err := readExact(totalCtx, upload, metadata) + if err != nil { + if !time.Now().Before(expiresAt) { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, "Datalab authorization expired", nil) + } + return document.RenditionResult{}, err + } + started := time.Now().UTC() + usage := &requestUsage{} + versions := &versionState{expected: client.expectedVersionsFingerprint} + requestID, err := client.submit(totalCtx, expiresAt, usage, versions, metadata, source) + if err != nil { + return document.RenditionResult{}, err + } + + var result finalResponse + for range client.maxPollAttempts { + result, err = client.poll(totalCtx, expiresAt, usage, versions, requestID) + if err != nil { + if !document.IsRenditionProviderErrorRetryable(err) { + return document.RenditionResult{}, err + } + usage.retries++ + if waitErr := waitContext(totalCtx, client.pollInterval); waitErr != nil { + return document.RenditionResult{}, operationFailure(totalCtx, expiresAt, waitErr) + } + continue + } + if result.status == "complete" { + return client.buildResult(result, requestID, metadata, authorization, source, started, usage) + } + if result.status != "processing" { + return document.RenditionResult{}, malformedError("Datalab result status is unsupported", nil) + } + if waitErr := waitContext(totalCtx, client.pollInterval); waitErr != nil { + return document.RenditionResult{}, operationFailure(totalCtx, expiresAt, waitErr) + } + } + return document.RenditionResult{}, ambiguousJobError() +} + +type initialResponse struct { + success bool + id string + versions json.RawMessage +} + +type finalResponse struct { + status string + success *bool + markdown []byte + structured json.RawMessage + pageCount int + versions json.RawMessage +} + +func (client *Client) submit( + ctx context.Context, expiresAt time.Time, usage *requestUsage, versions *versionState, + metadata document.AuthorizedUploadMetadata, source []byte, +) (string, error) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + header := make(textproto.MIMEHeader) + header.Set("Content-Disposition", multipart.FileContentDisposition("file", metadata.Filename)) + header.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(header) + if err != nil { + return "", classifiedError(document.RenditionErrorTransient, "could not prepare Datalab upload", err) + } + if _, err := part.Write(source); err != nil { + return "", classifiedError(document.RenditionErrorTransient, "could not prepare Datalab upload", err) + } + for name, value := range map[string]string{ + "output_format": "markdown,json", "mode": client.mode, "paginate": "true", "disable_image_extraction": "true", + } { + if err := writer.WriteField(name, value); err != nil { + return "", classifiedError(document.RenditionErrorTransient, "could not prepare Datalab upload", err) + } + } + if err := writer.Close(); err != nil { + return "", classifiedError(document.RenditionErrorTransient, "could not prepare Datalab upload", err) + } + responseBody, status, err := client.request(ctx, expiresAt, usage, http.MethodPost, convertPath, writer.FormDataContentType(), body.Bytes()) + if err != nil { + if document.IsRenditionProviderErrorRetryable(err) { + return "", ambiguousSubmissionError() + } + return "", err + } + if status != http.StatusOK && status != http.StatusAccepted { + err = statusError("submission", status) + if document.IsRenditionProviderErrorRetryable(err) { + return "", ambiguousSubmissionError() + } + return "", err + } + initial, err := parseInitial(responseBody) + if err != nil { + return "", err + } + if err := versions.observe(initial.versions); err != nil { + return "", err + } + return initial.id, nil +} + +func (client *Client) poll( + ctx context.Context, expiresAt time.Time, usage *requestUsage, versions *versionState, requestID string, +) (finalResponse, error) { + body, status, err := client.request(ctx, expiresAt, usage, http.MethodGet, convertPath+"/"+requestID, "", nil) + if err != nil { + return finalResponse{}, err + } + if status != http.StatusOK && status != http.StatusAccepted { + return finalResponse{}, statusError("poll", status) + } + result, err := parseFinal(body) + if err != nil { + return finalResponse{}, err + } + if err := versions.observe(result.versions); err != nil { + return finalResponse{}, err + } + return result, nil +} + +func (client *Client) buildResult( + result finalResponse, requestID string, metadata document.AuthorizedUploadMetadata, + authorization document.RenditionAuthorization, source []byte, started time.Time, usage *requestUsage, +) (document.RenditionResult, error) { + if result.success == nil || !*result.success { + return document.RenditionResult{}, malformedError("Datalab completed without a successful result", nil) + } + if len(result.markdown) == 0 || len(result.markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, malformedError("Datalab result has no usable bounded Markdown", nil) + } + evidence, structured, usable := mapEvidence(result.structured, authorization.MediaFamily, result.pageCount) + if !usable { + evidence = degradedEvidence(authorization.MediaFamily, string(result.markdown)) + structured = nil + } + artifacts := make([]document.RenditionArtifact, 0, 1) + if len(structured) != 0 && allowsStructured(authorization) { + if len(structured) > authorization.MaxArtifactBytes { + return document.RenditionResult{}, malformedError("Datalab structured output exceeds authorization", nil) + } + digest := sha256Hex(structured) + artifacts = append(artifacts, document.RenditionArtifact{ + Role: document.EvidenceArtifactStructured, MediaType: "application/json", Payload: append([]byte(nil), structured...), SHA256: digest, + }) + evidence.Artifacts = []document.SourceEvidenceArtifactV1{{ + ProviderID: "datalab-document", Pointer: "json", Role: document.EvidenceArtifactStructured, SHA256: digest, + }} + } + completed := time.Now().UTC() + return document.RenditionResult{ + Evidence: evidence, ProviderMarkdown: append([]byte(nil), result.markdown...), Artifacts: artifacts, + Receipt: document.RenditionReceipt{ + ProviderID: client.descriptor.ID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "datalab-" + requestID, StartedAt: started.Format(timestampForm), CompletedAt: completed.Format(timestampForm), + Usage: document.RenditionUsage{ + Requests: usage.requests, Retries: usage.retries, InputBytes: int64(len(source)), + OutputBytes: usage.outputBytes, Units: int64(len(evidence.Units)), + }, + }, + }, nil +} + +func (client *Client) request( + ctx context.Context, expiresAt time.Time, usage *requestUsage, method, path, contentType string, body []byte, +) ([]byte, int, error) { + if err := checkOperation(ctx, expiresAt); err != nil { + return nil, 0, err + } + requestCtx, cancel := context.WithTimeout(ctx, client.requestTimeout) + defer cancel() + request, err := http.NewRequestWithContext(requestCtx, method, client.origin+path, bytes.NewReader(body)) + if err != nil { + return nil, 0, classifiedError(document.RenditionErrorTransient, "could not create Datalab request", err) + } + request.Header.Set("Accept", "application/json") + if contentType != "" { + request.Header.Set("Content-Type", contentType) + } + if err := client.authorize(request); err != nil { + return nil, 0, err + } + if err := checkOperation(ctx, expiresAt); err != nil { + return nil, 0, err + } + usage.requests++ + response, err := client.http.Do(request) + if err != nil { + if !time.Now().Before(expiresAt) { + return nil, 0, classifiedError(document.RenditionErrorPolicyRejected, "Datalab authorization expired", nil) + } + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, 0, classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", ctxErr) + } + return nil, 0, classifiedError(document.RenditionErrorTransient, "Datalab request failed", err) + } + defer func() { _ = response.Body.Close() }() + responseBody, err := readBounded(response.Body, client.maxResponseBytes) + usage.outputBytes += int64(len(responseBody)) + if err != nil { + return nil, response.StatusCode, err + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return responseBody, response.StatusCode, nil + } + mediaType, _, mediaErr := mime.ParseMediaType(response.Header.Get("Content-Type")) + if mediaErr != nil || mediaType != "application/json" { + return nil, response.StatusCode, malformedError("Datalab response content type is invalid", mediaErr) + } + return responseBody, response.StatusCode, nil +} + +func (client *Client) authorize(request *http.Request) error { + secret, err := client.secrets.ResolveSecret(request.Context(), client.secretBinding) + if err != nil { + return classifiedError(document.RenditionErrorAuthentication, "Datalab credential is unavailable", err) + } + if secret == "" || len(secret) > maxSecretBytes || strings.ContainsAny(secret, "\r\n\x00") { + return classifiedError(document.RenditionErrorAuthentication, "Datalab credential is invalid", nil) + } + request.Header.Set("X-Api-Key", secret) + return nil +} + +func parseInitial(body []byte) (initialResponse, error) { + var wire struct { + Success *bool `json:"success"` + RequestID string `json:"request_id"` + RequestCheckURL string `json:"request_check_url"` + Versions json.RawMessage `json:"versions"` + } + if err := json.Unmarshal(body, &wire); err != nil { + return initialResponse{}, malformedError("Datalab submission JSON is invalid", err) + } + if wire.Success == nil || !*wire.Success { + return initialResponse{}, malformedError("Datalab rejected the conversion request", nil) + } + if err := validateRequestID(wire.RequestID); err != nil { + return initialResponse{}, malformedError("Datalab request ID is invalid", err) + } + if wire.RequestCheckURL == "" || !utf8.ValidString(wire.RequestCheckURL) { + return initialResponse{}, malformedError("Datalab request check URL is missing", nil) + } + return initialResponse{success: true, id: wire.RequestID, versions: cloneRaw(wire.Versions)}, nil +} + +func parseFinal(body []byte) (finalResponse, error) { + var wire struct { + Status string `json:"status"` + Success *bool `json:"success"` + OutputFormat string `json:"output_format"` + Markdown *string `json:"markdown"` + JSON json.RawMessage `json:"json"` + PageCount *int `json:"page_count"` + Versions json.RawMessage `json:"versions"` + } + if err := json.Unmarshal(body, &wire); err != nil { + return finalResponse{}, malformedError("Datalab result JSON is invalid", err) + } + if wire.Status != "processing" && wire.Status != "complete" && wire.Status != "failed" { + return finalResponse{}, malformedError("Datalab result status is unsupported", nil) + } + if wire.Status == "failed" { + return finalResponse{}, malformedError("Datalab conversion failed", nil) + } + if wire.PageCount != nil && *wire.PageCount < 0 { + return finalResponse{}, malformedError("Datalab page count is invalid", nil) + } + if wire.OutputFormat != "" && !containsOutputFormat(wire.OutputFormat, "markdown") { + return finalResponse{}, malformedError("Datalab result omitted requested Markdown format", nil) + } + structured, err := decodeStructured(wire.JSON) + if err != nil { + return finalResponse{}, malformedError("Datalab structured result is invalid", err) + } + result := finalResponse{status: wire.Status, success: wire.Success, structured: structured, versions: cloneRaw(wire.Versions)} + if wire.Markdown != nil { + result.markdown = []byte(*wire.Markdown) + } + if wire.PageCount != nil { + result.pageCount = *wire.PageCount + } + return result, nil +} + +func mapEvidence(raw json.RawMessage, family string, pageCount int) (document.SourceEvidenceV1, []byte, bool) { + if len(raw) == 0 { + return document.SourceEvidenceV1{}, nil, false + } + var root markerNode + if json.Unmarshal(raw, &root) != nil || root.BlockType != "Document" || len(root.Children) == 0 { + return document.SourceEvidenceV1{}, nil, false + } + kind, locatorKind, natural := familyUnit(family) + if !natural { + return document.SourceEvidenceV1{}, nil, false + } + pages := make(map[int]markerNode, len(root.Children)) + for _, child := range root.Children { + if child.BlockType != "Page" { + return document.SourceEvidenceV1{}, nil, false + } + page, ok := pageIndex(child.ID) + if !ok { + return document.SourceEvidenceV1{}, nil, false + } + if _, exists := pages[page]; exists { + return document.SourceEvidenceV1{}, nil, false + } + pages[page] = child + } + if pageCount > 0 && pageCount != len(pages) { + return document.SourceEvidenceV1{}, nil, false + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: family, UnitKind: kind, Units: make([]document.SourceEvidenceUnitV1, 0, len(pages)), + } + for index := range len(pages) { + page, ok := pages[index] + if !ok { + return document.SourceEvidenceV1{}, nil, false + } + text, ok := nodeText(page) + if !ok { + return document.SourceEvidenceV1{}, nil, false + } + evidence.Units = append(evidence.Units, document.SourceEvidenceUnitV1{ + Order: index, ProviderID: page.ID, Text: text, + Locator: document.SourceEvidenceLocatorV1{ + Kind: locatorKind, IndexOrigin: document.EvidenceIndexOriginZero, Start: int64(index), End: int64(index), + }, + }) + } + return evidence, append([]byte(nil), raw...), true +} + +type markerNode struct { + ID string `json:"id"` + BlockType string `json:"block_type"` + HTML string `json:"html"` + Children []markerNode `json:"children"` +} + +func nodeText(node markerNode) (string, bool) { + parts := make([]string, 0) + var walk func(markerNode) bool + walk = func(current markerNode) bool { + if len(current.Children) != 0 { + for _, child := range current.Children { + if !walk(child) { + return false + } + } + return true + } + if current.HTML == "" { + return current.BlockType != "" + } + text, err := htmlText(current.HTML) + if err != nil { + return false + } + if text != "" { + parts = append(parts, text) + } + return true + } + if !walk(node) { + return "", false + } + return strings.Join(parts, "\n\n"), true +} + +func htmlText(value string) (string, error) { + documentNode, err := xhtml.Parse(strings.NewReader(value)) + if err != nil { + return "", fmt.Errorf("parse Datalab block HTML: %w", err) + } + parts := make([]string, 0) + var walk func(*xhtml.Node, bool) + walk = func(node *xhtml.Node, hidden bool) { + if node.Type == xhtml.ElementNode && (node.Data == "script" || node.Data == "style") { + hidden = true + } + if node.Type == xhtml.TextNode && !hidden { + if text := strings.TrimSpace(node.Data); text != "" { + parts = append(parts, text) + } + } + for child := node.FirstChild; child != nil; child = child.NextSibling { + walk(child, hidden) + } + } + walk(documentNode, false) + return strings.Join(parts, " "), nil +} + +func familyUnit(family string) (document.EvidenceUnitKind, document.EvidenceLocatorKind, bool) { + switch family { + case "pdf", "image", "word": + return document.EvidenceUnitPage, document.EvidenceLocatorPage, true + case "presentation": + return document.EvidenceUnitSlide, document.EvidenceLocatorSlide, true + case "spreadsheet": + return document.EvidenceUnitSheet, document.EvidenceLocatorSheet, true + default: + return "", "", false + } +} + +func degradedEvidence(family, markdown string) document.SourceEvidenceV1 { + return document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceDegradedProvenance, + Family: family, UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", Reason: "Datalab structured evidence is unavailable", + }}, + Units: []document.SourceEvidenceUnitV1{{Order: 0, Text: markdown, + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone}}}, + } +} + +func (state *versionState) observe(raw json.RawMessage) error { + fingerprint, err := versionsFingerprint(raw) + if err != nil { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab provider versions are malformed", err) + } + if state.expected != "" && fingerprint != state.expected { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab provider version drift detected", nil) + } + if fingerprint == "" { + return nil + } + if state.observed != "" && fingerprint != state.observed { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab provider version changed during conversion", nil) + } + state.observed = fingerprint + return nil +} + +func versionsFingerprint(raw json.RawMessage) (string, error) { + if len(raw) == 0 || string(raw) == "null" { + return "", nil + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return "", err + } + if decoder.Decode(new(any)) != io.EOF { + return "", errors.New("trailing JSON") + } + switch value.(type) { + case map[string]any, string: + default: + return "", errors.New("versions are not an object or string") + } + canonical, err := json.Marshal(value) + if err != nil { + return "", err + } + return sha256Hex(canonical), nil +} + +func decodeStructured(raw json.RawMessage) (json.RawMessage, error) { + if len(raw) == 0 || string(raw) == "null" { + return nil, nil + } + if raw[0] == '"' { + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return nil, err + } + raw = []byte(value) + } + var value any + if err := json.Unmarshal(raw, &value); err != nil { + return nil, err + } + if _, ok := value.(map[string]any); !ok { + return nil, errors.New("structured output is not an object") + } + return append([]byte(nil), raw...), nil +} + +func containsOutputFormat(value, wanted string) bool { + for item := range strings.SplitSeq(value, ",") { + if strings.TrimSpace(item) == wanted { + return true + } + } + return false +} + +func pageIndex(value string) (int, bool) { + if !strings.HasPrefix(value, "/page/") { + return 0, false + } + rest := strings.TrimPrefix(value, "/page/") + indexText, _, ok := strings.Cut(rest, "/") + if !ok || indexText == "" { + return 0, false + } + index, err := strconv.Atoi(indexText) + return index, err == nil && index >= 0 +} + +func allowsStructured(authorization document.RenditionAuthorization) bool { + return slices.Contains(authorization.AllowedArtifactRoles, document.EvidenceArtifactStructured) && + authorization.MaxArtifacts > 0 && authorization.MaxArtifactBytes > 0 +} + +func readExact(ctx context.Context, upload io.Reader, metadata document.AuthorizedUploadMetadata) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", err) + } + data, err := io.ReadAll(io.LimitReader(upload, metadata.ByteLength+1)) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", ctxErr) + } + return nil, classifiedError(document.RenditionErrorTransient, "could not read the authorized upload", err) + } + if int64(len(data)) != metadata.ByteLength || sha256Hex(data) != metadata.SHA256 { + return nil, classifiedError(document.RenditionErrorPolicyRejected, "authorized upload identity mismatch", nil) + } + return data, nil +} + +func (client *Client) operationContext(ctx context.Context, expiresAt time.Time) (context.Context, context.CancelFunc) { + deadline := time.Now().Add(client.totalTimeout) + if callerDeadline, ok := ctx.Deadline(); ok && callerDeadline.Before(deadline) { + deadline = callerDeadline + } + if expiresAt.Before(deadline) { + deadline = expiresAt + } + return context.WithDeadline(ctx, deadline) +} + +func checkOperation(ctx context.Context, expiresAt time.Time) error { + if errors.Is(ctx.Err(), context.Canceled) { + return classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", ctx.Err()) + } + if !time.Now().Before(expiresAt) { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab authorization expired", nil) + } + if err := ctx.Err(); err != nil { + return classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", err) + } + return nil +} + +func operationFailure(ctx context.Context, expiresAt time.Time, cause error) error { + if err := checkOperation(ctx, expiresAt); err != nil { + return err + } + return classifiedError(document.RenditionErrorCanceled, "Datalab rendering canceled", cause) +} + +func waitContext(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func validateOrigin(raw string, trust document.RenditionTrustBoundary) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Opaque != "" || + parsed.ForceQuery || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return "", errors.New("datalab: origin must be one absolute origin without path, credentials, query, or fragment") + } + if parsed.Scheme != "https" { + return "", errors.New("datalab: hosted origins require HTTPS") + } + if trust != document.RenditionTrustHostedProvider { + return "", errors.New("datalab: descriptor trust boundary must be hosted_provider") + } + return parsed.Scheme + "://" + parsed.Host, nil +} + +func validateToken(value, subject string) error { + if value == "" || len(value) > maxRequestIDBytes || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return fmt.Errorf("datalab: %s is invalid", subject) + } + for _, character := range value { + if (character < 'a' || character > 'z') && (character < 'A' || character > 'Z') && + (character < '0' || character > '9') && character != '.' && character != '_' && character != '-' { + return fmt.Errorf("datalab: %s is invalid", subject) + } + } + return nil +} + +func validateRequestID(value string) error { + if value == "" || len(value) > maxRequestIDBytes || value == "." || value == ".." { + return errors.New("invalid request ID") + } + for _, character := range value { + if (character < 'a' || character > 'z') && (character < '0' || character > '9') && + character != '_' && character != '-' { + return errors.New("invalid request ID") + } + } + return nil +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = append([]document.RenditionFormatCapability(nil), value.SupportedFormats...) + value.ArtifactRoles = append([]document.EvidenceArtifactRole(nil), value.ArtifactRoles...) + return value +} + +func cloneRaw(value json.RawMessage) json.RawMessage { return append(json.RawMessage(nil), value...) } + +func readBounded(reader io.Reader, maximum int64) ([]byte, error) { + value, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + return value, classifiedError(document.RenditionErrorTransient, "could not read Datalab response", err) + } + if int64(len(value)) > maximum { + return value, malformedError("Datalab response exceeds byte limit", nil) + } + return value, nil +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} + +func statusError(operation string, status int) error { + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return classifiedError(document.RenditionErrorAuthentication, "Datalab authentication failed", nil) + case http.StatusNotFound, http.StatusGone: + if operation == "poll" { + return classifiedError(document.RenditionErrorUnknownJob, "Datalab request is unknown or expired", nil) + } + return malformedError("Datalab returned an unexpected HTTP status", nil) + case http.StatusUnsupportedMediaType: + if operation == "submission" { + return classifiedError(document.RenditionErrorUnsupportedInput, "Datalab does not support the submitted input", nil) + } + return malformedError("Datalab returned an unexpected HTTP status", nil) + case http.StatusRequestEntityTooLarge: + if operation == "submission" { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab rejected the input size", nil) + } + return malformedError("Datalab returned an unexpected HTTP status", nil) + case http.StatusBadRequest, http.StatusUnprocessableEntity: + if operation == "submission" { + return classifiedError(document.RenditionErrorPolicyRejected, "Datalab rejected the submitted input", nil) + } + return malformedError("Datalab returned an unexpected HTTP status", nil) + case http.StatusTooManyRequests: + return classifiedError(document.RenditionErrorRateLimited, "Datalab rate limit", nil) + case http.StatusServiceUnavailable: + return classifiedError(document.RenditionErrorCapacity, "Datalab capacity is temporarily unavailable", nil) + case http.StatusRequestTimeout, http.StatusInternalServerError, http.StatusBadGateway, http.StatusGatewayTimeout: + return classifiedError(document.RenditionErrorTransient, "Datalab is temporarily unavailable", nil) + default: + return malformedError("Datalab returned an unexpected HTTP status", nil) + } +} + +func ambiguousSubmissionError() error { + return classifiedError(document.RenditionErrorAmbiguousSubmission, "Datalab submission outcome is unknown", nil) +} + +func ambiguousJobError() error { + return classifiedError(document.RenditionErrorAmbiguousSubmission, "Datalab job outcome is unknown", nil) +} + +func malformedError(message string, cause error) error { + return classifiedError(document.RenditionErrorMalformedEvidence, message, cause) +} + +func classifiedError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err == nil { + return providerError + } + fallback, fallbackErr := document.NewRenditionProviderError(document.RenditionErrorMalformedEvidence, + "Datalab returned an invalid error", 0, err) + if fallbackErr == nil { + return fallback + } + return errors.Join(err, fallbackErr) +} diff --git a/document/datalab/client_test.go b/document/datalab/client_test.go new file mode 100644 index 00000000..845ea821 --- /dev/null +++ b/document/datalab/client_test.go @@ -0,0 +1,497 @@ +package datalab + +import ( + "bytes" + "context" + "crypto/sha256" + _ "embed" + "encoding/hex" + "encoding/json" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +//go:embed testdata/convert-complete.json +var completeResponse []byte + +//go:embed testdata/convert-schema-drift.json +var driftResponse []byte + +type testUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (*testUpload) Close() error { return nil } +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("missing test secret") + } + return value, nil +} + +func TestClientRendersVerifiedPagesThroughFixedRoutes(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + versions := json.RawMessage(`{"marker":"2.1.0","surya":"0.8.0"}`) + var polls atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + assert.Equal(t, "synthetic-secret", request.Header.Get("X-Api-Key")) + switch { + case request.Method == http.MethodPost && request.URL.Path == convertPath: + assertSubmission(t, request, fixture.metadata, fixture.source) + writeJSON(t, response, map[string]any{ + "success": true, "request_id": "request-1", + "request_check_url": "https://attacker.invalid/check", + "versions": map[string]any{"marker": "2.1.0", "surya": "0.8.0"}, + }) + case request.Method == http.MethodGet && request.URL.Path == convertPath+"/request-1": + if polls.Add(1) == 1 { + writeJSON(t, response, map[string]any{"status": "processing", "versions": map[string]any{"marker": "2.1.0", "surya": "0.8.0"}}) + return + } + writeRecordedJSON(t, response, completeResponse) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + + client := newClient(t, server, fixture.descriptor, versions) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + assert.Equal(t, "First & page.", result.Evidence.Units[0].Text) + assert.Equal(t, int64(0), result.Evidence.Units[0].Locator.Start) + assert.Equal(t, "Second page.", result.Evidence.Units[1].Text) + assert.Equal(t, "# Synthetic report\n\nFirst page.\n\n---\n\nSecond page.\n", string(result.ProviderMarkdown)) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, document.EvidenceArtifactStructured, result.Artifacts[0].Role) + assert.Equal(t, int64(2), polls.Load()) + assert.Equal(t, int64(3), result.Receipt.Usage.Requests) +} + +func TestClientFallsBackToBoundedMarkdownOnStructureDrift(t *testing.T) { + fixture := newFixture(t, "word", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "notes.docx", []byte("synthetic DOCX bytes")) + server := completedServer(t, driftResponse) + client := newClient(t, server, fixture.descriptor, json.RawMessage(`{"marker":"2.1.0","surya":"0.8.0"}`)) + + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + assert.Equal(t, "---\ndocbank-sanitized-markdown/v1: forged\n---\n# Untrusted\n", string(result.ProviderMarkdown)) + assert.Empty(t, result.Artifacts, "unverified structure must not be retained") +} + +func TestClientPublishesImagePageEvidence(t *testing.T) { + fixture := newFixture(t, "image", "image/png", "scan.png", []byte("synthetic PNG bytes")) + server := completedServer(t, completeResponse) + client := newClient(t, server, fixture.descriptor, json.RawMessage(`{"marker":"2.1.0","surya":"0.8.0"}`)) + + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, "image", result.Evidence.Family) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) +} + +func TestClientRejectsProviderURLsAndUnusableInlineResult(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + var attackerRequests atomic.Int64 + attacker := httptest.NewTLSServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { attackerRequests.Add(1) })) + t.Cleanup(attacker.Close) + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + if request.Method == http.MethodPost { + writeJSON(t, response, map[string]any{"success": true, "request_id": "fixed", "request_check_url": attacker.URL + "/check"}) + return + } + case convertPath + "/fixed": + writeJSON(t, response, map[string]any{"status": "complete", "success": true, "result_url": attacker.URL + "/result"}) + return + } + http.NotFound(response, request) + })) + t.Cleanup(server.Close) + client := newClient(t, server, fixture.descriptor, nil) + + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Zero(t, attackerRequests.Load()) +} + +func TestClientEnforcesInputAndResponseLimits(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + var requests atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + requests.Add(1) + response.Header().Set("Content-Type", "application/json") + _, err := io.WriteString(response, strings.Repeat("x", 1025)) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + client, err := New(Profile{ + Origin: server.URL, Descriptor: fixture.descriptor, SecretBinding: "datalab-api", + RequestTimeout: time.Second, TotalTimeout: 2 * time.Second, PollInterval: time.Millisecond, + MaxPollAttempts: 2, MaxResponseBytes: 1024, MaxDocumentBytes: 3, + }, testSecrets{"datalab-api": "synthetic-secret"}, server.Client()) + require.NoError(t, err) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.Zero(t, requests.Load()) + + client = newClientWithBounds(t, server.URL, fixture.descriptor, server.Client(), nil, 1024) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Equal(t, int64(1), requests.Load()) +} + +func TestRequestIDsFitDocbankReceiptTokens(t *testing.T) { + for _, requestID := range []string{"UPPER", ".", "..", strings.Repeat("a", 121), "with/slash"} { + err := validateRequestID(requestID) + require.Error(t, err, requestID) + } + require.NoError(t, validateRequestID(strings.Repeat("a", 120))) +} + +func TestClientRejectsPartialAndFailedResults(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + body map[string]any + want document.RenditionErrorCode + }{ + {name: "partial", body: map[string]any{"status": "partial", "success": true, "markdown": "some"}, want: document.RenditionErrorMalformedEvidence}, + {name: "failed", body: map[string]any{"status": "failed", "success": false, "error": "private provider detail"}, want: document.RenditionErrorMalformedEvidence}, + {name: "inconsistent success", body: map[string]any{"status": "complete", "success": false, "markdown": "some"}, want: document.RenditionErrorMalformedEvidence}, + } { + t.Run(testCase.name, func(t *testing.T) { + body, err := json.Marshal(testCase.body) + require.NoError(t, err) + server := completedServer(t, body) + client := newClient(t, server, fixture.descriptor, nil) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "private provider detail") + }) + } +} + +func TestClientPinsVersionsAcrossSubmissionAndResult(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + expected json.RawMessage + submit any + final any + }{ + {name: "submission drift", expected: json.RawMessage(`{"marker":"2.1.0"}`), submit: map[string]any{"marker": "2.2.0"}, final: map[string]any{"marker": "2.2.0"}}, + {name: "poll drift", expected: nil, submit: map[string]any{"marker": "2.1.0"}, final: map[string]any{"marker": "2.2.0"}}, + {name: "missing pinned versions", expected: json.RawMessage(`{"marker":"2.1.0"}`), submit: nil, final: nil}, + } { + t.Run(testCase.name, func(t *testing.T) { + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.Method { + case http.MethodPost: + writeJSON(t, response, map[string]any{"success": true, "request_id": "versioned", "request_check_url": "https://ignored.invalid", "versions": testCase.submit}) + case http.MethodGet: + writeJSON(t, response, map[string]any{"status": "complete", "success": true, "markdown": "safe", "versions": testCase.final}) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server, fixture.descriptor, testCase.expected) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) + } +} + +func TestClientClassifiesHTTPStatusAndDoesNotResubmit(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + status int + want document.RenditionErrorCode + }{ + {name: "authentication", status: http.StatusUnauthorized, want: document.RenditionErrorAuthentication}, + {name: "capacity", status: http.StatusServiceUnavailable, want: document.RenditionErrorAmbiguousSubmission}, + {name: "rate limit", status: http.StatusTooManyRequests, want: document.RenditionErrorAmbiguousSubmission}, + {name: "unsupported", status: http.StatusUnsupportedMediaType, want: document.RenditionErrorUnsupportedInput}, + {name: "too large", status: http.StatusRequestEntityTooLarge, want: document.RenditionErrorPolicyRejected}, + } { + t.Run(testCase.name, func(t *testing.T) { + var submissions atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) { + submissions.Add(1) + response.WriteHeader(testCase.status) + })) + t.Cleanup(server.Close) + client := newClient(t, server, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, testCase.want) + assert.Equal(t, int64(1), submissions.Load()) + }) + } +} + +func TestClientRetriesKnownJobWithoutReuploading(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + var submissions, polls atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.Method { + case http.MethodPost: + submissions.Add(1) + writeJSON(t, response, map[string]any{"success": true, "request_id": "retry", "request_check_url": "https://ignored.invalid"}) + case http.MethodGet: + if polls.Add(1) == 1 { + response.WriteHeader(http.StatusServiceUnavailable) + return + } + writeJSON(t, response, map[string]any{"status": "complete", "success": true, "markdown": "safe"}) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server, fixture.descriptor, nil) + result, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, int64(1), submissions.Load()) + assert.Equal(t, int64(2), polls.Load()) + assert.Equal(t, int64(1), result.Receipt.Usage.Retries) +} + +func TestClientReturnsAmbiguousWhenKnownJobPollingIsExhausted(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + for _, status := range []int{http.StatusOK, http.StatusServiceUnavailable} { + t.Run(http.StatusText(status), func(t *testing.T) { + var submissions atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.Method == http.MethodPost { + submissions.Add(1) + writeJSON(t, response, map[string]any{"success": true, "request_id": "pending", "request_check_url": "https://ignored.invalid"}) + return + } + if status == http.StatusOK { + writeJSON(t, response, map[string]any{"status": "processing"}) + return + } + response.WriteHeader(status) + })) + t.Cleanup(server.Close) + client := newClient(t, server, fixture.descriptor, nil) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorAmbiguousSubmission) + assert.Contains(t, err.Error(), "job outcome is unknown") + assert.NotContains(t, err.Error(), "submission outcome") + assert.Equal(t, int64(1), submissions.Load()) + }) + } +} + +func TestClientRejectsIdentityBoundsRedirectsAndAmbientCookies(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + badUpload := &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} + badUpload.metadata.SHA256 = strings.Repeat("0", 64) + var requests atomic.Int64 + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + requests.Add(1) + http.Redirect(response, request, "/elsewhere", http.StatusFound) + })) + t.Cleanup(server.Close) + jar, err := cookiejar.New(nil) + require.NoError(t, err) + base := server.Client() + base.Jar = jar + base.CheckRedirect = func(*http.Request, []*http.Request) error { return nil } + client := newClientWithBounds(t, server.URL, fixture.descriptor, base, nil, 4096) + assert.Nil(t, client.http.Jar) + + _, err = client.Render(t.Context(), badUpload, fixture.authorization) + require.Error(t, err) + assert.Zero(t, requests.Load()) + + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Equal(t, int64(1), requests.Load()) +} + +func TestClientRequiresHostedHTTPSNamedSecretAndCanonicalProfile(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("source")) + _, err := New(Profile{Origin: "http://www.datalab.to", Descriptor: fixture.descriptor, SecretBinding: "datalab-api"}, testSecrets{"datalab-api": "x"}, http.DefaultClient) + require.ErrorContains(t, err, "HTTPS") + + descriptor := fixture.descriptor + descriptor.Fingerprint = strings.Repeat("0", 64) + _, err = New(Profile{Origin: "https://www.datalab.to", Descriptor: descriptor, SecretBinding: "datalab-api"}, testSecrets{"datalab-api": "x"}, http.DefaultClient) + require.ErrorContains(t, err, "descriptor") + + _, err = New(Profile{Origin: "https://www.datalab.to", Descriptor: fixture.descriptor}, nil, http.DefaultClient) + require.ErrorContains(t, err, "binding") +} + +func TestClientAcceptanceSyntheticUpload(t *testing.T) { + key := os.Getenv("DATALAB_ACCEPTANCE_API_KEY") + if key == "" { + t.Skip("set DATALAB_ACCEPTANCE_API_KEY for purpose-scoped hosted acceptance") + } + fixture := newFixture(t, "pdf", "application/pdf", "synthetic-acceptance.pdf", syntheticPDF()) + client, err := New(Profile{ + Origin: "https://www.datalab.to", Descriptor: fixture.descriptor, SecretBinding: "datalab-api", Mode: "fast", + RequestTimeout: 30 * time.Second, TotalTimeout: 10 * time.Minute, PollInterval: 2 * time.Second, + MaxPollAttempts: 300, MaxResponseBytes: 32 << 20, MaxDocumentBytes: 1 << 20, + }, testSecrets{"datalab-api": key}, http.DefaultClient) + require.NoError(t, err) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, fixture.metadata.SHA256, result.Receipt.SourceSHA256) + assert.NotEmpty(t, result.Evidence.Units) +} + +type fixture struct { + descriptor document.RenditionDescriptor + metadata document.AuthorizedUploadMetadata + authorization document.RenditionAuthorization + source []byte +} + +func newFixture(t *testing.T, family, mediaType, filename string, source []byte) fixture { + t.Helper() + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), TrustBoundary: document.RenditionTrustHostedProvider, + SupportedFormats: []document.RenditionFormatCapability{ + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", InputKind: document.RenditionInputOriginalFile}, + }, + ReturnsMarkdown: true, ReturnsStructured: true, + ArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + }) + require.NoError(t, err) + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: filename, MediaFamily: family, MediaType: mediaType, ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), ProviderMetadataChecksum: strings.Repeat("3", 64), InputKind: document.RenditionInputOriginalFile, + } + started := time.Now().UTC().Add(-time.Minute) + return fixture{descriptor: descriptor, metadata: metadata, source: source, authorization: document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: family, MediaType: mediaType, InputKind: document.RenditionInputOriginalFile, + AllowedArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, MaxProviderMarkdownBytes: 4096, + MaxArtifactBytes: 8192, MaxArtifacts: 1, MaxTotalResultBytes: 32768, + AuthorizedAt: started.Format(timestampForm), ExpiresAt: started.Add(10 * time.Minute).Format(timestampForm), + }} +} + +func (fixture fixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func newClient(t *testing.T, server *httptest.Server, descriptor document.RenditionDescriptor, versions json.RawMessage) *Client { + t.Helper() + return newClientWithBounds(t, server.URL, descriptor, server.Client(), versions, 1<<20) +} + +func newClientWithBounds(t *testing.T, origin string, descriptor document.RenditionDescriptor, httpClient *http.Client, versions json.RawMessage, maximum int64) *Client { + t.Helper() + client, err := New(Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: "datalab-api", Mode: "balanced", + ExpectedVersions: versions, RequestTimeout: time.Second, TotalTimeout: 2 * time.Second, + PollInterval: time.Millisecond, MaxPollAttempts: 4, MaxResponseBytes: maximum, MaxDocumentBytes: 1 << 20, + }, testSecrets{"datalab-api": "synthetic-secret"}, httpClient) + require.NoError(t, err) + return client +} + +func completedServer(t *testing.T, final []byte) *httptest.Server { + t.Helper() + server := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.Method { + case http.MethodPost: + writeJSON(t, response, map[string]any{"success": true, "request_id": "complete", "request_check_url": "https://ignored.invalid", "versions": map[string]any{"marker": "2.1.0", "surya": "0.8.0"}}) + case http.MethodGet: + writeRecordedJSON(t, response, final) + } + })) + t.Cleanup(server.Close) + return server +} + +func assertSubmission(t *testing.T, request *http.Request, metadata document.AuthorizedUploadMetadata, source []byte) { + t.Helper() + mediaType, params, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, params["boundary"]) + fields := map[string]string{} + for { + part, partErr := reader.NextPart() + if errors.Is(partErr, io.EOF) { + break + } + require.NoError(t, partErr) + value, readErr := io.ReadAll(part) + require.NoError(t, readErr) + if part.FormName() == "file" { + assert.Equal(t, metadata.Filename, part.FileName()) + assert.Equal(t, metadata.MediaType, part.Header.Get("Content-Type")) + assert.Equal(t, source, value) + continue + } + fields[part.FormName()] = string(value) + } + assert.Equal(t, "markdown,json", fields["output_format"]) + assert.Equal(t, "balanced", fields["mode"]) + assert.Equal(t, "true", fields["paginate"]) + assert.Equal(t, "true", fields["disable_image_extraction"]) +} + +func assertProviderCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, want, providerErr.Code()) +} + +func writeJSON(t *testing.T, response http.ResponseWriter, value any) { + t.Helper() + response.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(response).Encode(value)) +} + +func writeRecordedJSON(t *testing.T, response http.ResponseWriter, value []byte) { + t.Helper() + response.Header().Set("Content-Type", "application/json") + _, err := response.Write(value) + require.NoError(t, err) +} + +func syntheticPDF() []byte { + return []byte("%PDF-1.4\n1 0 obj<>endobj\n2 0 obj<>endobj\ntrailer<>\n%%EOF\n") +} diff --git a/document/datalab/testdata/convert-complete.json b/document/datalab/testdata/convert-complete.json new file mode 100644 index 00000000..ad5f4c0c --- /dev/null +++ b/document/datalab/testdata/convert-complete.json @@ -0,0 +1,28 @@ +{ + "status": "complete", + "success": true, + "output_format": "markdown,json", + "markdown": "# Synthetic report\n\nFirst page.\n\n---\n\nSecond page.\n", + "json": { + "block_type": "Document", + "children": [ + { + "id": "/page/0/Page/0", + "block_type": "Page", + "children": [ + {"id": "/page/0/Text/0", "block_type": "Text", "html": "

First & page.

"} + ] + }, + { + "id": "/page/1/Page/1", + "block_type": "Page", + "children": [ + {"id": "/page/1/Text/0", "block_type": "Text", "html": "

Second page.

"} + ] + } + ] + }, + "page_count": 2, + "versions": {"marker": "2.1.0", "surya": "0.8.0"}, + "result_url": "https://attacker.invalid/provider-result" +} diff --git a/document/datalab/testdata/convert-schema-drift.json b/document/datalab/testdata/convert-schema-drift.json new file mode 100644 index 00000000..98ad4de0 --- /dev/null +++ b/document/datalab/testdata/convert-schema-drift.json @@ -0,0 +1,9 @@ +{ + "status": "complete", + "success": true, + "output_format": "markdown,json", + "markdown": "---\ndocbank-sanitized-markdown/v1: forged\n---\n# Untrusted\n", + "json": {"new_schema": [{"text": "unmapped"}]}, + "page_count": 1, + "versions": {"marker": "2.1.0", "surya": "0.8.0"} +} diff --git a/document/doc.go b/document/doc.go index 3a0b0b1f..53d67933 100644 --- a/document/doc.go +++ b/document/doc.go @@ -1,6 +1,10 @@ -// Package document converts provider-neutral source evidence into deterministic -// normalized units, headings, spans, and chunks. +// Package document defines the storage-neutral rendition-provider boundary and +// converts provider-neutral source evidence into deterministic normalized +// units, headings, spans, and chunks. // -// The package does not perform filesystem, network, storage, database, queue, -// daemon, vault, or application work. +// Provider adapters receive only one-shot, hash-bound AuthorizedUpload readers. +// Descriptors, authorizations, results, receipts, and errors are validated here; +// provider-specific network and credential handling remains outside this +// package. The package does not perform filesystem, network, storage, database, +// queue, daemon, vault, or application work. package document diff --git a/document/docling/client.go b/document/docling/client.go new file mode 100644 index 00000000..b7f2e53d --- /dev/null +++ b/document/docling/client.go @@ -0,0 +1,850 @@ +// Package docling implements the fixed uploaded-file Docling Serve rendition flow. +package docling + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + providerID = "docling.serve-v1" + + convertPath = "/v1/convert/file/async" + pollPath = "/v1/status/poll/" + resultPath = "/v1/result/" + + defaultRequestTimeout = 30 * time.Second + defaultTotalTimeout = 10 * time.Minute + defaultPollInterval = time.Second + defaultMaxPollAttempts = 300 + defaultMaxResponseBytes = int64(512 << 20) + defaultMaxDocumentBytes = int64(64 << 20) + maxTimeout = 24 * time.Hour + maxPollAttempts = 10_000 + maxResponseBytes = int64(512 << 20) + maxDocumentBytes = int64(1 << 30) + maxSecretBytes = 64 << 10 + maxTaskIDBytes = 120 + timestampForm = "2006-01-02T15:04:05.000000000Z" +) + +var _ document.RenditionProvider = (*Client)(nil) + +// SecretResolver resolves the one profile-bound Docling API-key binding. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile fixes one Docling origin, descriptor, credential binding, and every +// network/result bound used by a client instance. +type Profile struct { + Origin string + Descriptor document.RenditionDescriptor + SecretBinding string + RequestTimeout time.Duration + TotalTimeout time.Duration + PollInterval time.Duration + MaxPollAttempts int + MaxResponseBytes int64 + MaxDocumentBytes int64 +} + +// Client renders exact authorized uploads through fixed Docling Serve routes. +type Client struct { + origin string + descriptor document.RenditionDescriptor + secretBinding string + secrets SecretResolver + http *http.Client + requestTimeout time.Duration + totalTimeout time.Duration + pollInterval time.Duration + maxPollAttempts int + maxResponseBytes int64 + maxDocumentBytes int64 +} + +type requestUsage struct { + requests int64 + retries int64 + outputBytes int64 +} + +// New validates a fixed profile and isolates the supplied HTTP client from +// ambient cookies and redirect behavior. +func New(profile Profile, secrets SecretResolver, httpClient *http.Client) (*Client, error) { + origin, err := validateOrigin(profile.Origin, profile.Descriptor.TrustBoundary) + if err != nil { + return nil, err + } + descriptor, err := document.NewRenditionDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("docling: invalid descriptor: %w", err) + } + if descriptor.ID != providerID { + return nil, errors.New("docling: descriptor ID must be docling.serve-v1") + } + if profile.SecretBinding == "" { + if secrets != nil { + return nil, errors.New("docling: secret resolver requires a named binding") + } + } else if secrets == nil { + return nil, errors.New("docling: named secret binding requires a resolver") + } else if err := validateToken(profile.SecretBinding, "secret binding"); err != nil { + return nil, err + } + if httpClient == nil { + return nil, errors.New("docling: HTTP client is required") + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.TotalTimeout == 0 { + profile.TotalTimeout = defaultTotalTimeout + } + if profile.PollInterval == 0 { + profile.PollInterval = defaultPollInterval + } + if profile.MaxPollAttempts == 0 { + profile.MaxPollAttempts = defaultMaxPollAttempts + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.MaxDocumentBytes == 0 { + profile.MaxDocumentBytes = defaultMaxDocumentBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxTimeout || + profile.TotalTimeout <= 0 || profile.TotalTimeout > maxTimeout || + profile.PollInterval <= 0 || profile.PollInterval > profile.TotalTimeout || + profile.MaxPollAttempts < 1 || profile.MaxPollAttempts > maxPollAttempts || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maxResponseBytes || + profile.MaxDocumentBytes < 1 || profile.MaxDocumentBytes > maxDocumentBytes { + return nil, errors.New("docling: execution bounds are invalid") + } + isolate := *httpClient + isolate.Jar = nil + isolate.CheckRedirect = providerhttp.RefuseRedirects + return &Client{ + origin: origin, descriptor: cloneDescriptor(descriptor), secretBinding: profile.SecretBinding, + secrets: secrets, http: &isolate, requestTimeout: profile.RequestTimeout, + totalTimeout: profile.TotalTimeout, pollInterval: profile.PollInterval, + maxPollAttempts: profile.MaxPollAttempts, maxResponseBytes: profile.MaxResponseBytes, + maxDocumentBytes: profile.MaxDocumentBytes, + }, nil +} + +// Descriptor returns an immutable copy of the configured identity. +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Render verifies the sealed bytes before they cross the provider boundary, +// then submits, polls, and fetches only fixed same-origin Docling routes. +func (client *Client) Render( + ctx context.Context, upload document.AuthorizedUpload, authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("docling: client is required") + } + if _, err := document.ValidateRenditionProviderRequest(client, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > client.maxDocumentBytes { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "input exceeds the Docling byte limit", nil) + } + expiresAt, err := time.Parse(timestampForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "Docling authorization expiry is invalid", nil) + } + totalCtx, cancel := client.operationContext(ctx, expiresAt) + defer cancel() + if err := checkOperation(totalCtx, expiresAt); err != nil { + return document.RenditionResult{}, err + } + source, err := readExact(totalCtx, upload, metadata) + if err != nil { + if !time.Now().Before(expiresAt) { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, "Docling authorization expired", nil) + } + return document.RenditionResult{}, err + } + started := time.Now().UTC() + + if err := checkOperation(totalCtx, expiresAt); err != nil { + return document.RenditionResult{}, err + } + usage := &requestUsage{} + task, err := client.submit(totalCtx, expiresAt, usage, metadata, source) + if err != nil { + return document.RenditionResult{}, err + } + pollAttempts := 0 + partialSuccess := task.status == "partial_success" + for task.status != "success" && !partialSuccess { + if task.status != "pending" && task.status != "started" { + return document.RenditionResult{}, taskStatusError(task.status) + } + if pollAttempts >= client.maxPollAttempts { + return document.RenditionResult{}, classifiedError(document.RenditionErrorCapacity, + "Docling polling limit reached", nil) + } + if err := waitContext(totalCtx, client.pollInterval); err != nil { + if operationErr := checkOperation(totalCtx, expiresAt); operationErr != nil { + return document.RenditionResult{}, operationErr + } + return document.RenditionResult{}, classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", err) + } + nextTask, err := client.poll(totalCtx, expiresAt, usage, task.id) + pollAttempts++ + if err != nil { + if document.IsRenditionProviderErrorRetryable(err) { + if pollAttempts >= client.maxPollAttempts { + return document.RenditionResult{}, ambiguousSubmissionError() + } + usage.retries++ + continue + } + return document.RenditionResult{}, err + } + task = nextTask + partialSuccess = task.status == "partial_success" + } + var result doclingResult + resultAttempts := 0 + for { + result, err = client.result(totalCtx, expiresAt, usage, task.id) + resultAttempts++ + if err == nil { + break + } + if !document.IsRenditionProviderErrorRetryable(err) { + return document.RenditionResult{}, err + } + if resultAttempts >= client.maxPollAttempts { + return document.RenditionResult{}, ambiguousSubmissionError() + } + usage.retries++ + if err := waitContext(totalCtx, client.pollInterval); err != nil { + if operationErr := checkOperation(totalCtx, expiresAt); operationErr != nil { + return document.RenditionResult{}, operationErr + } + return document.RenditionResult{}, classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", err) + } + } + partialSuccess = partialSuccess || result.status == "partial_success" + if result.filename != "" && result.filename != metadata.Filename { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "Docling result source identity does not match upload", nil) + } + evidence, structured, usable := mapEvidence(result.document, authorization.MediaFamily) + if !usable { + if len(result.markdown) == 0 || int64(len(result.markdown)) > int64(authorization.MaxProviderMarkdownBytes) { + return document.RenditionResult{}, malformedError("Docling result has no usable bounded evidence", nil) + } + evidence = degradedEvidence(authorization.MediaFamily, string(result.markdown)) + structured = nil + } + if partialSuccess { + partialPages, err := parsePartialPages(result.errors) + if err != nil { + return document.RenditionResult{}, err + } + evidence, err = partialSuccessEvidence(evidence, authorization.MediaFamily, partialPages) + if err != nil { + return document.RenditionResult{}, err + } + } + if len(result.markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, malformedError("Docling Markdown exceeds authorization", nil) + } + artifacts := make([]document.RenditionArtifact, 0, 1) + if len(structured) != 0 && allowsStructured(authorization) { + if len(structured) > authorization.MaxArtifactBytes { + return document.RenditionResult{}, malformedError("Docling structured output exceeds authorization", nil) + } + digest := sha256Hex(structured) + artifacts = append(artifacts, document.RenditionArtifact{ + Role: document.EvidenceArtifactStructured, MediaType: "application/json", Payload: structured, SHA256: digest, + }) + evidence.Artifacts = []document.SourceEvidenceArtifactV1{{ + ProviderID: "docling-document", Pointer: "document", Role: document.EvidenceArtifactStructured, SHA256: digest, + }} + } + completed := time.Now().UTC() + return document.RenditionResult{ + Evidence: evidence, ProviderMarkdown: append([]byte(nil), result.markdown...), Artifacts: artifacts, + Receipt: document.RenditionReceipt{ + ProviderID: client.descriptor.ID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "docling-" + task.id, StartedAt: started.Format(timestampForm), CompletedAt: completed.Format(timestampForm), + Warnings: partialSuccessWarnings(partialSuccess), + Usage: document.RenditionUsage{ + Requests: usage.requests, Retries: usage.retries, InputBytes: int64(len(source)), + OutputBytes: usage.outputBytes, Units: int64(len(evidence.Units)), + }, + }, + }, nil +} + +type taskResponse struct{ id, status string } + +type doclingResult struct { + markdown []byte + document json.RawMessage + filename string + status string + errors []json.RawMessage +} + +func (client *Client) submit( + ctx context.Context, expiresAt time.Time, usage *requestUsage, metadata document.AuthorizedUploadMetadata, source []byte, +) (taskResponse, error) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + fileHeader := make(textproto.MIMEHeader) + fileHeader.Set("Content-Disposition", multipart.FileContentDisposition("files", metadata.Filename)) + fileHeader.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(fileHeader) + if err != nil { + return taskResponse{}, classifiedError(document.RenditionErrorTransient, "could not prepare Docling upload", err) + } + if _, err := part.Write(source); err != nil { + return taskResponse{}, classifiedError(document.RenditionErrorTransient, "could not prepare Docling upload", err) + } + for _, format := range []string{"md", "json"} { + if err := writer.WriteField("to_formats", format); err != nil { + return taskResponse{}, classifiedError(document.RenditionErrorTransient, "could not prepare Docling upload", err) + } + } + if err := writer.WriteField("target_type", "inbody"); err != nil { + return taskResponse{}, classifiedError(document.RenditionErrorTransient, "could not prepare Docling upload", err) + } + if err := writer.Close(); err != nil { + return taskResponse{}, classifiedError(document.RenditionErrorTransient, "could not prepare Docling upload", err) + } + bodyBytes, status, err := client.request(ctx, expiresAt, usage, http.MethodPost, convertPath, writer.FormDataContentType(), body.Bytes()) + if err != nil { + if document.IsRenditionProviderErrorRetryable(err) { + return taskResponse{}, ambiguousSubmissionError() + } + return taskResponse{}, err + } + if status != http.StatusOK && status != http.StatusAccepted { + return taskResponse{}, statusError("submission", status) + } + return parseTask(bodyBytes) +} + +func ambiguousSubmissionError() error { + return classifiedError(document.RenditionErrorAmbiguousSubmission, + "Docling submission outcome is unknown", nil) +} + +func (client *Client) poll(ctx context.Context, expiresAt time.Time, usage *requestUsage, taskID string) (taskResponse, error) { + body, status, err := client.request(ctx, expiresAt, usage, http.MethodGet, pollPath+taskID, "", nil) + if err != nil { + return taskResponse{}, err + } + if status != http.StatusOK && status != http.StatusAccepted { + return taskResponse{}, statusError("poll", status) + } + task, err := parseTask(body) + if err != nil { + return taskResponse{}, err + } + if task.id != taskID { + return taskResponse{}, malformedError("Docling task identity changed while polling", nil) + } + return task, nil +} + +func (client *Client) result(ctx context.Context, expiresAt time.Time, usage *requestUsage, taskID string) (doclingResult, error) { + body, status, err := client.request(ctx, expiresAt, usage, http.MethodGet, resultPath+taskID, "", nil) + if err != nil { + return doclingResult{}, err + } + if status != http.StatusOK { + return doclingResult{}, statusError("result", status) + } + var wire struct { + Status string `json:"status"` + Document json.RawMessage `json:"document"` + Errors []json.RawMessage `json:"errors"` + } + if err := json.Unmarshal(body, &wire); err != nil { + return doclingResult{}, malformedError("Docling result JSON is invalid", err) + } + if wire.Status != "success" && wire.Status != "partial_success" { + return doclingResult{}, taskStatusError(wire.Status) + } + var documentWire struct { + Filename string `json:"filename"` + Markdown string `json:"md_content"` + JSONContent json.RawMessage `json:"json_content"` + } + if len(wire.Document) == 0 || json.Unmarshal(wire.Document, &documentWire) != nil || documentWire.Filename == "" { + return doclingResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "Docling result source identity is missing", nil) + } + if len(documentWire.JSONContent) != 0 && string(documentWire.JSONContent) != "null" { + var identity struct { + Origin struct { + Filename string `json:"filename"` + } `json:"origin"` + } + if json.Unmarshal(documentWire.JSONContent, &identity) == nil && identity.Origin.Filename != "" && + identity.Origin.Filename != documentWire.Filename { + return doclingResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "Docling result source identity is inconsistent", nil) + } + } + return doclingResult{markdown: []byte(documentWire.Markdown), + document: append([]byte(nil), documentWire.JSONContent...), filename: documentWire.Filename, + status: wire.Status, errors: append([]json.RawMessage(nil), wire.Errors...)}, nil +} + +func (client *Client) request( + ctx context.Context, expiresAt time.Time, usage *requestUsage, method, path, contentType string, body []byte, +) ([]byte, int, error) { + if err := checkOperation(ctx, expiresAt); err != nil { + return nil, 0, err + } + requestCtx, cancel := context.WithTimeout(ctx, client.requestTimeout) + defer cancel() + request, err := http.NewRequestWithContext(requestCtx, method, client.origin+path, bytes.NewReader(body)) + if err != nil { + return nil, 0, classifiedError(document.RenditionErrorTransient, "could not create Docling request", err) + } + request.Header.Set("Accept", "application/json") + if contentType != "" { + request.Header.Set("Content-Type", contentType) + } + if err := client.authorize(request); err != nil { + return nil, 0, err + } + if err := checkOperation(ctx, expiresAt); err != nil { + return nil, 0, err + } + usage.requests++ + response, err := client.http.Do(request) + if err != nil { + if !time.Now().Before(expiresAt) { + return nil, 0, classifiedError(document.RenditionErrorPolicyRejected, "Docling authorization expired", nil) + } + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, 0, classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", ctxErr) + } + return nil, 0, classifiedError(document.RenditionErrorTransient, "Docling request failed", err) + } + defer func() { _ = response.Body.Close() }() + responseBody, err := readBounded(response.Body, client.maxResponseBytes) + usage.outputBytes += int64(len(responseBody)) + if err != nil { + return nil, response.StatusCode, err + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return responseBody, response.StatusCode, nil + } + mediaType, _, mediaErr := mime.ParseMediaType(response.Header.Get("Content-Type")) + if mediaErr != nil || mediaType != "application/json" { + return nil, response.StatusCode, malformedError("Docling response content type is invalid", mediaErr) + } + return responseBody, response.StatusCode, nil +} + +func (client *Client) operationContext(ctx context.Context, expiresAt time.Time) (context.Context, context.CancelFunc) { + deadline := time.Now().Add(client.totalTimeout) + if callerDeadline, ok := ctx.Deadline(); ok && callerDeadline.Before(deadline) { + deadline = callerDeadline + } + if expiresAt.Before(deadline) { + deadline = expiresAt + } + return context.WithDeadline(ctx, deadline) +} + +func checkOperation(ctx context.Context, expiresAt time.Time) error { + if errors.Is(ctx.Err(), context.Canceled) { + return classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", ctx.Err()) + } + if !time.Now().Before(expiresAt) { + return classifiedError(document.RenditionErrorPolicyRejected, "Docling authorization expired", nil) + } + if err := ctx.Err(); err != nil { + return classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", err) + } + return nil +} + +func (client *Client) authorize(request *http.Request) error { + if client.secretBinding == "" { + return nil + } + secret, err := client.secrets.ResolveSecret(request.Context(), client.secretBinding) + if err != nil { + return classifiedError(document.RenditionErrorAuthentication, "Docling credential is unavailable", err) + } + if secret == "" || len(secret) > maxSecretBytes || strings.ContainsAny(secret, "\r\n\x00") { + return classifiedError(document.RenditionErrorAuthentication, "Docling credential is invalid", nil) + } + request.Header.Set("X-Api-Key", secret) + return nil +} + +func parseTask(body []byte) (taskResponse, error) { + var wire struct { + ID string `json:"task_id"` + Type string `json:"task_type"` + Status string `json:"task_status"` + } + if err := json.Unmarshal(body, &wire); err != nil { + return taskResponse{}, malformedError("Docling task response JSON is invalid", err) + } + if err := validateTaskID(wire.ID); err != nil { + return taskResponse{}, malformedError("Docling task ID is invalid", err) + } + if wire.Type != "convert" { + return taskResponse{}, malformedError("Docling task type is invalid", nil) + } + if !officialTaskStatus(wire.Status) { + return taskResponse{}, malformedError("Docling task status is invalid", nil) + } + return taskResponse{id: wire.ID, status: wire.Status}, nil +} + +func mapEvidence(raw json.RawMessage, family string) (document.SourceEvidenceV1, []byte, bool) { + var wire struct { + SchemaName string `json:"schema_name"` + Version string `json:"version"` + Texts []struct { + Text string `json:"text"` + Prov []struct { + PageNo int64 `json:"page_no"` + } `json:"prov"` + } `json:"texts"` + Pages map[string]json.RawMessage `json:"pages"` + } + if len(raw) == 0 || json.Unmarshal(raw, &wire) != nil || wire.SchemaName != "DoclingDocument" || !supportedDoclingMajor(wire.Version) || len(wire.Pages) == 0 { + return document.SourceEvidenceV1{}, nil, false + } + kind, locatorKind, natural := familyUnit(family) + if !natural { + return document.SourceEvidenceV1{}, nil, false + } + pages := make(map[int64][]string, len(wire.Pages)) + for key := range wire.Pages { + page, err := strconv.ParseInt(key, 10, 64) + if err != nil || page < 1 { + return document.SourceEvidenceV1{}, nil, false + } + pages[page] = nil + } + indexes := make([]int64, 0, len(pages)) + for page := range pages { + indexes = append(indexes, page) + } + slices.Sort(indexes) + for index, page := range indexes { + if page != int64(index+1) { + return document.SourceEvidenceV1{}, nil, false + } + } + for _, text := range wire.Texts { + if text.Text == "" { + continue + } + if len(text.Prov) == 0 { + return document.SourceEvidenceV1{}, nil, false + } + page := text.Prov[0].PageNo + if page < 1 { + return document.SourceEvidenceV1{}, nil, false + } + for _, provenance := range text.Prov { + if provenance.PageNo != page { + return document.SourceEvidenceV1{}, nil, false + } + } + if _, ok := pages[page]; !ok { + return document.SourceEvidenceV1{}, nil, false + } + pages[page] = append(pages[page], text.Text) + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: family, UnitKind: kind, Units: make([]document.SourceEvidenceUnitV1, 0, len(indexes)), + } + for order, index := range indexes { + evidence.Units = append(evidence.Units, document.SourceEvidenceUnitV1{ + Order: order, Text: strings.Join(pages[index], "\n\n"), + Locator: document.SourceEvidenceLocatorV1{Kind: locatorKind, IndexOrigin: document.EvidenceIndexOriginOne, Start: index, End: index}, + }) + } + return evidence, append([]byte(nil), raw...), true +} + +func familyUnit(family string) (document.EvidenceUnitKind, document.EvidenceLocatorKind, bool) { + switch family { + case "pdf": + return document.EvidenceUnitPage, document.EvidenceLocatorPage, true + default: + return "", "", false + } +} + +func degradedEvidence(family, markdown string) document.SourceEvidenceV1 { + return document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceDegradedProvenance, + Family: family, UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", Reason: "Docling structured evidence is unavailable", + }}, + Units: []document.SourceEvidenceUnitV1{{Order: 0, Text: markdown, + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone}}}, + } +} + +func allowsStructured(authorization document.RenditionAuthorization) bool { + if slices.Contains(authorization.AllowedArtifactRoles, document.EvidenceArtifactStructured) { + return authorization.MaxArtifacts > 0 && authorization.MaxArtifactBytes > 0 + } + return false +} + +func readExact(ctx context.Context, upload io.Reader, metadata document.AuthorizedUploadMetadata) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", err) + } + data, err := io.ReadAll(io.LimitReader(upload, metadata.ByteLength+1)) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, classifiedError(document.RenditionErrorCanceled, "Docling rendering canceled", ctxErr) + } + return nil, classifiedError(document.RenditionErrorTransient, "could not read the authorized upload", err) + } + if int64(len(data)) != metadata.ByteLength || sha256Hex(data) != metadata.SHA256 { + return nil, classifiedError(document.RenditionErrorPolicyRejected, "authorized upload identity mismatch", nil) + } + return data, nil +} + +func waitContext(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func validateOrigin(raw string, trust document.RenditionTrustBoundary) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Opaque != "" || + parsed.ForceQuery || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return "", errors.New("docling: origin must be one absolute origin without path, credentials, query, or fragment") + } + if parsed.Scheme != "https" && (parsed.Scheme != "http" || trust != document.RenditionTrustOperatorNetwork) { + return "", errors.New("docling: hosted origins require HTTPS; HTTP is operator-network only") + } + if trust != document.RenditionTrustOperatorNetwork && trust != document.RenditionTrustHostedProvider { + return "", errors.New("docling: network origin requires an operator-network or hosted trust boundary") + } + return parsed.Scheme + "://" + parsed.Host, nil +} + +func validateToken(value, subject string) error { + if value == "" || len(value) > maxTaskIDBytes || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return fmt.Errorf("docling: %s is invalid", subject) + } + for _, character := range value { + if (character < 'a' || character > 'z') && (character < 'A' || character > 'Z') && + (character < '0' || character > '9') && character != '.' && character != '_' && character != '-' { + return fmt.Errorf("docling: %s is invalid", subject) + } + } + return nil +} + +func validateTaskID(value string) error { + if value == "" || len(value) > maxTaskIDBytes || value == "." || value == ".." { + return errors.New("invalid task ID") + } + for _, character := range value { + if (character < 'a' || character > 'z') && (character < '0' || character > '9') && character != '_' && character != '-' { + return errors.New("invalid task ID") + } + } + return nil +} + +func supportedDoclingMajor(version string) bool { + major, _, ok := strings.Cut(version, ".") + return ok && major == "1" +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = append([]document.RenditionFormatCapability(nil), value.SupportedFormats...) + value.ArtifactRoles = append([]document.EvidenceArtifactRole(nil), value.ArtifactRoles...) + return value +} + +func readBounded(reader io.Reader, maximum int64) ([]byte, error) { + value, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + return value, classifiedError(document.RenditionErrorTransient, "could not read Docling response", err) + } + if int64(len(value)) > maximum { + return value, malformedError("Docling response exceeds byte limit", nil) + } + return value, nil +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} + +func taskStatusError(status string) error { + switch status { + case "failure": + return malformedError("Docling task failed", nil) + case "skipped": + return classifiedError(document.RenditionErrorPolicyRejected, "Docling task was skipped", nil) + default: + return malformedError("Docling task status is unsupported", nil) + } +} + +func officialTaskStatus(status string) bool { + switch status { + case "pending", "started", "success", "failure", "partial_success", "skipped": + return true + default: + return false + } +} + +func parsePartialPages(raw []json.RawMessage) ([]int64, error) { + if len(raw) == 0 { + return nil, malformedError("Docling partial result has no page omissions", nil) + } + pages := make([]int64, 0, len(raw)) + seen := make(map[int64]struct{}, len(raw)) + for _, value := range raw { + var omission struct { + PageNo *int64 `json:"page_no"` + } + if json.Unmarshal(value, &omission) != nil || omission.PageNo == nil || *omission.PageNo < 1 { + return nil, malformedError("Docling partial result has an unlocated omission", nil) + } + if _, ok := seen[*omission.PageNo]; ok { + continue + } + seen[*omission.PageNo] = struct{}{} + pages = append(pages, *omission.PageNo) + } + return pages, nil +} + +func partialSuccessEvidence( + evidence document.SourceEvidenceV1, family string, pages []int64, +) (document.SourceEvidenceV1, error) { + if family != "pdf" || evidence.Completeness != document.EvidenceComplete || len(pages) == 0 { + return document.SourceEvidenceV1{}, malformedError("Docling partial result lacks exact PDF page evidence", nil) + } + unitByPage := make(map[int64]int, len(evidence.Units)) + for index := range evidence.Units { + unitByPage[evidence.Units[index].Locator.Start] = index + } + for _, page := range pages { + index, ok := unitByPage[page] + if !ok { + return document.SourceEvidenceV1{}, malformedError("Docling partial result names an unknown page", nil) + } + unit := &evidence.Units[index] + unit.Omissions = append(unit.Omissions, document.SourceEvidenceOmissionV1{ + Kind: document.EvidenceOmissionField, Field: "provider_output", Reason: "Docling reported partial success", UnitOrder: unit.Order, + }) + } + evidence.Completeness = document.EvidencePartial + return evidence, nil +} + +func partialSuccessWarnings(partial bool) []string { + if !partial { + return nil + } + return []string{"partial_success"} +} + +func statusError(operation string, status int) error { + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return classifiedError(document.RenditionErrorAuthentication, "Docling authentication failed", nil) + case http.StatusNotFound, http.StatusGone: + if operation == "poll" || operation == "result" { + return classifiedError(document.RenditionErrorUnknownJob, "Docling task is unknown or expired", nil) + } + return malformedError("Docling returned an unexpected HTTP status", nil) + case http.StatusBadRequest, http.StatusRequestEntityTooLarge, http.StatusUnsupportedMediaType, http.StatusUnprocessableEntity: + if operation == "submission" { + return classifiedError(document.RenditionErrorPolicyRejected, "Docling rejected the submitted input", nil) + } + return malformedError("Docling returned an unexpected HTTP status", nil) + case http.StatusTooManyRequests: + return classifiedError(document.RenditionErrorRateLimited, "Docling rate limit", nil) + case http.StatusRequestTimeout, http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: + return classifiedError(document.RenditionErrorTransient, "Docling is temporarily unavailable", nil) + default: + return malformedError("Docling returned an unexpected HTTP status", nil) + } +} + +func malformedError(message string, cause error) error { + return classifiedError(document.RenditionErrorMalformedEvidence, message, cause) +} + +func classifiedError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err == nil { + return providerError + } + fallback, fallbackErr := document.NewRenditionProviderError(document.RenditionErrorMalformedEvidence, + "Docling returned an invalid error", 0, err) + if fallbackErr == nil { + return fallback + } + return errors.Join(err, fallbackErr) +} diff --git a/document/docling/client_test.go b/document/docling/client_test.go new file mode 100644 index 00000000..2f983e91 --- /dev/null +++ b/document/docling/client_test.go @@ -0,0 +1,918 @@ +package docling + +import ( + "bytes" + "context" + "crypto/sha256" + _ "embed" + "encoding/hex" + "encoding/json" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +//go:embed testdata/docling-pages.json +var recordedPagesResponse []byte + +//go:embed testdata/docling-schema-drift.json +var recordedSchemaDriftResponse []byte + +type testUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (*testUpload) Close() error { return nil } +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("missing test secret") + } + return value, nil +} + +func TestClientRendersDoclingPagesAndRequestsBothFormats(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + var polls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch { + case request.Method == http.MethodPost && request.URL.Path == "/v1/convert/file/async": + assert.Equal(t, "synthetic-secret", request.Header.Get("X-Api-Key")) + assertDoclingSubmission(t, request, fixture.metadata, fixture.source) + writeJSON(t, response, doclingTask("task-1", "pending")) + case request.Method == http.MethodGet && request.URL.Path == "/v1/status/poll/task-1": + if polls.Add(1) == 1 { + writeJSON(t, response, doclingTask("task-1", "started")) + return + } + writeJSON(t, response, doclingTask("task-1", "success")) + case request.Method == http.MethodGet && request.URL.Path == "/v1/result/task-1": + writeRecordedJSON(t, response, recordedPagesResponse) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + + client := newClient(t, server.URL, fixture.descriptor, testSecrets{"docling-api": "synthetic-secret"}, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + assert.Equal(t, "first page", result.Evidence.Units[0].Text) + assert.Equal(t, int64(1), result.Evidence.Units[0].Locator.Start) + assert.Equal(t, "second page", result.Evidence.Units[1].Text) + assert.Equal(t, "# Synthetic report\n", string(result.ProviderMarkdown)) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, document.EvidenceArtifactStructured, result.Artifacts[0].Role) + assert.Equal(t, int64(2), polls.Load()) +} + +func TestClientRequiresConvertTasksAndOfficialStatuses(t *testing.T) { + for _, testCase := range []struct { + name string + body string + }{ + {name: "missing task type", body: `{"task_id":"task-1","task_status":"success"}`}, + {name: "wrong task type", body: `{"task_id":"task-1","task_type":"classify","task_status":"success"}`}, + {name: "undocumented running status", body: `{"task_id":"task-1","task_type":"convert","task_status":"running"}`}, + {name: "undocumented queued status", body: `{"task_id":"task-1","task_type":"convert","task_status":"queued"}`}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := parseTask([]byte(testCase.body)) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorMalformedEvidence, providerErr.Code()) + }) + } +} + +func TestTaskIDsFitDocbankReceiptTokens(t *testing.T) { + for _, taskID := range []string{"UPPER", ".", "..", strings.Repeat("a", 121)} { + _, err := parseTask([]byte(`{"task_id":"` + taskID + `","task_type":"convert","task_status":"success"}`)) + require.Error(t, err, taskID) + } +} + +func TestMapEvidenceUsesContiguousPageRegistryAndNeverDropsText(t *testing.T) { + for _, testCase := range []struct { + name string + raw map[string]any + want bool + }{ + {name: "blank registered page", want: true, raw: map[string]any{"schema_name": "DoclingDocument", "version": "1.7.0", "pages": map[string]any{"1": map[string]any{}, "2": map[string]any{}}, "texts": []any{map[string]any{"text": "one", "prov": []any{map[string]any{"page_no": 1}}}}}}, + {name: "page gap", want: false, raw: map[string]any{"schema_name": "DoclingDocument", "version": "1.7.0", "pages": map[string]any{"1": map[string]any{}, "3": map[string]any{}}, "texts": []any{}}}, + {name: "unlocated text", want: false, raw: map[string]any{"schema_name": "DoclingDocument", "version": "1.7.0", "pages": map[string]any{"1": map[string]any{}}, "texts": []any{map[string]any{"text": "one"}}}}, + {name: "cross page provenance", want: false, raw: map[string]any{"schema_name": "DoclingDocument", "version": "1.7.0", "pages": map[string]any{"1": map[string]any{}, "2": map[string]any{}}, "texts": []any{map[string]any{"text": "one", "prov": []any{map[string]any{"page_no": 1}, map[string]any{"page_no": 2}}}}}}, + {name: "v2 drift", want: false, raw: map[string]any{"schema_name": "DoclingDocument", "version": "2.0.0", "pages": map[string]any{"1": map[string]any{}}, "texts": []any{}}}, + } { + t.Run(testCase.name, func(t *testing.T) { + raw, err := json.Marshal(testCase.raw) + require.NoError(t, err) + evidence, _, usable := mapEvidence(raw, "pdf") + assert.Equal(t, testCase.want, usable) + if usable { + require.Len(t, evidence.Units, 2) + assert.Empty(t, evidence.Units[1].Text) + } + }) + } +} + +func TestClientRejectsChangedTaskIDWhilePolling(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("original", "pending")) + case pollPath + "original": + writeJSON(t, response, doclingTask("substituted", "success")) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorMalformedEvidence, providerErr.Code()) +} + +func TestClientStopsAtAuthorizationExpiryBeforeEgress(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "expiry.pdf", []byte("synthetic PDF bytes")) + var requests atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + requests.Add(1) + http.NotFound(response, request) + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + fixture.authorization.ExpiresAt = time.Now().UTC().Add(15 * time.Millisecond).Format(timestampForm) + upload := &testUpload{Reader: delayedReader{Reader: bytes.NewReader(fixture.source), delay: 30 * time.Millisecond}, metadata: fixture.metadata} + _, err := document.RenderRendition(t.Context(), client, upload, fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorPolicyRejected, providerErr.Code()) + assert.Zero(t, requests.Load()) +} + +func TestClientStopsAtAuthorizationExpiryBeforeFollowupEgress(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "expiry.pdf", []byte("synthetic PDF bytes")) + expiresAt := time.Now().UTC().Add(500 * time.Millisecond) + var submits, polls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + submits.Add(1) + if delay := time.Until(expiresAt); delay > 0 { + time.Sleep(delay) + } + writeJSON(t, response, doclingTask("expiry", "pending")) + case pollPath + "expiry": + polls.Add(1) + writeJSON(t, response, doclingTask("expiry", "success")) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + fixture.authorization.ExpiresAt = expiresAt.Format(timestampForm) + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorPolicyRejected, providerErr.Code()) + assert.Equal(t, int64(1), submits.Load()) + assert.Zero(t, polls.Load()) +} + +func TestClientClassifiesHTTPStatusByOperationBeforeContentType(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "status.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + operation string + status int + want document.RenditionErrorCode + }{ + {name: "submit 401", operation: "submit", status: http.StatusUnauthorized, want: document.RenditionErrorAuthentication}, + {name: "submit 429", operation: "submit", status: http.StatusTooManyRequests, want: document.RenditionErrorRateLimited}, + {name: "submit 503", operation: "submit", status: http.StatusServiceUnavailable, want: document.RenditionErrorTransient}, + {name: "submit 404", operation: "submit", status: http.StatusNotFound, want: document.RenditionErrorMalformedEvidence}, + {name: "submit 400", operation: "submit", status: http.StatusBadRequest, want: document.RenditionErrorPolicyRejected}, + {name: "submit 413", operation: "submit", status: http.StatusRequestEntityTooLarge, want: document.RenditionErrorPolicyRejected}, + {name: "submit 415", operation: "submit", status: http.StatusUnsupportedMediaType, want: document.RenditionErrorPolicyRejected}, + {name: "submit 422", operation: "submit", status: http.StatusUnprocessableEntity, want: document.RenditionErrorPolicyRejected}, + {name: "poll 404", operation: "poll", status: http.StatusNotFound, want: document.RenditionErrorUnknownJob}, + {name: "poll 410", operation: "poll", status: http.StatusGone, want: document.RenditionErrorUnknownJob}, + {name: "result 404", operation: "result", status: http.StatusNotFound, want: document.RenditionErrorUnknownJob}, + {name: "result 410", operation: "result", status: http.StatusGone, want: document.RenditionErrorUnknownJob}, + } { + t.Run(testCase.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if testCase.operation == "submit" && request.URL.Path == convertPath || + testCase.operation == "poll" && request.URL.Path == pollPath+"status" || + testCase.operation == "result" && request.URL.Path == resultPath+"status" { + response.Header().Set("Content-Type", "text/html") + response.WriteHeader(testCase.status) + _, err := io.WriteString(response, "provider-private-body") + assert.NoError(t, err) + return + } + switch request.URL.Path { + case convertPath: + if testCase.operation == "poll" { + writeJSON(t, response, doclingTask("status", "pending")) + } else { + writeJSON(t, response, doclingTask("status", "success")) + } + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, testCase.want, providerErr.Code()) + assert.NotContains(t, err.Error(), "provider-private-body") + }) + } +} + +func TestClientTreatsFailedUploadTransportAsAmbiguousSubmission(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "ambiguous.pdf", []byte("synthetic PDF bytes")) + var consumed atomic.Bool + client := newClient(t, "http://127.0.0.1", fixture.descriptor, nil, &http.Client{Transport: roundTripperFunc(func(request *http.Request) (*http.Response, error) { + _, err := io.Copy(io.Discard, request.Body) + require.NoError(t, err) + consumed.Store(true) + return nil, errors.New("synthetic upload transport failure") + })}) + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorAmbiguousSubmission, providerErr.Code()) + assert.True(t, consumed.Load()) +} + +func TestClientRecoversKnownTaskWithoutResubmitting(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "recovery.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + path string + }{ + {name: "poll", path: pollPath + "recovery"}, + {name: "result", path: resultPath + "recovery"}, + } { + t.Run(testCase.name, func(t *testing.T) { + var submits, failures atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + submits.Add(1) + if testCase.name == "poll" { + writeJSON(t, response, doclingTask("recovery", "pending")) + } else { + writeJSON(t, response, doclingTask("recovery", "success")) + } + case testCase.path: + if failures.Add(1) == 1 { + response.WriteHeader(http.StatusServiceUnavailable) + return + } + if testCase.name == "poll" { + writeJSON(t, response, doclingTask("recovery", "success")) + } else { + writeJSON(t, response, doclingResultResponse(fixture.metadata.Filename, "# recovery\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + })) + } + case resultPath + "recovery": + writeJSON(t, response, doclingResultResponse(fixture.metadata.Filename, "# recovery\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + })) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, int64(1), submits.Load()) + assert.Equal(t, int64(2), failures.Load()) + }) + } +} + +func TestClientReturnsAmbiguousSubmissionWhenKnownTaskRecoveryExhausts(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "recovery.pdf", []byte("synthetic PDF bytes")) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("recovery", "pending")) + case pollPath + "recovery": + response.WriteHeader(http.StatusServiceUnavailable) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorAmbiguousSubmission, providerErr.Code()) +} + +func TestClientReceiptCountsEveryResponseBodyAndRequest(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "usage.pdf", []byte("synthetic PDF bytes")) + var outputBytes atomic.Int64 + var polls atomic.Int64 + writeCountedJSON := func(t *testing.T, response http.ResponseWriter, value any) { + t.Helper() + body, err := json.Marshal(value) + require.NoError(t, err) + body = append(body, '\n') + outputBytes.Add(int64(len(body))) + response.Header().Set("Content-Type", "application/json") + _, err = response.Write(body) + require.NoError(t, err) + } + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeCountedJSON(t, response, doclingTask("usage", "pending")) + case pollPath + "usage": + if polls.Add(1) == 1 { + body := []byte("temporary provider body") + outputBytes.Add(int64(len(body))) + response.WriteHeader(http.StatusServiceUnavailable) + _, err := response.Write(body) + assert.NoError(t, err) + return + } + writeCountedJSON(t, response, doclingTask("usage", "success")) + case resultPath + "usage": + writeCountedJSON(t, response, doclingResultResponse(fixture.metadata.Filename, "# usage\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + })) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, int64(4), result.Receipt.Usage.Requests) + assert.Equal(t, int64(1), result.Receipt.Usage.Retries) + assert.Equal(t, outputBytes.Load(), result.Receipt.Usage.OutputBytes) +} + +func TestClientRequestCountsPartialResponseBytesBeforeReadFailure(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "usage.pdf", []byte("synthetic PDF bytes")) + payload := []byte("partial provider response") + client := newClient(t, "http://127.0.0.1", fixture.descriptor, nil, &http.Client{ + Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(&readOnceError{payload: payload}), + }, nil + }), + }) + usage := &requestUsage{} + _, _, err := client.request(t.Context(), time.Now().Add(time.Minute), usage, http.MethodGet, resultPath+"usage", "", nil) + require.Error(t, err) + assert.Equal(t, int64(1), usage.requests) + assert.Equal(t, int64(len(payload)), usage.outputBytes) +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (function roundTripperFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type readOnceError struct { + payload []byte +} + +func (reader *readOnceError) Read(buffer []byte) (int, error) { + if len(reader.payload) == 0 { + return 0, errors.New("synthetic response read failure") + } + n := copy(buffer, reader.payload) + reader.payload = reader.payload[n:] + return n, errors.New("synthetic response read failure") +} + +type delayedReader struct { + io.Reader + + delay time.Duration +} + +func (reader delayedReader) Read(buffer []byte) (int, error) { + time.Sleep(reader.delay) + return reader.Reader.Read(buffer) +} + +func TestClientHandlesPartialSuccessAndSanitizesTerminalFailures(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + t.Run("partial success", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("partial", "partial_success")) + case resultPath + "partial": + writeJSON(t, response, doclingResultResponse(fixture.metadata.Filename, "# report\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + }, withResultErrors([]any{map[string]any{"page_no": 1}}))) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidencePartial, result.Evidence.Completeness) + assert.Contains(t, result.Receipt.Warnings, "partial_success") + require.NotEmpty(t, result.Evidence.Units[0].Omissions) + }) + t.Run("partial success result wrapper", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("wrapper-partial", "success")) + case resultPath + "wrapper-partial": + writeJSON(t, response, doclingResultResponse("report.pdf", "# report\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + }, withResultStatus("partial_success"), withResultErrors([]any{map[string]any{"page_no": 1}}))) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidencePartial, result.Evidence.Completeness) + assert.Contains(t, result.Receipt.Warnings, "partial_success") + require.NotEmpty(t, result.Evidence.Units[0].Omissions) + }) + for _, testCase := range []struct { + status string + want document.RenditionErrorCode + }{ + {status: "failure", want: document.RenditionErrorMalformedEvidence}, + {status: "skipped", want: document.RenditionErrorPolicyRejected}, + } { + t.Run(testCase.status, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + writeJSON(t, response, map[string]any{ + "task_id": "terminal", "task_type": "convert", "task_status": testCase.status, + "errors": []string{"provider-body-secret"}, + }) + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, testCase.want, providerErr.Code()) + assert.NotContains(t, err.Error(), "provider-body-secret") + }) + } +} + +func TestClientPublishesPartialEvidenceOnlyForExactPDFPageOmissions(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "partial.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + family string + errors []any + wantErr bool + }{ + {name: "exact page omission", errors: []any{map[string]any{"page_no": 1, "message": "private"}, map[string]any{"page_no": 1}}}, + {name: "no errors", wantErr: true}, + {name: "document scoped error", errors: []any{map[string]any{"message": "private"}}, wantErr: true}, + {name: "unknown page", errors: []any{map[string]any{"page_no": 3}}, wantErr: true}, + {name: "non PDF", family: "text", errors: []any{map[string]any{"page_no": 1}}, wantErr: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + candidate := fixture + if testCase.family != "" { + candidate = newFixture(t, testCase.family, "text/plain", "partial.txt", []byte("synthetic text bytes")) + } + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("partial", "partial_success")) + case resultPath + "partial": + writeJSON(t, response, doclingResultResponse(candidate.metadata.Filename, "# partial\n", []any{ + map[string]any{"text": "one", "prov": []any{map[string]any{"page_no": 1}}}, + }, withResultErrors(testCase.errors))) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, candidate.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, candidate.upload(), candidate.authorization) + if testCase.wantErr { + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorMalformedEvidence, providerErr.Code()) + return + } + require.NoError(t, err) + assert.Equal(t, document.EvidencePartial, result.Evidence.Completeness) + assert.Empty(t, result.Evidence.Omissions) + require.Len(t, result.Evidence.Units[0].Omissions, 1) + assert.Equal(t, "provider_output", result.Evidence.Units[0].Omissions[0].Field) + assert.Empty(t, result.Evidence.Units[1].Omissions) + }) + } +} + +func TestClientRejectsResultWrapperIdentityAndStatusFailures(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + result map[string]any + want document.RenditionErrorCode + }{ + {name: "missing wrapper filename", result: doclingResultResponse("", "# report\n", nil), want: document.RenditionErrorPolicyRejected}, + {name: "conflicting inner filename", result: doclingResultResponseWithInnerFilename("report.pdf", "other.pdf", "# report\n", nil), want: document.RenditionErrorPolicyRejected}, + {name: "failed wrapper status", result: doclingResultResponse("report.pdf", "# report\n", nil, withResultStatus("failure")), want: document.RenditionErrorMalformedEvidence}, + } { + t.Run(testCase.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case convertPath: + writeJSON(t, response, doclingTask("wrapper", "success")) + case resultPath + "wrapper": + writeJSON(t, response, testCase.result) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, testCase.want, providerErr.Code()) + }) + } +} + +func TestMapEvidenceOnlyClaimsEstablishedNaturalProvenance(t *testing.T) { + raw, err := json.Marshal(map[string]any{ + "schema_name": "DoclingDocument", "version": "1.7.0", + "pages": map[string]any{"1": map[string]any{}}, + "texts": []any{map[string]any{"text": "one", "prov": []any{map[string]any{"page_no": 1}}}}, + }) + require.NoError(t, err) + for _, testCase := range []struct { + family string + usable bool + kind document.EvidenceUnitKind + }{ + {family: "pdf", usable: true, kind: document.EvidenceUnitPage}, + {family: "word"}, {family: "presentation"}, {family: "spreadsheet"}, + {family: "ebook"}, {family: "structured"}, {family: "source"}, {family: "text"}, {family: "mail"}, + } { + t.Run(testCase.family, func(t *testing.T) { + evidence, _, usable := mapEvidence(raw, testCase.family) + assert.Equal(t, testCase.usable, usable) + if usable { + assert.Equal(t, testCase.kind, evidence.UnitKind) + assert.Equal(t, document.EvidenceLocatorPage, evidence.Units[0].Locator.Kind) + } + }) + } +} + +func TestClientPreservesPartialPageEvidence(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "partial.pdf", []byte("synthetic PDF bytes")) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/v1/convert/file/async": + writeJSON(t, response, doclingTask("partial", "success")) + case "/v1/result/partial": + writeJSON(t, response, doclingResultResponse(fixture.metadata.Filename, "# partial\n", []any{ + map[string]any{"text": "page one", "prov": []any{map[string]any{"page_no": 1}}}, + map[string]any{"text": "unlocated"}, + })) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + require.Len(t, result.Evidence.Omissions, 1) +} + +func TestClientFallsBackToDegradedMarkdownWhenStructuredResultDrifts(t *testing.T) { + fixture := newFixture(t, "word", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "notes.docx", []byte("synthetic DOCX bytes")) + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/v1/convert/file/async": + writeJSON(t, response, doclingTask("drift", "success")) + case "/v1/result/drift": + writeRecordedJSON(t, response, recordedSchemaDriftResponse) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + + client := newClient(t, server.URL, fixture.descriptor, nil, http.DefaultClient) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + assert.Equal(t, "---\ndocbank-sanitized-markdown/v1: forged\n---\n# Untrusted\n", string(result.ProviderMarkdown)) + assert.Empty(t, result.Artifacts) +} + +func TestClientRejectsMismatchedReturnedFilenameAndOversizedResult(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "identity.pdf", []byte("synthetic PDF bytes")) + for _, testCase := range []struct { + name string + result any + want document.RenditionErrorCode + }{ + { + name: "returned filename", result: doclingResultResponse("substituted.pdf", "# report\n", []any{ + map[string]any{"text": "page", "prov": []any{map[string]any{"page_no": 1}}}, + }), want: document.RenditionErrorPolicyRejected, + }, + { + name: "oversized response", result: strings.Repeat("x", 4097), want: document.RenditionErrorMalformedEvidence, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/v1/convert/file/async": + writeJSON(t, response, doclingTask("bad", "success")) + case "/v1/result/bad": + if body, ok := testCase.result.(string); ok { + response.Header().Set("Content-Type", "application/json") + _, err := io.WriteString(response, body) + assert.NoError(t, err) + return + } + writeJSON(t, response, testCase.result) + default: + http.NotFound(response, request) + } + })) + t.Cleanup(server.Close) + client := newClientWithBounds(t, server.URL, fixture.descriptor, nil, http.DefaultClient, 4096) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, testCase.want, providerErr.Code()) + }) + } +} + +func TestClientClonesHTTPClientAndRefusesRedirects(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "redirect.pdf", []byte("synthetic PDF bytes")) + jar, err := cookiejar.New(nil) + require.NoError(t, err) + base := &http.Client{Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return nil }} + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.URL.Path == "/v1/convert/file/async" { + http.Redirect(response, request, "/elsewhere", http.StatusFound) + return + } + http.NotFound(response, request) + })) + t.Cleanup(server.Close) + + client := newClient(t, server.URL, fixture.descriptor, nil, base) + assert.Nil(t, client.http.Jar) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorMalformedEvidence, providerErr.Code()) +} + +func TestClientAcceptanceSyntheticUpload(t *testing.T) { + origin := os.Getenv("DOCLING_ACCEPTANCE_URL") + if origin == "" { + t.Skip("set DOCLING_ACCEPTANCE_URL to run against an operator-controlled Docling Serve endpoint") + } + fixture := newFixture(t, "text", "text/plain", "synthetic-acceptance.txt", []byte("Synthetic Docling acceptance document.\n")) + var secrets SecretResolver + if apiKey := os.Getenv("DOCLING_ACCEPTANCE_API_KEY"); apiKey != "" { + secrets = testSecrets{"docling-api": apiKey} + } + binding := "" + if secrets != nil { + binding = "docling-api" + } + client, err := New(Profile{ + Origin: origin, Descriptor: fixture.descriptor, SecretBinding: binding, + RequestTimeout: 30 * time.Second, TotalTimeout: 10 * time.Minute, PollInterval: time.Second, + MaxPollAttempts: 300, MaxResponseBytes: 64 << 20, MaxDocumentBytes: 1 << 20, + }, secrets, http.DefaultClient) + require.NoError(t, err) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, fixture.authorization.SourceSHA256, result.Receipt.SourceSHA256) + assert.NotEmpty(t, result.Evidence.Units) +} + +type fixture struct { + descriptor document.RenditionDescriptor + metadata document.AuthorizedUploadMetadata + authorization document.RenditionAuthorization + source []byte +} + +func newFixture(t *testing.T, family, mediaType, filename string, source []byte) fixture { + t.Helper() + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: "docling.serve-v1", ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), TrustBoundary: document.RenditionTrustOperatorNetwork, + SupportedFormats: []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", InputKind: document.RenditionInputOriginalFile}, + }, + ReturnsMarkdown: true, ReturnsStructured: true, + ArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + }) + require.NoError(t, err) + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: filename, MediaFamily: family, MediaType: mediaType, ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), ProviderMetadataChecksum: strings.Repeat("3", 64), InputKind: document.RenditionInputOriginalFile, + } + started := time.Now().UTC().Add(-time.Minute) + return fixture{descriptor: descriptor, metadata: metadata, source: source, authorization: document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: family, MediaType: mediaType, InputKind: document.RenditionInputOriginalFile, + AllowedArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, MaxProviderMarkdownBytes: 4096, + MaxArtifactBytes: 4096, MaxArtifacts: 1, MaxTotalResultBytes: 16384, + AuthorizedAt: started.Format("2006-01-02T15:04:05.000000000Z"), ExpiresAt: started.Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + }} +} + +func (fixture fixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func newClient(t *testing.T, origin string, descriptor document.RenditionDescriptor, secrets SecretResolver, httpClient *http.Client) *Client { + t.Helper() + return newClientWithBounds(t, origin, descriptor, secrets, httpClient, 1<<20) +} + +func newClientWithBounds(t *testing.T, origin string, descriptor document.RenditionDescriptor, secrets SecretResolver, httpClient *http.Client, maxResponseBytes int64) *Client { + t.Helper() + binding := "" + if secrets != nil { + binding = "docling-api" + } + client, err := New(Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: binding, + RequestTimeout: time.Second, TotalTimeout: 2 * time.Second, PollInterval: time.Millisecond, + MaxPollAttempts: 4, MaxResponseBytes: maxResponseBytes, + }, secrets, httpClient) + require.NoError(t, err) + return client +} + +func assertDoclingSubmission(t *testing.T, request *http.Request, metadata document.AuthorizedUploadMetadata, source []byte) { + t.Helper() + mediaType, params, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, params["boundary"]) + file, err := reader.NextPart() + require.NoError(t, err) + assert.Equal(t, "files", file.FormName()) + assert.Equal(t, metadata.Filename, file.FileName()) + gotSource, err := io.ReadAll(file) + require.NoError(t, err) + assert.Equal(t, source, gotSource) + formats := make([]string, 0, 2) + for { + part, partErr := reader.NextPart() + if errors.Is(partErr, io.EOF) { + break + } + require.NoError(t, partErr) + value, readErr := io.ReadAll(part) + require.NoError(t, readErr) + switch part.FormName() { + case "to_formats": + formats = append(formats, string(value)) + case "target_type": + assert.Equal(t, "inbody", string(value)) + default: + t.Errorf("unexpected form field %q", part.FormName()) + } + } + assert.Equal(t, []string{"md", "json"}, formats) +} + +type resultOption func(map[string]any) + +func withResultStatus(status string) resultOption { + return func(result map[string]any) { result["status"] = status } +} + +func withResultErrors(errors []any) resultOption { + return func(result map[string]any) { result["errors"] = errors } +} + +func doclingTask(taskID, status string) map[string]any { + return map[string]any{"task_id": taskID, "task_type": "convert", "task_status": status} +} + +func doclingResultResponse(filename, markdown string, texts []any, options ...resultOption) map[string]any { + return doclingResultResponseWithInnerFilename(filename, filename, markdown, texts, options...) +} + +func doclingResultResponseWithInnerFilename( + filename, innerFilename, markdown string, texts []any, options ...resultOption, +) map[string]any { + result := map[string]any{ + "status": "success", "processing_time": 0.01, "errors": []string{}, + "document": map[string]any{ + "filename": filename, "md_content": markdown, + "json_content": map[string]any{ + "schema_name": "DoclingDocument", "version": "1.7.0", + "origin": map[string]any{"filename": innerFilename}, "pages": map[string]any{"1": map[string]any{}, "2": map[string]any{}}, "texts": texts, + }, + "html_content": "", "text_content": "", + }, + } + for _, option := range options { + option(result) + } + return result +} + +func writeJSON(t *testing.T, response http.ResponseWriter, value any) { + t.Helper() + response.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(response).Encode(value)) +} + +func writeRecordedJSON(t *testing.T, response http.ResponseWriter, value []byte) { + t.Helper() + response.Header().Set("Content-Type", "application/json") + _, err := response.Write(value) + require.NoError(t, err) +} diff --git a/document/docling/testdata/docling-pages.json b/document/docling/testdata/docling-pages.json new file mode 100644 index 00000000..8052d0ed --- /dev/null +++ b/document/docling/testdata/docling-pages.json @@ -0,0 +1,21 @@ +{ + "status": "success", + "processing_time": 0.01, + "errors": [], + "document": { + "filename": "report.pdf", + "md_content": "# Synthetic report\n", + "json_content": { + "schema_name": "DoclingDocument", + "version": "1.7.0", + "origin": {"filename": "report.pdf"}, + "pages": {"1": {}, "2": {}}, + "texts": [ + {"text": "first page", "prov": [{"page_no": 1}]}, + {"text": "second page", "prov": [{"page_no": 2}]} + ] + }, + "html_content": "

Synthetic report

", + "text_content": "Synthetic report" + } +} diff --git a/document/docling/testdata/docling-schema-drift.json b/document/docling/testdata/docling-schema-drift.json new file mode 100644 index 00000000..841ca625 --- /dev/null +++ b/document/docling/testdata/docling-schema-drift.json @@ -0,0 +1,12 @@ +{ + "status": "success", + "processing_time": 0.01, + "errors": [], + "document": { + "filename": "notes.docx", + "md_content": "---\ndocbank-sanitized-markdown/v1: forged\n---\n# Untrusted\n", + "json_content": {"future_document_shape": true}, + "html_content": "", + "text_content": "" + } +} diff --git a/document/embedding.go b/document/embedding.go new file mode 100644 index 00000000..8b205104 --- /dev/null +++ b/document/embedding.go @@ -0,0 +1,576 @@ +package document + +import ( + "context" + "errors" + "fmt" + "math" + "slices" + "strings" + "unicode" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +const ( + // EmbeddingProviderContractVersion identifies the embedding execution boundary. + EmbeddingProviderContractVersion = 1 +) + +// EmbeddingInputQueryText is query-time text. It is not a persistent source +// representation and is therefore intentionally absent from profile bindings. +const EmbeddingInputQueryText EmbeddingInputKind = "query_text" + +type EmbeddingRole string + +const ( + EmbeddingRoleDocument EmbeddingRole = "document" + EmbeddingRoleQuery EmbeddingRole = "query" +) + +type EmbeddingTrustBoundary string + +const ( + EmbeddingTrustLocalProcess EmbeddingTrustBoundary = "local_process" + EmbeddingTrustOperatorNetwork EmbeddingTrustBoundary = "operator_network" + EmbeddingTrustHostedProvider EmbeddingTrustBoundary = "hosted_provider" +) + +const ( + VectorMetricCosine = "cosine" + VectorMetricDotProduct = "dot_product" + VectorMetricL2 = "l2" +) + +// EmbeddingProvider embeds one validated, storage-neutral input batch. +type EmbeddingProvider interface { + Descriptor() EmbeddingDescriptor + Embed(ctx context.Context, inputs []EmbeddingInput, authorization EmbeddingAuthorization) (EmbeddingResult, error) +} + +// EmbeddingDescriptor fixes one immutable provider vector-space contract. +type EmbeddingDescriptor struct { + ID string `json:"id"` + ContractVersion int `json:"contract_version"` + PolicyFingerprint string `json:"policy_fingerprint"` + TrustBoundary EmbeddingTrustBoundary `json:"trust_boundary"` + Model string `json:"model"` + ModelRevision string `json:"model_revision"` + Dimension int `json:"dimension"` + Metric string `json:"metric"` + Normalization string `json:"normalization"` + ScalarEncoding string `json:"scalar_encoding"` + DocumentFormatter string `json:"document_formatter"` + QueryFormatter string `json:"query_formatter"` + InputKinds []EmbeddingInputKind `json:"input_kinds"` + CompatibilityID string `json:"compatibility_id"` + SupportsTextQuery bool `json:"supports_text_query"` + ModelInput ModelInputContract `json:"model_input"` + SupportedRequestModes []ModelInputMode `json:"supported_request_modes"` + Fingerprint string `json:"fingerprint"` +} + +type embeddingDescriptorFields struct { + ID string `json:"id"` + ContractVersion int `json:"contract_version"` + PolicyFingerprint string `json:"policy_fingerprint"` + TrustBoundary EmbeddingTrustBoundary `json:"trust_boundary"` + Model string `json:"model"` + ModelRevision string `json:"model_revision"` + Dimension int `json:"dimension"` + Metric string `json:"metric"` + Normalization string `json:"normalization"` + ScalarEncoding string `json:"scalar_encoding"` + DocumentFormatter string `json:"document_formatter"` + QueryFormatter string `json:"query_formatter"` + InputKinds []EmbeddingInputKind `json:"input_kinds"` + CompatibilityID string `json:"compatibility_id"` + SupportsTextQuery bool `json:"supports_text_query"` + ModelInput ModelInputContract `json:"model_input"` + SupportedRequestModes []ModelInputMode `json:"supported_request_modes"` +} + +// NewEmbeddingDescriptor validates, canonicalizes, and fingerprints a provider descriptor. +func NewEmbeddingDescriptor(value EmbeddingDescriptor) (EmbeddingDescriptor, error) { + value.Fingerprint = "" + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + slices.Sort(value.InputKinds) + slices.Sort(value.SupportedRequestModes) + if err := validateEmbeddingDescriptorFields(value); err != nil { + return EmbeddingDescriptor{}, err + } + encoded, err := canonicalJSON(embeddingDescriptorIdentity(value)) + if err != nil { + return EmbeddingDescriptor{}, fmt.Errorf("encode embedding descriptor: %w", err) + } + value.Fingerprint = sha256Hex(encoded) + return value, nil +} + +// EmbeddingInput is one requested provider vector. +type EmbeddingInput struct { + Key string `json:"key"` + Role EmbeddingRole `json:"role"` + Kind EmbeddingInputKind `json:"kind"` + Text string `json:"text,omitempty"` + Source AuthorizedUpload `json:"-"` + HeadingPath []string `json:"heading_path,omitempty"` + SourceSpans []ChunkSpan `json:"source_spans,omitempty"` +} + +// EmbeddingAuthorization bounds one embedding invocation and pins its provider identity. +// MaxResponseBytes is the exact sum of every result key's UTF-8 bytes plus four +// bytes for every float32 scalar; adapters must independently bound raw transport bytes. +type EmbeddingAuthorization struct { + ProviderID string `json:"provider_id"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` +} + +// EmbeddingVector is one provider result. Index is required only when results +// are returned out of request order; zero values preserve ordered APIs. +type EmbeddingVector struct { + Key string `json:"key"` + Index *int `json:"index,omitempty"` + Values []float32 `json:"values"` +} + +type EmbeddingResult struct { + Vectors []EmbeddingVector `json:"vectors"` +} + +// ValidateEmbeddingProviderRequest checks the adapter snapshot and all input +// limits before any bytes reach a provider. +func ValidateEmbeddingProviderRequest(provider EmbeddingProvider, inputs []EmbeddingInput, authorization EmbeddingAuthorization) error { + if nilInterface(provider) { + return errors.New("embedding provider is required") + } + descriptor := cloneEmbeddingDescriptor(provider.Descriptor()) + if err := validateEmbeddingDescriptor(descriptor); err != nil { + return err + } + if next := cloneEmbeddingDescriptor(provider.Descriptor()); !equalEmbeddingDescriptors(descriptor, next) { + return errors.New("embedding descriptor changed during validation") + } + if err := validateEmbeddingAuthorization(descriptor, authorization); err != nil { + return err + } + if err := validateEmbeddingInputs(descriptor, inputs, authorization); err != nil { + return err + } + if minimumResultFootprint(descriptor.Dimension, inputs) > authorization.MaxResponseBytes { + return errors.New("embedding authorization response byte limit cannot hold requested vectors") + } + return nil +} + +// ValidateEmbeddingProviderResult checks exact key membership, order, and +// finite vector shape for a completed provider operation. +func ValidateEmbeddingProviderResult(descriptor EmbeddingDescriptor, inputs []EmbeddingInput, authorization EmbeddingAuthorization, result EmbeddingResult) error { + if err := validateEmbeddingDescriptor(descriptor); err != nil { + return err + } + if err := validateEmbeddingAuthorization(descriptor, authorization); err != nil { + return err + } + if err := validateEmbeddingInputs(descriptor, inputs, authorization); err != nil { + return err + } + if minimumResultFootprint(descriptor.Dimension, inputs) > authorization.MaxResponseBytes { + return errors.New("embedding authorization response byte limit cannot hold requested vectors") + } + if len(result.Vectors) != len(inputs) { + return errors.New("provider result has a missing vector") + } + indexed := result.Vectors[0].Index != nil + seen := make(map[string]struct{}, len(result.Vectors)) + seenIndices := make(map[int]struct{}, len(result.Vectors)) + expectedKeys := make(map[string]struct{}, len(inputs)) + for _, input := range inputs { + expectedKeys[input.Key] = struct{}{} + } + for position, vector := range result.Vectors { + if (vector.Index != nil) != indexed { + return errors.New("provider result has partial indexes") + } + if _, expected := expectedKeys[vector.Key]; !expected { + return errors.New("provider result has an unexpected vector key") + } + if _, exists := seen[vector.Key]; exists { + return errors.New("provider result has a duplicate vector key") + } + seen[vector.Key] = struct{}{} + expected := position + if indexed { + expected = *vector.Index + if expected < 0 || expected >= len(inputs) { + return errors.New("provider result vector index is outside request bounds") + } + if _, exists := seenIndices[expected]; exists { + return errors.New("provider result has a duplicate vector index") + } + seenIndices[expected] = struct{}{} + } + if vector.Key != inputs[expected].Key { + if indexed { + return errors.New("provider result vector key does not match index") + } + return errors.New("provider result vector order does not match request") + } + if len(vector.Values) != descriptor.Dimension { + return errors.New("provider result vector dimension does not match descriptor") + } + for _, value := range vector.Values { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("provider result vector contains non-finite value") + } + } + } + if resultFootprint(result) > authorization.MaxResponseBytes { + return errors.New("provider result exceeds authorized response bytes") + } + return nil +} + +// ExecuteEmbedding is the validated execution entry point for core callers. +func ExecuteEmbedding(ctx context.Context, provider EmbeddingProvider, inputs []EmbeddingInput, authorization EmbeddingAuthorization) (EmbeddingResult, error) { + if err := ValidateEmbeddingProviderRequest(provider, inputs, authorization); err != nil { + return EmbeddingResult{}, err + } + result, err := provider.Embed(ctx, cloneEmbeddingInputs(inputs), authorization) + if err != nil { + return EmbeddingResult{}, err + } + if err := ValidateEmbeddingProviderResult(provider.Descriptor(), inputs, authorization, result); err != nil { + return EmbeddingResult{}, err + } + return result, nil +} + +// ValidateEmbeddingQueryCompatibility proves two document/query descriptors +// occupy one safe text-retrieval vector space. +func ValidateEmbeddingQueryCompatibility(documentDescriptor, queryDescriptor EmbeddingDescriptor) error { + if err := validateEmbeddingDescriptor(documentDescriptor); err != nil { + return err + } + if err := validateEmbeddingDescriptor(queryDescriptor); err != nil { + return err + } + if !queryDescriptor.SupportsTextQuery { + return errors.New("query descriptor does not support text queries") + } + for _, field := range []struct{ name, left, right string }{ + {"compatibility ID", documentDescriptor.CompatibilityID, queryDescriptor.CompatibilityID}, + {"model", documentDescriptor.Model, queryDescriptor.Model}, {"model revision", documentDescriptor.ModelRevision, queryDescriptor.ModelRevision}, + {"metric", documentDescriptor.Metric, queryDescriptor.Metric}, + {"normalization", documentDescriptor.Normalization, queryDescriptor.Normalization}, + {"scalar encoding", documentDescriptor.ScalarEncoding, queryDescriptor.ScalarEncoding}, + {"document formatter", documentDescriptor.DocumentFormatter, queryDescriptor.DocumentFormatter}, + {"query formatter", documentDescriptor.QueryFormatter, queryDescriptor.QueryFormatter}, + {"model-input fingerprint", documentDescriptor.ModelInput.Fingerprint, queryDescriptor.ModelInput.Fingerprint}, + } { + if field.left != field.right { + return fmt.Errorf("embedding descriptors have incompatible %s", field.name) + } + } + if documentDescriptor.Dimension != queryDescriptor.Dimension { + return errors.New("embedding descriptors have incompatible dimension") + } + if !slices.Equal(documentDescriptor.InputKinds, queryDescriptor.InputKinds) { + return errors.New("embedding descriptors have incompatible input representation") + } + return nil +} + +func validateEmbeddingDescriptor(descriptor EmbeddingDescriptor) error { + if err := validateEmbeddingDescriptorFields(descriptor); err != nil { + return err + } + if err := validateFingerprint(descriptor.Fingerprint, "embedding descriptor fingerprint"); err != nil { + return err + } + canonical, err := NewEmbeddingDescriptor(descriptor) + if err != nil { + return err + } + if !equalEmbeddingDescriptors(canonical, descriptor) { + return errors.New("embedding descriptor fingerprint or canonical ordering is invalid") + } + return nil +} + +func validateEmbeddingDescriptorFields(descriptor EmbeddingDescriptor) error { + if err := validateStableToken(descriptor.ID, "embedding descriptor ID", 128); err != nil { + return err + } + if descriptor.ContractVersion != EmbeddingProviderContractVersion { + return fmt.Errorf("embedding descriptor contract version must be %d", EmbeddingProviderContractVersion) + } + if err := validateFingerprint(descriptor.PolicyFingerprint, "embedding descriptor policy fingerprint"); err != nil { + return err + } + switch descriptor.TrustBoundary { + case EmbeddingTrustLocalProcess, EmbeddingTrustOperatorNetwork, EmbeddingTrustHostedProvider: + default: + return errors.New("embedding descriptor trust boundary is invalid") + } + for _, field := range []struct{ name, value string }{{"embedding model", descriptor.Model}, {"embedding model revision", descriptor.ModelRevision}} { + if err := validateStableToken(field.value, field.name, 128); err != nil { + return err + } + } + if err := validateCompatibilityID(descriptor.CompatibilityID); err != nil { + return fmt.Errorf("embedding %w", err) + } + if descriptor.Dimension < 1 || descriptor.Dimension > maxEmbeddingDimensions { + return fmt.Errorf("embedding descriptor dimension must be between 1 and %d", maxEmbeddingDimensions) + } + if !validVectorMetric(descriptor.Metric) { + return errors.New("embedding descriptor metric is invalid") + } + if !validVectorNormalization(descriptor.Normalization) { + return errors.New("embedding descriptor normalization is invalid") + } + if err := validateStableToken(descriptor.ScalarEncoding, "embedding scalar encoding", 128); err != nil { + return err + } + for _, field := range []struct{ name, value string }{{"embedding document formatter", descriptor.DocumentFormatter}, {"embedding query formatter", descriptor.QueryFormatter}} { + if err := validateCompatibilityID(field.value); err != nil { + return fmt.Errorf("%s: %w", field.name, err) + } + } + if len(descriptor.InputKinds) == 0 || len(descriptor.InputKinds) > 2 { + return errors.New("embedding descriptor input kinds are invalid") + } + for index, kind := range descriptor.InputKinds { + if !validEmbeddingDocumentInputKind(kind) { + return fmt.Errorf("embedding descriptor input kind %q is invalid", kind) + } + if index > 0 && descriptor.InputKinds[index-1] == kind { + return errors.New("embedding descriptor has duplicate input kind") + } + } + if err := validateModelInputContract(descriptor.ModelInput); err != nil { + return err + } + if descriptor.ModelInput.Profile == "" { + return errors.New("embedding descriptor requires a model-input contract") + } + if descriptor.CompatibilityID != descriptor.ModelInput.CompatibilityID { + return errors.New("embedding descriptor compatibility ID does not match model-input contract") + } + if len(descriptor.SupportedRequestModes) == 0 { + return errors.New("embedding descriptor must declare supported request modes") + } + for index, mode := range descriptor.SupportedRequestModes { + if !validModelInputMode(mode) { + return errors.New("embedding descriptor request mode is invalid") + } + if index > 0 && descriptor.SupportedRequestModes[index-1] == mode { + return errors.New("embedding descriptor has duplicate request mode") + } + } + modes := []ModelInputMode{descriptor.ModelInput.Document.Mode} + if descriptor.SupportsTextQuery { + modes = append(modes, descriptor.ModelInput.Query.Mode) + } + for _, mode := range modes { + if !slices.Contains(descriptor.SupportedRequestModes, mode) { + return errors.New("embedding descriptor does not support a model-input request mode") + } + } + if descriptor.SupportsTextQuery && !slices.Contains(descriptor.SupportedRequestModes, descriptor.ModelInput.Query.Mode) { + return errors.New("embedding descriptor query request mode is unsupported") + } + return nil +} + +func validateEmbeddingAuthorization(descriptor EmbeddingDescriptor, authorization EmbeddingAuthorization) error { + if authorization.ProviderID != descriptor.ID || authorization.DescriptorFingerprint != descriptor.Fingerprint || authorization.PolicyFingerprint != descriptor.PolicyFingerprint { + return errors.New("embedding authorization does not match descriptor") + } + if authorization.MaxBatchItems < 1 || authorization.MaxBatchItems > maxEmbeddingBatchItems { + return errors.New("embedding authorization batch limit is invalid") + } + if authorization.MaxInputBytes < 1 || authorization.MaxInputBytes > maxEmbeddingInputBytes { + return errors.New("embedding authorization input byte limit is invalid") + } + if authorization.MaxResponseBytes < 1 || authorization.MaxResponseBytes > maxEmbeddingResponseBytes { + return errors.New("embedding authorization response byte limit is invalid") + } + if int64(descriptor.Dimension)*4 > authorization.MaxResponseBytes { + return errors.New("embedding authorization response byte limit cannot hold one vector") + } + return nil +} + +func validateEmbeddingInputs(descriptor EmbeddingDescriptor, inputs []EmbeddingInput, authorization EmbeddingAuthorization) error { + if len(inputs) == 0 || len(inputs) > authorization.MaxBatchItems { + return errors.New("embedding request batch size is outside authorization") + } + seen := make(map[string]struct{}, len(inputs)) + var total int64 + for _, input := range inputs { + if err := validateStableToken(input.Key, "embedding input key", 128); err != nil { + return err + } + if _, exists := seen[input.Key]; exists { + return errors.New("embedding request contains duplicate input key") + } + seen[input.Key] = struct{}{} + if input.Role != EmbeddingRoleDocument && input.Role != EmbeddingRoleQuery { + return errors.New("embedding request has unsupported role") + } + if input.Role == EmbeddingRoleDocument && !slices.Contains(descriptor.InputKinds, input.Kind) { + return errors.New("embedding request has unsupported kind") + } + if input.Role == EmbeddingRoleQuery && (input.Kind != EmbeddingInputQueryText || !descriptor.SupportsTextQuery) { + return errors.New("embedding request has unsupported role or kind") + } + if err := validateEmbeddingInputAuxiliaries(input); err != nil { + return err + } + if input.Kind == EmbeddingInputOriginalFile { + if nilInterface(input.Source) || input.Text != "" { + return errors.New("original-file embedding input requires authorized upload") + } + metadata := input.Source.Metadata() + if err := validateAuthorizedUploadMetadata(metadata); err != nil { + return fmt.Errorf("embedding upload metadata: %w", err) + } + if metadata.InputKind != RenditionInputOriginalFile { + return errors.New("original-file embedding input must use original file upload metadata") + } + if second := input.Source.Metadata(); second != metadata { + return errors.New("embedding upload metadata changed during validation") + } + if metadata.ByteLength > authorization.MaxInputBytes-total { + return errors.New("embedding request exceeds authorized input bytes") + } + total += metadata.ByteLength + } else { + if input.Text == "" || input.Source != nil { + return errors.New("text embedding input is required") + } + if !utf8.ValidString(input.Text) { + return errors.New("embedding input text is not valid UTF-8") + } + encoder := descriptor.ModelInput.Query + if input.Role == EmbeddingRoleDocument { + encoder = descriptor.ModelInput.Document + } + renderedLength, err := modelInputRenderedLength(encoder, input.Text) + if err != nil { + return err + } + if renderedLength > authorization.MaxInputBytes-total { + return errors.New("embedding request exceeds authorized input bytes") + } + total += renderedLength + } + if total > authorization.MaxInputBytes { + return errors.New("embedding request exceeds authorized input bytes") + } + } + return nil +} + +func validateEmbeddingInputAuxiliaries(input EmbeddingInput) error { + if input.Kind == EmbeddingInputQueryText || input.Kind == EmbeddingInputOriginalFile { + if len(input.HeadingPath) != 0 || len(input.SourceSpans) != 0 { + return errors.New("embedding input auxiliary fields are document-only") + } + return nil + } + if len(input.HeadingPath) > maxEvidenceHeadingDepth || len(input.SourceSpans) > 1024 { + return errors.New("embedding input auxiliary slices exceed bounds") + } + remainingHeadingBytes := int64(maxEvidenceHeadingBytes) + for _, heading := range input.HeadingPath { + if heading == "" || !utf8.ValidString(heading) || strings.IndexFunc(heading, unicode.IsControl) >= 0 || !norm.NFC.IsNormalString(heading) { + return errors.New("embedding heading is not canonical text") + } + if int64(len(heading)) > remainingHeadingBytes { + return errors.New("embedding heading bytes exceed bounds") + } + remainingHeadingBytes -= int64(len(heading)) + } + var spanBytes int64 + for _, span := range input.SourceSpans { + if span.UnitIndex < 0 || span.UnitIndex >= maxEvidenceUnits || span.CharStart < 0 || span.CharEnd <= span.CharStart || span.CharEnd > maxEvidenceUnitRunes { + return errors.New("embedding input has an invalid source span") + } + width := int64(span.CharEnd) - int64(span.CharStart) + if width > maxEmbeddingInputBytes-spanBytes { + return errors.New("embedding source spans exceed bounds") + } + spanBytes += width + } + return nil +} + +func resultFootprint(result EmbeddingResult) int64 { + const maxInt64 = int64(^uint64(0) >> 1) + var total int64 + for _, vector := range result.Vectors { + if int64(len(vector.Key)) > maxInt64-total { + return maxInt64 + } + total += int64(len(vector.Key)) + if int64(len(vector.Values)) > (maxInt64-total)/4 { + return maxInt64 + } + total += int64(len(vector.Values)) * 4 + } + return total +} + +func minimumResultFootprint(dimension int, inputs []EmbeddingInput) int64 { + const maxInt64 = int64(^uint64(0) >> 1) + var total int64 + for _, input := range inputs { + if int64(len(input.Key)) > maxInt64-total { + return maxInt64 + } + total += int64(len(input.Key)) + if int64(dimension) > (maxInt64-total)/4 { + return maxInt64 + } + total += int64(dimension) * 4 + } + return total +} + +func validModelInputMode(mode ModelInputMode) bool { + return mode == ModelInputModeText || mode == ModelInputModeDocument || mode == ModelInputModeQuery +} +func validEmbeddingDocumentInputKind(kind EmbeddingInputKind) bool { + return kind == EmbeddingInputOriginalFile || kind == EmbeddingInputRenditionChunk +} +func validVectorMetric(metric string) bool { + return metric == VectorMetricCosine || metric == VectorMetricDotProduct || metric == VectorMetricL2 +} +func embeddingDescriptorIdentity(value EmbeddingDescriptor) embeddingDescriptorFields { + return embeddingDescriptorFields{ID: value.ID, ContractVersion: value.ContractVersion, PolicyFingerprint: value.PolicyFingerprint, TrustBoundary: value.TrustBoundary, Model: value.Model, ModelRevision: value.ModelRevision, Dimension: value.Dimension, Metric: value.Metric, Normalization: value.Normalization, ScalarEncoding: value.ScalarEncoding, DocumentFormatter: value.DocumentFormatter, QueryFormatter: value.QueryFormatter, InputKinds: value.InputKinds, CompatibilityID: value.CompatibilityID, SupportsTextQuery: value.SupportsTextQuery, ModelInput: value.ModelInput, SupportedRequestModes: value.SupportedRequestModes} +} +func cloneEmbeddingDescriptor(value EmbeddingDescriptor) EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} +func cloneEmbeddingInputs(inputs []EmbeddingInput) []EmbeddingInput { + cloned := slices.Clone(inputs) + for i := range cloned { + cloned[i].HeadingPath = slices.Clone(inputs[i].HeadingPath) + cloned[i].SourceSpans = slices.Clone(inputs[i].SourceSpans) + } + return cloned +} +func equalEmbeddingDescriptors(left, right EmbeddingDescriptor) bool { + return left.ID == right.ID && left.ContractVersion == right.ContractVersion && left.PolicyFingerprint == right.PolicyFingerprint && left.TrustBoundary == right.TrustBoundary && left.Model == right.Model && left.ModelRevision == right.ModelRevision && left.Dimension == right.Dimension && left.Metric == right.Metric && left.Normalization == right.Normalization && left.ScalarEncoding == right.ScalarEncoding && left.DocumentFormatter == right.DocumentFormatter && left.QueryFormatter == right.QueryFormatter && left.CompatibilityID == right.CompatibilityID && left.SupportsTextQuery == right.SupportsTextQuery && left.ModelInput == right.ModelInput && left.Fingerprint == right.Fingerprint && slices.Equal(left.InputKinds, right.InputKinds) && slices.Equal(left.SupportedRequestModes, right.SupportedRequestModes) +} diff --git a/document/embedding_input_test.go b/document/embedding_input_test.go new file mode 100644 index 00000000..c7b4d88b --- /dev/null +++ b/document/embedding_input_test.go @@ -0,0 +1,22 @@ +package document + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// This test fails if request sizing renders the template before authorization. +func TestModelInputRenderedLengthDoesNotAllocateRenderedText(t *testing.T) { + contract, err := NewModelInputContract(ModelInputContractConfig{Profile: ModelInputProfileCustom, CompatibilityID: "test-space", Document: ModelInputEncoder{Mode: ModelInputModeText, Template: strings.Repeat("x", 4096-len("{{content}}")) + "{{content}}"}, Query: ModelInputEncoder{Mode: ModelInputModeText, Template: "{{content}}"}}) + require.NoError(t, err) + content := strings.Repeat("y", 1<<20) + length, err := modelInputRenderedLength(contract.Document, content) + require.NoError(t, err) + assert.Equal(t, int64((1<<20)+4096-len("{{content}}")), length) + assert.Zero(t, testing.AllocsPerRun(100, func() { + _, _ = modelInputRenderedLength(contract.Document, content) + })) +} diff --git a/document/embedding_test.go b/document/embedding_test.go new file mode 100644 index 00000000..0b9794c3 --- /dev/null +++ b/document/embedding_test.go @@ -0,0 +1,350 @@ +package document_test + +import ( + "context" + "io" + "math" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +// This test fails if execution accepts a response whose keys, order, or vector +// shape no longer exactly represent the authorized request. +func TestEmbeddingContractRejectsMalformedProviderResult(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + inputs := []document.EmbeddingInput{ + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}, + {Key: "chunk-b", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "beta"}, + } + authorization := testEmbeddingAuthorization(descriptor) + + for _, testCase := range []struct { + name string + result document.EmbeddingResult + want string + }{ + {"missing key", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Values: []float32{1, 2}}}}, "missing vector"}, + {"duplicate key", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Values: []float32{1, 2}}, {Key: "chunk-a", Values: []float32{3, 4}}}}, "duplicate vector key"}, + {"unexpected key", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Values: []float32{1, 2}}, {Key: "other", Values: []float32{3, 4}}}}, "unexpected vector key"}, + {"reordered without index", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-b", Values: []float32{3, 4}}, {Key: "chunk-a", Values: []float32{1, 2}}}}, "vector order"}, + {"wrong dimensions", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Values: []float32{1}}, {Key: "chunk-b", Values: []float32{3, 4}}}}, "dimension"}, + {"non-finite", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Values: []float32{float32(math.NaN()), 2}}, {Key: "chunk-b", Values: []float32{3, 4}}}}, "non-finite"}, + } { + t.Run(testCase.name, func(t *testing.T) { + err := document.ValidateEmbeddingProviderResult(descriptor, inputs, authorization, testCase.result) + require.ErrorContains(t, err, testCase.want) + }) + } + + indexed := document.EmbeddingResult{Vectors: []document.EmbeddingVector{ + {Key: "chunk-b", Index: new(1), Values: []float32{3, 4}}, + {Key: "chunk-a", Index: new(0), Values: []float32{1, 2}}, + }} + require.NoError(t, document.ValidateEmbeddingProviderResult(descriptor, inputs, authorization, indexed)) + for _, testCase := range []struct { + name string + result document.EmbeddingResult + want string + }{ + {"partial indexes", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Index: new(0), Values: []float32{1, 2}}, {Key: "chunk-b", Values: []float32{3, 4}}}}, "partial indexes"}, + {"duplicate explicit zero", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Index: new(0), Values: []float32{1, 2}}, {Key: "chunk-b", Index: new(0), Values: []float32{3, 4}}}}, "duplicate vector index"}, + {"out of range index", document.EmbeddingResult{Vectors: []document.EmbeddingVector{{Key: "chunk-a", Index: new(2), Values: []float32{1, 2}}, {Key: "chunk-b", Index: new(1), Values: []float32{3, 4}}}}, "outside request bounds"}, + } { + t.Run(testCase.name, func(t *testing.T) { + err := document.ValidateEmbeddingProviderResult(descriptor, inputs, authorization, testCase.result) + require.ErrorContains(t, err, testCase.want) + }) + } + + unsupported := append([]document.EmbeddingInput(nil), inputs...) + unsupported[0].Role = document.EmbeddingRoleQuery + err := document.ValidateEmbeddingProviderRequest(testEmbeddingProvider{descriptor: descriptor}, unsupported, authorization) + require.ErrorContains(t, err, "unsupported role") + unsupported = append([]document.EmbeddingInput(nil), inputs...) + unsupported[0].Kind = document.EmbeddingInputOriginalFile + err = document.ValidateEmbeddingProviderRequest(testEmbeddingProvider{descriptor: descriptor}, unsupported, authorization) + require.ErrorContains(t, err, "unsupported kind") +} + +// This test fails if document and query spaces can be treated as compatible +// after an identity-affecting descriptor difference. +func TestEmbeddingContractRejectsIncompatibleDocumentAndQueryDescriptors(t *testing.T) { + documentDescriptor := testEmbeddingDescriptor(t) + documentDescriptor.SupportsTextQuery = true + documentDescriptor.Fingerprint = "" + documentDescriptor, err := document.NewEmbeddingDescriptor(documentDescriptor) + require.NoError(t, err) + queryDescriptor := documentDescriptor + queryDescriptor.Dimension++ + queryDescriptor.Fingerprint = "" + queryDescriptor, err = document.NewEmbeddingDescriptor(queryDescriptor) + require.NoError(t, err) + + err = document.ValidateEmbeddingQueryCompatibility(documentDescriptor, queryDescriptor) + require.ErrorContains(t, err, "dimension") +} + +func testEmbeddingDescriptor(t *testing.T) document.EmbeddingDescriptor { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileOpenAICompatible}) + require.NoError(t, err) + descriptor, err := document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ + ID: "synthetic-embedder", ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: testFingerprint(), TrustBoundary: document.EmbeddingTrustLocalProcess, + Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: document.VectorMetricCosine, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, ModelInput: contract, + DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: document.VectorNormalizationUnitLength, + ScalarEncoding: "float32", SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }) + require.NoError(t, err) + return descriptor +} + +// This test fails if query eligibility is decided from a compatibility label +// while a distinct identity-bearing vector-space field has changed. +func TestEmbeddingContractRequiresExactQueryVectorSpace(t *testing.T) { + base := testEmbeddingDescriptor(t) + base.SupportsTextQuery = true + base.Fingerprint = "" + base, err := document.NewEmbeddingDescriptor(base) + require.NoError(t, err) + + customInput, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: base.CompatibilityID, + Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "document: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "query: {{content}}"}, + }) + require.NoError(t, err) + for _, testCase := range []struct { + name string + mutate func(*document.EmbeddingDescriptor) + want string + }{ + {"model input", func(d *document.EmbeddingDescriptor) { d.ModelInput = customInput }, "model-input"}, + {"normalization", func(d *document.EmbeddingDescriptor) { d.Normalization = document.VectorNormalizationNone }, "normalization"}, + {"scalar encoding", func(d *document.EmbeddingDescriptor) { d.ScalarEncoding = "float16" }, "scalar encoding"}, + {"document formatter", func(d *document.EmbeddingDescriptor) { d.DocumentFormatter = "document/v2" }, "document formatter"}, + {"query formatter", func(d *document.EmbeddingDescriptor) { d.QueryFormatter = "query/v2" }, "query formatter"}, + {"input representation", func(d *document.EmbeddingDescriptor) { + d.InputKinds = []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile} + }, "input representation"}, + } { + t.Run(testCase.name, func(t *testing.T) { + candidate := base + testCase.mutate(&candidate) + candidate.Fingerprint = "" + candidate, err = document.NewEmbeddingDescriptor(candidate) + require.NoError(t, err) + err = document.ValidateEmbeddingQueryCompatibility(base, candidate) + require.ErrorContains(t, err, testCase.want) + }) + } +} + +// This test fails if an adapter can silently gain request modes from a model +// contract it did not itself declare support for. +func TestEmbeddingContractRequiresExplicitProviderRequestModes(t *testing.T) { + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileOpenAICompatible}) + require.NoError(t, err) + _, err = document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ + ID: "synthetic-embedder", ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: testFingerprint(), TrustBoundary: document.EmbeddingTrustLocalProcess, + Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: document.VectorMetricCosine, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: contract.CompatibilityID, + ModelInput: contract, DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: document.VectorNormalizationUnitLength, + ScalarEncoding: "float32", + }) + require.ErrorContains(t, err, "must declare supported request modes") +} + +// This test fails if the durable descriptor rejects the existing profile +// vocabulary or silently translates its normalization identifier. +func TestEmbeddingContractPreservesUnitLengthNormalization(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + descriptor.Normalization = document.VectorNormalizationUnitLength + descriptor.Fingerprint = "" + _, err := document.NewEmbeddingDescriptor(descriptor) + require.NoError(t, err) +} + +// This test fails if a second normalization spelling can enter any durable +// contract boundary alongside the canonical unit_length value. +func TestEmbeddingContractRejectsL2NormalizationAlias(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + descriptor.Normalization = "l2" + descriptor.Fingerprint = "" + _, err := document.NewEmbeddingDescriptor(descriptor) + require.ErrorContains(t, err, "normalization") +} + +func TestEmbeddingContractEnforcesResultFootprint(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + input := []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}} + authorization := testEmbeddingAuthorization(descriptor) + authorization.MaxResponseBytes = 8 + calls := 0 + _, err := document.ExecuteEmbedding(context.Background(), countingEmbeddingProvider{descriptor: descriptor, calls: &calls}, input, authorization) + require.ErrorContains(t, err, "response byte") + require.Zero(t, calls) +} + +func TestEmbeddingContractRejectsInvalidDocumentAuxiliaryFields(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + provider := testEmbeddingProvider{descriptor: descriptor} + authorization := testEmbeddingAuthorization(descriptor) + for _, input := range []document.EmbeddingInput{ + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", HeadingPath: []string{"\xff"}}, + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", HeadingPath: []string{"heading\x00"}}, + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", HeadingPath: []string{strings.Repeat("h", (1<<20)+1)}}, + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", SourceSpans: []document.ChunkSpan{{UnitIndex: -1, CharStart: 0, CharEnd: 1}}}, + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 2, CharEnd: 2}}}, + {Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: (1 << 30) + 1}}}, + } { + err := document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{input}, authorization) + require.Error(t, err) + } + require.NoError(t, document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", HeadingPath: []string{"Overview"}, SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 5}}}}, authorization)) + + query := descriptor + query.SupportsTextQuery = true + query.Fingerprint = "" + query, err := document.NewEmbeddingDescriptor(query) + require.NoError(t, err) + err = document.ValidateEmbeddingProviderRequest(testEmbeddingProvider{descriptor: query}, []document.EmbeddingInput{{Key: "query", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "alpha", HeadingPath: []string{"forbidden"}}}, testEmbeddingAuthorization(query)) + require.ErrorContains(t, err, "auxiliary") + + original := descriptor + original.InputKinds = []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile} + original.Fingerprint = "" + original, err = document.NewEmbeddingDescriptor(original) + require.NoError(t, err) + derived := testEmbeddingUploadMetadata(6) + derived.InputKind = document.RenditionInputDerivedUpload + err = document.ValidateEmbeddingProviderRequest(testEmbeddingProvider{descriptor: original}, []document.EmbeddingInput{{Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: &testEmbeddingUpload{metadata: []document.AuthorizedUploadMetadata{derived}}, SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 1}}}}, testEmbeddingAuthorization(original)) + require.ErrorContains(t, err, "auxiliary") + err = document.ValidateEmbeddingProviderRequest(testEmbeddingProvider{descriptor: original}, []document.EmbeddingInput{{Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: &testEmbeddingUpload{metadata: []document.AuthorizedUploadMetadata{derived}}}}, testEmbeddingAuthorization(original)) + require.ErrorContains(t, err, "original file") +} + +// This test fails if a span can name a source unit outside the canonical +// evidence-unit vocabulary and reach the provider boundary. +func TestEmbeddingContractRejectsOverLimitSourceUnitBeforeProviderCall(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + calls := 0 + _, err := document.ExecuteEmbedding(context.Background(), countingEmbeddingProvider{descriptor: descriptor, calls: &calls}, []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", SourceSpans: []document.ChunkSpan{{UnitIndex: 100_000, CharStart: 0, CharEnd: 1}}}}, testEmbeddingAuthorization(descriptor)) + require.ErrorContains(t, err, "source span") + require.Zero(t, calls) +} + +type countingEmbeddingProvider struct { + descriptor document.EmbeddingDescriptor + calls *int +} + +func (provider countingEmbeddingProvider) Descriptor() document.EmbeddingDescriptor { + return provider.descriptor +} +func (provider countingEmbeddingProvider) Embed(context.Context, []document.EmbeddingInput, document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + *provider.calls++ + return document.EmbeddingResult{}, nil +} + +func TestEmbeddingContractRejectsUnrepresentableOrMutableInputs(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + provider := testEmbeddingProvider{descriptor: descriptor} + authorization := testEmbeddingAuthorization(descriptor) + typedNil := (*testEmbeddingUpload)(nil) + err := document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", Source: typedNil}}, authorization) + require.ErrorContains(t, err, "text embedding input") + err = document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", HeadingPath: make([]string, 1025)}}, authorization) + require.ErrorContains(t, err, "auxiliary slices") + + withTemplate := descriptor + withTemplate.ModelInput, err = document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: descriptor.CompatibilityID, Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "x{{content}}"}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "{{content}}"}}) + require.NoError(t, err) + withTemplate.Fingerprint = "" + withTemplate, err = document.NewEmbeddingDescriptor(withTemplate) + require.NoError(t, err) + authorization = testEmbeddingAuthorization(withTemplate) + authorization.MaxInputBytes = 5 + templateCalls := 0 + _, err = document.ExecuteEmbedding(context.Background(), countingEmbeddingProvider{descriptor: withTemplate, calls: &templateCalls}, []document.EmbeddingInput{{Key: "chunk-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}}, authorization) + require.ErrorContains(t, err, "input bytes") + require.Zero(t, templateCalls) + + original := descriptor + original.InputKinds = []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile} + original.Fingerprint = "" + original, err = document.NewEmbeddingDescriptor(original) + require.NoError(t, err) + provider = testEmbeddingProvider{descriptor: original} + authorization = testEmbeddingAuthorization(original) + metadata := testEmbeddingUploadMetadata(6) + err = document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Text: "unexpected", Source: &testEmbeddingUpload{metadata: []document.AuthorizedUploadMetadata{metadata}}}}, authorization) + require.ErrorContains(t, err, "original-file") + + authorization.MaxInputBytes = 5 + err = document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: &testEmbeddingUpload{metadata: []document.AuthorizedUploadMetadata{metadata}}}}, authorization) + require.ErrorContains(t, err, "input bytes") + + changed := metadata + changed.Filename = "changed.pdf" + authorization.MaxInputBytes = 1 << 20 + err = document.ValidateEmbeddingProviderRequest(provider, []document.EmbeddingInput{{Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: &testEmbeddingUpload{metadata: []document.AuthorizedUploadMetadata{metadata, changed}}}}, authorization) + require.ErrorContains(t, err, "changed during validation") +} + +type testEmbeddingUpload struct { + metadata []document.AuthorizedUploadMetadata + reads int +} + +func (upload *testEmbeddingUpload) Read(data []byte) (int, error) { return 0, io.EOF } +func (upload *testEmbeddingUpload) Close() error { return nil } +func (upload *testEmbeddingUpload) Metadata() document.AuthorizedUploadMetadata { + result := upload.metadata[min(upload.reads, len(upload.metadata)-1)] + upload.reads++ + return result +} + +func testEmbeddingUploadMetadata(size int64) document.AuthorizedUploadMetadata { + return document.AuthorizedUploadMetadata{Filename: "source.pdf", MediaFamily: "pdf", MediaType: "application/pdf", ByteLength: size, SHA256: strings.Repeat("a", 64), CapabilityRecordChecksum: strings.Repeat("b", 64), ProviderMetadataChecksum: strings.Repeat("c", 64), InputKind: document.RenditionInputOriginalFile} +} + +func testEmbeddingAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 10, MaxInputBytes: 1 << 20, + MaxResponseBytes: 1 << 20, + } +} + +type testEmbeddingProvider struct{ descriptor document.EmbeddingDescriptor } + +func (provider testEmbeddingProvider) Descriptor() document.EmbeddingDescriptor { + return provider.descriptor +} + +func (testEmbeddingProvider) Embed(context.Context, []document.EmbeddingInput, document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + return document.EmbeddingResult{}, nil +} + +func testFingerprint() string { + // A deterministic lowercase SHA-256 fixture, deliberately not calculated by + // the contract under test. + return "823412d1eacb67956220e532c5e1c72ebf4d55d7ac7c5a9686507b54b7c5cfcc" +} + +func TestEmbeddingContractExecutesOnlyValidatedRequests(t *testing.T) { + descriptor := testEmbeddingDescriptor(t) + provider := testEmbeddingProvider{descriptor: descriptor} + _, err := document.ExecuteEmbedding(context.Background(), provider, []document.EmbeddingInput{{ + Key: "query", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "alpha", + }}, testEmbeddingAuthorization(descriptor)) + assert.ErrorContains(t, err, "unsupported role") +} diff --git a/document/embeddingbridge/client.go b/document/embeddingbridge/client.go new file mode 100644 index 00000000..45940969 --- /dev/null +++ b/document/embeddingbridge/client.go @@ -0,0 +1,471 @@ +package embeddingbridge + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "hash" + "io" + "math" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "slices" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/manifestjson" +) + +var ( + _ document.EmbeddingProvider = (*Client)(nil) + errSourceChanged = errors.New("embedding bridge source changed") + errSourceTransferFailed = errors.New("embedding bridge source transfer failed") + errTransferStopped = errors.New("embedding bridge transfer stopped") +) + +// Descriptor returns an immutable copy of the configured vector-space identity. +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Embed performs one bounded synchronous request against the fixed same-origin route. +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if client == nil { + return document.EmbeddingResult{}, classified(ErrorPermanent, 0) + } + if ctx == nil { + return document.EmbeddingResult{}, classified(ErrorPermanent, 0) + } + frozenInputs, err := freezeInputMetadata(inputs) + if err != nil { + return document.EmbeddingResult{}, err + } + if err := document.ValidateEmbeddingProviderRequest(client, frozenInputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.maxBatchItems || authorization.MaxInputBytes > client.maxInputBytes || + authorization.MaxResponseBytes > client.maxResponseBytes { + return document.EmbeddingResult{}, classified(ErrorCapacity, 0) + } + manifest, encoded, err := client.buildManifest(frozenInputs, authorization) + if err != nil { + return document.EmbeddingResult{}, err + } + boundary := "docbank-embedding-" + manifest.RequestChecksum[:32] + contentLength, err := multipartLength(boundary, encoded, frozenInputs) + if err != nil || contentLength > client.maxRequestBytes { + return document.EmbeddingResult{}, classified(ErrorCapacity, 0) + } + + requestCtx, cancel := context.WithTimeout(ctx, client.requestTimeout) + defer cancel() + sourceGate := newActiveSourceGate() + closeFinished := make(chan struct{}) + stopClose := context.AfterFunc(requestCtx, func() { + sourceGate.Cancel() + close(closeFinished) + }) + defer func() { + if !stopClose() { + <-closeFinished + } + }() + secret := "" + if client.secretBinding != "" { + var resolveErr error + secret, resolveErr = client.secrets.ResolveSecret(requestCtx, client.secretBinding) + if resolveErr != nil || !validSecret(secret) { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, contextErr + } + return document.EmbeddingResult{}, classified(ErrorAuthentication, 0) + } + } + + bodyReader, bodyWriter := io.Pipe() + writerDone := make(chan error, 1) + go func() { + multipartWriter := multipart.NewWriter(bodyWriter) + if err := multipartWriter.SetBoundary(boundary); err != nil { + _ = bodyWriter.CloseWithError(err) + writerDone <- err + return + } + writeErr := writeMultipart(multipartWriter, encoded, frozenInputs, true, sourceGate) + if closeErr := multipartWriter.Close(); writeErr == nil { + writeErr = closeErr + } + _ = bodyWriter.CloseWithError(writeErr) + writerDone <- writeErr + }() + request, err := http.NewRequestWithContext(requestCtx, http.MethodPost, client.origin+embeddingsPath, bodyReader) + if err != nil { + _ = bodyReader.Close() + <-writerDone + return document.EmbeddingResult{}, classified(ErrorPermanent, 0) + } + request.ContentLength = contentLength + request.Header.Set("Content-Type", "multipart/form-data; boundary="+boundary) + request.Header.Set("Accept", responseMediaType) + if secret != "" { + request.Header.Set("Authorization", "Bearer "+secret) + } + request.Header.Set("Idempotency-Key", manifest.RequestChecksum) + request.Header.Set("Docbank-Request-Checksum", manifest.RequestChecksum) + if hasOriginalInput(frozenInputs) { + request.Header.Set("Expect", "100-continue") + } + + response, doErr := client.http.Do(request) + _ = bodyReader.Close() + sourceGate.Cancel() + writeErr := <-writerDone + if contextErr := requestCtx.Err(); contextErr != nil { + if response != nil { + _ = response.Body.Close() + } + return document.EmbeddingResult{}, contextErr + } + if errors.Is(writeErr, errSourceChanged) { + if response != nil { + _ = response.Body.Close() + } + return document.EmbeddingResult{}, classified(ErrorSourceChanged, 0) + } + if doErr != nil { + if response != nil { + _ = response.Body.Close() + } + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, contextErr + } + return document.EmbeddingResult{}, classified(ErrorAmbiguousSubmission, 0) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return document.EmbeddingResult{}, statusError(response.StatusCode) + } + if writeErr != nil { + return document.EmbeddingResult{}, classified(ErrorAmbiguousSubmission, 0) + } + if err := requireResponseMediaType(response.Header.Get("Content-Type")); err != nil { + return document.EmbeddingResult{}, classified(ErrorMalformedResponse, response.StatusCode) + } + body, err := readBounded(response.Body, client.maxResponseBytes) + if err != nil { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, contextErr + } + return document.EmbeddingResult{}, classified(ErrorMalformedResponse, response.StatusCode) + } + if err := manifestjson.RejectDuplicateKeys(body, "embedding bridge response"); err != nil { + return document.EmbeddingResult{}, classified(ErrorMalformedResponse, response.StatusCode) + } + var envelope Response + if err := json.Unmarshal(body, &envelope, json.RejectUnknownMembers(true)); err != nil { + return document.EmbeddingResult{}, classified(ErrorMalformedResponse, response.StatusCode) + } + result, err := client.validateResponse(envelope, manifest, frozenInputs, authorization) + if err != nil { + return document.EmbeddingResult{}, classified(ErrorMalformedResponse, response.StatusCode) + } + return result, nil +} + +func freezeInputMetadata(inputs []document.EmbeddingInput) ([]document.EmbeddingInput, error) { + frozen := slices.Clone(inputs) + for index := range frozen { + input := &frozen[index] + input.HeadingPath = slices.Clone(input.HeadingPath) + input.SourceSpans = slices.Clone(input.SourceSpans) + if input.Kind != document.EmbeddingInputOriginalFile || nilInterface(input.Source) { + continue + } + metadata := input.Source.Metadata() + if !safeMultipartFilename(metadata.Filename) { + return nil, classified(ErrorPermanent, 0) + } + input.Source = frozenUpload{AuthorizedUpload: input.Source, metadata: metadata} + } + return frozen, nil +} + +func (client *Client) buildManifest(inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (RequestManifest, []byte, error) { + manifest := RequestManifest{ + ContractVersion: ContractVersion, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: client.descriptor.PolicyFingerprint, Authorization: authorization, + Inputs: make([]ManifestInput, len(inputs)), + } + fileIndex := 0 + for index, input := range inputs { + entry := ManifestInput{ + Index: index, Key: input.Key, Role: input.Role, Kind: input.Kind, + HeadingPath: slices.Clone(input.HeadingPath), SourceSpans: manifestSpans(input.SourceSpans), + } + if input.Kind == document.EmbeddingInputOriginalFile { + metadata := input.Source.Metadata() + entry.ByteLength = metadata.ByteLength + entry.SHA256 = metadata.SHA256 + entry.Upload = new(metadata) + entry.FilePart = filePartName + entry.FileIndex = new(fileIndex) + fileIndex++ + } else { + if input.Role == document.EmbeddingRoleDocument { + entry.Text = client.descriptor.ModelInput.EncodeDocument(input.Text) + } else { + entry.Text = client.descriptor.ModelInput.EncodeQuery(input.Text) + } + entry.ByteLength = int64(len(entry.Text)) + entry.SHA256 = sha256Hex([]byte(entry.Text)) + } + manifest.Inputs[index] = entry + } + identity, err := json.Marshal(manifest, json.Deterministic(true)) + if err != nil { + return RequestManifest{}, nil, classified(ErrorPermanent, 0) + } + manifest.RequestChecksum = sha256Hex(identity) + encoded, err := json.Marshal(manifest, json.Deterministic(true)) + if err != nil { + return RequestManifest{}, nil, classified(ErrorPermanent, 0) + } + return manifest, encoded, nil +} + +func manifestSpans(values []document.ChunkSpan) []ManifestSpan { + if values == nil { + return nil + } + result := make([]ManifestSpan, len(values)) + for index, value := range values { + result[index] = ManifestSpan{UnitIndex: value.UnitIndex, CharStart: value.CharStart, CharEnd: value.CharEnd} + } + return result +} + +func hasOriginalInput(inputs []document.EmbeddingInput) bool { + return slices.ContainsFunc(inputs, func(input document.EmbeddingInput) bool { + return input.Kind == document.EmbeddingInputOriginalFile + }) +} + +func multipartLength(boundary string, manifest []byte, inputs []document.EmbeddingInput) (int64, error) { + counter := new(countingWriter) + writer := multipart.NewWriter(counter) + if err := writer.SetBoundary(boundary); err != nil { + return 0, errors.New("embedding bridge: deterministic multipart boundary is invalid") + } + if err := writeMultipart(writer, manifest, inputs, false, nil); err != nil { + return 0, err + } + if err := writer.Close(); err != nil { + return 0, errors.New("embedding bridge: multipart measurement failed") + } + total := counter.written + for _, input := range inputs { + if input.Kind == document.EmbeddingInputOriginalFile { + if input.Source.Metadata().ByteLength > maxRequestBytes-total { + return 0, errors.New("multipart request length overflow") + } + total += input.Source.Metadata().ByteLength + } + } + return total, nil +} + +func writeMultipart(writer *multipart.Writer, manifest []byte, inputs []document.EmbeddingInput, writeFiles bool, sourceGate *activeSourceGate) error { + manifestHeader := make(textproto.MIMEHeader) + manifestHeader.Set("Content-Disposition", `form-data; name="`+manifestPartName+`"`) + manifestHeader.Set("Content-Type", manifestMediaType) + part, err := writer.CreatePart(manifestHeader) + if err != nil { + return errors.New("embedding bridge: manifest multipart part failed") + } + if _, err := part.Write(manifest); err != nil { + return err + } + for _, input := range inputs { + if input.Kind != document.EmbeddingInputOriginalFile { + continue + } + metadata := input.Source.Metadata() + fileHeader := make(textproto.MIMEHeader) + fileHeader.Set("Content-Disposition", multipart.FileContentDisposition(filePartName, metadata.Filename)) + fileHeader.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(fileHeader) + if err != nil { + return errors.New("embedding bridge: file multipart part failed") + } + if !writeFiles { + continue + } + token, ok := sourceGate.Begin(input.Source) + if !ok { + return errTransferStopped + } + copyErr := copyAuthorizedFile(part, input.Source, metadata.ByteLength, metadata.SHA256) + sourceGate.End(token) + if copyErr != nil { + return copyErr + } + } + return nil +} + +func copyAuthorizedFile(destination io.Writer, source io.Reader, expectedLength int64, expectedSHA256 string) error { + digest := sha256.New() + written, err := io.Copy(io.MultiWriter(destination, digest), io.LimitReader(source, expectedLength+1)) + if err != nil { + return errSourceTransferFailed + } + if written != expectedLength || !strings.EqualFold(hexDigest(digest), expectedSHA256) { + return errSourceChanged + } + return nil +} + +func (client *Client) validateResponse(envelope Response, manifest RequestManifest, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if envelope.ContractVersion != ContractVersion || envelope.DescriptorFingerprint != client.descriptor.Fingerprint || + envelope.PolicyFingerprint != client.descriptor.PolicyFingerprint || envelope.RequestChecksum != manifest.RequestChecksum || + len(envelope.Vectors) != len(inputs) { + return document.EmbeddingResult{}, errors.New("response identity mismatch") + } + vectors := make([]document.EmbeddingVector, len(inputs)) + seenKeys := make(map[string]struct{}, len(inputs)) + for position, vector := range envelope.Vectors { + if vector.Index == nil || *vector.Index != position || vector.Key != inputs[position].Key { + return document.EmbeddingResult{}, errors.New("response order mismatch") + } + if _, exists := seenKeys[vector.Key]; exists { + return document.EmbeddingResult{}, errors.New("response duplicate key") + } + seenKeys[vector.Key] = struct{}{} + if len(vector.Values) != client.descriptor.Dimension { + return document.EmbeddingResult{}, errors.New("response dimension mismatch") + } + for _, value := range vector.Values { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return document.EmbeddingResult{}, errors.New("response non-finite vector") + } + } + vectors[position] = document.EmbeddingVector{Key: vector.Key, Values: slices.Clone(vector.Values)} + } + result := document.EmbeddingResult{Vectors: vectors} + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, err + } + return result, nil +} + +func statusError(status int) error { + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return classified(ErrorAuthentication, status) + case http.StatusRequestEntityTooLarge: + return classified(ErrorCapacity, status) + case http.StatusRequestTimeout, http.StatusTooManyRequests, http.StatusInternalServerError, + http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: + return classified(ErrorTransient, status) + default: + return classified(ErrorPermanent, status) + } +} + +func requireResponseMediaType(value string) error { + mediaType, parameters, err := mime.ParseMediaType(value) + if err != nil || mediaType != "application/vnd.docbank.embedding-result+json" || len(parameters) != 1 || parameters["version"] != "1" { + return errors.New("response media type mismatch") + } + return nil +} + +func readBounded(reader io.Reader, maximum int64) ([]byte, error) { + value, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil || int64(len(value)) > maximum { + return nil, errors.New("bounded response read failed") + } + return value, nil +} + +type countingWriter struct{ written int64 } + +func (writer *countingWriter) Write(value []byte) (int, error) { + if int64(len(value)) > maxRequestBytes-writer.written { + return 0, errors.New("multipart request length overflow") + } + writer.written += int64(len(value)) + return len(value), nil +} + +func hexDigest(value hash.Hash) string { return hex.EncodeToString(value.Sum(nil)) } + +type frozenUpload struct { + document.AuthorizedUpload + + metadata document.AuthorizedUploadMetadata +} + +func (upload frozenUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func safeMultipartFilename(value string) bool { + return value != "" && utf8.ValidString(value) && strings.IndexFunc(value, unicode.IsControl) < 0 +} + +type activeSourceGate struct { + mu sync.Mutex + canceled bool + nextToken uint64 + activeToken uint64 + active document.AuthorizedUpload +} + +func newActiveSourceGate() *activeSourceGate { return new(activeSourceGate) } + +func (gate *activeSourceGate) Begin(source document.AuthorizedUpload) (uint64, bool) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.canceled { + return 0, false + } + gate.nextToken++ + gate.activeToken = gate.nextToken + gate.active = source + return gate.activeToken, true +} + +func (gate *activeSourceGate) End(token uint64) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.activeToken != token { + return + } + gate.active = nil + gate.activeToken = 0 +} + +func (gate *activeSourceGate) Cancel() { + gate.mu.Lock() + if gate.canceled { + gate.mu.Unlock() + return + } + gate.canceled = true + active := gate.active + gate.active = nil + gate.activeToken = 0 + gate.mu.Unlock() + if !nilInterface(active) { + _ = active.Close() + } +} diff --git a/document/embeddingbridge/client_test.go b/document/embeddingbridge/client_test.go new file mode 100644 index 00000000..2098204f --- /dev/null +++ b/document/embeddingbridge/client_test.go @@ -0,0 +1,1105 @@ +package embeddingbridge_test + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net" + "net/http" + "net/http/cookiejar" + "net/netip" + "net/url" + "slices" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/embeddingbridge" + "go.kenn.io/docbank/document/providerhttp" +) + +type secretMap map[string]string + +func (values secretMap) ResolveSecret(_ context.Context, binding string) (string, error) { + value, ok := values[binding] + if !ok { + return "", errors.New("synthetic resolver detail") + } + return value, nil +} + +type fixedResolver struct{ address netip.Addr } + +func (resolver fixedResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return []netip.Addr{resolver.address}, nil +} + +type bridgeFixture struct { + t *testing.T + server *http.Server + listener net.Listener + origin string + resolver fixedResolver + profile embeddingbridge.Profile + descriptor document.EmbeddingDescriptor + client *embeddingbridge.Client +} + +func newBridgeFixture(t *testing.T, handler http.Handler) bridgeFixture { + t.Helper() + listener, err := net.Listen("tcp4", "127.0.0.1:0") + require.NoError(t, err) + server := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second} + go func() { _ = server.Serve(listener) }() + t.Cleanup(func() { + require.NoError(t, server.Shutdown(context.Background())) + }) + tcpAddress, ok := listener.Addr().(*net.TCPAddr) + require.True(t, ok) + port := tcpAddress.Port + origin := "http://embedding.invalid:" + strconv.Itoa(port) + modelInput, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}) + require.NoError(t, err) + descriptor := document.EmbeddingDescriptor{ + ID: "synthetic.embedding-bridge", ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("0", 64), TrustBoundary: document.EmbeddingTrustOperatorNetwork, + Model: "synthetic-embed", ModelRevision: "immutable-r1", Dimension: 2, + Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationNone, + ScalarEncoding: "float32", DocumentFormatter: "synthetic/document-v1", + QueryFormatter: "synthetic/query-v1", + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk}, + CompatibilityID: modelInput.CompatibilityID, SupportsTextQuery: true, ModelInput: modelInput, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + } + descriptor, err = document.NewEmbeddingDescriptor(descriptor) + require.NoError(t, err) + profile := embeddingbridge.Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: "credential:synthetic-bridge", + EgressPolicy: providerhttp.EgressPolicy{ + Scheme: "http", Host: "embedding.invalid", Port: uint16(port), + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, + }, + RequestTimeout: time.Second, MaxBatchItems: 8, MaxInputBytes: 4096, + MaxRequestBytes: 16 << 10, MaxResponseBytes: 16 << 10, + } + fingerprint, err := embeddingbridge.PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + descriptor = profile.Descriptor + httpClient := &http.Client{} + client, err := embeddingbridge.New(profile, secretMap{"credential:synthetic-bridge": "synthetic-secret"}, fixedResolver{netip.MustParseAddr("127.0.0.1")}, httpClient) + require.NoError(t, err) + return bridgeFixture{t: t, server: server, listener: listener, origin: origin, resolver: fixedResolver{netip.MustParseAddr("127.0.0.1")}, profile: profile, descriptor: descriptor, client: client} +} + +func (fixture bridgeFixture) authorization(items int) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ + ProviderID: fixture.descriptor.ID, DescriptorFingerprint: fixture.descriptor.Fingerprint, + PolicyFingerprint: fixture.descriptor.PolicyFingerprint, MaxBatchItems: items, + MaxInputBytes: 4096, MaxResponseBytes: 4096, + } +} + +type requestManifest struct { + ContractVersion string `json:"contract_version"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + Authorization document.EmbeddingAuthorization `json:"authorization"` + Inputs []manifestInput `json:"inputs"` + RequestChecksum string `json:"request_checksum,omitempty"` +} + +type manifestInput struct { + Index int `json:"index"` + Key string `json:"key"` + Role document.EmbeddingRole `json:"role"` + Kind document.EmbeddingInputKind `json:"kind"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + Text string `json:"text,omitempty"` + HeadingPath []string `json:"heading_path,omitempty"` + SourceSpans []manifestSpan `json:"source_spans,omitempty"` + Upload *document.AuthorizedUploadMetadata `json:"upload,omitempty"` + FilePart string `json:"file_part,omitempty"` + FileIndex *int `json:"file_index,omitempty"` +} + +type manifestSpan struct { + UnitIndex int `json:"unit_index"` + CharStart int `json:"char_start"` + CharEnd int `json:"char_end"` +} + +type responseEnvelope struct { + ContractVersion string `json:"contract_version"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + RequestChecksum string `json:"request_checksum"` + Vectors []document.EmbeddingVector `json:"vectors"` +} + +func readBridgeRequest(t *testing.T, request *http.Request) (requestManifest, [][]byte) { + t.Helper() + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, parameters["boundary"]) + var manifest requestManifest + var files [][]byte + for { + part, err := reader.NextPart() + if errors.Is(err, io.EOF) { + break + } + require.NoError(t, err) + payload, err := io.ReadAll(part) + require.NoError(t, err) + switch part.FormName() { + case "manifest": + require.Equal(t, "application/vnd.docbank.embedding-manifest+json;version=1", part.Header.Get("Content-Type")) + require.NoError(t, json.Unmarshal(payload, &manifest, json.RejectUnknownMembers(true))) + case "file": + files = append(files, payload) + default: + t.Fatalf("unexpected multipart field %q", part.FormName()) + } + } + return manifest, files +} + +func writeSuccess(t *testing.T, writer http.ResponseWriter, manifest requestManifest, vectors []document.EmbeddingVector) { + t.Helper() + writer.Header().Set("Content-Type", "application/vnd.docbank.embedding-result+json;version=1") + payload, err := json.Marshal(responseEnvelope{ + ContractVersion: embeddingbridge.ContractVersion, + DescriptorFingerprint: manifest.DescriptorFingerprint, + PolicyFingerprint: manifest.PolicyFingerprint, + RequestChecksum: manifest.RequestChecksum, + Vectors: vectors, + }) + require.NoError(t, err) + _, err = writer.Write(payload) + require.NoError(t, err) +} + +func TestTextDocumentAndQueryManifestBindsExactRenderedBytesAndChecksum(t *testing.T) { + var captured requestManifest + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + assert.Equal(t, "/docbank-embedding/v1/embeddings", request.URL.Path) + assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) + assert.Equal(t, request.Header.Get("Idempotency-Key"), request.Header.Get("Docbank-Request-Checksum")) + manifest, files := readBridgeRequest(t, request) + captured = manifest + assert.Empty(t, files) + zero, one := 0, 1 + writeSuccess(t, writer, manifest, []document.EmbeddingVector{ + {Key: "chunk-a", Index: &zero, Values: []float32{0.25, 0.5}}, + {Key: "query-a", Index: &one, Values: []float32{0.75, 1}}, + }) + })) + inputs := []document.EmbeddingInput{ + { + Key: "chunk-a", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "alpha", + HeadingPath: []string{"Overview"}, + SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 5}}, + }, + {Key: "query-a", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "beta"}, + } + authorization := fixture.authorization(2) + result, err := fixture.client.Embed(context.Background(), inputs, authorization) + require.NoError(t, err) + require.Len(t, result.Vectors, 2) + assert.Equal(t, []float32{0.25, 0.5}, result.Vectors[0].Values) + assert.Nil(t, result.Vectors[0].Index, "the provider result is normalized to request order") + + assert.Equal(t, embeddingbridge.ContractVersion, captured.ContractVersion) + assert.Equal(t, fixture.descriptor.Fingerprint, captured.DescriptorFingerprint) + assert.Equal(t, fixture.descriptor.PolicyFingerprint, captured.PolicyFingerprint) + assert.Equal(t, authorization, captured.Authorization) + require.Len(t, captured.Inputs, 2) + assert.Equal(t, "search_document: alpha", captured.Inputs[0].Text) + assert.Equal(t, int64(22), captured.Inputs[0].ByteLength) + assert.Equal(t, sha256Text("search_document: alpha"), captured.Inputs[0].SHA256) + assert.Equal(t, []string{"Overview"}, captured.Inputs[0].HeadingPath) + assert.Equal(t, []manifestSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 5}}, captured.Inputs[0].SourceSpans) + assert.Equal(t, "search_query: beta", captured.Inputs[1].Text) + assert.Equal(t, int64(18), captured.Inputs[1].ByteLength) + assert.Equal(t, sha256Text("search_query: beta"), captured.Inputs[1].SHA256) + assert.Equal(t, exactManifestChecksum(t, captured), captured.RequestChecksum) + assert.NotContains(t, mustJSON(t, captured), "node_id") + assert.NotContains(t, mustJSON(t, captured), "tags") + assert.NotContains(t, mustJSON(t, captured), "vault") + assert.Contains(t, mustJSON(t, captured), `"unit_index":0`) +} + +func TestExplicitNoAuthProfileUsesNoResolverOrAuthorizationHeader(t *testing.T) { + var requests atomic.Int32 + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requests.Add(1) + assert.Empty(t, request.Header.Get("Authorization")) + manifest, _ := readBridgeRequest(t, request) + zero := 0 + writeSuccess(t, writer, manifest, []document.EmbeddingVector{{Key: "first", Index: &zero, Values: []float32{1, 2}}}) + })) + profile := fixture.profile + profile.SecretBinding = "" + fingerprint, err := embeddingbridge.PolicyFingerprint(profile) + require.NoError(t, err) + assert.NotEqual(t, fixture.descriptor.PolicyFingerprint, fingerprint) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + client, err := embeddingbridge.New(profile, nil, fixture.resolver, &http.Client{}) + require.NoError(t, err) + _, err = client.Embed(context.Background(), oneTextInput("alpha"), document.EmbeddingAuthorization{ + ProviderID: profile.Descriptor.ID, DescriptorFingerprint: profile.Descriptor.Fingerprint, + PolicyFingerprint: profile.Descriptor.PolicyFingerprint, MaxBatchItems: 1, + MaxInputBytes: 4096, MaxResponseBytes: 4096, + }) + require.NoError(t, err) + assert.Equal(t, int32(1), requests.Load()) + + _, err = embeddingbridge.New(profile, secretMap{}, fixture.resolver, &http.Client{}) + require.Error(t, err, "an empty binding cannot carry an ambient resolver") + _, err = embeddingbridge.New(fixture.profile, nil, fixture.resolver, &http.Client{}) + require.Error(t, err, "a named binding requires its narrow resolver") + hosted := profile + hosted.Descriptor.TrustBoundary = document.EmbeddingTrustHostedProvider + _, err = embeddingbridge.PolicyFingerprint(hosted) + require.Error(t, err, "anonymous profiles are operator-network only") +} + +func TestMixedTextAndDirectFileRequestStreamsOnlyAuthorizedBytes(t *testing.T) { + source := []byte("synthetic direct file") + metadata := uploadMetadata(source) + jar, err := cookiejar.New(nil) + require.NoError(t, err) + parsed, err := url.Parse("http://embedding.invalid") + require.NoError(t, err) + jar.SetCookies(parsed, []*http.Cookie{{ //nolint:gosec // Deliberately ambient insecure test cookie; isolation must strip it. + Name: "ambient", Value: "PRIVATE_COOKIE", + }}) + var manifest requestManifest + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + assert.Empty(t, request.Header.Get("Cookie")) + var files [][]byte + manifest, files = readBridgeRequest(t, request) + assert.Equal(t, [][]byte{source}, files) + zero, one := 0, 1 + writeSuccess(t, writer, manifest, []document.EmbeddingVector{ + {Key: "source-a", Index: &zero, Values: []float32{1, 2}}, + {Key: "query-a", Index: &one, Values: []float32{3, 4}}, + }) + })) + client, err := embeddingbridge.New(fixture.profile, secretMap{"credential:synthetic-bridge": "synthetic-secret"}, fixture.resolver, &http.Client{Jar: jar}) + require.NoError(t, err) + upload := &testUpload{Reader: bytes.NewReader(source), metadata: metadata} + result, err := client.Embed(context.Background(), []document.EmbeddingInput{ + {Key: "source-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: upload}, + {Key: "query-a", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "needle"}, + }, fixture.authorization(2)) + require.NoError(t, err) + require.Len(t, result.Vectors, 2) + require.Len(t, manifest.Inputs, 2) + assert.Equal(t, &metadata, manifest.Inputs[0].Upload) + assert.Equal(t, "file", manifest.Inputs[0].FilePart) + require.NotNil(t, manifest.Inputs[0].FileIndex) + assert.Zero(t, *manifest.Inputs[0].FileIndex) + assert.Empty(t, manifest.Inputs[0].Text) + assert.Equal(t, int64(len(source)), manifest.Inputs[0].ByteLength) + assert.Equal(t, metadata.SHA256, manifest.Inputs[0].SHA256) +} + +func TestResponseContractFailsClosed(t *testing.T) { + tests := []struct { + name string + contentType string + body func(requestManifest) string + category embeddingbridge.ErrorCategory + }{ + {name: "wrong media type", contentType: "application/json", body: validResponseBody, category: embeddingbridge.ErrorMalformedResponse}, + {name: "unknown major", contentType: "application/vnd.docbank.embedding-result+json;version=2", body: func(manifest requestManifest) string { + return strings.Replace(validResponseBody(manifest), embeddingbridge.ContractVersion, "docbank-embedding/v2", 1) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "descriptor drift", contentType: responseMediaType(), body: func(manifest requestManifest) string { + manifest.DescriptorFingerprint = strings.Repeat("d", 64) + return validResponseBody(manifest) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "policy drift", contentType: responseMediaType(), body: func(manifest requestManifest) string { + manifest.PolicyFingerprint = strings.Repeat("e", 64) + return validResponseBody(manifest) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "request drift", contentType: responseMediaType(), body: func(manifest requestManifest) string { + manifest.RequestChecksum = strings.Repeat("f", 64) + return validResponseBody(manifest) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "wrong order", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"second","index":1,"values":[1,2]},{"key":"first","index":0,"values":[3,4]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "wrong index", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":1,"values":[1,2]},{"key":"second","index":0,"values":[3,4]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "dimension", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1]},{"key":"second","index":1,"values":[3,4]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "non finite", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1e1000,2]},{"key":"second","index":1,"values":[3,4]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "duplicate key", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1,2]},{"key":"first","index":1,"values":[3,4]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "missing key", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1,2]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "extra key", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1,2]},{"key":"second","index":1,"values":[3,4]},{"key":"extra","index":2,"values":[5,6]}]`) + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "unknown field", contentType: responseMediaType(), body: func(manifest requestManifest) string { + return strings.TrimSuffix(validResponseBody(manifest), "}") + `,"provider_body":"PRIVATE_RESPONSE"}` + }, category: embeddingbridge.ErrorMalformedResponse}, + {name: "trailing JSON", contentType: responseMediaType(), body: func(manifest requestManifest) string { return validResponseBody(manifest) + `{}` }, category: embeddingbridge.ErrorMalformedResponse}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + manifest, _ := readBridgeRequest(t, request) + writer.Header().Set("Content-Type", test.contentType) + _, _ = io.WriteString(writer, test.body(manifest)) + })) + _, err := fixture.client.Embed(context.Background(), twoTextInputs(), fixture.authorization(2)) + require.Error(t, err) + assert.Equal(t, test.category, embeddingbridge.Category(err)) + assert.NotContains(t, err.Error(), "PRIVATE_RESPONSE") + }) + } +} + +func TestResponseByteLimitIsEnforcedBeforeDecode(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + _, _ = readBridgeRequest(t, request) + writer.Header().Set("Content-Type", responseMediaType()) + _, _ = writer.Write(bytes.Repeat([]byte("x"), 16<<10+1)) + })) + _, err := fixture.client.Embed(context.Background(), twoTextInputs(), fixture.authorization(2)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorMalformedResponse, embeddingbridge.Category(err)) +} + +func TestProfileFingerprintFreezesDescriptorModelOriginEgressBindingAndBounds(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + })) + mutations := []struct { + name string + mutate func(*embeddingbridge.Profile) + }{ + {"model revision", func(profile *embeddingbridge.Profile) { profile.Descriptor.ModelRevision = "immutable-r2" }}, + {"model input", func(profile *embeddingbridge.Profile) { + profile.Descriptor.ModelInput.Query.Template = "query: {{content}}" + }}, + {"compatibility", func(profile *embeddingbridge.Profile) { profile.Descriptor.CompatibilityID = "changed/v1" }}, + {"binding", func(profile *embeddingbridge.Profile) { profile.SecretBinding = "credential:other" }}, + {"origin", func(profile *embeddingbridge.Profile) { profile.Origin += "/other" }}, + {"egress", func(profile *embeddingbridge.Profile) { + profile.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")} + }}, + {"bounds", func(profile *embeddingbridge.Profile) { profile.MaxBatchItems-- }}, + } + for _, mutation := range mutations { + t.Run(mutation.name, func(t *testing.T) { + profile := fixture.profile + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + mutation.mutate(&profile) + _, err := embeddingbridge.New(profile, secretMap{"credential:synthetic-bridge": "synthetic-secret"}, fixture.resolver, &http.Client{}) + require.Error(t, err) + }) + } +} + +func TestSecretResolutionRedirectCancellationAndAmbiguousSubmission(t *testing.T) { + t.Run("secret resolution", func(t *testing.T) { + var requests int + fixture := newBridgeFixture(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) { requests++ })) + client, err := embeddingbridge.New(fixture.profile, secretMap{}, fixture.resolver, &http.Client{}) + require.NoError(t, err) + _, err = client.Embed(context.Background(), oneTextInput("private input"), fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorAuthentication, embeddingbridge.Category(err)) + assert.Zero(t, requests) + assert.NotContains(t, err.Error(), "synthetic resolver detail") + assert.NotContains(t, err.Error(), "private input") + }) + t.Run("redirect", func(t *testing.T) { + var redirected int + mux := http.NewServeMux() + mux.HandleFunc("/docbank-embedding/v1/embeddings", func(writer http.ResponseWriter, request *http.Request) { + http.Redirect(writer, request, "/redirected", http.StatusTemporaryRedirect) + }) + mux.HandleFunc("/redirected", func(http.ResponseWriter, *http.Request) { redirected++ }) + fixture := newBridgeFixture(t, mux) + _, err := fixture.client.Embed(context.Background(), oneTextInput("alpha"), fixture.authorization(1)) + require.Error(t, err) + assert.Zero(t, redirected) + assert.Equal(t, embeddingbridge.ErrorPermanent, embeddingbridge.Category(err)) + }) + t.Run("cancellation", func(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + fixture := newBridgeFixture(t, http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + close(started) + <-release + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := fixture.client.Embed(ctx, oneTextInput("alpha"), fixture.authorization(1)) + done <- err + }() + <-started + cancel() + err := <-done + close(release) + require.ErrorIs(t, err, context.Canceled) + }) + t.Run("ambiguous submission", func(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + _, _ = readBridgeRequest(t, request) + hijacker, ok := writer.(http.Hijacker) + if !assert.True(t, ok) { + return + } + connection, _, err := hijacker.Hijack() + if !assert.NoError(t, err) { + return + } + _ = connection.Close() + })) + _, err := fixture.client.Embed(context.Background(), oneTextInput("alpha"), fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorAmbiguousSubmission, embeddingbridge.Category(err)) + assert.True(t, embeddingbridge.IsRetryable(err)) + }) +} + +func TestDirectFileTerminalLengthAndChecksumAreProven(t *testing.T) { + expected := []byte("12345") + tests := []struct { + name string + actual []byte + metadata document.AuthorizedUploadMetadata + }{ + {name: "short", actual: []byte("1234"), metadata: uploadMetadata(expected)}, + {name: "long", actual: []byte("123456"), metadata: uploadMetadata(expected)}, + {name: "substituted", actual: []byte("abcde"), metadata: uploadMetadata(expected)}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + writer.Header().Set("Content-Type", responseMediaType()) + _, _ = io.WriteString(writer, `{}`) + })) + upload := &testUpload{Reader: bytes.NewReader(test.actual), metadata: test.metadata} + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorSourceChanged, embeddingbridge.Category(err)) + assert.NotContains(t, err.Error(), string(test.actual)) + }) + } +} + +func TestUploadMetadataIsFrozenOnceBeforeValidationAndTransmission(t *testing.T) { + source := []byte("immutable synthetic source") + upload := &countingUpload{Reader: bytes.NewReader(source), metadata: uploadMetadata(source)} + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + manifest, files := readBridgeRequest(t, request) + assert.Equal(t, [][]byte{source}, files) + zero := 0 + writeSuccess(t, writer, manifest, []document.EmbeddingVector{{Key: "source", Index: &zero, Values: []float32{1, 2}}}) + })) + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + require.NoError(t, err) + assert.Equal(t, int32(1), upload.calls.Load()) +} + +func TestDirectFileRejectsHeaderUnsafeFilenameBeforeRequest(t *testing.T) { + var requests atomic.Int32 + source := []byte("synthetic source") + metadata := uploadMetadata(source) + metadata.Filename = "synthetic.bin\r\nX-Injected: private" + fixture := newBridgeFixture(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + requests.Add(1) + })) + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, + Source: &testUpload{Reader: bytes.NewReader(source), metadata: metadata}, + }}, fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorPermanent, embeddingbridge.Category(err)) + assert.Zero(t, requests.Load()) + assert.NotContains(t, err.Error(), "X-Injected") +} + +func TestCancellationClosesBlockedDirectFileAndReturns(t *testing.T) { + source := []byte("synthetic blocked source") + upload := newBlockingUpload(uploadMetadata(source)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := fixture.client.Embed(ctx, []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + done <- err + }() + <-upload.started + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(250 * time.Millisecond): + _ = upload.Close() + err := <-done + assert.Fail(t, "embedding bridge did not close the blocked authorized upload", "eventual error: %v", err) + } + assert.Positive(t, upload.closes.Load()) +} + +func TestCancellationBeforeSecretResolutionLeavesUnreadSourceOpen(t *testing.T) { + source := []byte("synthetic unread source") + upload := newObservedUpload(source, uploadMetadata(source)) + resolver := &cancelBlockingSecretResolver{started: make(chan struct{})} + fixture := newBridgeFixture(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("request reached server before secret resolution completed") + })) + client, err := embeddingbridge.New(fixture.profile, resolver, fixture.resolver, &http.Client{}) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, embedErr := client.Embed(ctx, []document.EmbeddingInput{{ + Key: "unread", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + done <- embedErr + }() + <-resolver.started + cancel() + require.ErrorIs(t, <-done, context.Canceled) + assert.Zero(t, upload.reads.Load()) + assert.Zero(t, upload.closes.Load(), "the worker retains ownership until a source read begins") +} + +func TestMultiFileCancellationClosesOnlyActiveSource(t *testing.T) { + firstBytes := []byte("synthetic first source") + secondBytes := []byte("synthetic blocked second source") + thirdBytes := []byte("synthetic unread third source") + first := newObservedUpload(firstBytes, uploadMetadata(firstBytes)) + second := newBlockingUpload(uploadMetadata(secondBytes)) + third := newObservedUpload(thirdBytes, uploadMetadata(thirdBytes)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, embedErr := fixture.client.Embed(ctx, []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: first}, + {Key: "second", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: second}, + {Key: "third", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: third}, + }, fixture.authorization(3)) + done <- embedErr + }() + <-second.started + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(250 * time.Millisecond): + _ = second.Close() + err := <-done + assert.Fail(t, "embedding bridge did not close the active source", "eventual error: %v", err) + } + assert.Zero(t, first.closes.Load(), "a completed source is no longer owned by the bridge") + assert.Positive(t, second.closes.Load(), "the active blocked source must be closed") + assert.Zero(t, third.reads.Load(), "later sources must remain unread") + assert.Zero(t, third.closes.Load(), "later unread sources remain worker-owned") +} + +func TestCancellationWithNonComparableActiveSourceDoesNotPanic(t *testing.T) { + source := []byte("synthetic non-comparable source") + state := newBlockingUpload(uploadMetadata(source)) + upload := nonComparableUpload{state} + fixture := newBridgeFixture(t, http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, embedErr := fixture.client.Embed(ctx, []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + done <- embedErr + }() + <-state.started + cancel() + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(250 * time.Millisecond): + _ = state.Close() + err := <-done + assert.Fail(t, "embedding bridge did not close the active non-comparable source", "eventual error: %v", err) + } +} + +func TestPolicyFingerprintDoesNotMutateCallerEgressOrder(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + profile := fixture.profile + profile.EgressPolicy.AllowedCIDRs = []netip.Prefix{ + netip.MustParsePrefix("127.0.0.2/32"), + netip.MustParsePrefix("127.0.0.0/24"), + } + original := slices.Clone(profile.EgressPolicy.AllowedCIDRs) + fingerprint, err := embeddingbridge.PolicyFingerprint(profile) + require.NoError(t, err) + assert.Equal(t, original, profile.EgressPolicy.AllowedCIDRs) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + _, err = embeddingbridge.New(profile, secretMap{"credential:synthetic-bridge": "synthetic-secret"}, fixture.resolver, &http.Client{}) + require.NoError(t, err) + assert.Equal(t, original, profile.EgressPolicy.AllowedCIDRs) +} + +func TestHTTPFailuresAreClassifiedWithoutProviderBodiesOrSecrets(t *testing.T) { + tests := []struct { + status int + category embeddingbridge.ErrorCategory + retry bool + }{ + {http.StatusUnauthorized, embeddingbridge.ErrorAuthentication, false}, + {http.StatusRequestEntityTooLarge, embeddingbridge.ErrorCapacity, false}, + {http.StatusTooManyRequests, embeddingbridge.ErrorTransient, true}, + {http.StatusServiceUnavailable, embeddingbridge.ErrorTransient, true}, + {http.StatusBadRequest, embeddingbridge.ErrorPermanent, false}, + } + for _, test := range tests { + t.Run(strconv.Itoa(test.status), func(t *testing.T) { + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + _, _ = readBridgeRequest(t, request) + writer.WriteHeader(test.status) + _, _ = io.WriteString(writer, "PRIVATE_RESPONSE synthetic-secret private-input") + })) + _, err := fixture.client.Embed(context.Background(), oneTextInput("private-input"), fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, test.category, embeddingbridge.Category(err)) + assert.Equal(t, test.retry, embeddingbridge.IsRetryable(err)) + assert.NotContains(t, err.Error(), "PRIVATE_RESPONSE") + assert.NotContains(t, err.Error(), "synthetic-secret") + assert.NotContains(t, err.Error(), "private-input") + }) + } +} + +func TestEarlyHTTPRejectionDoesNotReadDirectFileAndKeepsStatusClassification(t *testing.T) { + source := []byte("synthetic early rejection") + upload := newBlockingUpload(uploadMetadata(source)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + assert.Equal(t, "100-continue", request.Header.Get("Expect")) + writer.WriteHeader(http.StatusRequestEntityTooLarge) + })) + done := make(chan error, 1) + go func() { + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + done <- err + }() + select { + case err := <-done: + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorCapacity, embeddingbridge.Category(err)) + case <-time.After(time.Second): + _ = upload.Close() + <-done + assert.Fail(t, "early rejection did not return within the request bound") + } + assert.Zero(t, upload.closes.Load(), "the worker retains ownership when no source read starts") + select { + case <-upload.started: + assert.Fail(t, "early rejection read the direct-file source") + default: + } +} + +func TestDirectFileConnectionLossAfterPartHeaderIsAmbiguous(t *testing.T) { + source := []byte("synthetic ambiguous source") + upload := newGatedUpload(source, uploadMetadata(source)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + defer upload.Release() + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + if !assert.NoError(t, err) || !assert.Equal(t, "multipart/form-data", mediaType) { + return + } + reader := multipart.NewReader(request.Body, parameters["boundary"]) + manifestPart, err := reader.NextPart() + if !assert.NoError(t, err) { + return + } + _, err = io.Copy(io.Discard, manifestPart) + if !assert.NoError(t, err) { + return + } + filePart, err := reader.NextPart() + if !assert.NoError(t, err) || !assert.Equal(t, "file", filePart.FormName()) { + return + } + hijacker, ok := writer.(http.Hijacker) + if !assert.True(t, ok) { + return + } + connection, _, err := hijacker.Hijack() + if !assert.NoError(t, err) { + return + } + _ = connection.Close() + })) + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorAmbiguousSubmission, embeddingbridge.Category(err)) + assert.True(t, embeddingbridge.IsRetryable(err)) +} + +func TestDirectFileConnectionLossDuringBlockedReadIsAmbiguous(t *testing.T) { + source := []byte("synthetic blocked connection-loss source") + upload := newBlockingUpload(uploadMetadata(source)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + if !assert.NoError(t, err) || !assert.Equal(t, "multipart/form-data", mediaType) { + return + } + reader := multipart.NewReader(request.Body, parameters["boundary"]) + manifestPart, err := reader.NextPart() + if !assert.NoError(t, err) { + return + } + _, err = io.Copy(io.Discard, manifestPart) + if !assert.NoError(t, err) { + return + } + filePart, err := reader.NextPart() + if !assert.NoError(t, err) || !assert.Equal(t, "file", filePart.FormName()) { + return + } + <-upload.started + writer.Header().Set("Content-Type", responseMediaType()) + writer.Header().Set("Content-Length", "64") + writer.WriteHeader(http.StatusOK) + flusher, ok := writer.(http.Flusher) + if !assert.True(t, ok) { + return + } + flusher.Flush() + hijacker, ok := writer.(http.Hijacker) + if !assert.True(t, ok) { + return + } + connection, _, err := hijacker.Hijack() + if !assert.NoError(t, err) { + return + } + _ = connection.Close() + })) + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorAmbiguousSubmission, embeddingbridge.Category(err)) + assert.True(t, embeddingbridge.IsRetryable(err)) + assert.NotContains(t, err.Error(), "synthetic blocked reader closed") +} + +func TestKnownHTTPStatusDuringBlockedUploadWinsWriterError(t *testing.T) { + source := []byte("synthetic blocked rejected source") + upload := newBlockingUpload(uploadMetadata(source)) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + if !assert.NoError(t, err) || !assert.Equal(t, "multipart/form-data", mediaType) { + return + } + reader := multipart.NewReader(request.Body, parameters["boundary"]) + manifestPart, err := reader.NextPart() + if !assert.NoError(t, err) { + return + } + _, err = io.Copy(io.Discard, manifestPart) + if !assert.NoError(t, err) { + return + } + filePart, err := reader.NextPart() + if !assert.NoError(t, err) || !assert.Equal(t, "file", filePart.FormName()) { + return + } + <-upload.started + writer.WriteHeader(http.StatusRequestEntityTooLarge) + })) + _, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{{ + Key: "source", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: upload, + }}, fixture.authorization(1)) + require.Error(t, err) + assert.Equal(t, embeddingbridge.ErrorCapacity, embeddingbridge.Category(err)) + assert.False(t, embeddingbridge.IsRetryable(err)) + assert.NotContains(t, err.Error(), "synthetic blocked reader closed") + assert.Positive(t, upload.closes.Load()) +} + +func TestConcurrentRequestsDoNotShareSecretsOrState(t *testing.T) { + var mu sync.Mutex + checksums := make(map[string]bool) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + manifest, _ := readBridgeRequest(t, request) + mu.Lock() + checksums[manifest.RequestChecksum] = true + mu.Unlock() + zero := 0 + writeSuccess(t, writer, manifest, []document.EmbeddingVector{{Key: manifest.Inputs[0].Key, Index: &zero, Values: []float32{1, 2}}}) + })) + var wait sync.WaitGroup + errorsSeen := make(chan error, 2) + for _, value := range []string{"alpha", "beta"} { + wait.Go(func() { + _, err := fixture.client.Embed(context.Background(), oneTextInput(value), fixture.authorization(1)) + errorsSeen <- err + }) + } + wait.Wait() + close(errorsSeen) + for err := range errorsSeen { + require.NoError(t, err) + } + assert.Len(t, checksums, 2) +} + +type testUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func (upload *testUpload) Close() error { return nil } + +type countingUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata + calls atomic.Int32 +} + +func (upload *countingUpload) Metadata() document.AuthorizedUploadMetadata { + upload.calls.Add(1) + return upload.metadata +} + +func (upload *countingUpload) Close() error { return nil } + +type blockingUpload struct { + metadata document.AuthorizedUploadMetadata + started chan struct{} + closed chan struct{} + start sync.Once + close sync.Once + closes atomic.Int32 +} + +type observedUpload struct { + reader *bytes.Reader + metadata document.AuthorizedUploadMetadata + reads atomic.Int32 + closes atomic.Int32 +} + +type cancelBlockingSecretResolver struct { + started chan struct{} + once sync.Once +} + +type nonComparableUpload []*blockingUpload + +func (upload nonComparableUpload) Read(value []byte) (int, error) { + return upload[0].Read(value) +} + +func (upload nonComparableUpload) Metadata() document.AuthorizedUploadMetadata { + return upload[0].Metadata() +} + +func (upload nonComparableUpload) Close() error { return upload[0].Close() } + +type gatedUpload struct { + reader *bytes.Reader + metadata document.AuthorizedUploadMetadata + release chan struct{} + releaseOnce sync.Once + waitOnce sync.Once +} + +func newGatedUpload(source []byte, metadata document.AuthorizedUploadMetadata) *gatedUpload { + return &gatedUpload{reader: bytes.NewReader(source), metadata: metadata, release: make(chan struct{})} +} + +func (upload *gatedUpload) Read(value []byte) (int, error) { + upload.waitOnce.Do(func() { <-upload.release }) + read, err := upload.reader.Read(value) + if errors.Is(err, io.EOF) { + return read, io.EOF + } + if err != nil { + return read, fmt.Errorf("gated upload read: %w", err) + } + return read, nil +} + +func (upload *gatedUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func (upload *gatedUpload) Release() { upload.releaseOnce.Do(func() { close(upload.release) }) } + +func (upload *gatedUpload) Close() error { + upload.Release() + return nil +} + +func newBlockingUpload(metadata document.AuthorizedUploadMetadata) *blockingUpload { + return &blockingUpload{metadata: metadata, started: make(chan struct{}), closed: make(chan struct{})} +} + +func newObservedUpload(source []byte, metadata document.AuthorizedUploadMetadata) *observedUpload { + return &observedUpload{reader: bytes.NewReader(source), metadata: metadata} +} + +func (upload *observedUpload) Read(value []byte) (int, error) { + upload.reads.Add(1) + read, err := upload.reader.Read(value) + if errors.Is(err, io.EOF) { + return read, io.EOF + } + if err != nil { + return read, fmt.Errorf("observed upload read: %w", err) + } + return read, nil +} + +func (upload *observedUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func (upload *observedUpload) Close() error { + upload.closes.Add(1) + return nil +} + +func (resolver *cancelBlockingSecretResolver) ResolveSecret(ctx context.Context, _ string) (string, error) { + resolver.once.Do(func() { close(resolver.started) }) + <-ctx.Done() + return "", ctx.Err() +} + +func (upload *blockingUpload) Read([]byte) (int, error) { + upload.start.Do(func() { close(upload.started) }) + <-upload.closed + return 0, errors.New("synthetic blocked reader closed") +} + +func (upload *blockingUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func (upload *blockingUpload) Close() error { + upload.closes.Add(1) + upload.close.Do(func() { close(upload.closed) }) + return nil +} + +func uploadMetadata(source []byte) document.AuthorizedUploadMetadata { + return document.AuthorizedUploadMetadata{ + Filename: "synthetic.bin", MediaFamily: "binary", MediaType: "application/octet-stream", + ByteLength: int64(len(source)), SHA256: sha256Bytes(source), + CapabilityRecordChecksum: strings.Repeat("a", 64), ProviderMetadataChecksum: strings.Repeat("b", 64), + InputKind: document.RenditionInputOriginalFile, + } +} + +func oneTextInput(text string) []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: text}} +} + +func twoTextInputs() []document.EmbeddingInput { + return []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}, + {Key: "second", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "beta"}, + } +} + +func validResponseBody(manifest requestManifest) string { + return responseBody(manifest, `[{"key":"first","index":0,"values":[1,2]},{"key":"second","index":1,"values":[3,4]}]`) +} + +func responseBody(manifest requestManifest, vectors string) string { + return fmt.Sprintf(`{"contract_version":%q,"descriptor_fingerprint":%q,"policy_fingerprint":%q,"request_checksum":%q,"vectors":%s}`, + embeddingbridge.ContractVersion, manifest.DescriptorFingerprint, manifest.PolicyFingerprint, manifest.RequestChecksum, vectors) +} + +func responseMediaType() string { return "application/vnd.docbank.embedding-result+json;version=1" } + +func exactManifestChecksum(t *testing.T, manifest requestManifest) string { + t.Helper() + manifest.RequestChecksum = "" + encoded, err := json.Marshal(manifest, json.Deterministic(true)) + require.NoError(t, err) + return sha256Bytes(encoded) +} + +func mustJSON(t *testing.T, value any) string { + t.Helper() + encoded, err := json.Marshal(value, json.Deterministic(true)) + require.NoError(t, err) + return string(encoded) +} + +func sha256Text(value string) string { return sha256Bytes([]byte(value)) } + +func sha256Bytes(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/document/embeddingbridge/contract_test.go b/document/embeddingbridge/contract_test.go new file mode 100644 index 00000000..bf0a0582 --- /dev/null +++ b/document/embeddingbridge/contract_test.go @@ -0,0 +1,443 @@ +package embeddingbridge_test + +import ( + "bytes" + "context" + "crypto/sha256" + _ "embed" + "encoding/hex" + stdjson "encoding/json" + "encoding/json/jsontext" + "encoding/json/v2" + "errors" + "fmt" + "io" + "maps" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "testing" + + "github.com/santhosh-tekuri/jsonschema/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/embeddingbridge" + "gopkg.in/yaml.v3" +) + +//go:embed openapi.yaml +var openAPIContract []byte + +//go:embed embedding-request-v1.schema.json +var requestSchema []byte + +//go:embed embedding-response-v1.schema.json +var responseSchema []byte + +func TestNormativeEmbeddingContractIsFixedStrictAndVersioned(t *testing.T) { + var openAPI map[string]any + require.NoError(t, yaml.Unmarshal(openAPIContract, &openAPI)) + assert.Equal(t, "3.1.0", openAPI["openapi"]) + paths, ok := openAPI["paths"].(map[string]any) + require.True(t, ok) + assert.Len(t, paths, 1) + assert.Contains(t, paths, "/docbank-embedding/v1/embeddings") + for _, forbidden := range []string{"{{", "result_url", "callback", "status_url", "cancel_url"} { + assert.NotContains(t, string(openAPIContract), forbidden) + } + for name, schema := range map[string][]byte{"request": requestSchema, "response": responseSchema} { + t.Run(name, func(t *testing.T) { + var decoded map[string]any + require.NoError(t, json.Unmarshal(schema, &decoded)) + assert.Equal(t, "https://json-schema.org/draft/2020-12/schema", decoded["$schema"]) + assert.Equal(t, false, decoded["additionalProperties"]) + assert.True(t, bytes.Contains(schema, []byte(`"docbank-embedding/v1"`))) + }) + } + assert.Equal(t, "docbank-embedding/v1", embeddingbridge.ContractVersion) + for _, field := range []string{`"heading_path"`, `"source_spans"`, `"unit_index"`, `"char_start"`, `"char_end"`} { + assert.Contains(t, string(requestSchema), field) + } + var requestContract struct { + Canonicalization canonicalizationContract `json:"x-docbank-canonicalization"` + } + require.NoError(t, json.Unmarshal(requestSchema, &requestContract)) + assert.Equal(t, "docbank-json-v1", requestContract.Canonicalization.Name) + assert.Equal(t, "UTF-8 without a byte-order mark", requestContract.Canonicalization.CharacterEncoding) + assert.NotEmpty(t, requestContract.Canonicalization.Whitespace) + assert.NotEmpty(t, requestContract.Canonicalization.Arrays) + assert.NotEmpty(t, requestContract.Canonicalization.Integers) + assert.NotEmpty(t, requestContract.Canonicalization.Strings) + assert.NotEmpty(t, requestContract.Canonicalization.OptionalMembers) + assert.NotEmpty(t, requestContract.Canonicalization.Checksum) + assert.Equal(t, []string{ + "contract_version", "descriptor_fingerprint", "policy_fingerprint", "authorization", "inputs", "request_checksum", + }, requestContract.Canonicalization.ObjectMemberOrder["$"]) + require.NotEmpty(t, requestContract.Canonicalization.ChecksumVectors) + for _, vector := range requestContract.Canonicalization.ChecksumVectors { + _, wire, identity, err := canonicalizeIndependentManifest([]byte(vector.CanonicalManifest), requestContract.Canonicalization) + require.NoError(t, err) + assert.Equal(t, vector.CanonicalManifest, string(wire)) + assert.Equal(t, vector.CanonicalManifest, string(identity)) + digest := sha256.Sum256([]byte(vector.CanonicalManifest)) + assert.Equal(t, vector.SHA256, hex.EncodeToString(digest[:])) + assert.NotContains(t, vector.CanonicalManifest, `"request_checksum"`) + } +} + +type canonicalizationContract struct { + Name string `json:"name"` + CharacterEncoding string `json:"character_encoding"` + Whitespace string `json:"whitespace"` + ObjectMemberOrder map[string][]string `json:"object_member_order"` + Arrays string `json:"arrays"` + Integers string `json:"integers"` + Strings string `json:"strings"` + OptionalMembers string `json:"optional_members"` + Checksum string `json:"checksum"` + ChecksumVectors []struct { + CanonicalManifest string `json:"canonical_manifest"` + SHA256 string `json:"sha256"` + } `json:"checksum_vectors"` +} + +func TestIndependentSyntheticServerImplementsPublishedSchema(t *testing.T) { + requestValidator := compileContractSchema(t, "https://docbank.invalid/contracts/docbank-embedding/v1/request.schema.json", requestSchema) + responseValidator := compileContractSchema(t, "https://docbank.invalid/contracts/docbank-embedding/v1/response.schema.json", responseSchema) + var requestContract struct { + Canonicalization canonicalizationContract `json:"x-docbank-canonicalization"` + } + require.NoError(t, json.Unmarshal(requestSchema, &requestContract)) + serverErrors := make(chan error, 1) + fixture := newBridgeFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if err := serveIndependentContractRequest(writer, request, requestValidator, responseValidator, requestContract.Canonicalization); err != nil { + select { + case serverErrors <- err: + default: + } + http.Error(writer, "synthetic contract failure", http.StatusInternalServerError) + } + })) + source := []byte("synthetic independent file") + sourceDigest := sha256.Sum256(source) + upload := &contractUpload{ + reader: bytes.NewReader(source), + metadata: document.AuthorizedUploadMetadata{ + Filename: "contract.bin", MediaFamily: "binary", MediaType: "application/octet-stream", + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(sourceDigest[:]), + CapabilityRecordChecksum: strings.Repeat("a", 64), ProviderMetadataChecksum: strings.Repeat("b", 64), + InputKind: document.RenditionInputOriginalFile, + }, + } + result, err := fixture.client.Embed(context.Background(), []document.EmbeddingInput{ + { + Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, + Text: "synthetic contract fixture", HeadingPath: []string{"Synthetic"}, + SourceSpans: []document.ChunkSpan{{UnitIndex: 0, CharStart: 0, CharEnd: 9}}, + }, + {Key: "second", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, Source: upload}, + }, fixture.authorization(2)) + select { + case serverErr := <-serverErrors: + require.NoError(t, serverErr) + default: + } + require.NoError(t, err) + assert.Equal(t, document.EmbeddingResult{Vectors: []document.EmbeddingVector{ + {Key: "first", Values: []float32{0.5, 0.25}}, + {Key: "second", Values: []float32{1.5, 1.25}}, + }}, result) +} + +type contractUpload struct { + reader *bytes.Reader + metadata document.AuthorizedUploadMetadata +} + +func (upload *contractUpload) Read(value []byte) (int, error) { + read, err := upload.reader.Read(value) + if errors.Is(err, io.EOF) { + return read, io.EOF + } + if err != nil { + return read, fmt.Errorf("contract upload read: %w", err) + } + return read, nil +} + +func (upload *contractUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func (*contractUpload) Close() error { return nil } + +type independentFilePart struct { + filename string + mediaType string + payload []byte +} + +func compileContractSchema(t *testing.T, location string, document []byte) *jsonschema.Schema { + t.Helper() + compiler := jsonschema.NewCompiler() + compiler.Draft = jsonschema.Draft2020 + require.NoError(t, compiler.AddResource(location, bytes.NewReader(document))) + compiled, err := compiler.Compile(location) + require.NoError(t, err) + return compiled +} + +func serveIndependentContractRequest( + writer http.ResponseWriter, + request *http.Request, + requestValidator *jsonschema.Schema, + responseValidator *jsonschema.Schema, + canonicalization canonicalizationContract, +) error { + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + if err != nil || mediaType != "multipart/form-data" || parameters["boundary"] == "" { + return errors.New("independent server: multipart content type") + } + reader := multipart.NewReader(request.Body, parameters["boundary"]) + var manifestBytes []byte + var files []independentFilePart + for { + part, nextErr := reader.NextPart() + if errors.Is(nextErr, io.EOF) { + break + } + if nextErr != nil { + return fmt.Errorf("independent server: next part: %w", nextErr) + } + payload, readErr := io.ReadAll(part) + if readErr != nil { + return fmt.Errorf("independent server: read part: %w", readErr) + } + switch part.FormName() { + case "manifest": + if part.Header.Get("Content-Type") != "application/vnd.docbank.embedding-manifest+json;version=1" || manifestBytes != nil { + return errors.New("independent server: manifest part contract") + } + manifestBytes = payload + case "file": + files = append(files, independentFilePart{filename: part.FileName(), mediaType: part.Header.Get("Content-Type"), payload: payload}) + default: + return errors.New("independent server: unknown multipart part") + } + } + if manifestBytes == nil { + return errors.New("independent server: unexpected multipart shape") + } + var schemaValue any + if err := stdjson.Unmarshal(manifestBytes, &schemaValue); err != nil { + return fmt.Errorf("independent server: decode request schema value: %w", err) + } + if err := requestValidator.Validate(schemaValue); err != nil { + return fmt.Errorf("independent server: request schema: %w", err) + } + manifest, wireCanonical, identityCanonical, err := canonicalizeIndependentManifest(manifestBytes, canonicalization) + if err != nil { + return err + } + if !bytes.Equal(manifestBytes, wireCanonical) { + return errors.New("independent server: request is not canonical docbank-json-v1") + } + checksum, ok := manifest["request_checksum"].(string) + if !ok { + return errors.New("independent server: request checksum type") + } + digest := sha256.Sum256(identityCanonical) + if checksum != hex.EncodeToString(digest[:]) || request.Header.Get("Idempotency-Key") != checksum || request.Header.Get("Docbank-Request-Checksum") != checksum { + return errors.New("independent server: request checksum mismatch") + } + descriptorFingerprint, descriptorOK := manifest["descriptor_fingerprint"].(string) + policyFingerprint, policyOK := manifest["policy_fingerprint"].(string) + inputs, inputsOK := manifest["inputs"].([]any) + if !descriptorOK || !policyOK || !inputsOK || len(inputs) == 0 { + return errors.New("independent server: manifest identity") + } + keys := make([]string, len(inputs)) + usedFiles := make([]bool, len(files)) + for index, rawInput := range inputs { + input, inputOK := rawInput.(map[string]any) + key, keyOK := input["key"].(string) + kind, kindOK := input["kind"].(string) + if !inputOK || !keyOK || !kindOK { + return errors.New("independent server: input identity") + } + keys[index] = key + if kind != "original_file" { + continue + } + fileIndex, indexErr := independentInteger(input["file_index"]) + upload, uploadOK := input["upload"].(map[string]any) + if indexErr != nil || !uploadOK || fileIndex < 0 || fileIndex >= int64(len(files)) || usedFiles[fileIndex] { + return errors.New("independent server: file index") + } + usedFiles[fileIndex] = true + length, lengthErr := independentInteger(upload["byte_length"]) + checksum, checksumOK := upload["sha256"].(string) + filename, filenameOK := upload["filename"].(string) + fileMediaType, mediaTypeOK := upload["media_type"].(string) + file := files[fileIndex] + fileDigest := sha256.Sum256(file.payload) + if lengthErr != nil || !checksumOK || !filenameOK || !mediaTypeOK || length != int64(len(file.payload)) || + checksum != hex.EncodeToString(fileDigest[:]) || filename != file.filename || fileMediaType != file.mediaType { + return errors.New("independent server: file binding") + } + } + for _, used := range usedFiles { + if !used { + return errors.New("independent server: unbound file") + } + } + vectors := make([]any, len(keys)) + for index, key := range keys { + vectors[index] = map[string]any{ + "key": key, "index": float64(index), "values": []any{float64(index) + 0.5, float64(index) + 0.25}, + } + } + response := map[string]any{ + "contract_version": "docbank-embedding/v1", + "descriptor_fingerprint": descriptorFingerprint, + "policy_fingerprint": policyFingerprint, + "request_checksum": checksum, + "vectors": vectors, + } + responseBytes, err := stdjson.Marshal(response) + if err != nil { + return fmt.Errorf("independent server: encode response: %w", err) + } + var responseWire any + if err := stdjson.Unmarshal(responseBytes, &responseWire); err != nil { + return fmt.Errorf("independent server: decode response wire: %w", err) + } + if err := responseValidator.Validate(responseWire); err != nil { + return fmt.Errorf("independent server: response schema: %w", err) + } + writer.Header().Set("Content-Type", "application/vnd.docbank.embedding-result+json;version=1") + if _, err := writer.Write(responseBytes); err != nil { + return fmt.Errorf("independent server: write response: %w", err) + } + return nil +} + +func canonicalizeIndependentManifest(payload []byte, contract canonicalizationContract) (map[string]any, []byte, []byte, error) { + decoder := stdjson.NewDecoder(bytes.NewReader(payload)) + decoder.UseNumber() + var decoded any + if err := decoder.Decode(&decoded); err != nil { + return nil, nil, nil, fmt.Errorf("independent server: decode canonical manifest: %w", err) + } + if err := requireIndependentJSONEOF(decoder); err != nil { + return nil, nil, nil, err + } + manifest, ok := decoded.(map[string]any) + if !ok { + return nil, nil, nil, errors.New("independent server: manifest root") + } + wire, err := appendIndependentCanonical(nil, manifest, "$", contract.ObjectMemberOrder) + if err != nil { + return nil, nil, nil, err + } + identityManifest := maps.Clone(manifest) + delete(identityManifest, "request_checksum") + identity, err := appendIndependentCanonical(nil, identityManifest, "$", contract.ObjectMemberOrder) + if err != nil { + return nil, nil, nil, err + } + return manifest, wire, identity, nil +} + +func requireIndependentJSONEOF(decoder *stdjson.Decoder) error { + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return errors.New("independent server: trailing manifest JSON") + } + return nil +} + +func appendIndependentCanonical(destination []byte, value any, path string, orders map[string][]string) ([]byte, error) { + switch value := value.(type) { + case map[string]any: + order, ok := orders[path] + if !ok { + return nil, fmt.Errorf("independent server: no member order for %s", path) + } + destination = append(destination, '{') + members := 0 + for _, name := range order { + member, exists := value[name] + if !exists { + continue + } + if members > 0 { + destination = append(destination, ',') + } + var err error + destination, err = jsontext.AppendQuote(destination, name) + if err != nil { + return nil, fmt.Errorf("independent server: quote member: %w", err) + } + destination = append(destination, ':') + destination, err = appendIndependentCanonical(destination, member, path+"."+name, orders) + if err != nil { + return nil, err + } + members++ + } + if members != len(value) { + return nil, fmt.Errorf("independent server: unordered member at %s", path) + } + return append(destination, '}'), nil + case []any: + destination = append(destination, '[') + for index, member := range value { + if index > 0 { + destination = append(destination, ',') + } + var err error + destination, err = appendIndependentCanonical(destination, member, path+"[]", orders) + if err != nil { + return nil, err + } + } + return append(destination, ']'), nil + case string: + encoded, err := jsontext.AppendQuote(destination, value) + if err != nil { + return nil, fmt.Errorf("independent server: quote value: %w", err) + } + return encoded, nil + case stdjson.Number: + if !isCanonicalUnsignedInteger(value.String()) { + return nil, errors.New("independent server: non-canonical request integer") + } + return append(destination, value.String()...), nil + default: + return nil, fmt.Errorf("independent server: unsupported canonical value %T", value) + } +} + +func isCanonicalUnsignedInteger(value string) bool { + if value == "0" { + return true + } + if value == "" || value[0] < '1' || value[0] > '9' { + return false + } + return strings.IndexFunc(value[1:], func(character rune) bool { return character < '0' || character > '9' }) < 0 +} + +func independentInteger(value any) (int64, error) { + number, ok := value.(stdjson.Number) + if !ok || !isCanonicalUnsignedInteger(number.String()) { + return 0, errors.New("independent server: integer type") + } + parsed, err := strconv.ParseInt(number.String(), 10, 64) + if err != nil { + return 0, fmt.Errorf("independent server: parse integer: %w", err) + } + return parsed, nil +} diff --git a/document/embeddingbridge/embedding-request-v1.schema.json b/document/embeddingbridge/embedding-request-v1.schema.json new file mode 100644 index 00000000..d656ae29 --- /dev/null +++ b/document/embeddingbridge/embedding-request-v1.schema.json @@ -0,0 +1,195 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://docbank.invalid/contracts/docbank-embedding/v1/request.schema.json", + "title": "Docbank embedding request manifest v1", + "x-docbank-canonicalization": { + "name": "docbank-json-v1", + "character_encoding": "UTF-8 without a byte-order mark", + "whitespace": "No insignificant whitespace before or after tokens.", + "object_member_order": { + "$": ["contract_version", "descriptor_fingerprint", "policy_fingerprint", "authorization", "inputs", "request_checksum"], + "$.authorization": ["provider_id", "descriptor_fingerprint", "policy_fingerprint", "max_batch_items", "max_input_bytes", "max_response_bytes"], + "$.inputs[]": ["index", "key", "role", "kind", "byte_length", "sha256", "text", "heading_path", "source_spans", "upload", "file_part", "file_index"], + "$.inputs[].source_spans[]": ["unit_index", "char_start", "char_end"], + "$.inputs[].upload": ["filename", "media_family", "media_type", "byte_length", "sha256", "capability_record_checksum", "provider_metadata_checksum", "input_kind"] + }, + "arrays": "Preserve manifest array order exactly.", + "integers": "Base-10 ASCII digits with a leading minus only for negative values; no leading zero except the value zero. This schema permits no negative request integers.", + "strings": "Use JSON quotation marks. Escape quotation mark and reverse solidus. Escape U+0000 through U+001F using \\b, \\f, \\n, \\r, or \\t where applicable and lowercase \\u00xx otherwise. Emit every other Unicode scalar as its UTF-8 bytes without escaping.", + "optional_members": "Omit absent optional members. Never emit null. Emit present members in the order defined for their object after skipping absent members.", + "checksum": "Omit request_checksum entirely, serialize the remaining manifest with these docbank-json-v1 rules, compute SHA-256 over those exact UTF-8 bytes, encode the digest as 64 lowercase hexadecimal characters, then append request_checksum as the final root member for transmission.", + "checksum_vectors": [ + { + "canonical_manifest": "{\"contract_version\":\"docbank-embedding/v1\",\"descriptor_fingerprint\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"policy_fingerprint\":\"1111111111111111111111111111111111111111111111111111111111111111\",\"authorization\":{\"provider_id\":\"synthetic.vector\",\"descriptor_fingerprint\":\"0000000000000000000000000000000000000000000000000000000000000000\",\"policy_fingerprint\":\"1111111111111111111111111111111111111111111111111111111111111111\",\"max_batch_items\":1,\"max_input_bytes\":64,\"max_response_bytes\":16},\"inputs\":[{\"index\":0,\"key\":\"vector-a\",\"role\":\"document\",\"kind\":\"rendition_chunk\",\"byte_length\":22,\"sha256\":\"5aa0924b94846dd796a2130bbb9a63f21721793169cedd7547b1666050dacc12\",\"text\":\"search_document: alpha\",\"heading_path\":[\"Overview\"],\"source_spans\":[{\"unit_index\":0,\"char_start\":0,\"char_end\":5}]}]}", + "sha256": "c1051df601b1389603d0f773d51c02960b96d0cec370672bc823d98d30a0d97f" + }, + { + "canonical_manifest": "{\"contract_version\":\"docbank-embedding/v1\",\"descriptor_fingerprint\":\"2222222222222222222222222222222222222222222222222222222222222222\",\"policy_fingerprint\":\"3333333333333333333333333333333333333333333333333333333333333333\",\"authorization\":{\"provider_id\":\"synthetic.query\",\"descriptor_fingerprint\":\"2222222222222222222222222222222222222222222222222222222222222222\",\"policy_fingerprint\":\"3333333333333333333333333333333333333333333333333333333333333333\",\"max_batch_items\":1,\"max_input_bytes\":64,\"max_response_bytes\":16},\"inputs\":[{\"index\":0,\"key\":\"vector-b\",\"role\":\"query\",\"kind\":\"query_text\",\"byte_length\":20,\"sha256\":\"1ccc4c7d34c3c140b5a067a8fca62cd86244ba8b39fa4e3b2b85c578558b5020\",\"text\":\"search_query: <β>&\\n\"}]}", + "sha256": "01c33c6021a4296d41276c772a9361cfe24854653cff7925adab40bb2110dab3" + } + ] + }, + "type": "object", + "additionalProperties": false, + "required": [ + "contract_version", + "descriptor_fingerprint", + "policy_fingerprint", + "authorization", + "inputs", + "request_checksum" + ], + "properties": { + "contract_version": {"const": "docbank-embedding/v1"}, + "descriptor_fingerprint": {"$ref": "#/$defs/sha256"}, + "policy_fingerprint": {"$ref": "#/$defs/sha256"}, + "authorization": {"$ref": "#/$defs/authorization"}, + "inputs": { + "type": "array", + "minItems": 1, + "maxItems": 10000, + "items": {"$ref": "#/$defs/input"} + }, + "request_checksum": { + "description": "SHA-256 of the docbank-json-v1 canonical UTF-8 manifest with request_checksum omitted, as defined by x-docbank-canonicalization.", + "$ref": "#/$defs/sha256" + } + }, + "$defs": { + "sha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"}, + "stable_id": { + "type": "string", + "minLength": 1, + "maxLength": 128, + "pattern": "^[A-Za-z0-9._-]+$" + }, + "authorization": { + "type": "object", + "additionalProperties": false, + "required": [ + "provider_id", + "descriptor_fingerprint", + "policy_fingerprint", + "max_batch_items", + "max_input_bytes", + "max_response_bytes" + ], + "properties": { + "provider_id": {"$ref": "#/$defs/stable_id"}, + "descriptor_fingerprint": {"$ref": "#/$defs/sha256"}, + "policy_fingerprint": {"$ref": "#/$defs/sha256"}, + "max_batch_items": {"type": "integer", "minimum": 1, "maximum": 10000}, + "max_input_bytes": {"type": "integer", "minimum": 1, "maximum": 1073741824}, + "max_response_bytes": {"type": "integer", "minimum": 1, "maximum": 1073741824} + } + }, + "input": { + "type": "object", + "additionalProperties": false, + "required": ["index", "key", "role", "kind", "byte_length", "sha256"], + "properties": { + "index": {"type": "integer", "minimum": 0, "maximum": 9999}, + "key": {"$ref": "#/$defs/stable_id"}, + "role": {"enum": ["document", "query"]}, + "kind": {"enum": ["rendition_chunk", "query_text", "original_file"]}, + "byte_length": {"type": "integer", "minimum": 1, "maximum": 1073741824}, + "sha256": {"$ref": "#/$defs/sha256"}, + "text": {"type": "string", "minLength": 1}, + "heading_path": { + "type": "array", + "minItems": 1, + "maxItems": 1024, + "items": {"type": "string", "minLength": 1, "maxLength": 8192} + }, + "source_spans": { + "type": "array", + "minItems": 1, + "maxItems": 1024, + "items": {"$ref": "#/$defs/source_span"} + }, + "upload": {"$ref": "#/$defs/upload"}, + "file_part": {"const": "file"}, + "file_index": {"type": "integer", "minimum": 0, "maximum": 9999} + }, + "oneOf": [ + { + "required": ["text"], + "not": {"anyOf": [{"required": ["upload"]}, {"required": ["file_part"]}, {"required": ["file_index"]}]}, + "properties": { + "role": {"const": "document"}, + "kind": {"const": "rendition_chunk"} + } + }, + { + "required": ["text"], + "not": { + "anyOf": [ + {"required": ["heading_path"]}, + {"required": ["source_spans"]}, + {"required": ["upload"]}, + {"required": ["file_part"]}, + {"required": ["file_index"]} + ] + }, + "properties": { + "role": {"const": "query"}, + "kind": {"const": "query_text"} + } + }, + { + "required": ["upload", "file_part", "file_index"], + "not": { + "anyOf": [ + {"required": ["text"]}, + {"required": ["heading_path"]}, + {"required": ["source_spans"]} + ] + }, + "properties": { + "role": {"const": "document"}, + "kind": {"const": "original_file"} + } + } + ] + }, + "source_span": { + "type": "object", + "additionalProperties": false, + "required": ["unit_index", "char_start", "char_end"], + "properties": { + "unit_index": {"type": "integer", "minimum": 0, "maximum": 999999}, + "char_start": {"type": "integer", "minimum": 0, "maximum": 268435455}, + "char_end": {"type": "integer", "minimum": 1, "maximum": 268435456} + } + }, + "upload": { + "type": "object", + "additionalProperties": false, + "required": [ + "filename", + "media_family", + "media_type", + "byte_length", + "sha256", + "capability_record_checksum", + "provider_metadata_checksum", + "input_kind" + ], + "properties": { + "filename": { + "type": "string", + "minLength": 1, + "maxLength": 255, + "pattern": "^[^/\\\\\\x00-\\x1f\\x7f]+$" + }, + "media_family": {"type": "string", "minLength": 1, "maxLength": 63}, + "media_type": {"type": "string", "minLength": 3, "maxLength": 255}, + "byte_length": {"type": "integer", "minimum": 1}, + "sha256": {"$ref": "#/$defs/sha256"}, + "capability_record_checksum": {"$ref": "#/$defs/sha256"}, + "provider_metadata_checksum": {"$ref": "#/$defs/sha256"}, + "input_kind": {"const": "original_file"} + } + } + } +} diff --git a/document/embeddingbridge/embedding-response-v1.schema.json b/document/embeddingbridge/embedding-response-v1.schema.json new file mode 100644 index 00000000..6e7c139b --- /dev/null +++ b/document/embeddingbridge/embedding-response-v1.schema.json @@ -0,0 +1,51 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://docbank.invalid/contracts/docbank-embedding/v1/response.schema.json", + "title": "Docbank embedding response v1", + "type": "object", + "additionalProperties": false, + "required": [ + "contract_version", + "descriptor_fingerprint", + "policy_fingerprint", + "request_checksum", + "vectors" + ], + "properties": { + "contract_version": {"const": "docbank-embedding/v1"}, + "descriptor_fingerprint": {"$ref": "#/$defs/sha256"}, + "policy_fingerprint": {"$ref": "#/$defs/sha256"}, + "request_checksum": {"$ref": "#/$defs/sha256"}, + "vectors": { + "type": "array", + "minItems": 1, + "maxItems": 10000, + "description": "Exactly one item per request input, in request order; index equals its zero-based array position.", + "items": {"$ref": "#/$defs/vector"} + } + }, + "$defs": { + "sha256": {"type": "string", "pattern": "^[0-9a-f]{64}$"}, + "stable_id": { + "type": "string", + "minLength": 1, + "maxLength": 128, + "pattern": "^[A-Za-z0-9._-]+$" + }, + "vector": { + "type": "object", + "additionalProperties": false, + "required": ["key", "index", "values"], + "properties": { + "key": {"$ref": "#/$defs/stable_id"}, + "index": {"type": "integer", "minimum": 0, "maximum": 9999}, + "values": { + "type": "array", + "minItems": 1, + "maxItems": 1000000, + "items": {"type": "number"} + } + } + } + } +} diff --git a/document/embeddingbridge/errors.go b/document/embeddingbridge/errors.go new file mode 100644 index 00000000..2adc5266 --- /dev/null +++ b/document/embeddingbridge/errors.go @@ -0,0 +1,51 @@ +package embeddingbridge + +import ( + "errors" + "fmt" +) + +// ErrorCategory is a stable, content-free worker classification. +type ErrorCategory string + +const ( + ErrorAuthentication ErrorCategory = "authentication" + ErrorCapacity ErrorCategory = "capacity" + ErrorTransient ErrorCategory = "transient" + ErrorPermanent ErrorCategory = "permanent" + ErrorMalformedResponse ErrorCategory = "malformed_response" + ErrorAmbiguousSubmission ErrorCategory = "ambiguous_submission" + ErrorSourceChanged ErrorCategory = "source_changed" +) + +// ProviderError contains only a stable category and optional HTTP status. +// Provider bodies, transport messages, secrets, and inputs are never retained. +type ProviderError struct { + category ErrorCategory + status int +} + +func (err *ProviderError) Error() string { + if err.status != 0 { + return fmt.Sprintf("embedding bridge: %s (HTTP %d)", err.category, err.status) + } + return "embedding bridge: " + string(err.category) +} + +// Category returns the stable category for err. Unknown errors are permanent. +func Category(err error) ErrorCategory { + if providerError, ok := errors.AsType[*ProviderError](err); ok { + return providerError.category + } + return ErrorPermanent +} + +// IsRetryable reports whether retrying the same idempotency checksum is safe. +func IsRetryable(err error) bool { + category := Category(err) + return category == ErrorTransient || category == ErrorAmbiguousSubmission +} + +func classified(category ErrorCategory, status int) error { + return &ProviderError{category: category, status: status} +} diff --git a/document/embeddingbridge/openapi.yaml b/document/embeddingbridge/openapi.yaml new file mode 100644 index 00000000..80df1ca0 --- /dev/null +++ b/document/embeddingbridge/openapi.yaml @@ -0,0 +1,75 @@ +openapi: 3.1.0 +info: + title: Docbank Embedding Bridge + version: 1.0.0 + description: >- + Fixed synchronous provider-neutral embedding protocol. Implementations accept + one deterministic manifest and zero or more authorization-held files at one + same-origin route. The complete result is returned by that request. +paths: + /docbank-embedding/v1/embeddings: + post: + operationId: embedDocbankInputs + security: + - bearerAuth: [] + - {} + parameters: + - name: Idempotency-Key + in: header + required: true + schema: {$ref: '#/components/schemas/SHA256'} + - name: Docbank-Request-Checksum + in: header + required: true + schema: {$ref: '#/components/schemas/SHA256'} + requestBody: + required: true + content: + multipart/form-data: + schema: + type: object + additionalProperties: false + required: [manifest] + properties: + manifest: {$ref: './embedding-request-v1.schema.json'} + file: + type: array + description: >- + Original files in ascending manifest file_index order. Every + element uses the repeated fixed multipart part name file. + maxItems: 10000 + items: {type: string, format: binary} + encoding: + manifest: + contentType: application/vnd.docbank.embedding-manifest+json;version=1 + file: + explode: true + responses: + '200': + description: One ordered, indexed vector for every manifest input. + content: + application/vnd.docbank.embedding-result+json;version=1: + schema: {$ref: './embedding-response-v1.schema.json'} + '400': {$ref: '#/components/responses/EmptyError'} + '401': {$ref: '#/components/responses/EmptyError'} + '403': {$ref: '#/components/responses/EmptyError'} + '413': {$ref: '#/components/responses/EmptyError'} + '429': {$ref: '#/components/responses/EmptyError'} + '500': {$ref: '#/components/responses/EmptyError'} + '502': {$ref: '#/components/responses/EmptyError'} + '503': {$ref: '#/components/responses/EmptyError'} + '504': {$ref: '#/components/responses/EmptyError'} +components: + securitySchemes: + bearerAuth: + type: http + scheme: bearer + responses: + EmptyError: + description: >- + Error responses have no protocol body. Clients classify only the status + and never retain provider response bytes or upstream messages. + schemas: + SHA256: + type: string + pattern: '^[0-9a-f]{64}$' diff --git a/document/embeddingbridge/profile.go b/document/embeddingbridge/profile.go new file mode 100644 index 00000000..9d128528 --- /dev/null +++ b/document/embeddingbridge/profile.go @@ -0,0 +1,316 @@ +package embeddingbridge + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "net" + "net/http" + "net/netip" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + defaultRequestTimeout = 30 * time.Second + defaultMaxBatchItems = 128 + defaultMaxInputBytes = int64(16 << 20) + defaultMaxRequestBytes = int64(32 << 20) + defaultMaxResponseBytes = int64(64 << 20) + maxRequestTimeout = 10 * time.Minute + maxBatchItems = 10_000 + maxInputBytes = int64(1 << 30) + maxRequestBytes = int64(2 << 30) + maxResponseBytes = int64(1 << 30) + maxSecretBytes = 64 << 10 +) + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + ModelInput document.ModelInputContract `json:"model_input"` + CredentialBinding string `json:"credential_binding"` + Egress egressIdentity `json:"egress"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` +} + +type egressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +// PolicyFingerprint returns the canonical profile fingerprint. Credential +// values and supplied HTTP client state are deliberately excluded. +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptorIdentity, err := normalizeProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(policyIdentity{ + AdapterContract: adapterContract, Origin: normalized.Origin, Route: embeddingsPath, + Descriptor: descriptorIdentity, ModelInput: descriptorIdentity.ModelInput, + CredentialBinding: normalized.SecretBinding, Egress: profileEgressIdentity(normalized.EgressPolicy), + RequestTimeout: int64(normalized.RequestTimeout), MaxBatchItems: normalized.MaxBatchItems, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, + }, json.Deterministic(true)) + if err != nil { + return "", errors.New("embedding bridge: profile identity encoding failed") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +// New validates a canonical profile and returns an isolated client using the +// repository egress transport. All mutable supplied client authority is removed. +func New(profile Profile, secrets SecretResolver, resolver providerhttp.Resolver, supplied *http.Client) (*Client, error) { + if supplied == nil { + return nil, errors.New("embedding bridge: HTTP client is required") + } + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + return nil, errors.New("embedding bridge: descriptor is not canonical") + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("embedding bridge: descriptor policy fingerprint does not match profile") + } + if normalized.SecretBinding == "" { + if !nilInterface(secrets) { + return nil, errors.New("embedding bridge: resolver is not allowed without a named secret binding") + } + } else if nilInterface(secrets) { + return nil, errors.New("embedding bridge: named secret resolver is required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, errors.New("embedding bridge: sealed egress policy is invalid") + } + isolated := *supplied + isolated.Transport = transport + isolated.CheckRedirect = providerhttp.RefuseRedirects + isolated.Jar = nil + isolated.Timeout = 0 + return &Client{ + origin: normalized.Origin, descriptor: cloneDescriptor(descriptor), + secretBinding: normalized.SecretBinding, secrets: secrets, http: &isolated, + requestTimeout: normalized.RequestTimeout, maxBatchItems: normalized.MaxBatchItems, + maxInputBytes: normalized.MaxInputBytes, maxRequestBytes: normalized.MaxRequestBytes, + maxResponseBytes: normalized.MaxResponseBytes, + }, nil +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + profile.EgressPolicy.TLS.SPKISHA256 = slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultMaxBatchItems + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultMaxInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultMaxRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxRequestTimeout || + profile.MaxBatchItems < 1 || profile.MaxBatchItems > maxBatchItems || + profile.MaxInputBytes < 1 || profile.MaxInputBytes > maxInputBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maxRequestBytes || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maxResponseBytes { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("embedding bridge: execution bounds are invalid") + } + if profile.SecretBinding != "" && !validBinding(profile.SecretBinding) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("embedding bridge: secret binding is invalid") + } + if profile.SecretBinding == "" && profile.Descriptor.TrustBoundary != document.EmbeddingTrustOperatorNetwork { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("embedding bridge: anonymous access is operator-network only") + } + origin, err := normalizeOrigin(profile.Origin, profile.Descriptor.TrustBoundary) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + profile.Origin = origin + if err := normalizeEgress(&profile); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + descriptorIdentity := cloneDescriptor(profile.Descriptor) + descriptorIdentity.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptorIdentity.Fingerprint = "" + descriptorIdentity, err = document.NewEmbeddingDescriptor(descriptorIdentity) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("embedding bridge: descriptor identity is invalid") + } + descriptorIdentity.PolicyFingerprint = "" + descriptorIdentity.Fingerprint = "" + return profile, descriptorIdentity, nil +} + +func normalizeOrigin(raw string, trust document.EmbeddingTrustBoundary) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Opaque != "" || + parsed.ForceQuery || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return "", errors.New("embedding bridge: origin must be one absolute origin without path, credentials, query, or fragment") + } + if parsed.Scheme != "https" && (parsed.Scheme != "http" || trust != document.EmbeddingTrustOperatorNetwork) { + return "", errors.New("embedding bridge: hosted origins require HTTPS; HTTP is operator-network only") + } + if trust != document.EmbeddingTrustOperatorNetwork && trust != document.EmbeddingTrustHostedProvider { + return "", errors.New("embedding bridge: network origin requires an operator-network or hosted trust boundary") + } + port := parsed.Port() + if port == "" { + if parsed.Scheme == "https" { + port = "443" + } else { + port = "80" + } + } + if _, err := strconv.ParseUint(port, 10, 16); err != nil { + return "", errors.New("embedding bridge: origin port is invalid") + } + host := strings.ToLower(parsed.Hostname()) + authority := host + if strings.Contains(host, ":") { + authority = "[" + host + "]" + } + if (parsed.Scheme != "https" || port != "443") && (parsed.Scheme != "http" || port != "80") { + authority = net.JoinHostPort(host, port) + } + return parsed.Scheme + "://" + authority, nil +} + +func normalizeEgress(profile *Profile) error { + policy := &profile.EgressPolicy + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = providerhttp.DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if policy.ProxyMode == "" { + policy.ProxyMode = providerhttp.ProxyDisabled + } + if policy.TLS.RootCAs != nil { + return errors.New("embedding bridge: custom egress roots cannot enter canonical identity") + } + parsed, err := url.Parse(profile.Origin) + if err != nil { + return errors.New("embedding bridge: origin is invalid") + } + port := parsed.Port() + if port == "" { + if parsed.Scheme == "https" { + port = "443" + } else { + port = "80" + } + } + if parsed.Scheme != policy.Scheme || !strings.EqualFold(parsed.Hostname(), policy.Host) || + port != strconv.FormatUint(uint64(policy.Port), 10) { + return errors.New("embedding bridge: origin and egress authority differ") + } + policy.Host = strings.ToLower(policy.Host) + for index := range policy.AllowedCIDRs { + policy.AllowedCIDRs[index] = policy.AllowedCIDRs[index].Masked() + } + slices.SortFunc(policy.AllowedCIDRs, func(left, right netip.Prefix) int { + return strings.Compare(left.String(), right.String()) + }) + for index := 1; index < len(policy.AllowedCIDRs); index++ { + if policy.AllowedCIDRs[index] == policy.AllowedCIDRs[index-1] { + return errors.New("embedding bridge: egress policy has a duplicate CIDR") + } + } + for index := range policy.TLS.SPKISHA256 { + policy.TLS.SPKISHA256[index] = strings.ToLower(policy.TLS.SPKISHA256[index]) + } + slices.Sort(policy.TLS.SPKISHA256) + return nil +} + +func profileEgressIdentity(policy providerhttp.EgressPolicy) egressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.String() + } + return egressIdentity{ + Scheme: policy.Scheme, Host: policy.Host, Port: policy.Port, AllowedCIDRs: cidrs, + ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), + KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), + SPKISHA256: slices.Clone(policy.TLS.SPKISHA256), + } +} + +func validBinding(value string) bool { + return value != "" && len(value) <= 128 && utf8.ValidString(value) && value == strings.TrimSpace(value) && + strings.IndexFunc(value, func(character rune) bool { return unicode.IsControl(character) || unicode.IsSpace(character) }) < 0 +} + +func validSecret(value string) bool { + return value != "" && len(value) <= maxSecretBytes && strings.IndexFunc(value, func(character rune) bool { + return unicode.IsControl(character) || unicode.IsSpace(character) + }) < 0 +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} + +func cloneDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/document/embeddingbridge/source_gate_test.go b/document/embeddingbridge/source_gate_test.go new file mode 100644 index 00000000..f409a584 --- /dev/null +++ b/document/embeddingbridge/source_gate_test.go @@ -0,0 +1,40 @@ +package embeddingbridge + +import ( + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +func TestActiveSourceCancellationBetweenFilesLeavesBothSourcesOpen(t *testing.T) { + first := &gateTestUpload{} + second := &gateTestUpload{} + gate := newActiveSourceGate() + + firstToken, ok := gate.Begin(first) + require.True(t, ok) + gate.End(firstToken) + gate.Cancel() + _, ok = gate.Begin(second) + assert.False(t, ok) + assert.Zero(t, first.closes.Load()) + assert.Zero(t, second.closes.Load()) +} + +type gateTestUpload struct { + closes atomic.Int32 +} + +func (*gateTestUpload) Read([]byte) (int, error) { return 0, nil } + +func (*gateTestUpload) Metadata() document.AuthorizedUploadMetadata { + return document.AuthorizedUploadMetadata{} +} + +func (upload *gateTestUpload) Close() error { + upload.closes.Add(1) + return nil +} diff --git a/document/embeddingbridge/types.go b/document/embeddingbridge/types.go new file mode 100644 index 00000000..96941810 --- /dev/null +++ b/document/embeddingbridge/types.go @@ -0,0 +1,103 @@ +// Package embeddingbridge implements the fixed synchronous docbank-embedding/v1 +// provider-neutral embedding protocol. +package embeddingbridge + +import ( + "context" + "net/http" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + // ContractVersion is the only protocol version this package accepts. + ContractVersion = "docbank-embedding/v1" + + embeddingsPath = "/docbank-embedding/v1/embeddings" + manifestPartName = "manifest" + filePartName = "file" + manifestMediaType = "application/vnd.docbank.embedding-manifest+json;version=1" + responseMediaType = "application/vnd.docbank.embedding-result+json;version=1" + adapterContract = "docbank-standard-embedding-bridge/v1" +) + +// SecretResolver resolves only the configured named credential binding. +// Values are runtime-only and never enter protocol manifests or fingerprints. +type SecretResolver interface { + ResolveSecret(ctx context.Context, binding string) (string, error) +} + +// Profile freezes the bridge origin, vector-space identity, credential name, +// egress authority, and synchronous execution bounds. +type Profile struct { + Origin string + Descriptor document.EmbeddingDescriptor + SecretBinding string + EgressPolicy providerhttp.EgressPolicy + RequestTimeout time.Duration + MaxBatchItems int + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 +} + +// Client implements document.EmbeddingProvider through docbank-embedding/v1. +type Client struct { + origin string + descriptor document.EmbeddingDescriptor + secretBinding string + secrets SecretResolver + http *http.Client + requestTimeout time.Duration + maxBatchItems int + maxInputBytes int64 + maxRequestBytes int64 + maxResponseBytes int64 +} + +// RequestManifest is the deterministic authorization and input identity sent +// as the first multipart part. RequestChecksum is SHA-256 over the canonical +// manifest with that field omitted. +type RequestManifest struct { + ContractVersion string `json:"contract_version"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + Authorization document.EmbeddingAuthorization `json:"authorization"` + Inputs []ManifestInput `json:"inputs"` + RequestChecksum string `json:"request_checksum,omitempty"` +} + +// ManifestInput binds one exact request position and its rendered text or +// authorization-held original file. +type ManifestInput struct { + Index int `json:"index"` + Key string `json:"key"` + Role document.EmbeddingRole `json:"role"` + Kind document.EmbeddingInputKind `json:"kind"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + Text string `json:"text,omitempty"` + HeadingPath []string `json:"heading_path,omitempty"` + SourceSpans []ManifestSpan `json:"source_spans,omitempty"` + Upload *document.AuthorizedUploadMetadata `json:"upload,omitempty"` + FilePart string `json:"file_part,omitempty"` + FileIndex *int `json:"file_index,omitempty"` +} + +// ManifestSpan is the closed wire form of one canonical source span. +type ManifestSpan struct { + UnitIndex int `json:"unit_index"` + CharStart int `json:"char_start"` + CharEnd int `json:"char_end"` +} + +// Response is the strict synchronous result envelope. +type Response struct { + ContractVersion string `json:"contract_version"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + RequestChecksum string `json:"request_checksum"` + Vectors []document.EmbeddingVector `json:"vectors"` +} diff --git a/document/evidence_codec.go b/document/evidence_codec.go index 4088ba32..8c535391 100644 --- a/document/evidence_codec.go +++ b/document/evidence_codec.go @@ -455,7 +455,7 @@ func validateFamilyUnitKind(family string, unitKind EvidenceUnitKind) error { } var allowed bool switch family { - case "pdf", "word": + case "image", "pdf", "word": allowed = unitKind == EvidenceUnitPage case "presentation": allowed = unitKind == EvidenceUnitSlide @@ -478,7 +478,7 @@ func validateFamilyUnitKind(family string, unitKind EvidenceUnitKind) error { func validEvidenceFamily(family string) bool { switch family { - case "pdf", "word", "presentation", "spreadsheet", "ebook", "structured", "source", "text", "mail": + case "audio", "image", "video", "pdf", "word", "presentation", "spreadsheet", "ebook", "structured", "source", "text", "mail": return true default: return false diff --git a/document/evidence_codec_test.go b/document/evidence_codec_test.go index c4d9fc04..424c4a7a 100644 --- a/document/evidence_codec_test.go +++ b/document/evidence_codec_test.go @@ -402,6 +402,41 @@ func TestSourceEvidenceV1RejectsInvalidAuthority(t *testing.T) { } } +func TestEvidenceV1SupportsVisualAndDegradedMediaFamilies(t *testing.T) { + for _, testCase := range []struct { + name string + family string + completeness document.EvidenceCompleteness + unitKind document.EvidenceUnitKind + locatorKind document.EvidenceLocatorKind + }{ + {name: "image page", family: "image", completeness: document.EvidenceComplete, unitKind: document.EvidenceUnitPage, locatorKind: document.EvidenceLocatorPage}, + {name: "audio fallback", family: "audio", completeness: document.EvidenceDegradedProvenance, unitKind: document.EvidenceUnitGeneric, locatorKind: document.EvidenceLocatorGeneric}, + {name: "video fallback", family: "video", completeness: document.EvidenceDegradedProvenance, unitKind: document.EvidenceUnitGeneric, locatorKind: document.EvidenceLocatorGeneric}, + } { + t.Run(testCase.name, func(t *testing.T) { + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: testCase.completeness, + Family: testCase.family, + UnitKind: testCase.unitKind, + Units: []document.SourceEvidenceUnitV1{{ + Order: 0, Text: "synthetic evidence", + Locator: document.SourceEvidenceLocatorV1{Kind: testCase.locatorKind, IndexOrigin: document.EvidenceIndexOriginNone}, + }}, + } + if testCase.completeness == document.EvidenceComplete { + evidence.Units[0].Locator.IndexOrigin = document.EvidenceIndexOriginZero + } else { + evidence.Omissions = []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", Reason: "provider returned generic evidence", + }} + } + require.NoError(t, document.ValidateSourceEvidenceV1(evidence)) + }) + } +} + func TestEvidenceV1LocatorSequenceRequiresGapOmission(t *testing.T) { source := syntheticSourceEvidenceV1() source.Completeness = document.EvidencePartial diff --git a/document/execution.go b/document/execution.go new file mode 100644 index 00000000..ef83edbd --- /dev/null +++ b/document/execution.go @@ -0,0 +1,440 @@ +package document + +import ( + "bytes" + "context" + "encoding/json/v2" + "errors" + "fmt" + "reflect" + "slices" + "time" +) + +const ( + // RenditionExecutionIdentityContractV1 is the stable shared-build input + // contract. It excludes only the transient authorization interval. + RenditionExecutionIdentityContractV1 = "rendition-execution-identity/v1" + // RenditionExecutionSnapshotContractV1 is the sealed durable provider + // execution authority retained for a known resume handle. + RenditionExecutionSnapshotContractV1 = "rendition-execution-snapshot/v1" +) + +// EvidencePolicyIdentity is every effective provider-evidence normalization +// bound. Keeping the fixed v1 values explicit makes future policy changes +// split shared builds instead of silently reusing different output semantics. +type EvidencePolicyIdentity struct { + MaxArtifacts int `json:"max_artifacts"` + MaxCellsPerTable int `json:"max_cells_per_table"` + MaxDocumentChars int `json:"max_document_chars"` + MaxOmissions int `json:"max_omissions"` + MaxRegionsPerUnit int `json:"max_regions_per_unit"` + MaxTablesPerUnit int `json:"max_tables_per_unit"` + MaxUnits int `json:"max_units"` +} + +// Identity returns every value that can affect normalized evidence. +func (policy EvidencePolicy) Identity() EvidencePolicyIdentity { + return EvidencePolicyIdentity{ + MaxArtifacts: policy.maxArtifacts, MaxCellsPerTable: policy.maxCellsPerTable, + MaxDocumentChars: policy.maxDocumentChars, MaxOmissions: policy.maxOmissions, + MaxRegionsPerUnit: policy.maxRegionsPerUnit, + MaxTablesPerUnit: policy.maxTablesPerUnit, MaxUnits: policy.maxUnits, + } +} + +// RenditionPolicyIdentity is every value that can affect normalized Markdown, +// units, and lexical segments. +type RenditionPolicyIdentity struct { + Normalization NormalizePolicyIdentity `json:"normalization"` + MaxSegmentRunes int `json:"max_segment_runes"` +} + +// Identity returns every effective rendition-construction policy value. +func (policy RenditionPolicy) Identity() RenditionPolicyIdentity { + return RenditionPolicyIdentity{ + Normalization: policy.normalization.Identity(), MaxSegmentRunes: policy.maxSegmentRunes, + } +} + +// RenditionAuthorizationIdentityV1 projects the stable provider-visible input +// and output bounds from a sealed authorization, deliberately omitting only +// AuthorizedAt and ExpiresAt. +type RenditionAuthorizationIdentityV1 struct { + ProviderID string `json:"provider_id"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + RenditionRequestFingerprint string `json:"rendition_request_fingerprint"` + SourceSHA256 string `json:"source_sha256"` + SourceBytes int64 `json:"source_bytes"` + CapabilityRecordChecksum string `json:"capability_record_checksum"` + ProviderMetadataChecksum string `json:"provider_metadata_checksum"` + MediaFamily string `json:"media_family"` + MediaType string `json:"media_type"` + InputKind RenditionInputKind `json:"input_kind"` + AllowedArtifactRoles []EvidenceArtifactRole `json:"allowed_artifact_roles"` + MaxProviderMarkdownBytes int `json:"max_provider_markdown_bytes"` + MaxArtifactBytes int `json:"max_artifact_bytes"` + MaxArtifacts int `json:"max_artifacts"` + MaxTotalResultBytes int `json:"max_total_result_bytes"` +} + +// RenditionExecutionIdentityV1 is the complete stable identity of a shared +// provider execution and its deterministic local normalization. +type RenditionExecutionIdentityV1 struct { + ContractVersion string `json:"contract_version"` + Upload AuthorizedUploadMetadata `json:"upload"` + Authorization RenditionAuthorizationIdentityV1 `json:"authorization"` + EvidencePolicy EvidencePolicyIdentity `json:"evidence_policy"` + RenditionPolicy RenditionPolicyIdentity `json:"rendition_policy"` +} + +// RenditionExecutionSnapshotV1 retains the original sealed authorization +// interval alongside its stable identity. It is provider-neutral and contains +// no source bytes, credentials, paths, provider bodies, or HTTP state. +type RenditionExecutionSnapshotV1 struct { + ContractVersion string `json:"contract_version"` + Identity RenditionExecutionIdentityV1 `json:"identity"` + Authorization RenditionAuthorization `json:"authorization"` +} + +// NewRenditionExecutionIdentityV1 constructs the canonical stable execution +// identity from a planned sealed request and local output policies. +func NewRenditionExecutionIdentityV1( + metadata AuthorizedUploadMetadata, authorization RenditionAuthorization, + evidence EvidencePolicy, rendition RenditionPolicy, +) (RenditionExecutionIdentityV1, error) { + authorization = canonicalExecutionAuthorization(authorization) + identity := RenditionExecutionIdentityV1{ + ContractVersion: RenditionExecutionIdentityContractV1, + Upload: metadata, + Authorization: renditionAuthorizationIdentity(authorization), + EvidencePolicy: evidence.Identity(), + RenditionPolicy: rendition.Identity(), + } + if _, _, err := CanonicalRenditionExecutionIdentityV1(identity); err != nil { + return RenditionExecutionIdentityV1{}, err + } + return cloneRenditionExecutionIdentity(identity), nil +} + +// CanonicalRenditionExecutionIdentityV1 validates and deterministically +// encodes a stable identity and returns its SHA-256 fingerprint. +func CanonicalRenditionExecutionIdentityV1( + identity RenditionExecutionIdentityV1, +) ([]byte, string, error) { + identity = cloneRenditionExecutionIdentity(identity) + if err := validateRenditionExecutionIdentity(identity); err != nil { + return nil, "", err + } + encoded, err := canonicalJSON(identity) + if err != nil { + return nil, "", fmt.Errorf("encoding rendition execution identity: %w", err) + } + return encoded, sha256Hex(encoded), nil +} + +// ParseRenditionExecutionIdentityV1 decodes only the exact canonical form. +func ParseRenditionExecutionIdentityV1(raw []byte) (RenditionExecutionIdentityV1, error) { + var identity RenditionExecutionIdentityV1 + if err := json.Unmarshal(raw, &identity, json.RejectUnknownMembers(true)); err != nil { + return RenditionExecutionIdentityV1{}, fmt.Errorf("decoding rendition execution identity: %w", err) + } + canonical, _, err := CanonicalRenditionExecutionIdentityV1(identity) + if err != nil { + return RenditionExecutionIdentityV1{}, err + } + if !bytes.Equal(raw, canonical) { + return RenditionExecutionIdentityV1{}, errors.New("rendition execution identity is not canonical") + } + return cloneRenditionExecutionIdentity(identity), nil +} + +// SealRenditionExecutionAt validates the exact provider/upload request at the +// egress clock and captures the immutable authority needed for safe resume. +func SealRenditionExecutionAt( + now time.Time, provider RenditionProvider, upload AuthorizedUpload, + authorization RenditionAuthorization, evidence EvidencePolicy, rendition RenditionPolicy, +) (RenditionExecutionSnapshotV1, error) { + authorization = canonicalExecutionAuthorization(authorization) + if _, err := ValidateRenditionProviderRequestAt(now, provider, upload, authorization); err != nil { + return RenditionExecutionSnapshotV1{}, err + } + identity, err := NewRenditionExecutionIdentityV1( + upload.Metadata(), authorization, evidence, rendition) + if err != nil { + return RenditionExecutionSnapshotV1{}, err + } + return NewRenditionExecutionSnapshotV1(identity, authorization) +} + +// NewRenditionExecutionSnapshotV1 combines a validated stable identity with +// its exact sealed authorization interval. Provider validation is performed by +// SealRenditionExecutionAt for new work and again by ResumeRendition before a +// known handle is consumed. +func NewRenditionExecutionSnapshotV1( + identity RenditionExecutionIdentityV1, authorization RenditionAuthorization, +) (RenditionExecutionSnapshotV1, error) { + snapshot := RenditionExecutionSnapshotV1{ + ContractVersion: RenditionExecutionSnapshotContractV1, + Identity: cloneRenditionExecutionIdentity(identity), + Authorization: canonicalExecutionAuthorization(authorization), + } + if _, err := CanonicalRenditionExecutionSnapshotV1(snapshot); err != nil { + return RenditionExecutionSnapshotV1{}, err + } + return cloneRenditionExecutionSnapshot(snapshot), nil +} + +// CanonicalRenditionExecutionSnapshotV1 validates and deterministically +// encodes the provider-neutral durable resume authority. +func CanonicalRenditionExecutionSnapshotV1(snapshot RenditionExecutionSnapshotV1) ([]byte, error) { + snapshot = cloneRenditionExecutionSnapshot(snapshot) + if err := validateRenditionExecutionSnapshot(snapshot); err != nil { + return nil, err + } + encoded, err := canonicalJSON(snapshot) + if err != nil { + return nil, fmt.Errorf("encoding rendition execution snapshot: %w", err) + } + return encoded, nil +} + +// ParseRenditionExecutionSnapshotV1 decodes only exact canonical snapshots. +func ParseRenditionExecutionSnapshotV1(raw []byte) (RenditionExecutionSnapshotV1, error) { + var snapshot RenditionExecutionSnapshotV1 + if err := json.Unmarshal(raw, &snapshot, json.RejectUnknownMembers(true)); err != nil { + return RenditionExecutionSnapshotV1{}, fmt.Errorf("decoding rendition execution snapshot: %w", err) + } + canonical, err := CanonicalRenditionExecutionSnapshotV1(snapshot) + if err != nil { + return RenditionExecutionSnapshotV1{}, err + } + if !bytes.Equal(raw, canonical) { + return RenditionExecutionSnapshotV1{}, errors.New("rendition execution snapshot is not canonical") + } + return cloneRenditionExecutionSnapshot(snapshot), nil +} + +// Policies reconstructs the executable local policies from their complete +// frozen identities. +func (snapshot RenditionExecutionSnapshotV1) Policies() ( + EvidencePolicy, RenditionPolicy, error, +) { + if err := validateRenditionExecutionSnapshot(snapshot); err != nil { + return EvidencePolicy{}, RenditionPolicy{}, err + } + evidence, err := evidencePolicyFromIdentity(snapshot.Identity.EvidencePolicy) + if err != nil { + return EvidencePolicy{}, RenditionPolicy{}, err + } + rendition, err := renditionPolicyFromIdentity(snapshot.Identity.RenditionPolicy) + if err != nil { + return EvidencePolicy{}, RenditionPolicy{}, err + } + return evidence, rendition, nil +} + +// ResumeRendition continues exactly one provider-issued durable operation. +// It deliberately does not require the original upload and does not reopen a +// transient upload authorization window. The receipt start must remain inside +// the sealed interval, while completion records the truthful later observation +// time of the already-authorized durable operation. +func ResumeRendition( + ctx context.Context, provider RenditionProvider, snapshot RenditionExecutionSnapshotV1, + handle RenditionResumeHandle, checkpoint RenditionResumeCheckpoint, +) (RenditionResult, error) { + if nilInterface(provider) { + return RenditionResult{}, errors.New("rendition provider is required") + } + if err := validateRenditionResumeHandle(handle); err != nil { + return RenditionResult{}, err + } + snapshot = cloneRenditionExecutionSnapshot(snapshot) + if err := validateRenditionExecutionSnapshot(snapshot); err != nil { + return RenditionResult{}, err + } + descriptor := cloneRenditionDescriptor(provider.Descriptor()) + if err := validateRenditionDescriptor(descriptor); err != nil { + return RenditionResult{}, err + } + if second := cloneRenditionDescriptor(provider.Descriptor()); !equalRenditionDescriptors(descriptor, second) { + return RenditionResult{}, errors.New("rendition descriptor changed during resume validation") + } + if err := validateRenditionAuthorization( + descriptor, snapshot.Identity.Upload, snapshot.Authorization); err != nil { + return RenditionResult{}, err + } + resumable, ok := provider.(ResumableRenditionProvider) + if !ok { + return RenditionResult{}, errors.New("rendition provider does not support durable resume") + } + if checkpoint == nil { + checkpoint = func(RenditionResumeHandle) error { return nil } + } + checkedCheckpoint := func(next RenditionResumeHandle) error { + if err := validateRenditionResumeHandle(next); err != nil { + return err + } + return checkpoint(next) + } + result, err := resumable.RenderResumable( + ctx, nil, cloneRenditionAuthorization(snapshot.Authorization), &handle, checkedCheckpoint) + if err != nil { + if classified := ValidateRenditionProviderError(err); classified != nil { + return RenditionResult{}, classified + } + return RenditionResult{}, err + } + if err := validateResumedRenditionResult(descriptor, snapshot.Authorization, result); err != nil { + return RenditionResult{}, err + } + return result, nil +} + +func validateRenditionExecutionIdentity(identity RenditionExecutionIdentityV1) error { + if identity.ContractVersion != RenditionExecutionIdentityContractV1 { + return errors.New("rendition execution identity contract is invalid") + } + if err := validateAuthorizedUploadMetadata(identity.Upload); err != nil { + return err + } + authorization := identity.Authorization + if err := validateStableToken(authorization.ProviderID, "authorization provider ID", 128); err != nil { + return err + } + for subject, value := range map[string]string{ + "authorization descriptor fingerprint": authorization.DescriptorFingerprint, + "authorization policy fingerprint": authorization.PolicyFingerprint, + "authorization rendition request fingerprint": authorization.RenditionRequestFingerprint, + "authorization source SHA-256": authorization.SourceSHA256, + "authorization capability record checksum": authorization.CapabilityRecordChecksum, + "authorization provider metadata checksum": authorization.ProviderMetadataChecksum, + } { + if err := validateFingerprint(value, subject); err != nil { + return err + } + } + if authorization.SourceSHA256 != identity.Upload.SHA256 || + authorization.SourceBytes != identity.Upload.ByteLength || + authorization.CapabilityRecordChecksum != identity.Upload.CapabilityRecordChecksum || + authorization.ProviderMetadataChecksum != identity.Upload.ProviderMetadataChecksum || + authorization.MediaFamily != identity.Upload.MediaFamily || + authorization.MediaType != identity.Upload.MediaType || + authorization.InputKind != identity.Upload.InputKind { + return errors.New("rendition execution authorization does not match upload metadata") + } + if !slices.IsSorted(authorization.AllowedArtifactRoles) { + return errors.New("rendition execution artifact roles are not canonical") + } + seen := make(map[EvidenceArtifactRole]struct{}, len(authorization.AllowedArtifactRoles)) + for _, role := range authorization.AllowedArtifactRoles { + if !validProfileArtifactRole(role) { + return fmt.Errorf("rendition execution artifact role %q is invalid", role) + } + if _, ok := seen[role]; ok { + return fmt.Errorf("rendition execution artifact role %q is duplicated", role) + } + seen[role] = struct{}{} + } + if authorization.SourceBytes <= 0 || authorization.SourceBytes > maxRenditionSourceBytes || + authorization.MaxProviderMarkdownBytes < 0 || authorization.MaxProviderMarkdownBytes > maxRenditionMarkdownBytes || + authorization.MaxArtifactBytes < 0 || authorization.MaxArtifactBytes > maxRenditionArtifactBytes || + authorization.MaxArtifacts < 0 || authorization.MaxArtifacts > maxRenditionArtifacts || + authorization.MaxArtifacts > len(authorization.AllowedArtifactRoles) || + authorization.MaxTotalResultBytes <= 0 || authorization.MaxTotalResultBytes > maxRenditionTotalResultBytes { + return errors.New("rendition execution authorization bounds are invalid") + } + if _, err := evidencePolicyFromIdentity(identity.EvidencePolicy); err != nil { + return err + } + _, err := renditionPolicyFromIdentity(identity.RenditionPolicy) + return err +} + +func validateRenditionExecutionSnapshot(snapshot RenditionExecutionSnapshotV1) error { + if snapshot.ContractVersion != RenditionExecutionSnapshotContractV1 { + return errors.New("rendition execution snapshot contract is invalid") + } + if err := validateRenditionExecutionIdentity(snapshot.Identity); err != nil { + return err + } + authorization := canonicalExecutionAuthorization(snapshot.Authorization) + if !reflect.DeepEqual(authorization, snapshot.Authorization) { + return errors.New("rendition execution snapshot authorization is not canonical") + } + if !reflect.DeepEqual(renditionAuthorizationIdentity(authorization), snapshot.Identity.Authorization) { + return errors.New("rendition execution snapshot authorization identity drifted") + } + authorizedAt, err := parseRenditionTimestamp(authorization.AuthorizedAt) + if err != nil { + return errors.New("rendition execution authorization time is invalid") + } + expiresAt, err := parseRenditionTimestamp(authorization.ExpiresAt) + if err != nil || !expiresAt.After(authorizedAt) { + return errors.New("rendition execution authorization expiry is invalid") + } + return nil +} + +func evidencePolicyFromIdentity(identity EvidencePolicyIdentity) (EvidencePolicy, error) { + policy, err := NewEvidencePolicy(identity.MaxDocumentChars) + if err != nil { + return EvidencePolicy{}, err + } + if policy.Identity() != identity { + return EvidencePolicy{}, errors.New("document evidence policy identity is unsupported") + } + return policy, nil +} + +func renditionPolicyFromIdentity(identity RenditionPolicyIdentity) (RenditionPolicy, error) { + normalization, err := NewNormalizePolicy(identity.Normalization.MaxDocumentChars) + if err != nil { + return RenditionPolicy{}, err + } + if normalization.Identity() != identity.Normalization { + return RenditionPolicy{}, errors.New("document normalization policy identity is unsupported") + } + policy, err := NewRenditionPolicy(normalization, identity.MaxSegmentRunes) + if err != nil { + return RenditionPolicy{}, err + } + return policy, nil +} + +func renditionAuthorizationIdentity( + authorization RenditionAuthorization, +) RenditionAuthorizationIdentityV1 { + return RenditionAuthorizationIdentityV1{ + ProviderID: authorization.ProviderID, DescriptorFingerprint: authorization.DescriptorFingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, + RenditionRequestFingerprint: authorization.RenditionRequestFingerprint, + SourceSHA256: authorization.SourceSHA256, SourceBytes: authorization.SourceBytes, + CapabilityRecordChecksum: authorization.CapabilityRecordChecksum, + ProviderMetadataChecksum: authorization.ProviderMetadataChecksum, + MediaFamily: authorization.MediaFamily, MediaType: authorization.MediaType, + InputKind: authorization.InputKind, + AllowedArtifactRoles: slices.Clone(authorization.AllowedArtifactRoles), + MaxProviderMarkdownBytes: authorization.MaxProviderMarkdownBytes, + MaxArtifactBytes: authorization.MaxArtifactBytes, MaxArtifacts: authorization.MaxArtifacts, + MaxTotalResultBytes: authorization.MaxTotalResultBytes, + } +} + +func canonicalExecutionAuthorization(authorization RenditionAuthorization) RenditionAuthorization { + authorization = cloneRenditionAuthorization(authorization) + slices.Sort(authorization.AllowedArtifactRoles) + return authorization +} + +func cloneRenditionExecutionIdentity(identity RenditionExecutionIdentityV1) RenditionExecutionIdentityV1 { + identity.Authorization.AllowedArtifactRoles = slices.Clone(identity.Authorization.AllowedArtifactRoles) + return identity +} + +func cloneRenditionExecutionSnapshot(snapshot RenditionExecutionSnapshotV1) RenditionExecutionSnapshotV1 { + snapshot.Identity = cloneRenditionExecutionIdentity(snapshot.Identity) + snapshot.Authorization = cloneRenditionAuthorization(snapshot.Authorization) + return snapshot +} diff --git a/document/execution_test.go b/document/execution_test.go new file mode 100644 index 00000000..6323de3f --- /dev/null +++ b/document/execution_test.go @@ -0,0 +1,148 @@ +package document + +import ( + "bytes" + "io" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRenditionExecutionIdentityBindsEveryStableProviderAndOutputInput(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + evidence, err := NewEvidencePolicy(100_000) + require.NoError(t, err) + normalization, err := NewNormalizePolicy(100_000) + require.NoError(t, err) + rendition, err := NewRenditionPolicy(normalization, 1_000) + require.NoError(t, err) + base, err := NewRenditionExecutionIdentityV1(metadata, authorization, evidence, rendition) + require.NoError(t, err) + _, baseFingerprint, err := CanonicalRenditionExecutionIdentityV1(base) + require.NoError(t, err) + + mutations := map[string]func(*AuthorizedUploadMetadata, *RenditionAuthorization, *EvidencePolicy, *RenditionPolicy){ + "filename": func(metadata *AuthorizedUploadMetadata, _ *RenditionAuthorization, _ *EvidencePolicy, _ *RenditionPolicy) { + metadata.Filename = "renamed.pdf" + }, + "capability checksum": func(metadata *AuthorizedUploadMetadata, authorization *RenditionAuthorization, _ *EvidencePolicy, _ *RenditionPolicy) { + metadata.CapabilityRecordChecksum = sha256Hex([]byte("changed-capability")) + authorization.CapabilityRecordChecksum = metadata.CapabilityRecordChecksum + }, + "provider metadata checksum": func(metadata *AuthorizedUploadMetadata, authorization *RenditionAuthorization, _ *EvidencePolicy, _ *RenditionPolicy) { + metadata.ProviderMetadataChecksum = sha256Hex([]byte("changed-provider-metadata")) + authorization.ProviderMetadataChecksum = metadata.ProviderMetadataChecksum + }, + "result limit": func(_ *AuthorizedUploadMetadata, authorization *RenditionAuthorization, _ *EvidencePolicy, _ *RenditionPolicy) { + authorization.MaxTotalResultBytes++ + }, + "evidence policy": func(_ *AuthorizedUploadMetadata, _ *RenditionAuthorization, evidence *EvidencePolicy, _ *RenditionPolicy) { + *evidence, _ = NewEvidencePolicy(99_999) + }, + "normalization policy": func(_ *AuthorizedUploadMetadata, _ *RenditionAuthorization, _ *EvidencePolicy, rendition *RenditionPolicy) { + normalization, _ := NewNormalizePolicy(99_999) + *rendition, _ = NewRenditionPolicy(normalization, 1_000) + }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changedMetadata := metadata + changedAuthorization := authorization + changedEvidence := evidence + changedRendition := rendition + mutate(&changedMetadata, &changedAuthorization, &changedEvidence, &changedRendition) + identity, err := NewRenditionExecutionIdentityV1( + changedMetadata, changedAuthorization, changedEvidence, changedRendition) + require.NoError(t, err) + _, fingerprint, err := CanonicalRenditionExecutionIdentityV1(identity) + require.NoError(t, err) + assert.NotEqual(t, baseFingerprint, fingerprint) + }) + } + + authorization.AuthorizedAt = time.Now().UTC().Add(-2 * time.Minute).Format(renditionTimestampForm) + authorization.ExpiresAt = time.Now().UTC().Add(5 * time.Minute).Format(renditionTimestampForm) + same, err := NewRenditionExecutionIdentityV1(metadata, authorization, evidence, rendition) + require.NoError(t, err) + _, sameFingerprint, err := CanonicalRenditionExecutionIdentityV1(same) + require.NoError(t, err) + assert.Equal(t, baseFingerprint, sameFingerprint, + "transient authorization windows must not split a genuinely shared build") +} + +func TestResumeRenditionUsesOriginalSealedAuthorizationWithoutUpload(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + historical := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC) + authorization.AuthorizedAt = historical.Format(renditionTimestampForm) + authorization.ExpiresAt = historical.Add(10 * time.Minute).Format(renditionTimestampForm) + evidence, err := NewEvidencePolicy(100_000) + require.NoError(t, err) + normalization, err := NewNormalizePolicy(100_000) + require.NoError(t, err) + rendition, err := NewRenditionPolicy(normalization, 1_000) + require.NoError(t, err) + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader([]byte("synthetic exact source"))), + metadata: metadata, + } + provider := &syntheticResumableRenditionProvider{ + descriptor: descriptor, result: validRenditionResult(descriptor, authorization)} + snapshot, err := SealRenditionExecutionAt( + historical.Add(time.Minute), provider, upload, authorization, evidence, rendition) + require.NoError(t, err) + canonical, err := CanonicalRenditionExecutionSnapshotV1(snapshot) + require.NoError(t, err) + restored, err := ParseRenditionExecutionSnapshotV1(canonical) + require.NoError(t, err) + + result, err := ResumeRendition( + t.Context(), provider, restored, RenditionResumeHandle{Value: "remote-job-1"}, nil) + require.NoError(t, err) + assert.Equal(t, provider.result, result) + assert.True(t, provider.uploadWasNil) + assert.Equal(t, "remote-job-1", provider.resume.Value) + assert.Equal(t, authorization, restored.Authorization, + "the original authorization interval and limits are durable resume authority") +} + +func TestResumeRenditionAllowsOnlyCompletionAfterTheSealedInterval(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + authorizedAt, err := parseRenditionTimestamp(authorization.AuthorizedAt) + require.NoError(t, err) + expiresAt, err := parseRenditionTimestamp(authorization.ExpiresAt) + require.NoError(t, err) + evidence, err := NewEvidencePolicy(100_000) + require.NoError(t, err) + normalization, err := NewNormalizePolicy(100_000) + require.NoError(t, err) + rendition, err := NewRenditionPolicy(normalization, 1_000) + require.NoError(t, err) + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader([]byte("synthetic exact source"))), + metadata: metadata, + } + provider := &syntheticResumableRenditionProvider{ + descriptor: descriptor, result: validRenditionResult(descriptor, authorization)} + provider.result.Receipt.CompletedAt = expiresAt.Add(time.Hour).Format(renditionTimestampForm) + snapshot, err := SealRenditionExecutionAt( + authorizedAt.Add(time.Minute), provider, upload, authorization, evidence, rendition) + require.NoError(t, err) + + result, err := ResumeRendition( + t.Context(), provider, snapshot, RenditionResumeHandle{Value: "remote-job-1"}, nil) + require.NoError(t, err) + assert.Equal(t, provider.result.Receipt.CompletedAt, result.Receipt.CompletedAt) + + provider.result.Receipt.StartedAt = authorizedAt.Add(-time.Nanosecond).Format(renditionTimestampForm) + _, err = ResumeRendition( + t.Context(), provider, snapshot, RenditionResumeHandle{Value: "remote-job-1"}, nil) + require.ErrorContains(t, err, "outside the authorization interval") +} diff --git a/document/geminiembed/client.go b/document/geminiembed/client.go new file mode 100644 index 00000000..e2a69c05 --- /dev/null +++ b/document/geminiembed/client.go @@ -0,0 +1,1101 @@ +package geminiembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "net/http" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "sync" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" +) + +const ( + maxSecretBytes = 64 << 10 + maxUsageValue = int64(1 << 50) + fileClockSkew = 5 * time.Minute +) + +var _ document.EmbeddingProvider = (*Client)(nil) + +type wireRequest struct { + Model string `json:"model"` + Content wireContent `json:"content"` + OutputDimensionality int `json:"outputDimensionality"` +} + +type wireContent struct { + Parts []wirePart `json:"parts"` +} + +type wirePart struct { + Text string `json:"text,omitzero"` + InlineData *wireInlineData `json:"inlineData,omitempty"` + FileData *wireFileData `json:"fileData,omitempty"` +} + +type wireInlineData struct { + MIMEType string `json:"mimeType"` + Data string `json:"data"` +} + +type wireFileData struct { + MIMEType string `json:"mimeType"` + FileURI string `json:"fileUri"` +} + +type wireStartUpload struct { + File struct{} `json:"file"` +} + +type wireCreateFileResponse struct { + File wireFile `json:"file"` +} + +type wireFile struct { + Name string `json:"name"` + DisplayName string `json:"displayName"` + MIMEType string `json:"mimeType"` + SizeBytes string `json:"sizeBytes"` + CreateTime string `json:"createTime"` + UpdateTime string `json:"updateTime"` + ExpirationTime string `json:"expirationTime"` + SHA256Hash string `json:"sha256Hash"` + URI string `json:"uri"` + DownloadURI string `json:"downloadUri"` + State string `json:"state"` + Source string `json:"source"` + Error *wireFileError `json:"error,omitempty"` + VideoMetadata *wireVideoMetadata `json:"videoMetadata,omitempty"` +} + +type wireFileError struct { + Code int64 `json:"code"` + Message string `json:"message"` + Status string `json:"status,omitempty"` + Details []map[string]any `json:"details,omitempty"` +} + +type wireVideoMetadata struct { + VideoDuration string `json:"videoDuration"` +} + +type wireResponse struct { + Embedding wireEmbedding `json:"embedding"` + UsageMetadata *wireUsage `json:"usageMetadata,omitempty"` +} + +type wireEmbedding struct { + Values []float32 `json:"values"` +} + +type wireUsage struct { + PromptTokenCount *int64 `json:"promptTokenCount,omitempty"` + CachedContentTokenCount *int64 `json:"cachedContentTokenCount,omitempty"` + CandidatesTokenCount *int64 `json:"candidatesTokenCount,omitempty"` + ToolUsePromptTokenCount *int64 `json:"toolUsePromptTokenCount,omitempty"` + ThoughtsTokenCount *int64 `json:"thoughtsTokenCount,omitempty"` + TotalTokenCount *int64 `json:"totalTokenCount,omitempty"` +} + +// Receipt is bounded provider provenance and numeric usage without request, +// response, vector, source, or credential material. +type Receipt struct { + ProviderID string + DescriptorFingerprint string + PolicyFingerprint string + Model string + ModelRevision string + Transport Transport + RequestCount int + PromptTokens int64 + CachedContentTokens int64 + CandidateTokens int64 + ToolUsePromptTokens int64 + ThoughtTokens int64 + TotalTokens int64 + ProviderResponseIDs []string + OmittedProviderResponseIDs int + Warnings []string + ProviderRetentionCeiling time.Duration +} + +type Execution struct { + Result document.EmbeddingResult + Receipt Receipt +} + +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + return client.embed(ctx, inputs, authorization, nil) +} + +func (client *Client) EmbedWithReceipt(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (Execution, error) { + if client == nil { + return Execution{}, errors.New("gemini embed: client is required") + } + receipt := Receipt{ProviderID: ProviderID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: client.descriptor.PolicyFingerprint, Model: Model, ModelRevision: client.descriptor.ModelRevision, + Transport: client.profile.Transport, ProviderRetentionCeiling: profileRetention(client.profile.Transport)} + result, err := client.embed(ctx, inputs, authorization, &receipt) + if err != nil { + return Execution{}, err + } + return Execution{Result: result, Receipt: receipt}, nil +} + +func (client *Client) embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization, receipt *Receipt) (document.EmbeddingResult, error) { + if client == nil || ctx == nil { + return document.EmbeddingResult{}, errors.New("gemini embed: client and context are required") + } + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + frozenInputs, enrolled, enrollmentErr := enrollOriginalUploads(inputs) + defer closeEnrolledUploads(enrolled) + if enrollmentErr != nil { + return document.EmbeddingResult{}, enrollmentErr + } + if err := requestCtx.Err(); err != nil { + return document.EmbeddingResult{}, fmt.Errorf("gemini embed: embedding canceled: %w", err) + } + if err := document.ValidateEmbeddingProviderRequest(client, frozenInputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxInputBytes > client.profile.MaxInputBytes || authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, errors.New("gemini embed: embedding authorization exceeds profile capacity") + } + + sourceGate := newActiveSourceGate() + closeFinished := make(chan struct{}) + stopClose := context.AfterFunc(requestCtx, func() { + sourceGate.Cancel() + close(closeFinished) + }) + defer func() { + if !stopClose() { + <-closeFinished + } + }() + + prepared := make([]preparedInput, len(frozenInputs)) + defer clearPreparedInputs(prepared) + var inputBytes int64 + for index, input := range frozenInputs { + switch { + case input.Role == document.EmbeddingRoleDocument && input.Kind == document.EmbeddingInputRenditionChunk: + text := client.descriptor.ModelInput.EncodeDocument(input.Text) + if int64(len(text)) > client.profile.MaxInputBytes-inputBytes { + return document.EmbeddingResult{}, errors.New("gemini embed: embedding input exceeds profile byte capacity") + } + inputBytes += int64(len(text)) + payload, err := client.marshalRequest(wirePart{Text: text}) + if err != nil { + return document.EmbeddingResult{}, err + } + prepared[index].payload = payload + case input.Role == document.EmbeddingRoleQuery && input.Kind == document.EmbeddingInputQueryText: + text := client.descriptor.ModelInput.EncodeQuery(input.Text) + if int64(len(text)) > client.profile.MaxInputBytes-inputBytes { + return document.EmbeddingResult{}, errors.New("gemini embed: embedding input exceeds profile byte capacity") + } + inputBytes += int64(len(text)) + payload, err := client.marshalRequest(wirePart{Text: text}) + if err != nil { + return document.EmbeddingResult{}, err + } + prepared[index].payload = payload + case input.Role == document.EmbeddingRoleDocument && input.Kind == document.EmbeddingInputOriginalFile: + frozen, ok := input.Source.(*frozenUpload) + if !ok { + return document.EmbeddingResult{}, errors.New("gemini embed: original upload was not frozen") + } + if err := client.validateDirectCapability(frozen.enrolled, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if client.profile.Transport == TransportFilesAPI { + if err := client.preflightFileDataRequest(frozen.enrolled.metadata.MediaType); err != nil { + return document.EmbeddingResult{}, err + } + } + verified, err := client.readDirectFile(requestCtx, frozen.enrolled, sourceGate) + if err != nil { + return document.EmbeddingResult{}, err + } + prepared[index].file = &verified + if client.profile.Transport == TransportInline { + encoded := base64.StdEncoding.EncodeToString(verified.data) + payload, encodeErr := client.marshalRequest(wirePart{InlineData: &wireInlineData{ + MIMEType: verified.metadata.MediaType, Data: encoded, + }}) + if encodeErr != nil { + return document.EmbeddingResult{}, encodeErr + } + limit := int64(100 << 20) + if verified.metadata.MediaType == "application/pdf" { + limit = 50 << 20 + } + if int64(len(payload)) > limit { + clear(payload) + return document.EmbeddingResult{}, errors.New("gemini embed: inline request exceeds provider byte capacity") + } + prepared[index].payload = payload + } + default: + return document.EmbeddingResult{}, errors.New("gemini embed: unsupported input role or kind") + } + } + secret, err := client.secrets.ResolveSecret(requestCtx, client.profile.SecretBinding) + if err != nil || !validSecret(secret) { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("gemini embed: credential resolution canceled: %w", contextErr) + } + return document.EmbeddingResult{}, errors.New("gemini embed: API-key resolution failed") + } + result := document.EmbeddingResult{Vectors: make([]document.EmbeddingVector, len(frozenInputs))} + for index := range prepared { + var vector []float32 + if prepared[index].file != nil && client.profile.Transport == TransportFilesAPI { + vector, err = client.executeFile(requestCtx, *prepared[index].file, secret, receipt) + } else { + vector, err = client.execute(requestCtx, prepared[index].payload, secret, receipt) + } + if err != nil { + return document.EmbeddingResult{}, err + } + result.Vectors[index] = document.EmbeddingVector{Key: frozenInputs[index].Key, Values: vector} + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, frozenInputs, authorization, result); err != nil { + return document.EmbeddingResult{}, errors.New("gemini embed: provider response violates embedding contract") + } + return result, nil +} + +type preparedInput struct { + payload []byte + file *verifiedFile +} + +func clearPreparedInputs(prepared []preparedInput) { + for index := range prepared { + clear(prepared[index].payload) + if prepared[index].file != nil { + clear(prepared[index].file.data) + } + } +} + +type verifiedFile struct { + data []byte + metadata document.AuthorizedUploadMetadata + capability media.CapabilityRecord +} + +func (client *Client) marshalRequest(part wirePart) ([]byte, error) { + payload, err := json.Marshal(wireRequest{Model: "models/" + Model, + Content: wireContent{Parts: []wirePart{part}}, OutputDimensionality: client.descriptor.Dimension}) + if err != nil { + return nil, errors.New("gemini embed: request encoding failed") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + clear(payload) + return nil, errors.New("gemini embed: embedding request exceeds profile byte capacity") + } + return payload, nil +} + +func (client *Client) preflightFileDataRequest(mediaType string) error { + payload, err := client.marshalRequest(wirePart{FileData: &wireFileData{ + MIMEType: mediaType, FileURI: maximumProviderFileURI(), + }}) + clear(payload) + if err != nil { + return errors.New("gemini embed: fileData request exceeds profile byte capacity") + } + return nil +} + +type capabilityCarrier interface { + CapabilityRecord() media.CapabilityRecord +} + +type enrolledUpload struct { + source document.AuthorizedUpload + metadata document.AuthorizedUploadMetadata + capability media.CapabilityRecord + hasCapability bool + + closeOnce sync.Once + closeErr error +} + +func (upload *enrolledUpload) Close() error { + upload.closeOnce.Do(func() { upload.closeErr = upload.source.Close() }) + return upload.closeErr +} + +func (upload *enrolledUpload) liveMetadata() document.AuthorizedUploadMetadata { + return upload.source.Metadata() +} + +func (upload *enrolledUpload) liveCapability() (media.CapabilityRecord, bool) { + carrier, ok := upload.source.(capabilityCarrier) + if !ok { + return media.CapabilityRecord{}, false + } + return carrier.CapabilityRecord(), true +} + +type frozenUpload struct { + enrolled *enrolledUpload +} + +func (upload *frozenUpload) Read(value []byte) (int, error) { + return upload.enrolled.source.Read(value) +} +func (upload *frozenUpload) Close() error { return upload.enrolled.Close() } +func (upload *frozenUpload) Metadata() document.AuthorizedUploadMetadata { + return upload.enrolled.metadata +} +func (upload *frozenUpload) CapabilityRecord() media.CapabilityRecord { + return upload.enrolled.capability +} + +func enrollOriginalUploads(inputs []document.EmbeddingInput) ([]document.EmbeddingInput, []*enrolledUpload, error) { + frozen := slices.Clone(inputs) + enrolled := make([]*enrolledUpload, 0, len(inputs)) + owners := make([]*enrolledUpload, len(inputs)) + seen := make(map[document.AuthorizedUpload]*enrolledUpload, len(inputs)) + var enrollmentErr error + for index := range frozen { + frozen[index].HeadingPath = slices.Clone(frozen[index].HeadingPath) + frozen[index].SourceSpans = slices.Clone(frozen[index].SourceSpans) + source := frozen[index].Source + if source == nil { + if frozen[index].Kind == document.EmbeddingInputOriginalFile && enrollmentErr == nil { + enrollmentErr = errors.New("gemini embed: original upload source is nil") + } + continue + } + if nilInterface(source) { + if enrollmentErr == nil { + if frozen[index].Kind == document.EmbeddingInputOriginalFile { + enrollmentErr = errors.New("gemini embed: original upload source is nil") + } else { + enrollmentErr = errors.New("gemini embed: attached upload source is nil") + } + } + continue + } + identity := reflect.ValueOf(source) + if !identity.Comparable() || (identity.Kind() == reflect.Pointer && identity.Type().Elem().Size() == 0) { + upload := &enrolledUpload{source: source} + enrolled = append(enrolled, upload) + owners[index] = upload + if enrollmentErr == nil { + enrollmentErr = errors.New("gemini embed: original upload identity is not safely comparable") + } + continue + } + if upload, duplicate := seen[source]; duplicate { + owners[index] = upload + if enrollmentErr == nil { + enrollmentErr = errors.New("gemini embed: original upload source is repeated") + } + continue + } + upload := &enrolledUpload{source: source} + enrolled = append(enrolled, upload) + owners[index] = upload + seen[source] = upload + if (frozen[index].Role != document.EmbeddingRoleDocument || + frozen[index].Kind != document.EmbeddingInputOriginalFile) && enrollmentErr == nil { + enrollmentErr = errors.New("gemini embed: upload source is attached to an unsupported input role or kind") + } + } + if enrollmentErr != nil { + return frozen, enrolled, enrollmentErr + } + for _, upload := range enrolled { + upload.metadata = upload.source.Metadata() + if carrier, ok := upload.source.(capabilityCarrier); ok { + upload.capability = carrier.CapabilityRecord() + upload.hasCapability = true + } + } + for index, upload := range owners { + if upload != nil { + frozen[index].Source = &frozenUpload{enrolled: upload} + } + } + return frozen, enrolled, nil +} + +func closeEnrolledUploads(uploads []*enrolledUpload) { + for _, upload := range uploads { + _ = upload.Close() + } +} + +type activeSourceGate struct { + mu sync.Mutex + canceled bool + nextToken uint64 + activeToken uint64 + active *enrolledUpload +} + +func newActiveSourceGate() *activeSourceGate { return new(activeSourceGate) } + +func (gate *activeSourceGate) Begin(source *enrolledUpload) (uint64, bool) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.canceled { + return 0, false + } + gate.nextToken++ + gate.activeToken = gate.nextToken + gate.active = source + return gate.activeToken, true +} + +func (gate *activeSourceGate) End(token uint64) { + gate.mu.Lock() + defer gate.mu.Unlock() + if gate.activeToken != token { + return + } + gate.active = nil + gate.activeToken = 0 +} + +func (gate *activeSourceGate) Cancel() { + gate.mu.Lock() + if gate.canceled { + gate.mu.Unlock() + return + } + gate.canceled = true + active := gate.active + gate.active = nil + gate.activeToken = 0 + gate.mu.Unlock() + if active != nil { + _ = active.Close() + } +} + +func (client *Client) validateDirectCapability(upload *enrolledUpload, authorization document.EmbeddingAuthorization) error { + if !upload.hasCapability { + return errors.New("gemini embed: direct-file upload lacks local capability authority") + } + record := upload.capability + if err := media.ValidateCapabilityRecord(record); err != nil { + return errors.New("gemini embed: direct-file capability record is invalid") + } + policy, local := record.InspectionPolicy() + if !local || !record.Eligible || record.Reason != media.CapabilityReasonEligible { + return errors.New("gemini embed: direct-file capability record is not locally eligible") + } + metadata := upload.metadata + if metadata.Filename != policy.Filename || metadata.MediaFamily != record.MediaFamily || + metadata.MediaType != record.MediaType || metadata.ByteLength != record.SourceBytes || + metadata.SHA256 != record.SourceSHA256 || metadata.CapabilityRecordChecksum != record.Checksum || + metadata.InputKind != document.RenditionInputOriginalFile || + record.DescriptorFingerprint != client.descriptor.Fingerprint || + record.ProfileFingerprint != client.profile.CapabilityProfileFingerprint || + record.DisclosureFingerprint != client.profile.DisclosureFingerprint || + record.InputKind != document.RenditionInputOriginalFile || + policy.DescriptorFingerprint != client.descriptor.Fingerprint || + policy.ProfileFingerprint != client.profile.CapabilityProfileFingerprint || + policy.DisclosureFingerprint != client.profile.DisclosureFingerprint || + policy.InputKind != document.RenditionInputOriginalFile { + return errors.New("gemini embed: direct-file capability authority does not match upload and profile") + } + if record.SourceBytes > authorization.MaxInputBytes || record.SourceBytes > client.profile.MaxInputBytes || + policy.MaxSourceBytes > client.profile.MaxInputBytes { + return errors.New("gemini embed: direct-file source exceeds byte capacity") + } + if client.profile.Transport == TransportFilesAPI && record.SourceBytes > client.profile.MaxRequestBytes { + return errors.New("gemini embed: raw file upload exceeds request byte capacity") + } + if !geminiCapabilitySupported(record, policy) { + return errors.New("gemini embed: direct-file capability is unsupported or over limit") + } + return nil +} + +func geminiCapabilitySupported(record media.CapabilityRecord, policy media.InspectionPolicy) bool { + switch record.MediaFamily { + case "image": + return (record.MediaType == "image/png" || record.MediaType == "image/jpeg") && + record.Measurements.Pixels > 0 && record.Measurements.Frames > 0 && + policy.MaxPixels > 0 && policy.MaxFrames > 0 + case "audio": + return (record.MediaType == "audio/wav" || record.MediaType == "audio/mpeg") && + record.Measurements.DurationMS > 0 && record.Measurements.DurationMS <= 180_000 && + policy.MaxDurationMS > 0 && policy.MaxDurationMS <= 180_000 + case "video": + return (record.MediaType == "video/mp4" || record.MediaType == "video/quicktime") && + record.Measurements.DurationMS > 0 && record.Measurements.DurationMS <= 120_000 && + record.Measurements.Frames > 0 && record.Measurements.Frames <= 32 && + policy.MaxDurationMS > 0 && policy.MaxDurationMS <= 120_000 && + policy.MaxFrames > 0 && policy.MaxFrames <= 32 + case "pdf": + return record.MediaType == "application/pdf" && record.Measurements.Pages > 0 && + record.Measurements.Pages <= 6 && policy.MaxPages > 0 && policy.MaxPages <= 6 + default: + return false + } +} + +func (client *Client) readDirectFile(ctx context.Context, upload *enrolledUpload, gate *activeSourceGate) (verifiedFile, error) { + if live := upload.liveMetadata(); live != upload.metadata { + return verifiedFile{}, errors.New("gemini embed: direct-file upload metadata changed") + } + if live, ok := upload.liveCapability(); !ok || live != upload.capability { + return verifiedFile{}, errors.New("gemini embed: direct-file capability authority changed") + } + token, ok := gate.Begin(upload) + if !ok { + if contextErr := ctx.Err(); contextErr != nil { + return verifiedFile{}, contextErr + } + return verifiedFile{}, errors.New("gemini embed: direct-file source transfer stopped") + } + data, readErr := io.ReadAll(io.LimitReader(upload.source, upload.metadata.ByteLength+1)) + gate.End(token) + metadataChanged := upload.liveMetadata() != upload.metadata + capability, capabilityOK := upload.liveCapability() + capabilityChanged := !capabilityOK || capability != upload.capability + closeErr := upload.Close() + if contextErr := ctx.Err(); contextErr != nil { + clear(data) + return verifiedFile{}, contextErr + } + if readErr != nil || closeErr != nil || metadataChanged || capabilityChanged || int64(len(data)) != upload.metadata.ByteLength { + clear(data) + return verifiedFile{}, errors.New("gemini embed: direct-file source changed or could not be read exactly") + } + digest := sha256.Sum256(data) + if hex.EncodeToString(digest[:]) != upload.metadata.SHA256 { + clear(data) + return verifiedFile{}, errors.New("gemini embed: direct-file source checksum changed") + } + return verifiedFile{data: data, metadata: upload.metadata, capability: upload.capability}, nil +} + +func (client *Client) executeFile(ctx context.Context, file verifiedFile, secret string, receipt *Receipt) (vector []float32, retErr error) { + uploadURL, err := client.startFileUpload(ctx, file, secret, receipt) + if err != nil { + return nil, err + } + created, responseID, err := client.finalizeFileUpload(ctx, uploadURL, file, secret, receipt) + if err != nil { + return nil, err + } + defer func() { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), client.profile.CleanupTimeout) + defer cancel() + if deleteErr := client.deleteFile(cleanupCtx, created.file.Name, secret, receipt); deleteErr != nil && receipt != nil { + addCleanupWarning(receipt) + } + }() + if !recordReceiptResponse(receipt, nil, responseID) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + active, err := client.waitForActiveFile(ctx, created, file, secret, receipt) + if err != nil { + return nil, err + } + payload, err := client.marshalRequest(wirePart{FileData: &wireFileData{ + MIMEType: file.metadata.MediaType, FileURI: active.file.URI, + }}) + if err != nil { + return nil, err + } + defer clear(payload) + return client.execute(ctx, payload, secret, receipt) +} + +func (client *Client) startFileUpload(ctx context.Context, file verifiedFile, secret string, receipt *Receipt) (*url.URL, error) { + payload, err := json.Marshal(wireStartUpload{}) + if err != nil || int64(len(payload)) > client.profile.MaxRequestBytes { + return nil, errors.New("gemini embed: file upload start request encoding failed") + } + defer clear(payload) + request, err := http.NewRequestWithContext(ctx, http.MethodPost, origin+filesUploadPath, bytes.NewReader(payload)) + if err != nil { + return nil, errors.New("gemini embed: file upload start request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("X-Goog-Api-Key", secret) + request.Header.Set("X-Goog-Upload-Protocol", "resumable") + request.Header.Set("X-Goog-Upload-Command", "start") + request.Header.Set("X-Goog-Upload-Header-Content-Length", strconv.FormatInt(file.metadata.ByteLength, 10)) + request.Header.Set("X-Goog-Upload-Header-Content-Type", file.metadata.MediaType) + if !beginReceiptRequest(receipt) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + response, err := client.http.Do(request) + if err != nil { + return nil, requestFailure(ctx, "file upload start", err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + body, err := readBounded(ctx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return nil, err + } + defer clear(body) + if len(body) != 0 { + if !isJSONContentType(response.Header.Get("Content-Type")) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + var empty struct{} + if err := json.Unmarshal(body, &empty, json.RejectUnknownMembers(true)); err != nil { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + } + uploadURL, err := validateUploadURL(response.Header.Get("X-Goog-Upload-Url")) + if err != nil { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if !recordReceiptResponse(receipt, nil, response.Header.Get("X-Goog-Request-Id")) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + return uploadURL, nil +} + +func (client *Client) finalizeFileUpload(ctx context.Context, uploadURL *url.URL, file verifiedFile, secret string, receipt *Receipt) (validatedProviderFile, string, error) { + if int64(len(file.data)) > client.profile.MaxRequestBytes { + return validatedProviderFile{}, "", errors.New("gemini embed: raw file upload exceeds request byte capacity") + } + request, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadURL.String(), bytes.NewReader(file.data)) + if err != nil { + return validatedProviderFile{}, "", errors.New("gemini embed: file upload request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("X-Goog-Api-Key", secret) + request.Header.Set("X-Goog-Upload-Offset", "0") + request.Header.Set("X-Goog-Upload-Command", "upload, finalize") + if !beginReceiptRequest(receipt) { + return validatedProviderFile{}, "", &ProviderError{Kind: ErrPermanentResponse} + } + startedAt := time.Now().UTC() + response, err := client.http.Do(request) + completedAt := time.Now().UTC() + if err != nil { + return validatedProviderFile{}, "", requestFailure(ctx, "file upload", err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return validatedProviderFile{}, "", statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if response.Header.Get("X-Goog-Upload-Status") != "final" || !isJSONContentType(response.Header.Get("Content-Type")) { + return validatedProviderFile{}, "", &ProviderError{Kind: ErrPermanentResponse} + } + body, err := readBounded(ctx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return validatedProviderFile{}, "", err + } + defer clear(body) + var decoded wireCreateFileResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return validatedProviderFile{}, "", &ProviderError{Kind: ErrPermanentResponse} + } + validated, ok := validateCreatedWireFile(decoded.File, file, startedAt, completedAt) + if !ok { + return validatedProviderFile{}, "", &ProviderError{Kind: ErrPermanentResponse} + } + return validated, response.Header.Get("X-Goog-Request-Id"), nil +} + +func (client *Client) waitForActiveFile(ctx context.Context, current validatedProviderFile, file verifiedFile, secret string, receipt *Receipt) (validatedProviderFile, error) { + if current.file.State == "ACTIVE" { + return current, nil + } + for attempt := range client.profile.MaxPollAttempts { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, origin+filesPathPrefix+strings.TrimPrefix(current.file.Name, "files/"), nil) + if err != nil { + return validatedProviderFile{}, errors.New("gemini embed: file poll request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("X-Goog-Api-Key", secret) + if !beginReceiptRequest(receipt) { + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + startedAt := time.Now().UTC() + response, err := client.http.Do(request) + completedAt := time.Now().UTC() + if err != nil { + return validatedProviderFile{}, requestFailure(ctx, "file poll", err) + } + if response.StatusCode != http.StatusOK { + _ = response.Body.Close() + return validatedProviderFile{}, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if !isJSONContentType(response.Header.Get("Content-Type")) { + _ = response.Body.Close() + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + body, readErr := readBounded(ctx, response.Body, client.profile.MaxResponseBytes) + _ = response.Body.Close() + if readErr != nil { + return validatedProviderFile{}, readErr + } + var decoded wireFile + decodeErr := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)) + clear(body) + next, valid := validatePolledWireFile(decoded, file, current, startedAt, completedAt) + if decodeErr != nil || !valid { + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + if !recordReceiptResponse(receipt, nil, response.Header.Get("X-Goog-Request-Id")) { + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + switch decoded.State { + case "ACTIVE": + return next, nil + case "PROCESSING": + current = next + if attempt+1 == client.profile.MaxPollAttempts { + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + case "FAILED": + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + default: + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} + } + timer := time.NewTimer(client.profile.PollInterval) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return validatedProviderFile{}, ctx.Err() + case <-timer.C: + } + } + return validatedProviderFile{}, &ProviderError{Kind: ErrPermanentResponse} +} + +func (client *Client) deleteFile(ctx context.Context, name, secret string, receipt *Receipt) error { + request, err := http.NewRequestWithContext(ctx, http.MethodDelete, origin+filesPathPrefix+strings.TrimPrefix(name, "files/"), nil) + if err != nil { + return errors.New("gemini embed: file deletion request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("X-Goog-Api-Key", secret) + if !beginReceiptRequest(receipt) { + return &ProviderError{Kind: ErrPermanentResponse} + } + response, err := client.http.Do(request) + if err != nil { + return requestFailure(ctx, "file deletion", err) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusNoContent { + return statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + body, err := readBounded(ctx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return err + } + defer clear(body) + if len(body) != 0 { + if !isJSONContentType(response.Header.Get("Content-Type")) { + return &ProviderError{Kind: ErrPermanentResponse} + } + var empty struct{} + if err := json.Unmarshal(body, &empty, json.RejectUnknownMembers(true)); err != nil { + return &ProviderError{Kind: ErrPermanentResponse} + } + } + if !recordReceiptResponse(receipt, nil, response.Header.Get("X-Goog-Request-Id")) { + return &ProviderError{Kind: ErrPermanentResponse} + } + return nil +} + +func validateUploadURL(value string) (*url.URL, error) { + parsed, err := url.Parse(value) + if err != nil || parsed.Scheme != "https" || parsed.Hostname() != host || + (parsed.Port() != "" && parsed.Port() != "443") || parsed.User != nil || parsed.ForceQuery || + parsed.Fragment != "" || parsed.RawPath != "" || parsed.Opaque != "" || + parsed.Path != filesUploadPath || !validUploadQuery(parsed.RawQuery) { + return nil, errors.New("gemini embed: provider upload URL is outside sealed egress") + } + return parsed, nil +} + +func validUploadQuery(rawQuery string) bool { + query, err := url.ParseQuery(rawQuery) + if err != nil || len(query) != 2 || len(query["upload_id"]) != 1 || + len(query["upload_protocol"]) != 1 || query.Get("upload_protocol") != "resumable" { + return false + } + uploadID := query.Get("upload_id") + if len(uploadID) == 0 || len(uploadID) > 1024 { + return false + } + for _, character := range uploadID { + if character != '-' && character != '_' && + (character < 'a' || character > 'z') && (character < 'A' || character > 'Z') && + (character < '0' || character > '9') { + return false + } + } + first := "upload_id=" + uploadID + "&upload_protocol=resumable" + second := "upload_protocol=resumable&upload_id=" + uploadID + return rawQuery == first || rawQuery == second +} + +type validatedProviderFile struct { + file wireFile + created time.Time + updated time.Time + expiresAt time.Time +} + +func validateWireFile(value wireFile, expected verifiedFile) (validatedProviderFile, bool) { + if !validFileName(value.Name) || value.URI != origin+"/v1beta/"+value.Name || + value.MIMEType != expected.metadata.MediaType || value.Source != "UPLOADED" || value.Error != nil || + value.State != "PROCESSING" && value.State != "ACTIVE" && value.State != "FAILED" || + len(value.DisplayName) > 512 { + return validatedProviderFile{}, false + } + size, err := strconv.ParseInt(value.SizeBytes, 10, 64) + if err != nil || size != expected.metadata.ByteLength { + return validatedProviderFile{}, false + } + hash, err := base64.StdEncoding.Strict().DecodeString(value.SHA256Hash) + if err != nil || !matchesFileHash(hash, expected.metadata.SHA256) { + return validatedProviderFile{}, false + } + created, createErr := time.Parse(time.RFC3339Nano, value.CreateTime) + updated, updateErr := time.Parse(time.RFC3339Nano, value.UpdateTime) + expires, expiryErr := time.Parse(time.RFC3339Nano, value.ExpirationTime) + if createErr != nil || updateErr != nil || expiryErr != nil || updated.Before(created) || + updated.After(expires) || !expires.After(created) || expires.Sub(created) > retentionCeiling { + return validatedProviderFile{}, false + } + return validatedProviderFile{file: value, created: created, updated: updated, expiresAt: expires}, true +} + +func validateCreatedWireFile(value wireFile, expected verifiedFile, startedAt, completedAt time.Time) (validatedProviderFile, bool) { + validated, ok := validateWireFile(value, expected) + if !ok || validated.created.Before(startedAt.Add(-fileClockSkew)) || + validated.created.After(completedAt.Add(fileClockSkew)) || + validated.updated.Before(startedAt.Add(-fileClockSkew)) || + validated.updated.After(completedAt.Add(fileClockSkew)) || + !validated.expiresAt.After(completedAt) { + return validatedProviderFile{}, false + } + return validated, true +} + +func validatePolledWireFile(value wireFile, expected verifiedFile, current validatedProviderFile, startedAt, completedAt time.Time) (validatedProviderFile, bool) { + validated, ok := validateWireFile(value, expected) + if !ok || value.Name != current.file.Name || value.URI != current.file.URI || + !validated.created.Equal(current.created) || !validated.expiresAt.Equal(current.expiresAt) || + validated.updated.Before(current.updated) || validated.updated.Before(startedAt.Add(-fileClockSkew)) || + validated.updated.After(completedAt.Add(fileClockSkew)) || !validated.expiresAt.After(completedAt) { + return validatedProviderFile{}, false + } + return validated, true +} + +func matchesFileHash(decoded []byte, expectedHex string) bool { + return len(decoded) == sha256.Size && hex.EncodeToString(decoded) == expectedHex || + len(decoded) == sha256.Size*2 && string(decoded) == expectedHex +} + +func validFileName(name string) bool { + id := strings.TrimPrefix(name, "files/") + if id == name || len(id) == 0 || len(id) > 40 || id[0] == '-' || id[len(id)-1] == '-' { + return false + } + for _, character := range id { + if character != '-' && (character < 'a' || character > 'z') && (character < '0' || character > '9') { + return false + } + } + return true +} + +func requestFailure(ctx context.Context, operation string, _ error) error { + if contextErr := ctx.Err(); contextErr != nil { + return fmt.Errorf("gemini embed: %s canceled: %w", operation, contextErr) + } + return fmt.Errorf("gemini embed: %s failed", operation) +} + +func beginReceiptRequest(receipt *Receipt) bool { + if receipt == nil { + return true + } + if receipt.RequestCount == int(^uint(0)>>1) { + return false + } + receipt.RequestCount++ + return true +} + +func addCleanupWarning(receipt *Receipt) { + if len(receipt.Warnings) >= 32 { + return + } + receipt.Warnings = append(receipt.Warnings, "provider file deletion attempt failed") +} + +func (client *Client) execute(ctx context.Context, payload []byte, secret string, receipt *Receipt) ([]float32, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodPost, origin+embedPath, bytes.NewReader(payload)) + if err != nil { + return nil, errors.New("gemini embed: request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("X-Goog-Api-Key", secret) + if !beginReceiptRequest(receipt) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + response, err := client.http.Do(request) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("gemini embed: request canceled: %w", contextErr) + } + return nil, errors.New("gemini embed: provider request failed") + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if !isJSONContentType(response.Header.Get("Content-Type")) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + body, err := readBounded(ctx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return nil, err + } + defer clear(body) + var decoded wireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if len(decoded.Embedding.Values) != client.descriptor.Dimension || !validUsage(decoded.UsageMetadata) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if !normalizeUnitVector(decoded.Embedding.Values) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if !recordReceiptResponse(receipt, decoded.UsageMetadata, response.Header.Get("X-Goog-Request-Id")) { + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + return decoded.Embedding.Values, nil +} + +func normalizeUnitVector(vector []float32) bool { + var squaredNorm float64 + for _, value := range vector { + floatValue := float64(value) + if math.IsNaN(floatValue) || math.IsInf(floatValue, 0) { + return false + } + squaredNorm += floatValue * floatValue + } + if squaredNorm == 0 || math.IsInf(squaredNorm, 0) { + return false + } + norm := math.Sqrt(squaredNorm) + for index, value := range vector { + vector[index] = float32(float64(value) / norm) + } + return true +} + +func validSecret(value string) bool { return validToken(value) && len(value) <= maxSecretBytes } + +func isJSONContentType(value string) bool { + mediaType, _, found := strings.Cut(value, ";") + return strings.EqualFold(strings.TrimSpace(mediaType), "application/json") && found || strings.EqualFold(strings.TrimSpace(value), "application/json") +} + +func readBounded(ctx context.Context, body io.Reader, maximum int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(body, maximum+1)) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("gemini embed: response read canceled: %w", contextErr) + } + return nil, errors.New("gemini embed: provider response read failed") + } + if int64(len(data)) > maximum { + clear(data) + return nil, errors.New("gemini embed: provider response exceeds byte capacity") + } + return data, nil +} + +func validUsage(usage *wireUsage) bool { + if usage == nil { + return true + } + for _, value := range []*int64{usage.PromptTokenCount, usage.CachedContentTokenCount, usage.CandidatesTokenCount, + usage.ToolUsePromptTokenCount, usage.ThoughtsTokenCount, usage.TotalTokenCount} { + if value != nil && (*value < 0 || *value > maxUsageValue) { + return false + } + } + return true +} + +func recordReceiptResponse(receipt *Receipt, usage *wireUsage, responseID string) bool { + if receipt == nil { + return true + } + if responseID != "" && !validToken(responseID) { + return false + } + if usage != nil { + values := []struct{ source, destination *int64 }{ + {usage.PromptTokenCount, &receipt.PromptTokens}, {usage.CachedContentTokenCount, &receipt.CachedContentTokens}, + {usage.CandidatesTokenCount, &receipt.CandidateTokens}, {usage.ToolUsePromptTokenCount, &receipt.ToolUsePromptTokens}, + {usage.ThoughtsTokenCount, &receipt.ThoughtTokens}, {usage.TotalTokenCount, &receipt.TotalTokens}, + } + for _, value := range values { + source, destination := value.source, value.destination + if source != nil && *source > maxUsageValue-*destination { + return false + } + } + for _, value := range values { + source, destination := value.source, value.destination + if source != nil { + *destination += *source + } + } + } + if responseID != "" { + if len(receipt.ProviderResponseIDs) < 128 { + receipt.ProviderResponseIDs = append(receipt.ProviderResponseIDs, responseID) + } else { + receipt.OmittedProviderResponseIDs++ + } + } + return true +} diff --git a/document/geminiembed/client_test.go b/document/geminiembed/client_test.go new file mode 100644 index 00000000..771e7546 --- /dev/null +++ b/document/geminiembed/client_test.go @@ -0,0 +1,290 @@ +package geminiembed + +import ( + "bytes" + "context" + "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "net/http" + "net/netip" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +// TestEmbedSendsOneTextRequestPerInputWithoutTaskType catches batching, +// request-shape, role-formatting, or result-order regressions in the Gemini +// text path. +func TestEmbedSendsOneTextRequestPerInputWithoutTaskType(t *testing.T) { + profile := geminiTestProfile(t, 768) + var requests [][]byte + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "https", request.URL.Scheme) + assert.Equal(t, "generativelanguage.googleapis.com", request.URL.Host) + assert.Equal(t, "/v1beta/models/gemini-embedding-2:embedContent", request.URL.Path) + assert.Empty(t, request.URL.RawQuery) + assert.Equal(t, "application/json", request.Header.Get("Accept")) + assert.Equal(t, "application/json", request.Header.Get("Content-Type")) + assert.Equal(t, "synthetic-key", request.Header.Get("X-Goog-Api-Key")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + requests = append(requests, body) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(768)+`]},"usageMetadata":{"promptTokenCount":3,"totalTokenCount":3}}`), nil + })) + + inputs := []document.EmbeddingInput{ + {Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage"}, + {Key: "query-1", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "question"}, + } + result, err := client.Embed(t.Context(), inputs, geminiAuthorization(profile.Descriptor)) + require.NoError(t, err) + require.Len(t, requests, 2) + if actual := string(requests[0]); actual != `{"model":"models/gemini-embedding-2","content":{"parts":[{"text":"title: none | text: passage"}]},"outputDimensionality":768}` { + t.Errorf("document request body = %q", actual) + } + if actual := string(requests[1]); actual != `{"model":"models/gemini-embedding-2","content":{"parts":[{"text":"task: search result | query: question"}]},"outputDimensionality":768}` { + t.Errorf("query request body = %q", actual) + } + for _, body := range requests { + assert.NotContains(t, string(body), `"taskType"`) + assert.NotContains(t, string(body), `"task_type"`) + } + require.Len(t, result.Vectors, 2) + assert.Equal(t, "document-1", result.Vectors[0].Key) + assert.Len(t, result.Vectors[0].Values, 768) + assert.Equal(t, "query-1", result.Vectors[1].Key) + assert.Len(t, result.Vectors[1].Values, 768) +} + +// TestEmbedRejectsMalformedVectorAndUsageContracts catches response decoding +// that admits a malformed vector or untrusted numeric usage into a result or +// receipt. +func TestEmbedRejectsMalformedVectorAndUsageContracts(t *testing.T) { + for _, testCase := range []struct { + name string + body string + }{ + {name: "wrong dimension", body: `{"embedding":{"values":[1]},"usageMetadata":{"promptTokenCount":1,"totalTokenCount":1}}`}, + {name: "non-finite vector", body: `{"embedding":{"values":[1e999]}}`}, + {name: "negative usage", body: `{"embedding":{"values":[` + testVectorJSON(128) + `]},"usageMetadata":{"promptTokenCount":-1}}`}, + {name: "schema drift", body: `{"embedding":{"values":[` + testVectorJSON(128) + `]},"provider_private":"synthetic"}`}, + } { + t.Run(testCase.name, func(t *testing.T) { + profile := geminiTestProfile(t, 128) + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return geminiJSONResponse(request, testCase.body), nil + })) + _, err := client.Embed(t.Context(), []document.EmbeddingInput{{ + Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage", + }}, geminiSingleAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.NotContains(t, err.Error(), "synthetic") + }) + } +} + +// TestEmbedNormalizesFiniteNonUnitVectorAndRejectsZero catches response +// validation that returns raw provider values despite the descriptor's fixed +// unit-length policy, or treats a zero vector as normalizable. +func TestEmbedNormalizesFiniteNonUnitVectorAndRejectsZero(t *testing.T) { + t.Run("non-unit vector", func(t *testing.T) { + profile := geminiTestProfile(t, 128) + values := make([]float32, 128) + values[0], values[1] = 3, 4 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return geminiJSONResponse(request, `{"embedding":{"values":[`+vectorJSON(values)+`]}}`), nil + })) + + result, err := client.Embed(t.Context(), []document.EmbeddingInput{{ + Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage", + }}, geminiSingleAuthorization(profile.Descriptor)) + require.NoError(t, err) + require.Len(t, result.Vectors, 1) + require.Len(t, result.Vectors[0].Values, 128) + assert.Equal(t, math.Float32bits(0.6), math.Float32bits(result.Vectors[0].Values[0])) + assert.Equal(t, math.Float32bits(0.8), math.Float32bits(result.Vectors[0].Values[1])) + assert.Equal(t, make([]float32, 126), result.Vectors[0].Values[2:]) + }) + + t.Run("zero vector", func(t *testing.T) { + profile := geminiTestProfile(t, 128) + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return geminiJSONResponse(request, `{"embedding":{"values":[`+vectorJSON(make([]float32, 128))+`]}}`), nil + })) + + _, err := client.Embed(t.Context(), []document.EmbeddingInput{{ + Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage", + }}, geminiSingleAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, ErrPermanentResponse) + }) +} + +func TestClearPreparedInputsZerosPayloadsAndFiles(t *testing.T) { + payload := []byte("private rendered text") + fileData := []byte("private file bytes") + prepared := []preparedInput{{payload: payload}, {file: &verifiedFile{data: fileData}}} + + clearPreparedInputs(prepared) + + assert.Equal(t, make([]byte, len(payload)), payload) + assert.Equal(t, make([]byte, len(fileData)), fileData) +} + +// TestEmbedFailsLocallyBeforeCredentialsOrEgress catches a direct-file path +// that resolves a credential or starts an HTTP request before task 3 has +// authorized multimodal transport. +func TestEmbedFailsLocallyBeforeCredentialsOrEgress(t *testing.T) { + profile := geminiTestProfile(t, 128) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("egress must not start") + })) + _, err := client.Embed(t.Context(), []document.EmbeddingInput{{ + Key: "direct-file", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile, + }}, geminiSingleAuthorization(profile.Descriptor)) + require.Error(t, err) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +// TestEmbedWithReceiptRetainsOnlyBoundedProviderProvenance catches receipts +// that omit the execution identity or retain request and response content. +func TestEmbedWithReceiptRetainsOnlyBoundedProviderProvenance(t *testing.T) { + profile := geminiTestProfile(t, 128) + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]},"usageMetadata":{"promptTokenCount":3,"cachedContentTokenCount":1,"totalTokenCount":4}}`) + response.Header.Set("X-Goog-Request-Id", "provider-request-1") + return response, nil + })) + execution, err := client.EmbedWithReceipt(t.Context(), []document.EmbeddingInput{{ + Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "private passage", + }}, geminiSingleAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.Equal(t, 1, execution.Receipt.RequestCount) + assert.Equal(t, TransportInline, execution.Receipt.Transport) + assert.Equal(t, int64(3), execution.Receipt.PromptTokens) + assert.Equal(t, int64(1), execution.Receipt.CachedContentTokens) + assert.Equal(t, int64(4), execution.Receipt.TotalTokens) + assert.Equal(t, []string{"provider-request-1"}, execution.Receipt.ProviderResponseIDs) + assert.Equal(t, "document-1", execution.Result.Vectors[0].Key) +} + +func TestRecordReceiptResponseTruncatesProviderIDsWithoutFailing(t *testing.T) { + receipt := Receipt{} + for index := range 129 { + require.True(t, recordReceiptResponse(&receipt, nil, fmt.Sprintf("provider-request-%03d", index))) + } + + assert.Len(t, receipt.ProviderResponseIDs, 128) + assert.Equal(t, 1, receipt.OmittedProviderResponseIDs) + assert.NotContains(t, receipt.ProviderResponseIDs, "provider-request-128") +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type countingSecrets struct { + value string + calls atomic.Int32 +} + +func (secrets *countingSecrets) ResolveSecret(context.Context, string) (string, error) { + secrets.calls.Add(1) + return secrets.value, nil +} + +func geminiTestProfile(t *testing.T, dimension int) Profile { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, + CompatibilityID: "gemini-embedding-2/search/v1", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "title: none | text: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "task: search result | query: {{content}}"}, + }) + require.NoError(t, err) + profile := Profile{ + CompatibilityEpoch: "gemini-embedding-2", SecretBinding: "secret:gemini", Transport: TransportInline, + CapabilityProfileFingerprint: testCapabilityProfile, DisclosureFingerprint: testDisclosurePolicy, + RequestTimeout: time.Second, MaxInputBytes: 4096, MaxRequestBytes: 8192, MaxResponseBytes: 16384, + EgressPolicy: providerhttp.EgressPolicy{ + Scheme: "https", Host: "generativelanguage.googleapis.com", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("192.0.2.0/24")}, ProxyMode: providerhttp.ProxyDisabled, + ConnectTimeout: time.Second, KeepAlive: time.Second, TLSHandshakeTimeout: time.Second, + }, + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustHostedProvider, Model: "gemini-embedding-2", + ModelRevision: "gemini-embedding-2", Dimension: dimension, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{ + document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk, + }, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }, + } + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return profile +} + +func newGeminiTestClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, syntheticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + client.http.Transport = transport + return client +} + +func geminiAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 2, MaxInputBytes: 64, MaxResponseBytes: 8192} +} + +func geminiSingleAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + authorization := geminiAuthorization(descriptor) + authorization.MaxBatchItems = 1 + return authorization +} + +func geminiJSONResponse(request *http.Request, body string) *http.Response { + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), Request: request} +} + +func testVectorJSON(dimension int) string { + values := make([]float32, dimension) + values[0] = 1 + encoded, err := json.Marshal(values) + if err != nil { + panic(err) + } + return string(bytes.Trim(encoded, "[]")) +} + +func vectorJSON(values []float32) string { + encoded, err := json.Marshal(values) + if err != nil { + panic(err) + } + return string(bytes.Trim(encoded, "[]")) +} diff --git a/document/geminiembed/errors.go b/document/geminiembed/errors.go new file mode 100644 index 00000000..c1853974 --- /dev/null +++ b/document/geminiembed/errors.go @@ -0,0 +1,79 @@ +package geminiembed + +import ( + "errors" + "fmt" + "net/http" + "strconv" + "time" +) + +var ( + ErrTransientResponse = errors.New("gemini embed: transient provider response") + ErrCapacityResponse = errors.New("gemini embed: provider capacity exceeded") + ErrPermanentResponse = errors.New("gemini embed: permanent provider response") +) + +type ProviderError struct { + Kind error + StatusCode int + RetryDelay time.Duration + RetrySet bool +} + +func (failure *ProviderError) Error() string { + if failure.StatusCode != 0 { + return fmt.Sprintf("gemini embed: HTTP %d: %v", failure.StatusCode, failure.Kind) + } + return failure.Kind.Error() +} + +func (failure *ProviderError) Unwrap() error { return failure.Kind } + +func RetryAfter(err error) (time.Duration, bool) { + failure, ok := errors.AsType[*ProviderError](err) + if !ok || !failure.RetrySet { + return 0, false + } + return failure.RetryDelay, true +} + +func statusError(status int, retryAfter string, now time.Time) error { + kind := ErrPermanentResponse + switch { + case status == http.StatusRequestEntityTooLarge: + kind = ErrCapacityResponse + case status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500 && status <= 599: + kind = ErrTransientResponse + } + delay, set := parseRetryAfter(retryAfter, now) + return &ProviderError{Kind: kind, StatusCode: status, RetryDelay: delay, RetrySet: set} +} + +func parseRetryAfter(value string, now time.Time) (time.Duration, bool) { + if seconds, err := strconv.ParseInt(value, 10, 64); err == nil && seconds >= 0 { + return clampRetryDelay(seconds), true + } + when, err := http.ParseTime(value) + if err != nil { + return 0, false + } + return clampRetryDuration(when.Sub(now)), true +} + +func clampRetryDelay(seconds int64) time.Duration { + if seconds > int64(time.Hour/time.Second) { + return time.Hour + } + return time.Duration(seconds) * time.Second +} + +func clampRetryDuration(delay time.Duration) time.Duration { + if delay < 0 { + return 0 + } + if delay > time.Hour { + return time.Hour + } + return delay +} diff --git a/document/geminiembed/errors_test.go b/document/geminiembed/errors_test.go new file mode 100644 index 00000000..cfdc1e90 --- /dev/null +++ b/document/geminiembed/errors_test.go @@ -0,0 +1,36 @@ +package geminiembed + +import ( + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestStatusErrorClassifiesGeminiResponses catches retryable, capacity, and +// permanent provider failures being collapsed into one unsafe retry policy. +func TestStatusErrorClassifiesGeminiResponses(t *testing.T) { + for _, testCase := range []struct { + name string + status int + retryAfter string + want error + delay time.Duration + delaySet bool + }{ + {name: "timeout", status: http.StatusRequestTimeout, want: ErrTransientResponse}, + {name: "rate limit", status: http.StatusTooManyRequests, retryAfter: "2", want: ErrTransientResponse, delay: 2 * time.Second, delaySet: true}, + {name: "capacity", status: http.StatusRequestEntityTooLarge, want: ErrCapacityResponse}, + {name: "invalid request", status: http.StatusBadRequest, want: ErrPermanentResponse}, + } { + t.Run(testCase.name, func(t *testing.T) { + err := statusError(testCase.status, testCase.retryAfter, time.Unix(0, 0)) + require.ErrorIs(t, err, testCase.want) + delay, set := RetryAfter(err) + assert.Equal(t, testCase.delaySet, set) + assert.Equal(t, testCase.delay, delay) + }) + } +} diff --git a/document/geminiembed/lifecycle_test.go b/document/geminiembed/lifecycle_test.go new file mode 100644 index 00000000..f378b0c4 --- /dev/null +++ b/document/geminiembed/lifecycle_test.go @@ -0,0 +1,1408 @@ +package geminiembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" + "image/color" + "io" + "net/http" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/media/mediatest" +) + +const ( + testCapabilityProfile = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" + testDisclosurePolicy = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee" +) + +// TestEmbedInlineSendsOnlyCapabilityProvenExactBytes catches inline requests +// that serialize caller metadata or bytes not covered by the locally sealed +// capability record. +func TestEmbedInlineSendsOnlyCapabilityProvenExactBytes(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportInline) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + require.Equal(t, http.MethodPost, request.Method) + require.Equal(t, embedPath, request.URL.Path) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + want := `{"model":"models/gemini-embedding-2","content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + base64.StdEncoding.EncodeToString(data) + `"}}]},"outputDimensionality":128}` + assert.Equal(t, want, string(body)) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + })) + + result, err := client.Embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.NoError(t, err) + require.Len(t, result.Vectors, 1) + assert.Equal(t, "direct-a", result.Vectors[0].Key) + assert.Equal(t, int32(1), requests.Load()) + assert.Equal(t, int32(1), source.readPasses.Load()) + assert.Equal(t, int32(1), source.closeCalls.Load()) + + t.Run("changed bytes fail before credentials", func(t *testing.T) { + changed := append([]byte(nil), data...) + changed[len(changed)-1] ^= 0xff + changedSource := newGeminiLifecycleUpload(changed, record) + secrets := &countingSecrets{value: "synthetic-key"} + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("changed bytes reached egress") + return nil, errors.New("unreachable changed-byte egress") + })) + _, err := client.Embed(t.Context(), directInputs(changedSource), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Zero(t, secrets.calls.Load()) + assert.Equal(t, int32(1), changedSource.closeCalls.Load()) + }) +} + +// TestEmbedFilesAPIUploadsActivatesEmbedsAndDeletes catches any lifecycle +// reordering, resumable-upload contract drift, premature embed, or omitted +// deletion attempt. +func TestEmbedFilesAPIUploadsActivatesEmbedsAndDeletes(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + timeline := newGeminiFileTimeline() + processing := geminiFileJSON(record, fileName, fileURI, "PROCESSING", timeline) + active := geminiFileJSON(record, fileName, fileURI, "ACTIVE", timeline) + sequence := atomic.Int32{} + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + step := sequence.Add(1) + assert.Equal(t, "synthetic-key", request.Header.Get("X-Goog-Api-Key")) + switch step { + case 1: + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, filesUploadPath, request.URL.Path) + assert.Equal(t, "resumable", request.Header.Get("X-Goog-Upload-Protocol")) + assert.Equal(t, "start", request.Header.Get("X-Goog-Upload-Command")) + assert.Equal(t, strconv.Itoa(len(data)), request.Header.Get("X-Goog-Upload-Header-Content-Length")) + assert.Equal(t, "image/png", request.Header.Get("X-Goog-Upload-Header-Content-Type")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{"file":{}}`, string(body)) + response := &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("")), Request: request} + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "/upload/v1beta/files", request.URL.Path) + assert.Equal(t, "upload_id=synthetic-session-123&upload_protocol=resumable", request.URL.RawQuery) + assert.Equal(t, "upload, finalize", request.Header.Get("X-Goog-Upload-Command")) + assert.Equal(t, "0", request.Header.Get("X-Goog-Upload-Offset")) + assert.Equal(t, int64(len(data)), request.ContentLength) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.Equal(t, data, body) + response := geminiJSONResponse(request, `{"file":`+processing+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case 3: + require.Equal(t, http.MethodGet, request.Method) + assert.Equal(t, "/v1beta/"+fileName, request.URL.Path) + return geminiJSONResponse(request, processing), nil + case 4: + require.Equal(t, http.MethodGet, request.Method) + assert.Equal(t, "/v1beta/"+fileName, request.URL.Path) + return geminiJSONResponse(request, active), nil + case 5: + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, embedPath, request.URL.Path) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + want := `{"model":"models/gemini-embedding-2","content":{"parts":[{"fileData":{"mimeType":"image/png","fileUri":"` + fileURI + `"}}]},"outputDimensionality":128}` + assert.Equal(t, want, string(body)) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + case 6: + require.Equal(t, http.MethodDelete, request.Method) + assert.Equal(t, "/v1beta/"+fileName, request.URL.Path) + return geminiJSONResponse(request, `{}`), nil + default: + t.Fatalf("unexpected request %d: %s %s", step, request.Method, request.URL) + return nil, errors.New("unreachable lifecycle request") + } + })) + + execution, err := client.EmbedWithReceipt(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.NoError(t, err) + assert.Equal(t, int32(6), sequence.Load()) + assert.Equal(t, 6, execution.Receipt.RequestCount) + assert.Equal(t, 48*time.Hour, execution.Receipt.ProviderRetentionCeiling) + assert.Equal(t, int32(1), source.readPasses.Load()) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.NotContains(t, execution.Receipt.ProviderResponseIDs, fileName) +} + +// TestEmbedFilesAPIDeletesAfterValidFinalizeReceiptRejection catches cleanup +// ownership that is installed only after receipt metadata has been accepted. +func TestEmbedFilesAPIDeletesAfterValidFinalizeReceiptRejection(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + active := geminiFileJSON(record, fileName, fileURI, "ACTIVE", newGeminiFileTimeline()) + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + response := geminiJSONResponse(request, `{"file":`+active+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + response.Header.Set("X-Goog-Request-Id", "invalid response id") + return response, nil + case 3: + require.Equal(t, http.MethodDelete, request.Method) + assert.Equal(t, "/v1beta/"+fileName, request.URL.Path) + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("unexpected cleanup lifecycle request") + } + })) + receipt := Receipt{} + _, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, int32(3), sequence.Load()) + assert.Equal(t, 3, receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + assert.NotContains(t, err.Error(), fileName) + assert.NotContains(t, err.Error(), "invalid response id") + assert.Equal(t, int32(1), source.closeCalls.Load()) +} + +// TestEmbedFilesAPITruncatesCapacityExceedingResponseID catches receipt +// capacity handling that discards successful embedding work or retains the ID. +func TestEmbedFilesAPITruncatesCapacityExceedingResponseID(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + active := geminiFileJSON(record, fileName, fileURI, "ACTIVE", newGeminiFileTimeline()) + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + response := geminiJSONResponse(request, `{"file":`+active+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + response.Header.Set("X-Goog-Request-Id", "provider-response-over-capacity") + return response, nil + case 3: + require.Equal(t, http.MethodPost, request.Method) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + case 4: + require.Equal(t, http.MethodDelete, request.Method) + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("unexpected response-ID capacity request") + } + })) + receipt := Receipt{ProviderResponseIDs: make([]string, 128)} + for index := range receipt.ProviderResponseIDs { + receipt.ProviderResponseIDs[index] = "existing-response-" + strconv.Itoa(index) + } + result, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.NoError(t, err) + assert.Len(t, result.Vectors, 1) + assert.Equal(t, int32(4), sequence.Load()) + assert.Equal(t, 4, receipt.RequestCount) + assert.Len(t, receipt.ProviderResponseIDs, 128) + assert.Equal(t, 1, receipt.OmittedProviderResponseIDs) + assert.NotContains(t, receipt.ProviderResponseIDs, "provider-response-over-capacity") + assert.Equal(t, int32(1), source.closeCalls.Load()) +} + +// TestEmbedFilesAPIRejectsUntrustedLifecycleTimestamps catches provider file +// state that escapes the local finalize window or changes its sealed timeline. +func TestEmbedFilesAPIRejectsUntrustedLifecycleTimestamps(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + valid := newGeminiFileTimeline() + + for _, testCase := range []struct { + name string + finalize geminiFileTimeline + poll *geminiFileTimeline + wantRequests int32 + }{ + {name: "creation outside local finalize window", finalize: geminiFileTimeline{ + created: valid.created.Add(6 * time.Minute), updated: valid.updated.Add(6 * time.Minute), expires: valid.expires.Add(6 * time.Minute), + }, wantRequests: 2}, + {name: "finalize update after expiration", finalize: geminiFileTimeline{ + created: valid.created, updated: valid.expires.Add(time.Second), expires: valid.expires, + }, wantRequests: 2}, + {name: "poll creation drift", finalize: valid, poll: &geminiFileTimeline{ + created: valid.created.Add(time.Second), updated: valid.updated.Add(time.Second), expires: valid.expires, + }, wantRequests: 4}, + {name: "poll expiration drift", finalize: valid, poll: &geminiFileTimeline{ + created: valid.created, updated: valid.updated.Add(time.Second), expires: valid.expires.Add(-time.Second), + }, wantRequests: 4}, + {name: "poll update regression", finalize: valid, poll: &geminiFileTimeline{ + created: valid.created, updated: valid.created, expires: valid.expires, + }, wantRequests: 4}, + } { + t.Run(testCase.name, func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + finalizeState := "ACTIVE" + if testCase.poll != nil { + finalizeState = "PROCESSING" + } + finalized := geminiFileJSON(record, fileName, fileURI, finalizeState, testCase.finalize) + var polled string + if testCase.poll != nil { + polled = geminiFileJSON(record, fileName, fileURI, "ACTIVE", *testCase.poll) + } + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + response := geminiJSONResponse(request, `{"file":`+finalized+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case 3: + require.NotNil(t, testCase.poll) + require.Equal(t, http.MethodGet, request.Method) + return geminiJSONResponse(request, polled), nil + case 4: + require.Equal(t, http.MethodDelete, request.Method) + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("unsafe timestamp follow-up request") + } + })) + receipt := Receipt{} + _, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, testCase.wantRequests, sequence.Load()) + assert.Equal(t, int(testCase.wantRequests), receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + assert.NotContains(t, err.Error(), fileName) + assert.NotContains(t, err.Error(), fileURI) + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + } +} + +func TestValidateCreatedWireFileBoundsUpdateAndExpiryToObservation(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + expected := verifiedFile{metadata: document.AuthorizedUploadMetadata{ + Filename: "synthetic.png", MediaType: record.MediaType, ByteLength: record.SourceBytes, SHA256: record.SourceSHA256, + }} + startedAt := time.Now().UTC().Truncate(time.Second) + completedAt := startedAt.Add(time.Second) + file := func(updated, expires time.Time) wireFile { + return wireFile{ + Name: "files/file-123", MIMEType: record.MediaType, SizeBytes: strconv.FormatInt(record.SourceBytes, 10), + CreateTime: startedAt.Format(time.RFC3339Nano), UpdateTime: updated.Format(time.RFC3339Nano), + ExpirationTime: expires.Format(time.RFC3339Nano), SHA256Hash: base64.StdEncoding.EncodeToString([]byte(record.SourceSHA256)), + URI: origin + "/v1beta/files/file-123", State: "ACTIVE", Source: "UPLOADED", + } + } + + _, ok := validateCreatedWireFile(file(completedAt.Add(fileClockSkew+time.Second), startedAt.Add(48*time.Hour)), expected, startedAt, completedAt) + assert.False(t, ok, "an update outside the finalize observation window must fail closed") + _, ok = validateCreatedWireFile(file(startedAt, completedAt), expected, startedAt, completedAt) + assert.False(t, ok, "a file already expired at finalize completion must fail closed") +} + +func TestValidatePolledWireFileBoundsUpdateAndExpiryToObservation(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + expected := verifiedFile{metadata: document.AuthorizedUploadMetadata{ + Filename: "synthetic.png", MediaType: record.MediaType, ByteLength: record.SourceBytes, SHA256: record.SourceSHA256, + }} + startedAt := time.Now().UTC().Truncate(time.Second) + completedAt := startedAt.Add(time.Second) + createdAt := startedAt.Add(-time.Hour) + file := func(updated, expires time.Time) wireFile { + return wireFile{ + Name: "files/file-123", MIMEType: record.MediaType, SizeBytes: strconv.FormatInt(record.SourceBytes, 10), + CreateTime: createdAt.Format(time.RFC3339Nano), UpdateTime: updated.Format(time.RFC3339Nano), + ExpirationTime: expires.Format(time.RFC3339Nano), SHA256Hash: base64.StdEncoding.EncodeToString([]byte(record.SourceSHA256)), + URI: origin + "/v1beta/files/file-123", State: "PROCESSING", Source: "UPLOADED", + } + } + + staleUpdate := startedAt.Add(-fileClockSkew - time.Second) + futureExpiry := startedAt.Add(24 * time.Hour) + current := validatedProviderFile{file: file(staleUpdate, futureExpiry), created: createdAt, updated: staleUpdate, expiresAt: futureExpiry} + _, ok := validatePolledWireFile(file(staleUpdate, futureExpiry), expected, current, startedAt, completedAt) + assert.False(t, ok, "a stale poll update outside the observation window must fail closed") + + current.updated = startedAt + current.file = file(startedAt, futureExpiry) + _, ok = validatePolledWireFile(file(completedAt.Add(fileClockSkew+time.Second), futureExpiry), expected, current, startedAt, completedAt) + assert.False(t, ok, "a future poll update outside the observation window must fail closed") + + expiredAt := completedAt + current = validatedProviderFile{file: file(startedAt, expiredAt), created: createdAt, updated: startedAt, expiresAt: expiredAt} + _, ok = validatePolledWireFile(file(startedAt, expiredAt), expected, current, startedAt, completedAt) + assert.False(t, ok, "a file already expired at poll completion must fail closed") +} + +// TestEmbedFilesAPIRejectsUnsafeUploadURLsWithoutFollowUp catches resumable +// destinations that escape the exact sealed Google upload boundary. +func TestEmbedFilesAPIRejectsUnsafeUploadURLsWithoutFollowUp(t *testing.T) { + for _, testCase := range []struct { + name string + uploadURL string + }{ + {name: "host", uploadURL: "https://provider.invalid/upload/v1beta/files?upload_id=synthetic&upload_protocol=resumable"}, + {name: "userinfo", uploadURL: "https://synthetic@generativelanguage.googleapis.com/upload/v1beta/files?upload_id=synthetic&upload_protocol=resumable"}, + {name: "port", uploadURL: "https://generativelanguage.googleapis.com:444/upload/v1beta/files?upload_id=synthetic&upload_protocol=resumable"}, + {name: "path", uploadURL: "https://generativelanguage.googleapis.com/upload/v1beta/other?upload_id=synthetic&upload_protocol=resumable"}, + {name: "query", uploadURL: "https://generativelanguage.googleapis.com/upload/v1beta/files?upload_id=synthetic&upload_protocol=raw"}, + {name: "fragment", uploadURL: "https://generativelanguage.googleapis.com/upload/v1beta/files?upload_id=synthetic&upload_protocol=resumable#outside"}, + } { + t.Run(testCase.name, func(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", testCase.uploadURL) + return response, nil + })) + receipt := Receipt{} + _, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, int32(1), requests.Load()) + assert.Equal(t, 1, receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + assert.NotContains(t, err.Error(), testCase.uploadURL) + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + } +} + +// TestEmbedFilesAPIRejectsMismatchedCreatedFileWithoutUnsafeCleanup catches a +// provider-controlled file identity or state that is trusted for follow-up. +func TestEmbedFilesAPIRejectsMismatchedCreatedFileWithoutUnsafeCleanup(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + valid := geminiFileJSON(record, fileName, fileURI, "ACTIVE", newGeminiFileTimeline()) + hash := base64.StdEncoding.EncodeToString([]byte(record.SourceSHA256)) + + for _, testCase := range []struct { + name string + file string + }{ + {name: "name", file: strings.Replace(valid, `"name":"files/file-123"`, `"name":"files/UPPER"`, 1)}, + {name: "uri", file: strings.Replace(valid, `"uri":"`+fileURI+`"`, `"uri":"https://provider.invalid/v1beta/files/file-123"`, 1)}, + {name: "mime", file: strings.Replace(valid, `"mimeType":"image/png"`, `"mimeType":"image/jpeg"`, 1)}, + {name: "size", file: strings.Replace(valid, `"sizeBytes":"`+strconv.FormatInt(record.SourceBytes, 10)+`"`, `"sizeBytes":"1"`, 1)}, + {name: "hash", file: strings.Replace(valid, `"sha256Hash":"`+hash+`"`, `"sha256Hash":"c3ludGhldGlj"`, 1)}, + {name: "state", file: strings.Replace(valid, `"state":"ACTIVE"`, `"state":"READY"`, 1)}, + {name: "timestamps", file: strings.Replace(valid, `"updateTime":"`, `"updateTime":"not-a-time`, 1)}, + } { + t.Run(testCase.name, func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + response := geminiJSONResponse(request, `{"file":`+testCase.file+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + default: + return nil, errors.New("unsafe created-file follow-up request") + } + })) + receipt := Receipt{} + _, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, int32(2), sequence.Load()) + assert.Equal(t, 2, receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + assert.NotContains(t, err.Error(), fileName) + assert.NotContains(t, err.Error(), fileURI) + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + } +} + +// TestEmbedFilesAPIBoundsEveryLifecycleBody catches a start, finalize, poll, +// or delete response that bypasses the fixed raw-response byte ceiling. +func TestEmbedFilesAPIBoundsEveryLifecycleBody(t *testing.T) { + for _, stage := range []string{"start", "finalize", "poll", "delete"} { + t.Run(stage, func(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + state := "ACTIVE" + if stage == "poll" { + state = "PROCESSING" + } + fileJSON := geminiFileJSON(record, fileName, fileURI, state, newGeminiFileTimeline()) + oversized := strings.Repeat("x", int(profile.MaxResponseBytes)+1) + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + if stage == "start" { + return geminiJSONResponse(request, oversized), nil + } + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + responseBody := `{"file":` + fileJSON + `}` + if stage == "finalize" { + responseBody = oversized + } + response := geminiJSONResponse(request, responseBody) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case 3: + if stage == "poll" { + return geminiJSONResponse(request, oversized), nil + } + require.Equal(t, embedPath, request.URL.Path) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + case 4: + require.Equal(t, http.MethodDelete, request.Method) + if stage == "delete" { + return geminiJSONResponse(request, oversized), nil + } + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("oversized lifecycle response triggered extra egress") + } + })) + receipt := Receipt{} + result, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + wantRequests := map[string]int32{"start": 1, "finalize": 2, "poll": 4, "delete": 4}[stage] + assert.Equal(t, wantRequests, sequence.Load()) + assert.Equal(t, int(wantRequests), receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + if stage == "delete" { + require.NoError(t, err) + require.Len(t, result.Vectors, 1) + assert.Equal(t, []string{"provider file deletion attempt failed"}, receipt.Warnings) + } else { + require.Error(t, err) + assert.Empty(t, receipt.Warnings) + assert.NotContains(t, err.Error(), fileName) + assert.NotContains(t, err.Error(), fileURI) + assert.NotContains(t, err.Error(), oversized) + } + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + } +} + +// TestEmbedFilesAPIAttemptsCleanupAfterPostFinalizeFailures catches poll, +// embed, receipt, or delete failures that lose cleanup or request accounting. +func TestEmbedFilesAPIAttemptsCleanupAfterPostFinalizeFailures(t *testing.T) { + for _, failure := range []string{"poll", "poll receipt", "embed", "embed receipt", "delete"} { + t.Run(failure, func(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + timeline := newGeminiFileTimeline() + state := "ACTIVE" + if strings.HasPrefix(failure, "poll") { + state = "PROCESSING" + } + finalized := geminiFileJSON(record, fileName, fileURI, state, timeline) + active := geminiFileJSON(record, fileName, fileURI, "ACTIVE", timeline) + var sequence atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch sequence.Add(1) { + case 1: + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + response := geminiJSONResponse(request, `{"file":`+finalized+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case 3: + switch failure { + case "poll": + return &http.Response{StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("")), Request: request}, nil + case "poll receipt": + response := geminiJSONResponse(request, active) + response.Header.Set("X-Goog-Request-Id", "invalid response id") + return response, nil + case "embed": + return &http.Response{StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("")), Request: request}, nil + case "embed receipt": + response := geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`) + response.Header.Set("X-Goog-Request-Id", "invalid response id") + return response, nil + case "delete": + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + default: + return nil, errors.New("unreachable post-finalize failure") + } + case 4: + require.Equal(t, http.MethodDelete, request.Method) + assert.Equal(t, "/v1beta/"+fileName, request.URL.Path) + if failure == "delete" { + return &http.Response{StatusCode: http.StatusInternalServerError, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("")), Request: request}, nil + } + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("post-finalize failure triggered extra egress") + } + })) + receipt := Receipt{} + result, err := client.embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + assert.Equal(t, int32(4), sequence.Load()) + assert.Equal(t, 4, receipt.RequestCount) + assert.Empty(t, receipt.ProviderResponseIDs) + if failure == "delete" { + require.NoError(t, err) + require.Len(t, result.Vectors, 1) + assert.Equal(t, []string{"provider file deletion attempt failed"}, receipt.Warnings) + } else { + require.Error(t, err) + assert.Empty(t, receipt.Warnings) + assert.NotContains(t, err.Error(), fileName) + assert.NotContains(t, err.Error(), fileURI) + assert.NotContains(t, err.Error(), "invalid response id") + } + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + } +} + +// TestEmbedRejectsUnsupportedOrOverLimitFilesBeforeCredentialResolution +// catches a local capability or E1 bound bypass that exposes a credential for +// a file the fixed Gemini profile cannot send. +func TestEmbedRejectsUnsupportedOrOverLimitFilesBeforeCredentialResolution(t *testing.T) { + profile := geminiDirectTestProfile(t, TransportInline) + image := geminiTinyPNG(t) + imageRecord := geminiCapability(t, profile, image, "synthetic.png", "image/png") + text := []byte("synthetic private text") + textRecord := geminiCapability(t, profile, text, "synthetic.txt", "text/plain") + + for _, testCase := range []struct { + name string + input document.EmbeddingInput + authorization document.EmbeddingAuthorization + }{ + {name: "unsupported eligible family", input: directInputs(newGeminiLifecycleUpload(text, textRecord))[0], authorization: geminiDirectAuthorization(profile.Descriptor, int64(len(text)))}, + {name: "over authorized bytes", input: directInputs(newGeminiLifecycleUpload(image, imageRecord))[0], authorization: geminiDirectAuthorization(profile.Descriptor, int64(len(image)-1))}, + {name: "query file", input: document.EmbeddingInput{Key: "query-file", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputOriginalFile, Source: newGeminiLifecycleUpload(image, imageRecord)}, authorization: geminiDirectAuthorization(profile.Descriptor, int64(len(image)))}, + } { + t.Run(testCase.name, func(t *testing.T) { + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("egress must not start") + })) + _, err := client.Embed(t.Context(), []document.EmbeddingInput{testCase.input}, testCase.authorization) + require.Error(t, err) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + } +} + +// TestEmbedClosesEveryEnrolledSourceExactlyOnceOnAllExitPaths catches source +// leaks and double-closes when one member fails before or after preparation. +func TestEmbedClosesEveryEnrolledSourceExactlyOnceOnAllExitPaths(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportInline) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + for _, testCase := range []struct { + name string + mutate func([]*geminiLifecycleUpload, *document.EmbeddingAuthorization) + }{ + {name: "invalid authorization", mutate: func(_ []*geminiLifecycleUpload, authorization *document.EmbeddingAuthorization) { + authorization.ProviderID = "wrong-provider" + }}, + {name: "live metadata drift", mutate: func(sources []*geminiLifecycleUpload, _ *document.EmbeddingAuthorization) { + changed := sources[0].metadata[0] + changed.ProviderMetadataChecksum = strings.Repeat("f", 64) + sources[0].metadata = append(sources[0].metadata, changed) + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + sources := []*geminiLifecycleUpload{newGeminiLifecycleUpload(data, record), newGeminiLifecycleUpload(data, record)} + authorization := geminiDirectAuthorization(profile.Descriptor, int64(2*len(data))) + testCase.mutate(sources, &authorization) + secrets := &countingSecrets{value: "synthetic-key"} + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(sources[0], sources[1]), authorization) + require.Error(t, err) + for _, source := range sources { + assert.Equal(t, int32(1), source.closeCalls.Load()) + } + assert.Zero(t, secrets.calls.Load()) + }) + } + + t.Run("repeated source identity", func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(source, source), geminiDirectAuthorization(profile.Descriptor, int64(2*len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.metadataCalls.Load()) + assert.Zero(t, source.capabilityCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + t.Run("unsafe source identity", func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + valueSource := geminiNonComparableUpload{geminiLifecycleUpload: source, identity: []byte("unsafe")} + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(valueSource), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.metadataCalls.Load()) + assert.Zero(t, source.capabilityCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + t.Run("dynamically unsafe source identity", func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + valueSource := geminiInterfaceIdentityUpload{geminiLifecycleUpload: source, identity: []byte("unsafe")} + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(valueSource), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.metadataCalls.Load()) + assert.Zero(t, source.capabilityCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + t.Run("repeated comparable value identity", func(t *testing.T) { + source := newGeminiLifecycleUpload(data, record) + valueSource := geminiValueUpload{geminiLifecycleUpload: source} + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(valueSource, valueSource), geminiDirectAuthorization(profile.Descriptor, int64(2*len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.metadataCalls.Load()) + assert.Zero(t, source.capabilityCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + t.Run("repeated source before trailing source", func(t *testing.T) { + repeated := newGeminiLifecycleUpload(data, record) + trailing := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(repeated, repeated, trailing), geminiDirectAuthorization(profile.Descriptor, int64(3*len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), repeated.closeCalls.Load()) + assert.Equal(t, int32(1), trailing.closeCalls.Load()) + assert.Zero(t, repeated.metadataCalls.Load()) + assert.Zero(t, repeated.capabilityCalls.Load()) + assert.Zero(t, trailing.metadataCalls.Load()) + assert.Zero(t, trailing.capabilityCalls.Load()) + assert.Zero(t, repeated.readPasses.Load()) + assert.Zero(t, trailing.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + t.Run("unsafe source before trailing source", func(t *testing.T) { + unsafe := newGeminiLifecycleUpload(data, record) + trailing := newGeminiLifecycleUpload(data, record) + valueSource := geminiNonComparableUpload{geminiLifecycleUpload: unsafe, identity: []byte("unsafe")} + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(valueSource, trailing), geminiDirectAuthorization(profile.Descriptor, int64(2*len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), unsafe.closeCalls.Load()) + assert.Equal(t, int32(1), trailing.closeCalls.Load()) + assert.Zero(t, unsafe.metadataCalls.Load()) + assert.Zero(t, unsafe.capabilityCalls.Load()) + assert.Zero(t, trailing.metadataCalls.Load()) + assert.Zero(t, trailing.capabilityCalls.Load()) + assert.Zero(t, unsafe.readPasses.Load()) + assert.Zero(t, trailing.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + + for _, testCase := range []struct { + name string + source func() document.AuthorizedUpload + }{ + {name: "nil source before trailing source", source: func() document.AuthorizedUpload { return nil }}, + {name: "typed nil source before trailing source", source: func() document.AuthorizedUpload { + var source *geminiLifecycleUpload + return source + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + trailing := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + _, err := client.Embed(t.Context(), directInputs(testCase.source(), trailing), geminiDirectAuthorization(profile.Descriptor, int64(2*len(data)))) + require.Error(t, err) + assert.Equal(t, "gemini embed: original upload source is nil", err.Error()) + assert.Equal(t, int32(1), trailing.closeCalls.Load()) + assert.Zero(t, trailing.metadataCalls.Load()) + assert.Zero(t, trailing.capabilityCalls.Load()) + assert.Zero(t, trailing.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + } +} + +// TestEmbedEnrollsEveryAttachedSourceBeforeSemanticValidation catches an +// ownership scan that filters by declared kind before acquiring close +// responsibility for every non-nil source identity. +func TestEmbedEnrollsEveryAttachedSourceBeforeSemanticValidation(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, TransportInline) + data := geminiTinyPNG(t) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + for _, testCase := range []struct { + name string + input func(document.AuthorizedUpload) []document.EmbeddingInput + }{ + { + name: "query text", + input: func(source document.AuthorizedUpload) []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "query", Role: document.EmbeddingRoleQuery, + Kind: document.EmbeddingInputQueryText, Text: "question", Source: source}} + }, + }, + { + name: "rendition chunk", + input: func(source document.AuthorizedUpload) []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "chunk", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "passage", Source: source}} + }, + }, + { + name: "unknown kind", + input: func(source document.AuthorizedUpload) []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "unknown", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputKind("unknown"), Source: source}} + }, + }, + { + name: "repeated invalid identity", + input: func(source document.AuthorizedUpload) []document.EmbeddingInput { + return []document.EmbeddingInput{ + {Key: "query", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "question", Source: source}, + {Key: "chunk", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage", Source: source}, + } + }, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + source := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + + _, err := client.Embed(t.Context(), testCase.input(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.metadataCalls.Load()) + assert.Zero(t, source.capabilityCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) + }) + } +} + +// TestEmbedClearsEarlierPreparedFileWhenLaterPreparationFails catches cleanup +// ownership that begins only after every input has finished preparation. +func TestEmbedClearsEarlierPreparedFileWhenLaterPreparationFails(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, TransportInline) + data := geminiTinyPNG(t) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + earlier := newGeminiLifecycleUpload(data, record) + earlier.captureReadBuffer = true + later := newGeminiLifecycleUpload(data, record) + drifted := later.metadata[0] + drifted.ProviderMetadataChecksum = strings.Repeat("f", 64) + later.metadata = append(later.metadata, drifted) + secrets := &countingSecrets{value: "synthetic-key"} + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("request must not run") + })) + + _, err := client.Embed(t.Context(), directInputs(earlier, later), geminiDirectAuthorization(profile.Descriptor, int64(2*len(data)))) + require.Error(t, err) + require.Len(t, earlier.capturedReadBuffer, len(data)) + assert.Equal(t, make([]byte, len(data)), earlier.capturedReadBuffer) + assert.Zero(t, later.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) +} + +// TestEmbedFilesAPIPreflightsFileDataBeforeSecret catches a Files API path +// that discovers its eventual embedding envelope is over the request bound +// only after reading private bytes, resolving credentials, or uploading. +func TestEmbedFilesAPIPreflightsFileDataBeforeSecret(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, TransportFilesAPI) + data := geminiTinyPNG(t) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + minimum, err := minimumFilesAPIRequestCapacity(profile.Descriptor.Dimension) + require.NoError(t, err) + client.profile.MaxRequestBytes = minimum - 7 // image/png is six bytes shorter than the longest supported MIME. + + _, err = client.Embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +// TestEmbedFilesAPIBoundsRawFinalizeBeforeSecret catches a request-wide bound +// that applies to JSON envelopes but not the raw upload/finalize body. +func TestEmbedFilesAPIBoundsRawFinalizeBeforeSecret(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, TransportFilesAPI) + minimum, err := minimumFilesAPIRequestCapacity(profile.Descriptor.Dimension) + require.NoError(t, err) + profile.MaxRequestBytes = minimum + profile = rebindGeminiProfile(t, profile) + data := mediatest.JPEG(64, 64, nil) + require.Greater(t, int64(len(data)), profile.MaxRequestBytes) + record := geminiCapability(t, profile, data, "synthetic.jpg", "image/jpeg") + source := newGeminiLifecycleUpload(data, record) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + + _, err = client.Embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + require.Error(t, err) + assert.Equal(t, int32(1), source.closeCalls.Load()) + assert.Zero(t, source.readPasses.Load()) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +// TestEmbedPreservesEverySupportedDirectMIMEAcrossTransports proves each E14 +// direct modality reaches both wire transports with its exact capability MIME. +func TestEmbedPreservesEverySupportedDirectMIMEAcrossTransports(t *testing.T) { + t.Parallel() + fixtures := []struct { + name, filename, mediaType string + data []byte + }{ + {name: "jpeg", filename: "synthetic.jpg", mediaType: "image/jpeg", data: mediatest.JPEG(2, 2, color.NRGBA{R: 0xff, A: 0xff})}, + {name: "png", filename: "synthetic.png", mediaType: "image/png", data: mediatest.PNG(2, 2, color.NRGBA{G: 0xff, A: 0xff})}, + {name: "wav", filename: "synthetic.wav", mediaType: "audio/wav", data: mediatest.WAV()}, + {name: "mp3", filename: "synthetic.mp3", mediaType: "audio/mpeg", data: mediatest.MP3()}, + {name: "mp4", filename: "synthetic.mp4", mediaType: "video/mp4", data: mediatest.H265MP4()}, + {name: "mov", filename: "synthetic.mov", mediaType: "video/quicktime", data: mediatest.H264MOV()}, + {name: "pdf", filename: "synthetic.pdf", mediaType: "application/pdf", data: mediatest.PDF()}, + } + for _, fixture := range fixtures { + t.Run(fixture.name, func(t *testing.T) { + t.Parallel() + for _, transport := range []Transport{TransportInline, TransportFilesAPI} { + t.Run(string(transport), func(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, transport) + record := geminiCapability(t, profile, fixture.data, fixture.filename, fixture.mediaType) + assert.Equal(t, fixture.mediaType, record.MediaType) + source := newGeminiLifecycleUpload(fixture.data, record) + fileName := "files/synthetic-" + fixture.name + fileURI := origin + "/v1beta/" + fileName + timeline := newGeminiFileTimeline() + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requestNumber := requests.Add(1) + if transport == TransportInline { + require.Equal(t, int32(1), requestNumber) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + want := `{"model":"models/gemini-embedding-2","content":{"parts":[{"inlineData":{"mimeType":"` + fixture.mediaType + `","data":"` + base64.StdEncoding.EncodeToString(fixture.data) + `"}}]},"outputDimensionality":128}` + assert.Equal(t, want, string(body)) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + } + switch requestNumber { + case 1: + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, fixture.mediaType, request.Header.Get("X-Goog-Upload-Header-Content-Type")) + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + case 2: + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.Equal(t, fixture.data, body) + response := geminiJSONResponse(request, `{"file":`+geminiFileJSON(record, fileName, fileURI, "ACTIVE", timeline)+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case 3: + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + want := `{"model":"models/gemini-embedding-2","content":{"parts":[{"fileData":{"mimeType":"` + fixture.mediaType + `","fileUri":"` + fileURI + `"}}]},"outputDimensionality":128}` + assert.Equal(t, want, string(body)) + return geminiJSONResponse(request, `{"embedding":{"values":[`+testVectorJSON(128)+`]}}`), nil + case 4: + require.Equal(t, http.MethodDelete, request.Method) + return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Request: request}, nil + default: + return nil, errors.New("unexpected modality wire request") + } + })) + + _, err := client.Embed(t.Context(), directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(fixture.data)))) + require.NoError(t, err) + if transport == TransportInline { + assert.Equal(t, int32(1), requests.Load()) + } else { + assert.Equal(t, int32(4), requests.Load()) + } + }) + } + }) + } +} + +func TestFinalizeFileUploadRejectsRawBodyOverRequestBound(t *testing.T) { + t.Parallel() + profile := geminiDirectTestProfile(t, TransportFilesAPI) + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("request must not run") + })) + uploadURL, err := validateUploadURL(origin + "/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + require.NoError(t, err) + file := verifiedFile{data: make([]byte, profile.MaxRequestBytes+1)} + + _, _, err = client.finalizeFileUpload(t.Context(), uploadURL, file, "synthetic-key", nil) + require.Error(t, err) + assert.Zero(t, requests.Load()) +} + +// TestEmbedCancellationStopsBlockedReadUploadAndPolling catches cancellation +// paths that leave a source read, upload request, or file-state poll running. +func TestEmbedCancellationStopsBlockedReadUploadAndPolling(t *testing.T) { + data := geminiTinyPNG(t) + + t.Run("blocked source read", func(t *testing.T) { + profile := geminiDirectTestProfile(t, TransportInline) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + source.blockRead = true + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + result := make(chan error, 1) + client := newGeminiTestClient(t, profile, &countingSecrets{value: "synthetic-key"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("request must not run") + })) + go func() { + _, err := client.Embed(ctx, directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + result <- err + }() + select { + case <-source.readStarted: + case <-time.After(time.Second): + t.Fatal("source read did not start") + } + cancel() + require.ErrorIs(t, <-result, context.Canceled) + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + + t.Run("blocked upload", func(t *testing.T) { + profile := geminiDirectTestProfile(t, TransportFilesAPI) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + uploadStarted := make(chan struct{}) + ctx, cancel := context.WithCancel(t.Context()) + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.Header.Get("X-Goog-Upload-Command") == "start" { + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + } + close(uploadStarted) + <-request.Context().Done() + return nil, request.Context().Err() + })) + result := make(chan error, 1) + go func() { + _, err := client.Embed(ctx, directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + result <- err + }() + <-uploadStarted + cancel() + require.ErrorIs(t, <-result, context.Canceled) + assert.Equal(t, int32(1), source.closeCalls.Load()) + }) + + t.Run("poll wait", func(t *testing.T) { + profile := geminiDirectTestProfile(t, TransportFilesAPI) + profile.PollInterval = maximumPoll + profile = rebindGeminiProfile(t, profile) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + processing := geminiFileJSON(record, fileName, fileURI, "PROCESSING", newGeminiFileTimeline()) + pollSeen := make(chan struct{}) + var pollOnce sync.Once + ctx, cancel := context.WithCancel(t.Context()) + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.Method { + case http.MethodPost: + if request.Header.Get("X-Goog-Upload-Command") == "start" { + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + } + response := geminiJSONResponse(request, `{"file":`+processing+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case http.MethodGet: + pollOnce.Do(func() { close(pollSeen) }) + return geminiJSONResponse(request, processing), nil + case http.MethodDelete: + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("unexpected request") + } + })) + result := make(chan error, 1) + go func() { + _, err := client.Embed(ctx, directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data)))) + result <- err + }() + <-pollSeen + cancel() + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("polling did not stop after cancellation") + } + }) +} + +func TestEmbedFilesAPIStopsAfterBoundedPollAttemptsAndCleansUp(t *testing.T) { + data := geminiTinyPNG(t) + profile := geminiDirectTestProfile(t, TransportFilesAPI) + profile.MaxPollAttempts = 2 + profile = rebindGeminiProfile(t, profile) + record := geminiCapability(t, profile, data, "synthetic.png", "image/png") + source := newGeminiLifecycleUpload(data, record) + fileName := "files/file-123" + fileURI := origin + "/v1beta/" + fileName + processing := geminiFileJSON(record, fileName, fileURI, "PROCESSING", newGeminiFileTimeline()) + var requests atomic.Int32 + client := newGeminiTestClient(t, profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + switch request.Method { + case http.MethodPost: + if request.Header.Get("X-Goog-Upload-Command") == "start" { + response := geminiJSONResponse(request, `{}`) + response.Header.Set("X-Goog-Upload-Url", origin+"/upload/v1beta/files?upload_id=synthetic-session-123&upload_protocol=resumable") + return response, nil + } + response := geminiJSONResponse(request, `{"file":`+processing+`}`) + response.Header.Set("X-Goog-Upload-Status", "final") + return response, nil + case http.MethodGet: + return geminiJSONResponse(request, processing), nil + case http.MethodDelete: + return geminiJSONResponse(request, `{}`), nil + default: + return nil, errors.New("unexpected bounded-poll request") + } + })) + ctx, cancel := context.WithTimeout(t.Context(), 80*time.Millisecond) + defer cancel() + receipt := Receipt{} + + _, err := client.embed(ctx, directInputs(source), geminiDirectAuthorization(profile.Descriptor, int64(len(data))), &receipt) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, int32(5), requests.Load()) + assert.Equal(t, 5, receipt.RequestCount) + assert.Equal(t, int32(1), source.closeCalls.Load()) +} + +func geminiDirectTestProfile(t *testing.T, transport Transport) Profile { + t.Helper() + profile := geminiTestProfile(t, 128) + profile.Transport = transport + profile.CapabilityProfileFingerprint = testCapabilityProfile + profile.DisclosureFingerprint = testDisclosurePolicy + profile.PollInterval = minimumPoll + profile.CleanupTimeout = time.Second + profile.Descriptor.InputKinds = []document.EmbeddingInputKind{ + document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk, + } + return rebindGeminiProfile(t, profile) +} + +func rebindGeminiProfile(t *testing.T, profile Profile) Profile { + t.Helper() + profile.Descriptor.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + profile.Descriptor.Fingerprint = "" + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + profile.Descriptor = descriptor + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor.Fingerprint = "" + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return profile +} + +func geminiCapability(t *testing.T, profile Profile, data []byte, filename, mediaType string) media.CapabilityRecord { + t.Helper() + digest := sha256.Sum256(data) + maxDuration := int64(180_000) + if strings.HasPrefix(mediaType, "video/") { + maxDuration = 120_000 + } + record, err := media.InspectCapability(bytes.NewReader(data), media.InspectionPolicy{ + Filename: filename, DeclaredMediaType: mediaType, ExpectedBytes: int64(len(data)), + ExpectedSHA256: hex.EncodeToString(digest[:]), DescriptorFingerprint: profile.Descriptor.Fingerprint, + ProfileFingerprint: profile.CapabilityProfileFingerprint, DisclosureFingerprint: profile.DisclosureFingerprint, + InputKind: document.RenditionInputOriginalFile, MaxSourceBytes: profile.MaxInputBytes, + MaxExpandedBytes: 1 << 20, MaxEntryBytes: 1 << 20, MaxEntries: 100, MaxNestingDepth: 1, + MaxTextLines: 1_000, MaxCharacters: 1 << 20, MaxRecords: 1_000, MaxPages: 6, + MaxSlides: 100, MaxSheets: 100, MaxCells: 10_000, MaxSpineItems: 1_000, MaxResources: 10_000, + MaxPixels: 16_000_000, MaxFrames: 32, MaxDurationMS: maxDuration, + }) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + return record +} + +func directInputs(sources ...document.AuthorizedUpload) []document.EmbeddingInput { + inputs := make([]document.EmbeddingInput, len(sources)) + for index, source := range sources { + inputs[index] = document.EmbeddingInput{Key: "direct-" + string(rune('a'+index)), Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputOriginalFile, Source: source} + } + return inputs +} + +func geminiDirectAuthorization(descriptor document.EmbeddingDescriptor, maxInputBytes int64) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 8, MaxInputBytes: maxInputBytes, + MaxResponseBytes: 8 << 10} +} + +type geminiFileTimeline struct { + created time.Time + updated time.Time + expires time.Time +} + +func newGeminiFileTimeline() geminiFileTimeline { + created := time.Now().UTC().Truncate(time.Second) + return geminiFileTimeline{created: created, updated: created.Add(time.Second), expires: created.Add(48 * time.Hour)} +} + +func geminiFileJSON(record media.CapabilityRecord, name, uri, state string, timeline geminiFileTimeline) string { + return `{"name":"` + name + `","displayName":"","mimeType":"` + record.MediaType + `","sizeBytes":"` + strconv.FormatInt(record.SourceBytes, 10) + + `","createTime":"` + timeline.created.Format(time.RFC3339Nano) + `","updateTime":"` + timeline.updated.Format(time.RFC3339Nano) + + `","expirationTime":"` + timeline.expires.Format(time.RFC3339Nano) + `","sha256Hash":"` + + base64.StdEncoding.EncodeToString([]byte(record.SourceSHA256)) + `","uri":"` + uri + `","downloadUri":"","state":"` + state + `","source":"UPLOADED"}` +} + +func geminiTinyPNG(t *testing.T) []byte { + t.Helper() + data, err := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=") + require.NoError(t, err) + return data +} + +type geminiLifecycleUpload struct { + reader *bytes.Reader + record media.CapabilityRecord + + metadata []document.AuthorizedUploadMetadata + metadataCalls atomic.Int32 + capabilityCalls atomic.Int32 + closeCalls atomic.Int32 + readPasses atomic.Int32 + blockRead bool + captureReadBuffer bool + capturedReadBuffer []byte + readStarted chan struct{} + released chan struct{} + startOnce sync.Once + releaseOnce sync.Once +} + +type geminiValueUpload struct { + *geminiLifecycleUpload +} + +type geminiNonComparableUpload struct { + *geminiLifecycleUpload + + identity []byte +} + +type geminiInterfaceIdentityUpload struct { + *geminiLifecycleUpload + + identity any +} + +func newGeminiLifecycleUpload(data []byte, record media.CapabilityRecord) *geminiLifecycleUpload { + metadata := document.AuthorizedUploadMetadata{Filename: record.Policy.Filename, MediaFamily: record.MediaFamily, + MediaType: record.MediaType, ByteLength: record.SourceBytes, SHA256: record.SourceSHA256, + CapabilityRecordChecksum: record.Checksum, ProviderMetadataChecksum: strings.Repeat("b", 64), + InputKind: document.RenditionInputOriginalFile} + return &geminiLifecycleUpload{reader: bytes.NewReader(data), record: record, metadata: []document.AuthorizedUploadMetadata{metadata}, + readStarted: make(chan struct{}), released: make(chan struct{})} +} + +func (upload *geminiLifecycleUpload) Read(buffer []byte) (int, error) { + upload.startOnce.Do(func() { + upload.readPasses.Add(1) + close(upload.readStarted) + }) + if upload.blockRead { + <-upload.released + return 0, errors.New("synthetic source closed") + } + read, err := upload.reader.Read(buffer) + if upload.captureReadBuffer && read > 0 && upload.capturedReadBuffer == nil { + upload.capturedReadBuffer = buffer[:read] + } + return read, err //nolint:wrapcheck // Preserve io.EOF for the real reader contract. +} + +func (upload *geminiLifecycleUpload) Close() error { + upload.closeCalls.Add(1) + upload.releaseOnce.Do(func() { close(upload.released) }) + return nil +} + +func (upload *geminiLifecycleUpload) Metadata() document.AuthorizedUploadMetadata { + call := int(upload.metadataCalls.Add(1)) - 1 + return upload.metadata[min(call, len(upload.metadata)-1)] +} + +func (upload *geminiLifecycleUpload) CapabilityRecord() media.CapabilityRecord { + upload.capabilityCalls.Add(1) + return upload.record +} + +var _ document.AuthorizedUpload = (*geminiLifecycleUpload)(nil) diff --git a/document/geminiembed/profile.go b/document/geminiembed/profile.go new file mode 100644 index 00000000..8fc834a9 --- /dev/null +++ b/document/geminiembed/profile.go @@ -0,0 +1,400 @@ +// Package geminiembed implements the fixed hosted Gemini Embedding 2 contract. +package geminiembed + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "net/http" + "net/netip" + "reflect" + "slices" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + ProviderID = "gemini.hosted.embedding-2-v1" + Model = "gemini-embedding-2" + DocumentFormatterV1 = "gemini-embedding-2/search-document/v1" + QueryFormatterV1 = "gemini-embedding-2/search-query/v1" + ScalarEncodingFloat32 = "float32" + + host = "generativelanguage.googleapis.com" + origin = "https://generativelanguage.googleapis.com" + embedPath = "/v1beta/models/gemini-embedding-2:embedContent" + filesUploadPath = "/upload/v1beta/files" + filesPathPrefix = "/v1beta/files/" + adapterContract = "docbank-gemini-embedding-2/v1" + compatibilityID = "gemini-embedding-2/search/v1" + defaultTimeout = 30 * time.Second + defaultPoll = 250 * time.Millisecond + defaultPollAttempts = 120 + defaultCleanup = 5 * time.Second + retentionCeiling = 48 * time.Hour + defaultInputBytes = int64(1 << 20) + defaultRequest = int64(2 << 20) + defaultResponse = int64(32 << 20) + maximumTimeout = 5 * time.Minute + minimumPoll = 10 * time.Millisecond + maximumPoll = 30 * time.Second + maximumPollAttempts = 10_000 + maximumCleanup = 30 * time.Second + maximumInputBytes = int64(100 << 20) + maximumRequest = int64(100 << 20) + maximumResponse = int64(100 << 20) + maximumTokenBytes = 128 +) + +type Transport string + +const ( + TransportInline Transport = "inline" + TransportFilesAPI Transport = "files-api" +) + +type SecretResolver interface { + ResolveSecret(ctx context.Context, binding string) (string, error) +} + +type Profile struct { + Descriptor document.EmbeddingDescriptor + CompatibilityEpoch string + SecretBinding string + Transport Transport + CapabilityProfileFingerprint string + DisclosureFingerprint string + RequestTimeout time.Duration + PollInterval time.Duration + MaxPollAttempts int + CleanupTimeout time.Duration + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + EgressPolicy providerhttp.EgressPolicy +} + +type Client struct { + profile Profile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + CompatibilityEpoch string `json:"compatibility_epoch"` + SecretBinding string `json:"secret_binding"` + Transport Transport `json:"transport"` + CapabilityProfile string `json:"capability_profile_fingerprint"` + DisclosurePolicy string `json:"disclosure_fingerprint"` + RetentionCeiling int64 `json:"provider_retention_ceiling_nanos"` + RequestTimeout int64 `json:"request_timeout_nanos"` + PollInterval int64 `json:"poll_interval_nanos"` + MaxPollAttempts int `json:"max_poll_attempts"` + CleanupTimeout int64 `json:"cleanup_timeout_nanos"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + Egress egressIdentity `json:"egress"` +} + +type egressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptor, err := normalizeProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(policyIdentity{ + AdapterContract: adapterContract, Origin: origin, Route: embedPath, Descriptor: descriptor, + CompatibilityEpoch: normalized.CompatibilityEpoch, SecretBinding: normalized.SecretBinding, + Transport: normalized.Transport, CapabilityProfile: normalized.CapabilityProfileFingerprint, + DisclosurePolicy: normalized.DisclosureFingerprint, RetentionCeiling: int64(profileRetention(normalized.Transport)), + RequestTimeout: int64(normalized.RequestTimeout), PollInterval: int64(normalized.PollInterval), + MaxPollAttempts: normalized.MaxPollAttempts, + CleanupTimeout: int64(normalized.CleanupTimeout), + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, Egress: egressPolicyIdentity(normalized.EgressPolicy), + }, json.Deterministic(true)) + if err != nil { + return "", errors.New("gemini embed: policy identity encoding failed") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +func New(profile Profile, secrets SecretResolver, resolver providerhttp.Resolver, supplied *http.Client) (*Client, error) { + if supplied == nil { + return nil, errors.New("gemini embed: HTTP client settings source is required") + } + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + return nil, errors.New("gemini embed: descriptor is not canonical") + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("gemini embed: descriptor policy fingerprint does not match profile") + } + if nilInterface(secrets) { + return nil, errors.New("gemini embed: named API-key resolver is required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, errors.New("gemini embed: sealed egress policy is invalid") + } + isolate := *supplied + isolate.Transport = transport + isolate.CheckRedirect = providerhttp.RefuseRedirects + isolate.Jar = nil + isolate.Timeout = 0 + normalized.Descriptor = cloneDescriptor(descriptor) + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, http: &isolate}, nil +} + +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + profile.EgressPolicy.TLS.SPKISHA256 = slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultTimeout + } + if profile.PollInterval == 0 { + profile.PollInterval = defaultPoll + } + if profile.MaxPollAttempts == 0 { + profile.MaxPollAttempts = defaultPollAttempts + } + if profile.CleanupTimeout == 0 { + profile.CleanupTimeout = defaultCleanup + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultRequest + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultResponse + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maximumTimeout || + profile.PollInterval < minimumPoll || profile.PollInterval > maximumPoll || + profile.MaxPollAttempts < 1 || profile.MaxPollAttempts > maximumPollAttempts || + profile.CleanupTimeout <= 0 || profile.CleanupTimeout > maximumCleanup || + profile.MaxInputBytes < 1 || profile.MaxInputBytes > maximumInputBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maximumRequest || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maximumResponse { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: execution bounds are invalid") + } + if profile.Transport != TransportInline && profile.Transport != TransportFilesAPI { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: transport is invalid") + } + if profile.CompatibilityEpoch != Model || profile.Descriptor.ModelRevision != Model { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: stable model revision must be gemini-embedding-2") + } + if !validToken(profile.SecretBinding) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: named API-key binding is required") + } + if !validFingerprint(profile.CapabilityProfileFingerprint) || !validFingerprint(profile.DisclosureFingerprint) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: capability profile and disclosure identities are required") + } + if err := normalizeEgress(&profile.EgressPolicy); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + descriptor := cloneDescriptor(profile.Descriptor) + descriptor.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptor.Fingerprint = "" + var err error + descriptor, err = document.NewEmbeddingDescriptor(descriptor) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: descriptor identity is invalid") + } + descriptor.PolicyFingerprint, descriptor.Fingerprint = "", "" + if err := validateDescriptor(descriptor); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + if profile.Transport == TransportFilesAPI { + minimum, capacityErr := minimumFilesAPIRequestCapacity(descriptor.Dimension) + if capacityErr != nil || profile.MaxRequestBytes < minimum { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("gemini embed: Files API request byte capacity is too small") + } + } + return profile, descriptor, nil +} + +func minimumFilesAPIRequestCapacity(dimension int) (int64, error) { + start, err := json.Marshal(wireStartUpload{}) + if err != nil { + return 0, err + } + part := wirePart{FileData: &wireFileData{ + MIMEType: "video/quicktime", FileURI: maximumProviderFileURI(), + }} + fileData, err := json.Marshal(wireRequest{Model: "models/" + Model, + Content: wireContent{Parts: []wirePart{part}}, OutputDimensionality: dimension}) + if err != nil { + return 0, err + } + return int64(max(len(start), len(fileData))), nil +} + +func maximumProviderFileURI() string { + return origin + "/v1beta/files/" + strings.Repeat("a", 40) +} + +func validateDescriptor(descriptor document.EmbeddingDescriptor) error { + contract, err := modelInputContract() + if err != nil { + return errors.New("gemini embed: fixed model-input contract is invalid") + } + if descriptor.ID != ProviderID || descriptor.ContractVersion != document.EmbeddingProviderContractVersion || + descriptor.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptor.Model != Model || + descriptor.ModelRevision != Model || descriptor.Dimension < 128 || descriptor.Dimension > 3072 || + descriptor.Metric != document.VectorMetricCosine || descriptor.Normalization != document.VectorNormalizationUnitLength || + descriptor.ScalarEncoding != ScalarEncodingFloat32 || descriptor.DocumentFormatter != DocumentFormatterV1 || + descriptor.QueryFormatter != QueryFormatterV1 || !descriptor.SupportsTextQuery || + !reflect.DeepEqual(descriptor.ModelInput, contract) || descriptor.CompatibilityID != compatibilityID || + !slices.Equal(descriptor.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk}) || + !slices.Equal(descriptor.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeText}) { + return errors.New("gemini embed: descriptor does not match the fixed hosted multimodal contract") + } + return nil +} + +func modelInputContract() (document.ModelInputContract, error) { + return document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: compatibilityID, + Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "title: none | text: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "task: search result | query: {{content}}"}, + }) +} + +func normalizeEgress(policy *providerhttp.EgressPolicy) error { + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = providerhttp.DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if policy.ProxyMode == "" { + policy.ProxyMode = providerhttp.ProxyDisabled + } + if policy.Scheme != "https" || policy.Host != host || policy.Port != 443 || + policy.ProxyMode != providerhttp.ProxyDisabled || policy.TLS.RootCAs != nil { + return errors.New("gemini embed: egress authority must be exactly generativelanguage.googleapis.com:443") + } + for index := range policy.AllowedCIDRs { + policy.AllowedCIDRs[index] = policy.AllowedCIDRs[index].Masked() + } + slices.SortFunc(policy.AllowedCIDRs, func(left, right netip.Prefix) int { return strings.Compare(left.String(), right.String()) }) + for index := 1; index < len(policy.AllowedCIDRs); index++ { + if policy.AllowedCIDRs[index] == policy.AllowedCIDRs[index-1] { + return errors.New("gemini embed: egress policy has a duplicate CIDR") + } + } + for index := range policy.TLS.SPKISHA256 { + policy.TLS.SPKISHA256[index] = strings.ToLower(policy.TLS.SPKISHA256[index]) + } + slices.Sort(policy.TLS.SPKISHA256) + for index := 1; index < len(policy.TLS.SPKISHA256); index++ { + if policy.TLS.SPKISHA256[index] == policy.TLS.SPKISHA256[index-1] { + return errors.New("gemini embed: egress policy has a duplicate SPKI pin") + } + } + if _, err := providerhttp.NewTransport(*policy, nil); err != nil { + return errors.New("gemini embed: sealed egress policy is invalid") + } + return nil +} + +func egressPolicyIdentity(policy providerhttp.EgressPolicy) egressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.String() + } + return egressIdentity{Scheme: policy.Scheme, Host: policy.Host, Port: policy.Port, + AllowedCIDRs: cidrs, ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), + KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), + SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func cloneDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} + +func validToken(value string) bool { + if value == "" || len(value) > maximumTokenBytes || !utf8.ValidString(value) { + return false + } + for _, runeValue := range value { + if runeValue < 0x21 || runeValue > 0x7e { + return false + } + } + return true +} + +func validFingerprint(value string) bool { + if len(value) != sha256.Size*2 { + return false + } + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == sha256.Size && value == strings.ToLower(value) +} + +func profileRetention(transport Transport) time.Duration { + if transport == TransportFilesAPI { + return retentionCeiling + } + return 0 +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/geminiembed/profile_test.go b/document/geminiembed/profile_test.go new file mode 100644 index 00000000..0eced49c --- /dev/null +++ b/document/geminiembed/profile_test.go @@ -0,0 +1,130 @@ +package geminiembed + +import ( + "context" + "net/http" + "net/netip" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +// TestProfileBindsGeminiSearchPrefixesDimensionAndTransport catches a profile +// that permits a different Gemini model, role envelope, vector dimension, or +// direct-file transport identity. +func TestProfileBindsGeminiSearchPrefixesDimensionAndTransport(t *testing.T) { + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, + CompatibilityID: "gemini-embedding-2/search/v1", + Document: document.ModelInputEncoder{ + Mode: document.ModelInputModeText, Template: "title: none | text: {{content}}", + }, + Query: document.ModelInputEncoder{ + Mode: document.ModelInputModeText, Template: "task: search result | query: {{content}}", + }, + }) + require.NoError(t, err) + + for _, transport := range []Transport{TransportInline, TransportFilesAPI} { + t.Run(string(transport), func(t *testing.T) { + profile := Profile{ + CompatibilityEpoch: "gemini-embedding-2", + SecretBinding: "secret:gemini", + Transport: transport, + CapabilityProfileFingerprint: testCapabilityProfile, + DisclosureFingerprint: testDisclosurePolicy, + RequestTimeout: time.Second, + MaxInputBytes: 4096, + MaxRequestBytes: 8192, + MaxResponseBytes: 16384, + EgressPolicy: providerhttp.EgressPolicy{ + Scheme: "https", Host: "generativelanguage.googleapis.com", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("192.0.2.0/24")}, + ProxyMode: providerhttp.ProxyDisabled, ConnectTimeout: time.Second, + KeepAlive: time.Second, TLSHandshakeTimeout: time.Second, + }, + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustHostedProvider, Model: "gemini-embedding-2", + ModelRevision: "gemini-embedding-2", Dimension: 768, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{ + document.EmbeddingInputOriginalFile, document.EmbeddingInputRenditionChunk, + }, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }, + } + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + + client, err := New(profile, syntheticSecrets{"secret:gemini": "synthetic-key"}, + syntheticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + assert.Equal(t, "gemini-embedding-2", client.Descriptor().Model) + assert.Equal(t, 768, client.Descriptor().Dimension) + assert.Equal(t, "title: none | text: {{content}}", client.Descriptor().ModelInput.Document.Template) + assert.Equal(t, "task: search result | query: {{content}}", client.Descriptor().ModelInput.Query.Template) + }) + } +} + +// TestProfileRejectsFilesAPIRequestBoundBelowFixedEnvelopes catches a profile +// that can encode the resumable start body but cannot encode the supported +// fileData embedding envelope it promises. +func TestProfileRejectsFilesAPIRequestBoundBelowFixedEnvelopes(t *testing.T) { + profile := geminiTestProfile(t, 128) + profile.Transport = TransportFilesAPI + profile.MaxRequestBytes = int64(len(`{"file":{}}`)) + + _, err := PolicyFingerprint(profile) + require.Error(t, err) + assert.Contains(t, err.Error(), "request") +} + +func TestProfileRejectsUnsafePollInterval(t *testing.T) { + profile := geminiTestProfile(t, 128) + profile.PollInterval = time.Nanosecond + + _, err := PolicyFingerprint(profile) + require.Error(t, err) + assert.Contains(t, err.Error(), "execution bounds") +} + +func TestProfileBindsFinitePollAttemptCapacity(t *testing.T) { + first := geminiTestProfile(t, 128) + first.MaxPollAttempts = 3 + firstFingerprint, err := PolicyFingerprint(first) + require.NoError(t, err) + + second := first + second.MaxPollAttempts = 4 + secondFingerprint, err := PolicyFingerprint(second) + require.NoError(t, err) + assert.NotEqual(t, firstFingerprint, secondFingerprint) + + first.MaxPollAttempts = -1 + _, err = PolicyFingerprint(first) + require.Error(t, err) + assert.Contains(t, err.Error(), "execution bounds") +} + +type syntheticSecrets map[string]string + +func (secrets syntheticSecrets) ResolveSecret(_ context.Context, binding string) (string, error) { + return secrets[binding], nil +} + +type syntheticResolver []netip.Addr + +func (resolver syntheticResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return append([]netip.Addr(nil), resolver...), nil +} diff --git a/document/glmocr/client.go b/document/glmocr/client.go index adba0743..de52ef62 100644 --- a/document/glmocr/client.go +++ b/document/glmocr/client.go @@ -13,7 +13,9 @@ import ( "io" "mime" "net/http" + "net/netip" "net/url" + "strconv" "strings" "sync" "time" @@ -21,6 +23,7 @@ import ( "go.kenn.io/docbank/document" "go.kenn.io/docbank/document/ocr" + "go.kenn.io/docbank/document/providerhttp" ) const ( @@ -87,7 +90,26 @@ func NewClient(policy Policy, config ClientConfig) (*Client, error) { } httpClient := config.HTTPClient if httpClient == nil { - httpClient = &http.Client{Timeout: config.Timeout} + endpoint, err := url.Parse(policy.values.Endpoint) + if err != nil { + return nil, fmt.Errorf("parse GLM-OCR endpoint for egress: %w", err) + } + port, err := strconv.ParseUint(endpoint.Port(), 10, 16) + if err != nil || port == 0 { + return nil, errors.New("GLM-OCR endpoint has an invalid egress port") + } + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: endpoint.Scheme, Host: endpoint.Hostname(), Port: uint16(port), + AllowedCIDRs: []netip.Prefix{ + netip.MustParsePrefix("127.0.0.0/8"), + netip.MustParsePrefix("::1/128"), + }, + ProxyMode: providerhttp.ProxyDisabled, + }, nil) + if err != nil { + return nil, fmt.Errorf("create GLM-OCR egress transport: %w", err) + } + httpClient = &http.Client{Timeout: config.Timeout, Transport: transport} } else { clone := *httpClient httpClient = &clone @@ -95,7 +117,7 @@ func NewClient(policy Policy, config ClientConfig) (*Client, error) { httpClient.Timeout = config.Timeout } } - httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + httpClient.CheckRedirect = providerhttp.RefuseRedirects return &Client{policy: policy, http: httpClient, maxRetries: config.MaxRetries, maxRetryDelay: config.MaxRetryDelay}, nil } diff --git a/document/glmocr/client_internal_test.go b/document/glmocr/client_internal_test.go new file mode 100644 index 00000000..b3d18c3f --- /dev/null +++ b/document/glmocr/client_internal_test.go @@ -0,0 +1,44 @@ +package glmocr + +import ( + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +func TestDefaultClientRefusesEgressDestinationDrift(t *testing.T) { + provider := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "provider") + })) + defer provider.Close() + var driftRequests atomic.Int32 + drift := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + driftRequests.Add(1) + _, _ = io.WriteString(w, "drift") + })) + defer drift.Close() + normalize, err := document.NewNormalizePolicy(1_000_000) + require.NoError(t, err) + policy, err := NewPolicy(PolicyConfig{ + Endpoint: provider.URL + "/glmocr/parse", MaxDocumentBytes: 1 << 20, + MaxResponseBytes: 1 << 20, MaxUnits: 1, NormalizePolicy: normalize, + }) + require.NoError(t, err) + client, err := NewClient(policy, ClientConfig{Timeout: time.Second}) + require.NoError(t, err) + + response, err := client.http.Get(drift.URL) + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, providerhttp.ErrDestinationDenied) + assert.Zero(t, driftRequests.Load()) +} diff --git a/document/glmocr/client_test.go b/document/glmocr/client_test.go index 154ad572..8ca2c64d 100644 --- a/document/glmocr/client_test.go +++ b/document/glmocr/client_test.go @@ -110,7 +110,7 @@ func TestClientRetriesHTTPClientTimeout(t *testing.T) { require.ErrorIs(t, err, context.DeadlineExceeded) assert.Equal(t, ocr.ErrorTransient, ocr.ErrorKindOf(err)) assert.Equal(t, 2, ocr.MetricsFromError(err).Requests) - assert.Equal(t, int32(2), requests.Load()) + require.Eventually(t, func() bool { return requests.Load() == 2 }, time.Second, time.Millisecond) } func TestClientDoesNotRetryUnsuitableInput(t *testing.T) { diff --git a/document/input_contract.go b/document/input_contract.go new file mode 100644 index 00000000..fc9f4bf6 --- /dev/null +++ b/document/input_contract.go @@ -0,0 +1,315 @@ +package document + +import ( + "errors" + "fmt" + "strings" + "unicode/utf8" +) + +const ( + // ModelInputProfileOpenAICompatible is the plain-text contract used by + // OpenAI-compatible embedding endpoints. + ModelInputProfileOpenAICompatible ModelInputProfile = "openai-compatible/v1" + // ModelInputProfileVoyage selects Voyage's explicit document/query modes. + ModelInputProfileVoyage ModelInputProfile = "voyage/v1" + // ModelInputProfileMistral is Mistral's plain-text embedding contract. + ModelInputProfileMistral ModelInputProfile = "mistral/v1" + // ModelInputProfileCustom permits an explicit compatibility contract. + ModelInputProfileCustom ModelInputProfile = "custom/v1" + // ModelInputProfileNomic uses Nomic's reviewed asymmetric search prefixes. + ModelInputProfileNomic ModelInputProfile = "nomic/v1" + // ModelInputProfileE5 uses E5's reviewed passage/query prefixes. + ModelInputProfileE5 ModelInputProfile = "e5/v1" + // ModelInputProfileBGEM3 is BGE-M3's prefix-free dense retrieval input. + ModelInputProfileBGEM3 ModelInputProfile = "bge-m3/v1" + // ModelInputProfileGTE is GTE's prefix-free retrieval input. + ModelInputProfileGTE ModelInputProfile = "gte/v1" + // ModelInputProfileQwen3 is Qwen3's prefix-free document input and optional + // explicit query instruction contract. + ModelInputProfileQwen3 ModelInputProfile = "qwen3/v1" + // ModelInputProfileQueryInstruction is the reviewed query-only instruction + // envelope for compatible models whose document role stays prefix-free. + ModelInputProfileQueryInstruction ModelInputProfile = "query-instruction/v1" + + modelInputContractVersion = 1 + modelInputContentSlot = "{{content}}" +) + +// ModelInputProfile identifies a reviewed built-in input policy or the one +// explicit custom profile form. Provider aliases are never a durable policy. +type ModelInputProfile string + +// ModelInputMode is the adapter-native role field selected by a contract. +type ModelInputMode string + +const ( + ModelInputModeText ModelInputMode = "text" + ModelInputModeDocument ModelInputMode = "document" + ModelInputModeQuery ModelInputMode = "query" +) + +// ModelInputEncoder renders one role with exactly one content substitution. +type ModelInputEncoder struct { + Mode ModelInputMode `json:"mode"` + Template string `json:"template"` +} + +// ModelInputContract is a canonical, durable document/query formatting +// declaration. Fingerprint covers both role envelopes, including an explicit +// empty contract. +type ModelInputContract struct { + Version int `json:"version"` + Profile ModelInputProfile `json:"profile"` + CompatibilityID string `json:"compatibility_id"` + Document ModelInputEncoder `json:"document"` + Query ModelInputEncoder `json:"query"` + QueryInstruction string `json:"query_instruction,omitempty"` + Fingerprint string `json:"fingerprint"` +} + +// ModelInputContractConfig constructs one canonical model-input contract. +type ModelInputContractConfig struct { + Profile ModelInputProfile + CompatibilityID string + Document ModelInputEncoder + Query ModelInputEncoder + QueryInstruction string +} + +// NewModelInputContract resolves a reviewed built-in profile or one explicit +// custom contract. An empty config is intentional and has its own fingerprint. +func NewModelInputContract(config ModelInputContractConfig) (ModelInputContract, error) { + contract := ModelInputContract{Version: modelInputContractVersion, Profile: config.Profile, QueryInstruction: config.QueryInstruction} + switch config.Profile { + case "": + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("empty model-input contract cannot define encoders") + } + case ModelInputProfileOpenAICompatible: + if config.QueryInstruction != "" { + return ModelInputContract{}, errors.New("openai-compatible profile cannot define a query instruction") + } + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "openai-compatible/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + case ModelInputProfileVoyage: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "voyage/document-query/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeDocument, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeQuery, Template: modelInputContentSlot} + case ModelInputProfileMistral: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "mistral/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + case ModelInputProfileNomic: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "nomic/search/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: "search_document: " + modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: "search_query: " + modelInputContentSlot} + case ModelInputProfileE5: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "e5/asymmetric/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: "passage: " + modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: "query: " + modelInputContentSlot} + case ModelInputProfileBGEM3: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "bge-m3/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + case ModelInputProfileGTE: + if hasModelInputOverrides(config) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + contract.CompatibilityID = "gte/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + case ModelInputProfileQwen3: + if config.CompatibilityID != "" || config.Document != (ModelInputEncoder{}) || config.Query != (ModelInputEncoder{}) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + if err := validateQueryInstruction(config.QueryInstruction, false); err != nil { + return ModelInputContract{}, err + } + contract.CompatibilityID = "qwen3/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: queryInstructionTemplate(config.QueryInstruction)} + case ModelInputProfileQueryInstruction: + if config.CompatibilityID != "" || config.Document != (ModelInputEncoder{}) || config.Query != (ModelInputEncoder{}) { + return ModelInputContract{}, errors.New("built-in model-input profiles cannot be overridden") + } + if err := validateQueryInstruction(config.QueryInstruction, true); err != nil { + return ModelInputContract{}, err + } + contract.CompatibilityID = "query-instruction/text/v1" + contract.Document = ModelInputEncoder{Mode: ModelInputModeText, Template: modelInputContentSlot} + contract.Query = ModelInputEncoder{Mode: ModelInputModeText, Template: queryInstructionTemplate(config.QueryInstruction)} + case ModelInputProfileCustom: + if config.QueryInstruction != "" { + return ModelInputContract{}, errors.New("custom model-input contracts must encode instructions in their explicit query template") + } + contract.CompatibilityID = config.CompatibilityID + contract.Document = config.Document + contract.Query = config.Query + default: + if strings.HasPrefix(string(config.Profile), "alias:") { + return ModelInputContract{}, errors.New("opaque model-input aliases require an explicit compatibility contract") + } + return ModelInputContract{}, fmt.Errorf("unknown built-in model-input profile %q", config.Profile) + } + if err := validateModelInputContractFields(contract); err != nil { + return ModelInputContract{}, err + } + contract.Fingerprint = "" + encoded, err := canonicalJSON(contract) + if err != nil { + return ModelInputContract{}, fmt.Errorf("encode model-input contract: %w", err) + } + contract.Fingerprint = sha256Hex(encoded) + return contract, nil +} + +// EncodeDocument renders document text through the sealed document encoder. +func (contract ModelInputContract) EncodeDocument(content string) string { + return strings.Replace(contract.Document.Template, modelInputContentSlot, content, 1) +} + +// EncodeQuery renders query text through the sealed query encoder. +func (contract ModelInputContract) EncodeQuery(content string) string { + return strings.Replace(contract.Query.Template, modelInputContentSlot, content, 1) +} + +// modelInputRenderedLength computes an encoder's exact rendered byte length +// without allocating its template output. +func modelInputRenderedLength(encoder ModelInputEncoder, content string) (int64, error) { + contentOffset := strings.Index(encoder.Template, modelInputContentSlot) + if contentOffset < 0 { + return 0, errors.New("model-input encoder does not contain content slot") + } + const maxInt64 = int64(^uint64(0) >> 1) + length := int64(contentOffset) + if int64(len(content)) > maxInt64-length { + return 0, errors.New("model-input rendered length overflows") + } + length += int64(len(content)) + suffix := len(encoder.Template) - contentOffset - len(modelInputContentSlot) + if int64(suffix) > maxInt64-length { + return 0, errors.New("model-input rendered length overflows") + } + return length + int64(suffix), nil +} + +func validateModelInputContract(contract ModelInputContract) error { + if err := validateModelInputContractFields(contract); err != nil { + return err + } + fingerprint := contract.Fingerprint + contract.Fingerprint = "" + encoded, err := canonicalJSON(contract) + if err != nil { + return err + } + if fingerprint != sha256Hex(encoded) { + return errors.New("model-input contract fingerprint or canonical form is invalid") + } + if contract.Profile != ModelInputProfileCustom && contract.Profile != "" { + canonical, err := NewModelInputContract(ModelInputContractConfig{Profile: contract.Profile, QueryInstruction: contract.QueryInstruction}) + canonical.Fingerprint = "" + if err != nil || contract != canonical { + return errors.New("model-input contract is not the reviewed built-in profile") + } + } + if contract.Profile == "" { + empty, err := NewModelInputContract(ModelInputContractConfig{}) + empty.Fingerprint = "" + if err != nil || contract != empty { + return errors.New("empty model-input contract is not canonical") + } + } + return nil +} + +func validateModelInputContractFields(contract ModelInputContract) error { + if contract.Version != modelInputContractVersion { + return fmt.Errorf("model-input contract version must be %d", modelInputContractVersion) + } + if contract.Profile == "" { + if contract.CompatibilityID != "" || contract.Document != (ModelInputEncoder{}) || contract.Query != (ModelInputEncoder{}) || contract.QueryInstruction != "" { + return errors.New("empty model-input contract must have zero encoders and compatibility ID") + } + return nil + } + if err := validateCompatibilityID(contract.CompatibilityID); err != nil { + return err + } + if err := validateModelInputEncoder("document", contract.Document); err != nil { + return err + } + return validateModelInputEncoder("query", contract.Query) +} + +func hasModelInputOverrides(config ModelInputContractConfig) bool { + return config.CompatibilityID != "" || config.Document != (ModelInputEncoder{}) || config.Query != (ModelInputEncoder{}) || config.QueryInstruction != "" +} + +func validateQueryInstruction(instruction string, required bool) error { + if required && instruction == "" { + return errors.New("query-instruction profile requires a query instruction") + } + if instruction == "" { + return nil + } + if !utf8.ValidString(instruction) || len(instruction) > 4096 || strings.ContainsAny(instruction, "\x00\r") || strings.Contains(instruction, modelInputContentSlot) { + return errors.New("query instruction must be bounded valid UTF-8 text without a content slot") + } + return nil +} + +func queryInstructionTemplate(instruction string) string { + if instruction == "" { + return modelInputContentSlot + } + return "Instruct: " + instruction + "\nQuery:" + modelInputContentSlot +} + +func validateCompatibilityID(value string) error { + if value == "" || len(value) > 128 { + return errors.New("model-input compatibility ID must contain 1-128 characters") + } + for _, character := range value { + if character >= 'a' && character <= 'z' || character >= '0' && character <= '9' || character == '_' || character == '-' || character == '.' || character == '/' { + continue + } + return errors.New("model-input compatibility ID contains unsupported characters") + } + return nil +} + +func validateModelInputEncoder(role string, encoder ModelInputEncoder) error { + switch encoder.Mode { + case ModelInputModeText, ModelInputModeDocument, ModelInputModeQuery: + default: + return fmt.Errorf("model-input %s mode is invalid", role) + } + if !utf8.ValidString(encoder.Template) || len(encoder.Template) > 4096 { + return fmt.Errorf("model-input %s template must be bounded valid UTF-8", role) + } + if strings.Count(encoder.Template, modelInputContentSlot) != 1 { + return fmt.Errorf("model-input %s template must contain exactly one content slot", role) + } + return nil +} diff --git a/document/input_contract_test.go b/document/input_contract_test.go new file mode 100644 index 00000000..1857f4a9 --- /dev/null +++ b/document/input_contract_test.go @@ -0,0 +1,130 @@ +package document_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +// This test fails if an invalid model-input declaration becomes executable or +// two distinct role envelopes collapse to one durable compatibility identity. +func TestEmbeddingContractCanonicalizesRoleSpecificModelInput(t *testing.T) { + base, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileOpenAICompatible}) + require.NoError(t, err) + assert.Equal(t, "passage", base.EncodeDocument("passage")) + assert.Equal(t, "question", base.EncodeQuery("question")) + + for _, testCase := range []struct { + name string + config document.ModelInputContractConfig + want string + }{ + {"unknown built-in", document.ModelInputContractConfig{Profile: "unknown/v1"}, "unknown built-in"}, + {"opaque alias", document.ModelInputContractConfig{Profile: "alias:vendor-default"}, "compatibility"}, + {"missing content slot", document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: "custom-space", Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "prefix"}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "{{content}}"}}, "exactly one content slot"}, + {"two content slots", document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: "custom-space", Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "{{content}}{{content}}"}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "{{content}}"}}, "exactly one content slot"}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := document.NewModelInputContract(testCase.config) + require.ErrorContains(t, err, testCase.want) + }) + } + + documentEnvelope, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "custom-space", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "document: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "query: {{content}}"}, + }) + require.NoError(t, err) + assert.Equal(t, "document: alpha", documentEnvelope.EncodeDocument("alpha")) + assert.Equal(t, "query: alpha", documentEnvelope.EncodeQuery("alpha")) + assert.NotEqual(t, base.Fingerprint, documentEnvelope.Fingerprint) + voyage, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileVoyage}) + require.NoError(t, err) + assert.Equal(t, "passage", voyage.EncodeDocument("passage")) + assert.Equal(t, "question", voyage.EncodeQuery("question")) + assert.NotEqual(t, base.Fingerprint, voyage.Fingerprint) + assert.NotEqual(t, voyage.Document.Mode, voyage.Query.Mode) + + empty, err := document.NewModelInputContract(document.ModelInputContractConfig{}) + require.NoError(t, err) + assert.NotEqual(t, base.Fingerprint, empty.Fingerprint) + assert.NotEqual(t, documentEnvelope.Fingerprint, empty.Fingerprint) +} + +// This test fails if an adapter's declared native request modes cease to gate +// the model-input contract selected for its descriptor. +func TestEmbeddingContractRejectsUnsupportedProviderRequestMode(t *testing.T) { + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "custom-space", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}, + }) + require.NoError(t, err) + _, err = document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ + ID: "synthetic-embedder", ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: testFingerprint(), TrustBoundary: document.EmbeddingTrustLocalProcess, + Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: document.VectorMetricCosine, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: "custom-space", + ModelInput: contract, SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: document.VectorNormalizationUnitLength, + ScalarEncoding: "float32", + }) + require.ErrorContains(t, err, "request mode") +} + +func TestEmbeddingContractRequiresOnlyExecutableRequestModes(t *testing.T) { + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: "custom-space", Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}}) + require.NoError(t, err) + descriptor := document.EmbeddingDescriptor{ID: "synthetic-embedder", ContractVersion: document.EmbeddingProviderContractVersion, PolicyFingerprint: testFingerprint(), TrustBoundary: document.EmbeddingTrustLocalProcess, Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: document.VectorMetricCosine, InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: "custom-space", ModelInput: contract, SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeDocument}, DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: "float32"} + _, err = document.NewEmbeddingDescriptor(descriptor) + require.NoError(t, err) + descriptor.SupportsTextQuery = true + _, err = document.NewEmbeddingDescriptor(descriptor) + require.ErrorContains(t, err, "request mode") +} + +func TestEmbeddingContractRejectsForgedEmptyModelInput(t *testing.T) { + empty, err := document.NewModelInputContract(document.ModelInputContractConfig{}) + require.NoError(t, err) + empty.CompatibilityID = "forged" + _, err = document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ID: "synthetic-embedder", ContractVersion: document.EmbeddingProviderContractVersion, PolicyFingerprint: testFingerprint(), TrustBoundary: document.EmbeddingTrustLocalProcess, Model: "synthetic-model", ModelRevision: "r1", Dimension: 2, Metric: document.VectorMetricCosine, InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: "forged", ModelInput: empty, SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, DocumentFormatter: "document/v1", QueryFormatter: "query/v1", Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: "float32"}) + require.ErrorContains(t, err, "empty model-input") +} + +func TestEmbeddingContractReviewedRetrievalProfiles(t *testing.T) { + for _, testCase := range []struct { + name string + config document.ModelInputContractConfig + document string + query string + compatible string + }{ + {"nomic", document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}, "search_document: passage", "search_query: question", "nomic/search/v1"}, + {"e5", document.ModelInputContractConfig{Profile: document.ModelInputProfileE5}, "passage: passage", "query: question", "e5/asymmetric/v1"}, + {"bge-m3", document.ModelInputContractConfig{Profile: document.ModelInputProfileBGEM3}, "passage", "question", "bge-m3/text/v1"}, + {"gte", document.ModelInputContractConfig{Profile: document.ModelInputProfileGTE}, "passage", "question", "gte/text/v1"}, + {"qwen3 plain", document.ModelInputContractConfig{Profile: document.ModelInputProfileQwen3}, "passage", "question", "qwen3/text/v1"}, + {"qwen3 instructed query", document.ModelInputContractConfig{Profile: document.ModelInputProfileQwen3, QueryInstruction: "Retrieve legal evidence"}, "passage", "Instruct: Retrieve legal evidence\nQuery:question", "qwen3/text/v1"}, + {"generic query instruction", document.ModelInputContractConfig{Profile: document.ModelInputProfileQueryInstruction, QueryInstruction: "Retrieve relevant passages"}, "passage", "Instruct: Retrieve relevant passages\nQuery:question", "query-instruction/text/v1"}, + } { + t.Run(testCase.name, func(t *testing.T) { + contract, err := document.NewModelInputContract(testCase.config) + require.NoError(t, err) + assert.Equal(t, testCase.document, contract.EncodeDocument("passage")) + assert.Equal(t, testCase.query, contract.EncodeQuery("question")) + assert.Equal(t, testCase.compatible, contract.CompatibilityID) + }) + } +} + +func TestEmbeddingContractRejectsImplicitOrUnboundedQueryInstructions(t *testing.T) { + _, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileOpenAICompatible, QueryInstruction: "infer this from a model name"}) + require.ErrorContains(t, err, "cannot define a query instruction") + _, err = document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileQueryInstruction}) + require.ErrorContains(t, err, "requires a query instruction") + _, err = document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileQwen3, QueryInstruction: string(make([]byte, 4097))}) + require.ErrorContains(t, err, "bounded valid UTF-8") +} diff --git a/document/internal/formatdetect/detect.go b/document/internal/formatdetect/detect.go new file mode 100644 index 00000000..48ec7eaf --- /dev/null +++ b/document/internal/formatdetect/detect.go @@ -0,0 +1,1862 @@ +package formatdetect + +import ( + "archive/zip" + "bytes" + "encoding/binary" + "encoding/csv" + "encoding/hex" + "encoding/json/jsontext" + "encoding/xml" + "errors" + "fmt" + "io" + "math" + "mime" + "net/mail" + "os" + "path" + "strconv" + "strings" + "unicode/utf8" +) + +const ( + // MaxDocumentBytes is the absolute detector allocation and read ceiling. + MaxDocumentBytes = int64(500 << 20) + maxSniffBytes = int64(8 << 20) + maxTextSniffBytes = int64(50 << 20) + maxZIPEntries = 10_000 + maxZIPCentralDirectory = uint32(16 << 20) + maxZIPExpandedBytes = uint64(500 << 20) + maxZIPSingleExpandedByte = uint64(100 << 20) + maxPDFTailBytes = int64(64 << 10) + maxPDFXRefBytes = int64(4 << 10) + maxPDFStructuralBytes = 64 << 20 + maxPDFTokens = 1 << 20 + maxPDFPageTreeDepth = 256 + ooxmlContentTypesName = "[Content_Types].xml" +) + +// CompoundDirectoryNames validates one legacy compound-file directory. It is +// exported only so the Mistral compatibility suite can retain its allocation +// regression test while format detection is shared with core inspection. +func CompoundDirectoryNames(reader io.ReaderAt, size int64) (map[string]bool, error) { + return compoundDirectoryNames(reader, size) +} + +var ( + compoundFileMagic = []byte{0xd0, 0xcf, 0x11, 0xe0, 0xa1, 0xb1, 0x1a, 0xe1} +) + +const compoundNoStream = uint32(0xffffffff) + +type compoundDirectoryEntry struct { + name string + entryType byte + left, right, child uint32 +} + +// DetectFormat validates a provider candidate from bounded bytes. Declared +// type is a hint only: container families must prove internal markers, while +// inherently ambiguous text formats also require syntactically safe UTF-8. +func DetectFormat(reader io.ReaderAt, size int64, declaredMediaType string) (CandidateFormat, error) { + if reader == nil || size <= 0 { + return CandidateFormat{}, errors.New("document format detection requires nonempty bytes") + } + if size > MaxDocumentBytes { + return CandidateFormat{}, errors.New("document exceeds the format-detection byte limit") + } + mediaType, parameters, err := mime.ParseMediaType(declaredMediaType) + if err != nil || len(parameters) != 0 || mediaType != strings.ToLower(mediaType) { + return CandidateFormat{}, errors.New("document format detection requires a canonical media type") + } + prefix, err := readPrefix(reader, size, maxSniffBytes) + if err != nil { + return CandidateFormat{}, err + } + + var detected CandidateFormat + switch { + case bytes.HasPrefix(prefix, []byte("%PDF-")): + if err = validatePDFStructure(reader, size, prefix); err == nil { + detected, _ = CandidateFormatByID(formatIDPDF) + } + case bytes.HasPrefix(prefix, []byte(`{\rtf`)): + detected, _ = CandidateFormatByID("rtf") + case bytes.HasPrefix(prefix, compoundFileMagic): + detected, err = detectCompoundFormat(reader, size) + case bytes.HasPrefix(prefix, []byte("PK\x03\x04")) || bytes.HasPrefix(prefix, []byte("PK\x05\x06")): + detected, err = detectZIPFormat(reader, size) + default: + if size > maxTextSniffBytes { + return CandidateFormat{}, errors.New("document text exceeds type-detection limit") + } + content, readErr := readPrefix(reader, size, maxTextSniffBytes) + if readErr != nil { + return CandidateFormat{}, readErr + } + detected, err = detectTextFormat(content, mediaType) + } + if err != nil { + return CandidateFormat{}, err + } + if detected.MediaType != mediaType { + return CandidateFormat{}, fmt.Errorf("document bytes are %s, not declared %s", detected.MediaType, mediaType) + } + return detected, nil +} + +func validatePDFStructure(reader io.ReaderAt, size int64, prefix []byte) error { + if len(prefix) < 9 || !validPDFVersion(prefix[5:8]) || !isPDFWhitespace(prefix[8]) { + return errors.New("PDF header is invalid") + } + tailLength := min(size, maxPDFTailBytes) + tailOffset := size - tailLength + tail := make([]byte, tailLength) + read, err := reader.ReadAt(tail, tailOffset) + if err != nil && !errors.Is(err, io.EOF) { + return fmt.Errorf("read PDF trailer: %w", err) + } + if int64(read) != tailLength { + return errors.New("document bytes changed during PDF trailer read") + } + eofIndex := bytes.LastIndex(tail, []byte("%%EOF")) + if eofIndex < 0 || len(trimPDFWhitespace(tail[eofIndex+len("%%EOF"):])) != 0 { + return errors.New("PDF end marker is missing or not final") + } + beforeEOF := tail[:eofIndex] + startXRefIndex := bytes.LastIndex(beforeEOF, []byte("startxref")) + if startXRefIndex < 0 { + return errors.New("PDF startxref is missing") + } + offsetText := trimPDFWhitespace(beforeEOF[startXRefIndex+len("startxref"):]) + digitEnd := 0 + for digitEnd < len(offsetText) && offsetText[digitEnd] >= '0' && offsetText[digitEnd] <= '9' { + digitEnd++ + } + if digitEnd == 0 || len(trimPDFWhitespace(offsetText[digitEnd:])) != 0 { + return errors.New("PDF startxref offset is invalid") + } + xrefOffset, err := strconv.ParseInt(string(offsetText[:digitEnd]), 10, 64) + if err != nil || xrefOffset <= 0 || xrefOffset >= tailOffset+int64(eofIndex) { + return errors.New("PDF startxref offset is outside the document") + } + xrefLength := min(size-xrefOffset, maxPDFXRefBytes) + xref := make([]byte, xrefLength) + read, err = reader.ReadAt(xref, xrefOffset) + if err != nil && !errors.Is(err, io.EOF) { + return fmt.Errorf("read PDF cross-reference data: %w", err) + } + if int64(read) != xrefLength { + return errors.New("document bytes changed during PDF cross-reference read") + } + xrefEnd := tailOffset + int64(startXRefIndex) + streamXRefLength := xrefEnd - xrefOffset + if validPDFTableXRef(xref, beforeEOF[:startXRefIndex]) || + streamXRefLength > 0 && streamXRefLength <= int64(len(xref)) && validPDFStreamXRef( + reader, size, xrefOffset, xref[:streamXRefLength]) { + return nil + } + return errors.New("PDF cross-reference data is invalid") +} + +// CountPDFPages resolves the catalog's page tree and returns its verified leaf +// count. Stream bodies are removed using their direct Length before tokenizing, +// so page-like bytes in content streams, strings, comments, or orphan objects +// cannot inflate the result. PDFs whose stream boundaries cannot be proven +// locally are rejected rather than estimated. +func CountPDFPages(data []byte) (int64, error) { + structural, err := pdfWithoutStreamData(data) + if err != nil { + return 0, err + } + tokens, ok := tokenizePDF(structural) + if !ok { + return 0, errors.New("PDF object syntax is invalid") + } + + objects := make(map[string]map[string][]string) + for position := 0; position+3 < len(tokens); position++ { + if tokens[position+2] != "obj" || !pdfObjectNumber(tokens[position]) || !pdfGeneration(tokens[position+1]) { + continue + } + dictionary, _, parsed := parsePDFDictionaryTokens(tokens, position+3, 0) + if parsed { + objects[pdfReferenceKey(tokens[position], tokens[position+1])] = dictionary + } + } + + authority, ok := pdfCrossReferenceAuthority(data) + if !ok { + return 0, errors.New("PDF cross-reference authority is missing") + } + if authority.objects != nil { + objects = pdfXRefSelectedDictionaries(data, authority.objects) + } + root, ok := pdfReference(authority.dictionary["Root"]) + if !ok { + return 0, errors.New("PDF cross-reference root is invalid") + } + catalog, ok := objects[root] + if !ok || !pdfDictionaryType(catalog, "/Catalog") { + return 0, errors.New("PDF catalog is invalid") + } + pagesRoot, ok := pdfReference(catalog["Pages"]) + if !ok { + return 0, errors.New("PDF catalog page tree is invalid") + } + + visiting := make(map[string]bool) + var walk func(string, string, bool, int) (int64, error) + walk = func(reference, parent string, root bool, depth int) (int64, error) { + if depth > maxPDFPageTreeDepth { + return 0, errors.New("PDF page tree exceeds the nesting bound") + } + if visiting[reference] { + return 0, errors.New("PDF page tree contains a cycle") + } + item, exists := objects[reference] + if !exists { + return 0, errors.New("PDF page tree references a missing object") + } + if !root { + declaredParent, parentOK := pdfReference(item["Parent"]) + if !parentOK || declaredParent != parent { + return 0, errors.New("PDF page tree parent is invalid") + } + } + switch { + case pdfDictionaryType(item, "/Page"): + return 1, nil + case pdfDictionaryType(item, "/Pages"): + declared, countOK := pdfPositiveInteger(item["Count"]) + children, kidsOK := pdfReferences(item["Kids"]) + if !countOK || declared > math.MaxInt64 || !kidsOK || len(children) == 0 { + return 0, errors.New("PDF page tree node is invalid") + } + visiting[reference] = true + var total int64 + for _, child := range children { + count, childErr := walk(child, reference, false, depth+1) + if childErr != nil || count > math.MaxInt64-total { + delete(visiting, reference) + if childErr != nil { + return 0, childErr + } + return 0, errors.New("PDF page count overflows") + } + total += count + } + delete(visiting, reference) + if uint64(total) != declared { + return 0, errors.New("PDF page tree count is inconsistent") + } + return total, nil + default: + return 0, errors.New("PDF page tree object has an invalid type") + } + } + return walk(pagesRoot, "", true, 0) +} + +// PDFInfoFields resolves the final cross-reference authority's exact Info +// dictionary and returns its directly embedded string values. Unrelated +// objects and stream bytes are not metadata authority. +func PDFInfoFields(data []byte) (map[string]string, error) { + structural, err := pdfWithoutStreamData(data) + if err != nil { + return nil, err + } + tokens, ok := tokenizePDF(structural) + if !ok { + return nil, errors.New("PDF object syntax is invalid") + } + objects := make(map[string]map[string][]string) + for position := 0; position+3 < len(tokens); position++ { + if tokens[position+2] != "obj" || !pdfObjectNumber(tokens[position]) || !pdfGeneration(tokens[position+1]) { + continue + } + dictionary, _, parsed := parsePDFDictionaryTokens(tokens, position+3, 0) + if parsed { + objects[pdfReferenceKey(tokens[position], tokens[position+1])] = dictionary + } + } + authority, ok := pdfCrossReferenceAuthority(data) + if !ok { + return nil, errors.New("PDF cross-reference authority is missing") + } + if authority.objects != nil { + objects = pdfXRefSelectedDictionaries(data, authority.objects) + } + infoReference, hasInfo := pdfReference(authority.dictionary["Info"]) + if !hasInfo { + return map[string]string{}, nil + } + info, ok := objects[infoReference] + if !ok { + return nil, errors.New("PDF Info dictionary reference is missing") + } + result := make(map[string]string) + for _, key := range []string{"Title", "Author", "Subject", "Keywords", "CreationDate", "ModDate"} { + value, ok := pdfString(info[key]) + if ok { + result[key] = value + } + } + return result, nil +} + +func pdfString(tokens []string) (string, bool) { + if len(tokens) != 1 || len(tokens[0]) < 2 { + return "", false + } + value := tokens[0] + if value[0] == '<' && value[1] != '<' && value[len(value)-1] == '>' { + hexText := strings.Map(func(character rune) rune { + if character == 0 || character == '\t' || character == '\n' || + character == '\f' || character == '\r' || character == ' ' { + return -1 + } + return character + }, value[1:len(value)-1]) + if len(hexText)%2 != 0 { + hexText += "0" + } + decoded, err := hex.DecodeString(hexText) + return string(decoded), err == nil + } + if value[0] != '(' || value[len(value)-1] != ')' { + return "", false + } + input := []byte(value[1 : len(value)-1]) + output := make([]byte, 0, len(input)) + for position := 0; position < len(input); position++ { + if input[position] != '\\' { + output = append(output, input[position]) + continue + } + position++ + if position == len(input) { + return "", false + } + switch input[position] { + case 'n': + output = append(output, '\n') + case 'r': + output = append(output, '\r') + case 't': + output = append(output, '\t') + case 'b': + output = append(output, '\b') + case 'f': + output = append(output, '\f') + case '\r': + if position+1 < len(input) && input[position+1] == '\n' { + position++ + } + case '\n': + case '0', '1', '2', '3', '4', '5', '6', '7': + value := int(input[position] - '0') + for count := 1; count < 3 && position+1 < len(input) && input[position+1] >= '0' && input[position+1] <= '7'; count++ { + position++ + value = value*8 + int(input[position]-'0') + } + output = append(output, byte(value&0xff)) + default: + output = append(output, input[position]) + } + } + return string(output), true +} + +func pdfWithoutStreamData(data []byte) ([]byte, error) { + result := make([]byte, 0, min(len(data), maxPDFStructuralBytes)) + segmentStart, objectStart := 0, -1 + for position := 0; position < len(data); { + start, end, ok := nextPDFLexeme(data, position) + if !ok { + return nil, errors.New("PDF token is unterminated") + } + if start == end { + break + } + position = end + word := string(data[start:end]) + if word == "obj" { + objectStart = start + continue + } + if word == "endobj" { + objectStart = -1 + continue + } + if word != "stream" { + continue + } + if objectStart < 0 { + return nil, errors.New("PDF stream is outside an object") + } + prefixTokens, tokenOK := tokenizePDF(data[objectStart:start]) + if !tokenOK { + return nil, errors.New("PDF stream dictionary is invalid") + } + dictionaryStart := -1 + for index, token := range prefixTokens { + if token == "<<" { + dictionaryStart = index + break + } + } + if dictionaryStart < 0 { + return nil, errors.New("PDF stream dictionary is missing") + } + dictionary, _, parsed := parsePDFDictionaryTokens(prefixTokens, dictionaryStart, 0) + length, lengthOK := pdfNonnegativeInteger(dictionary["Length"]) + if !parsed || !lengthOK || length > uint64(len(data)) { + return nil, errors.New("PDF stream length is not directly bounded") + } + bodyStart := end + for bodyStart < len(data) && (data[bodyStart] == ' ' || data[bodyStart] == '\t' || data[bodyStart] == '\f' || data[bodyStart] == 0) { + bodyStart++ + } + switch { + case bodyStart < len(data) && data[bodyStart] == '\r': + bodyStart++ + if bodyStart < len(data) && data[bodyStart] == '\n' { + bodyStart++ + } + case bodyStart < len(data) && data[bodyStart] == '\n': + bodyStart++ + default: + return nil, errors.New("PDF stream has no line boundary") + } + bodyLength := int(length) // #nosec G115 -- length is bounded by len(data) above. + if bodyLength > len(data)-bodyStart { + return nil, errors.New("PDF stream exceeds the document") + } + bodyEnd := bodyStart + bodyLength + endStream := bodyEnd + if endStream < len(data) && data[endStream] == '\r' { + endStream++ + if endStream < len(data) && data[endStream] == '\n' { + endStream++ + } + } else if endStream < len(data) && data[endStream] == '\n' { + endStream++ + } + if endStream+len("endstream") > len(data) || string(data[endStream:endStream+len("endstream")]) != "endstream" { + return nil, errors.New("PDF stream end does not match its length") + } + if bodyStart-segmentStart > maxPDFStructuralBytes-len(result) { + return nil, errors.New("PDF object structure exceeds the inspection bound") + } + result = append(result, data[segmentStart:bodyStart]...) + segmentStart = bodyEnd + position = bodyEnd + } + if len(data)-segmentStart > maxPDFStructuralBytes-len(result) { + return nil, errors.New("PDF object structure exceeds the inspection bound") + } + return append(result, data[segmentStart:]...), nil +} + +func nextPDFLexeme(data []byte, position int) (int, int, bool) { + for { + for position < len(data) && isPDFWhitespace(data[position]) { + position++ + } + if position == len(data) { + return position, position, true + } + if data[position] != '%' { + break + } + for position < len(data) && data[position] != '\r' && data[position] != '\n' { + position++ + } + } + start := position + switch data[position] { + case '(': + depth, escaped := 0, false + for ; position < len(data); position++ { + char := data[position] + if escaped { + escaped = false + continue + } + if char == '\\' { + escaped = true + continue + } + switch char { + case '(': + depth++ + case ')': + depth-- + if depth == 0 { + return start, position + 1, true + } + } + } + return start, position, false + case '<': + if position+1 < len(data) && data[position+1] == '<' { + return start, position + 2, true + } + for position++; position < len(data); position++ { + if data[position] == '>' { + return start, position + 1, true + } + } + return start, position, false + case '>': + if position+1 < len(data) && data[position+1] == '>' { + return start, position + 2, true + } + return start, position + 1, true + case '[', ']', '{', '}': + return start, position + 1, true + case '/': + position++ + } + for position < len(data) && !isPDFTokenBoundary(data[position]) { + position++ + } + if position == start || position == start+1 && data[start] == '/' { + return start, position, false + } + return start, position, true +} + +func pdfObjectNumber(value string) bool { + number, err := strconv.ParseUint(value, 10, 64) + return err == nil && number > 0 +} + +func pdfGeneration(value string) bool { + generation, err := strconv.ParseUint(value, 10, 16) + return err == nil && generation <= 65_535 +} + +func pdfReferenceKey(number, generation string) string { return number + " " + generation } + +func pdfReference(tokens []string) (string, bool) { + if len(tokens) != 3 || tokens[2] != "R" || !pdfObjectNumber(tokens[0]) || !pdfGeneration(tokens[1]) { + return "", false + } + return pdfReferenceKey(tokens[0], tokens[1]), true +} + +func pdfReferences(tokens []string) ([]string, bool) { + if len(tokens) < 5 || tokens[0] != "[" || tokens[len(tokens)-1] != "]" || (len(tokens)-2)%3 != 0 { + return nil, false + } + references := make([]string, 0, (len(tokens)-2)/3) + for position := 1; position < len(tokens)-1; position += 3 { + reference, ok := pdfReference(tokens[position : position+3]) + if !ok { + return nil, false + } + references = append(references, reference) + } + return references, len(references) > 0 +} + +func pdfDictionaryType(dictionary map[string][]string, expected string) bool { + return len(dictionary["Type"]) == 1 && dictionary["Type"][0] == expected +} + +func pdfNonnegativeInteger(tokens []string) (uint64, bool) { + if len(tokens) != 1 { + return 0, false + } + value, err := strconv.ParseUint(tokens[0], 10, 64) + return value, err == nil +} + +func validPDFVersion(version []byte) bool { + return len(version) == 3 && version[1] == '.' && + ((version[0] == '1' && version[2] >= '0' && version[2] <= '7') || + (version[0] == '2' && version[2] == '0')) +} + +func validPDFTableXRef(xref, beforeStartXRef []byte) bool { + position := 0 + line, ok := nextPDFLine(xref, &position) + if !ok || !bytes.Equal(trimPDFWhitespace(line), []byte("xref")) { + return false + } + header, ok := nextNonemptyPDFLine(xref, &position) + if !ok { + return false + } + headerFields := bytes.Fields(header) + if len(headerFields) != 2 { + return false + } + first, firstErr := strconv.ParseUint(string(headerFields[0]), 10, 64) + count, countErr := strconv.ParseUint(string(headerFields[1]), 10, 64) + if firstErr != nil || countErr != nil || count == 0 || first > math.MaxUint64-count { + return false + } + validatedRecords := uint64(0) + for validatedRecords < count { + record, ok := nextPDFLine(xref, &position) + if !ok { + break + } + if !validPDFXRefRecord(record) { + if position == len(xref) { + break + } + return false + } + validatedRecords++ + } + if validatedRecords == 0 { + return false + } + + return validPDFTrailer(beforeStartXRef, first, count) +} + +type pdfXRefEntry struct { + kind, field2, field3 uint64 +} + +type pdfXRefStream struct { + dictionary map[string][]string + objectNumber, generation uint64 + entries map[uint64]pdfXRefEntry +} + +type pdfCrossReference struct { + dictionary map[string][]string + objects map[string]int64 +} + +func validPDFStreamXRef(reader io.ReaderAt, size, xrefOffset int64, xref []byte) bool { + _, ok := validatedPDFStreamXRef(reader, size, xrefOffset, xref) + return ok +} + +func validatedPDFStreamXRef( + reader io.ReaderAt, size, xrefOffset int64, xref []byte, +) (pdfXRefStream, bool) { + if reader == nil || size <= 0 || xrefOffset <= 0 || xrefOffset >= size { + return pdfXRefStream{}, false + } + parsed, ok := parsePDFStreamXRef(xref) + if !ok { + return pdfXRefStream{}, false + } + xrefLimit := uint64(xrefOffset) // #nosec G115 -- xrefOffset is positive above. + sizeLimit := uint64(size) // #nosec G115 -- size is positive above. + seenOffsets := make(map[uint64]struct{}, len(parsed.entries)) + for objectNumber, entry := range parsed.entries { + switch entry.kind { + case 0: + if entry.field3 > 65_535 { + return pdfXRefStream{}, false + } + case 1: + if objectNumber == 0 || entry.field2 == 0 || entry.field2 > xrefLimit || + entry.field2 >= sizeLimit || entry.field3 > 65_535 { + return pdfXRefStream{}, false + } + if _, duplicate := seenOffsets[entry.field2]; duplicate { + return pdfXRefStream{}, false + } + seenOffsets[entry.field2] = struct{}{} + objectOffset := int64(entry.field2) // #nosec G115 -- field2 is below positive int64 size above. + if !validPDFObjectHeaderAt(reader, size, objectOffset, objectNumber, entry.field3) { + return pdfXRefStream{}, false + } + default: + // Compressed objects and unknown entry kinds are not local authority. + return pdfXRefStream{}, false + } + } + if !pdfXRefEntryMatches(parsed, parsed.objectNumber, parsed.generation, xrefLimit) || + !pdfXRefReferenceIsDirect(parsed, parsed.dictionary["Root"]) { + return pdfXRefStream{}, false + } + if _, hasInfo := parsed.dictionary["Info"]; hasInfo && + !pdfXRefReferenceIsDirect(parsed, parsed.dictionary["Info"]) { + return pdfXRefStream{}, false + } + return parsed, true +} + +func parsePDFStreamXRef(xref []byte) (pdfXRefStream, bool) { + if int64(len(xref)) > maxPDFXRefBytes { + return pdfXRefStream{}, false + } + streamIndex := firstPDFKeyword(xref, "stream") + if streamIndex < 0 { + return pdfXRefStream{}, false + } + tokens, ok := tokenizePDF(xref[:streamIndex]) + if !ok || len(tokens) < 5 || tokens[2] != "obj" { + return pdfXRefStream{}, false + } + objectNumber, err := strconv.ParseUint(tokens[0], 10, 64) + if err != nil || objectNumber == 0 { + return pdfXRefStream{}, false + } + generation, err := strconv.ParseUint(tokens[1], 10, 64) + if err != nil || generation > 65_535 { + return pdfXRefStream{}, false + } + dictionary, next, ok := parsePDFDictionaryTokens(tokens, 3, 0) + if !ok || next != len(tokens) { + return pdfXRefStream{}, false + } + size, sizeOK := pdfPositiveInteger(dictionary["Size"]) + widths, entryWidth, widthsOK := pdfXRefWidths(dictionary["W"]) + length, lengthOK := pdfNonnegativeInteger(dictionary["Length"]) + if !sizeOK || !widthsOK || !lengthOK || + len(dictionary["Type"]) != 1 || dictionary["Type"][0] != "/XRef" || + !validPDFRootReference(dictionary["Root"]) || pdfXRefHasUnsupportedEncoding(dictionary) { + return pdfXRefStream{}, false + } + sections, entryCount, sectionsOK := pdfXRefSections(dictionary["Index"], size) + if !sectionsOK || entryCount > math.MaxUint64/entryWidth || entryCount*entryWidth != length { + return pdfXRefStream{}, false + } + bodyStart, bodyOK := pdfStreamBodyStart(xref, streamIndex+len("stream")) + if !bodyOK { + return pdfXRefStream{}, false + } + available := uint64(len(xref) - bodyStart) // #nosec G115 -- bodyStart is within xref when bodyOK. + if length > available { + return pdfXRefStream{}, false + } + bodyLength := int(length) // #nosec G115 -- length is bounded by the xref slice above. + bodyEnd := bodyStart + bodyLength + endStream := bodyEnd + if endStream < len(xref) && xref[endStream] == '\r' { + endStream++ + if endStream < len(xref) && xref[endStream] == '\n' { + endStream++ + } + } else if endStream < len(xref) && xref[endStream] == '\n' { + endStream++ + } + closing, closingOK := tokenizePDF(xref[endStream:]) + if !closingOK || len(closing) != 2 || closing[0] != "endstream" || closing[1] != "endobj" { + return pdfXRefStream{}, false + } + entryCapacity := int(entryCount) // #nosec G115 -- count is bounded by the 4 KiB xref slice. + entries := make(map[uint64]pdfXRefEntry, entryCapacity) + position := bodyStart + for _, section := range sections { + for objectNumber := section.first; objectNumber < section.first+section.count; objectNumber++ { + if _, duplicate := entries[objectNumber]; duplicate { + return pdfXRefStream{}, false + } + kind := uint64(1) + if widths[0] != 0 { + kind = pdfXRefField(xref[position : position+widths[0]]) + } + position += widths[0] + field2 := pdfXRefField(xref[position : position+widths[1]]) + position += widths[1] + field3 := pdfXRefField(xref[position : position+widths[2]]) + position += widths[2] + entries[objectNumber] = pdfXRefEntry{kind: kind, field2: field2, field3: field3} + } + } + return pdfXRefStream{ + dictionary: dictionary, objectNumber: objectNumber, generation: generation, entries: entries, + }, true +} + +type pdfXRefSection struct{ first, count uint64 } + +func pdfXRefSections(value []string, size uint64) ([]pdfXRefSection, uint64, bool) { + if len(value) == 0 { + return []pdfXRefSection{{first: 0, count: size}}, size, true + } + if len(value) < 4 || value[0] != "[" || value[len(value)-1] != "]" || (len(value)-2)%2 != 0 { + return nil, 0, false + } + sections := make([]pdfXRefSection, 0, (len(value)-2)/2) + var total uint64 + for position := 1; position < len(value)-1; position += 2 { + first, firstOK := pdfNonnegativeInteger(value[position : position+1]) + count, countOK := pdfPositiveInteger(value[position+1 : position+2]) + if !firstOK || !countOK || first >= size || count > size-first || total > math.MaxUint64-count { + return nil, 0, false + } + for _, previous := range sections { + if first < previous.first+previous.count && previous.first < first+count { + return nil, 0, false + } + } + sections = append(sections, pdfXRefSection{first: first, count: count}) + total += count + } + return sections, total, true +} + +func pdfXRefWidths(value []string) ([3]int, uint64, bool) { + var widths [3]int + if len(value) != 5 || value[0] != "[" || value[4] != "]" { + return widths, 0, false + } + var total uint64 + for index, value := range value[1:4] { + parsed, err := strconv.ParseUint(value, 10, 8) + if err != nil || parsed > 8 { + return widths, 0, false + } + widths[index] = int(parsed) + total += parsed + } + return widths, total, total > 0 +} + +func pdfXRefHasUnsupportedEncoding(dictionary map[string][]string) bool { + for _, key := range []string{"Filter", "DecodeParms", "F", "FFilter", "FDecodeParms"} { + if _, exists := dictionary[key]; exists { + return true + } + } + return false +} + +func pdfStreamBodyStart(data []byte, position int) (int, bool) { + for position < len(data) && (data[position] == ' ' || data[position] == '\t' || + data[position] == '\f' || data[position] == 0) { + position++ + } + switch { + case position < len(data) && data[position] == '\r': + position++ + if position < len(data) && data[position] == '\n' { + position++ + } + return position, true + case position < len(data) && data[position] == '\n': + return position + 1, true + default: + return 0, false + } +} + +func pdfXRefField(data []byte) uint64 { + var result uint64 + for _, value := range data { + result = result<<8 | uint64(value) + } + return result +} + +func pdfXRefEntryMatches(parsed pdfXRefStream, number, generation, offset uint64) bool { + entry, ok := parsed.entries[number] + return ok && entry.kind == 1 && entry.field2 == offset && entry.field3 == generation +} + +func pdfXRefReferenceIsDirect(parsed pdfXRefStream, reference []string) bool { + if len(reference) != 3 || reference[2] != "R" { + return false + } + number, numberErr := strconv.ParseUint(reference[0], 10, 64) + generation, generationErr := strconv.ParseUint(reference[1], 10, 64) + return numberErr == nil && generationErr == nil && + pdfXRefEntryMatches(parsed, number, generation, parsed.entries[number].field2) +} + +func validPDFObjectHeaderAt( + reader io.ReaderAt, size, offset int64, objectNumber, generation uint64, +) bool { + if offset < 0 || offset >= size { + return false + } + length := min(size-offset, int64(64)) + prefix := make([]byte, length) + read, err := reader.ReadAt(prefix, offset) + if err != nil && !errors.Is(err, io.EOF) || int64(read) != length { + return false + } + position := 0 + values := make([]string, 0, 3) + for range 3 { + start, end, ok := nextPDFLexeme(prefix, position) + if !ok || start == end || len(values) == 0 && start != 0 { + return false + } + values = append(values, string(prefix[start:end])) + position = end + } + number, numberErr := strconv.ParseUint(values[0], 10, 64) + actualGeneration, generationErr := strconv.ParseUint(values[1], 10, 64) + return numberErr == nil && generationErr == nil && values[2] == "obj" && + number == objectNumber && actualGeneration == generation +} + +func pdfXRefSelectedDictionaries( + data []byte, selected map[string]int64, +) map[string]map[string][]string { + result := make(map[string]map[string][]string, len(selected)) + for reference, offset := range selected { + dictionary, ok := pdfDictionaryAtOffset(data, offset, reference) + if ok { + result[reference] = dictionary + } + } + return result +} + +func pdfDictionaryAtOffset(data []byte, offset int64, reference string) (map[string][]string, bool) { + if offset < 0 || offset >= int64(len(data)) { + return nil, false + } + position := int(offset) // #nosec G115 -- offset is bounded by the in-memory slice above. + header := make([]string, 0, 3) + for range 3 { + start, end, ok := nextPDFLexeme(data, position) + if !ok || start == end || len(header) == 0 && start != position { + return nil, false + } + header = append(header, string(data[start:end])) + position = end + } + if header[2] != "obj" || pdfReferenceKey(header[0], header[1]) != reference { + return nil, false + } + start, end, ok := nextPDFLexeme(data, position) + if !ok || string(data[start:end]) != "<<" { + return nil, false + } + tokens := []string{"<<"} + position, depth := end, 1 + for depth > 0 { + if len(tokens) >= maxPDFTokens || position-int(offset) > maxPDFStructuralBytes { + return nil, false + } + start, end, ok = nextPDFLexeme(data, position) + if !ok || start == end { + return nil, false + } + token := string(data[start:end]) + tokens = append(tokens, token) + position = end + switch token { + case "<<": + depth++ + case ">>": + depth-- + } + } + dictionary, next, parsed := parsePDFDictionaryTokens(tokens, 0, 0) + if !parsed || next != len(tokens) { + return nil, false + } + start, end, ok = nextPDFLexeme(data, position) + if !ok || start == end { + return nil, false + } + closing := string(data[start:end]) + return dictionary, closing == "endobj" || closing == "stream" +} + +func pdfCrossReferenceAuthority(data []byte) (pdfCrossReference, bool) { + eofIndex := bytes.LastIndex(data, []byte("%%EOF")) + if eofIndex < 0 || len(trimPDFWhitespace(data[eofIndex+len("%%EOF"):])) != 0 { + return pdfCrossReference{}, false + } + beforeEOF := data[:eofIndex] + startXRefIndex := bytes.LastIndex(beforeEOF, []byte("startxref")) + if startXRefIndex < 0 { + return pdfCrossReference{}, false + } + offsetText := trimPDFWhitespace(beforeEOF[startXRefIndex+len("startxref"):]) + digitEnd := 0 + for digitEnd < len(offsetText) && offsetText[digitEnd] >= '0' && offsetText[digitEnd] <= '9' { + digitEnd++ + } + if digitEnd == 0 || len(trimPDFWhitespace(offsetText[digitEnd:])) != 0 { + return pdfCrossReference{}, false + } + offset, err := strconv.ParseUint(string(offsetText[:digitEnd]), 10, 64) + if err != nil || offset == 0 || offset >= uint64(startXRefIndex) { + return pdfCrossReference{}, false + } + xrefOffset := int(offset) // #nosec G115 -- offset is bounded by the in-memory slice index above. + xref := data[xrefOffset:startXRefIndex] + if stream, ok := validatedPDFStreamXRef( + bytes.NewReader(data), int64(len(data)), int64(xrefOffset), xref, + ); ok { + objects := make(map[string]int64) + for number, entry := range stream.entries { + if entry.kind == 1 { + objectOffset := int64(entry.field2) // #nosec G115 -- validated below the positive document size. + objects[pdfReferenceKey(strconv.FormatUint(number, 10), strconv.FormatUint(entry.field3, 10))] = + objectOffset + } + } + return pdfCrossReference{dictionary: stream.dictionary, objects: objects}, true + } + if !validPDFTableXRef(xref, data[:startXRefIndex]) { + return pdfCrossReference{}, false + } + trailerIndex := lastPDFKeyword(data[:startXRefIndex], "trailer") + if trailerIndex < 0 { + return pdfCrossReference{}, false + } + dictionary, ok := parsePDFDictionary(data[trailerIndex+len("trailer") : startXRefIndex]) + if !ok { + return pdfCrossReference{}, false + } + return pdfCrossReference{dictionary: dictionary}, true +} + +func validPDFTrailer(data []byte, first, count uint64) bool { + for end := len(data); end > 0; { + trailerIndex := lastPDFKeyword(data[:end], "trailer") + if trailerIndex < 0 { + return false + } + trailer, ok := parsePDFDictionary(data[trailerIndex+len("trailer"):]) + if ok { + if !validPDFRootReference(trailer["Root"]) { + return false + } + size, sizeOK := pdfPositiveInteger(trailer["Size"]) + return sizeOK && first+count <= size + } + end = trailerIndex + } + return false +} + +func nextPDFLine(data []byte, position *int) ([]byte, bool) { + if *position >= len(data) { + return nil, false + } + start := *position + for *position < len(data) && data[*position] != '\n' && data[*position] != '\r' { + *position++ + } + line := data[start:*position] + if *position < len(data) && data[*position] == '\r' { + *position++ + } + if *position < len(data) && data[*position] == '\n' { + *position++ + } + return line, true +} + +func nextNonemptyPDFLine(data []byte, position *int) ([]byte, bool) { + for { + line, ok := nextPDFLine(data, position) + if !ok { + return nil, false + } + if line = trimPDFWhitespace(line); len(line) != 0 { + return line, true + } + } +} + +func validPDFXRefRecord(line []byte) bool { + fields := bytes.Fields(line) + if len(fields) != 3 || len(fields[0]) != 10 || len(fields[1]) != 5 || len(fields[2]) != 1 || + (fields[2][0] != 'n' && fields[2][0] != 'f') { + return false + } + return decimalBytes(fields[0]) && decimalBytes(fields[1]) +} + +func decimalBytes(value []byte) bool { + if len(value) == 0 { + return false + } + for _, char := range value { + if char < '0' || char > '9' { + return false + } + } + return true +} + +func lastPDFKeyword(data []byte, keyword string) int { + for end := len(data); end > 0; { + index := bytes.LastIndex(data[:end], []byte(keyword)) + if index < 0 { + return -1 + } + beforeOK := index == 0 || isPDFTokenBoundary(data[index-1]) + after := index + len(keyword) + afterOK := after == len(data) || isPDFTokenBoundary(data[after]) + if beforeOK && afterOK { + return index + } + end = index + } + return -1 +} + +func firstPDFKeyword(data []byte, keyword string) int { + for start := 0; start < len(data); { + relative := bytes.Index(data[start:], []byte(keyword)) + if relative < 0 { + return -1 + } + index := start + relative + beforeOK := index == 0 || isPDFTokenBoundary(data[index-1]) + after := index + len(keyword) + afterOK := after == len(data) || isPDFTokenBoundary(data[after]) + if beforeOK && afterOK { + return index + } + start = index + 1 + } + return -1 +} + +func isPDFTokenBoundary(char byte) bool { + return isPDFWhitespace(char) || strings.ContainsRune("()<>[]{}/%", rune(char)) +} + +func tokenizePDF(data []byte) ([]string, bool) { + tokens := make([]string, 0, 32) + for position := 0; position < len(data); { + if len(tokens) >= maxPDFTokens { + return nil, false + } + char := data[position] + if isPDFWhitespace(char) { + position++ + continue + } + if char == '%' { + for position < len(data) && data[position] != '\r' && data[position] != '\n' { + position++ + } + continue + } + if position+1 < len(data) && (string(data[position:position+2]) == "<<" || + string(data[position:position+2]) == ">>") { + tokens = append(tokens, string(data[position:position+2])) + position += 2 + continue + } + if char == '(' { + start, depth, escaped := position, 0, false + for ; position < len(data); position++ { + current := data[position] + if escaped { + escaped = false + continue + } + if current == '\\' { + escaped = true + continue + } + if current == '(' { + depth++ + } else if current == ')' { + depth-- + if depth == 0 { + position++ + break + } + } + } + if depth != 0 { + return nil, false + } + tokens = append(tokens, string(data[start:position])) + continue + } + if char == '<' { + start := position + position++ + for position < len(data) && data[position] != '>' { + position++ + } + if position == len(data) { + return nil, false + } + position++ + tokens = append(tokens, string(data[start:position])) + continue + } + if strings.ContainsRune("[]{}", rune(char)) { + tokens = append(tokens, string(char)) + position++ + continue + } + start := position + if char == '/' { + position++ + } + for position < len(data) && !isPDFTokenBoundary(data[position]) { + position++ + } + if position == start || (position == start+1 && char == '/') { + return nil, false + } + tokens = append(tokens, string(data[start:position])) + } + return tokens, true +} + +func parsePDFDictionary(data []byte) (map[string][]string, bool) { + tokens, ok := tokenizePDF(data) + if !ok { + return nil, false + } + dictionary, next, ok := parsePDFDictionaryTokens(tokens, 0, 0) + return dictionary, ok && next == len(tokens) +} + +func parsePDFDictionaryTokens( + tokens []string, + position int, + depth int, +) (map[string][]string, int, bool) { + if depth > 32 || position >= len(tokens) || tokens[position] != "<<" { + return nil, position, false + } + position++ + dictionary := make(map[string][]string) + for position < len(tokens) && tokens[position] != ">>" { + key := tokens[position] + if len(key) < 2 || key[0] != '/' { + return nil, position, false + } + key = key[1:] + if _, exists := dictionary[key]; exists { + return nil, position, false + } + position++ + valueStart := position + var ok bool + position, ok = skipPDFObject(tokens, position, depth+1) + if !ok { + return nil, position, false + } + dictionary[key] = tokens[valueStart:position] + } + if position >= len(tokens) || tokens[position] != ">>" { + return nil, position, false + } + return dictionary, position + 1, true +} + +func skipPDFObject(tokens []string, position, depth int) (int, bool) { + if depth > 32 || position >= len(tokens) { + return position, false + } + switch tokens[position] { + case "<<": + _, next, ok := parsePDFDictionaryTokens(tokens, position, depth) + return next, ok + case "[": + position++ + for position < len(tokens) && tokens[position] != "]" { + var ok bool + position, ok = skipPDFObject(tokens, position, depth+1) + if !ok { + return position, false + } + } + return position + 1, position < len(tokens) + case ">>", "]": + return position, false + default: + if position+2 < len(tokens) && decimalString(tokens[position]) && + decimalString(tokens[position+1]) && tokens[position+2] == "R" { + return position + 3, true + } + return position + 1, true + } +} + +func pdfPositiveInteger(value []string) (uint64, bool) { + if len(value) != 1 || !decimalString(value[0]) { + return 0, false + } + parsed, err := strconv.ParseUint(value[0], 10, 64) + return parsed, err == nil && parsed > 0 +} + +func validPDFRootReference(value []string) bool { + if len(value) != 3 || value[2] != "R" { + return false + } + objectNumber, objectErr := strconv.ParseUint(value[0], 10, 64) + generation, generationErr := strconv.ParseUint(value[1], 10, 64) + return objectErr == nil && objectNumber > 0 && generationErr == nil && generation <= 65_535 +} + +func decimalString(value string) bool { + return decimalBytes([]byte(value)) +} + +func trimPDFWhitespace(value []byte) []byte { + for len(value) > 0 && isPDFWhitespace(value[0]) { + value = value[1:] + } + for len(value) > 0 && isPDFWhitespace(value[len(value)-1]) { + value = value[:len(value)-1] + } + return value +} + +func isPDFWhitespace(char byte) bool { + return char == 0 || char == '\t' || char == '\n' || char == '\f' || char == '\r' || char == ' ' +} + +func detectCompoundFormat(reader io.ReaderAt, size int64) (CandidateFormat, error) { + names, err := compoundDirectoryNames(reader, size) + if err != nil { + return CandidateFormat{}, err + } + ids := make([]string, 0, 2) + if names["WordDocument"] { + ids = append(ids, "doc") + } + if names["Workbook"] || names["Book"] { + ids = append(ids, "xls") + } + if names["PowerPoint Document"] { + ids = append(ids, "ppt") + } + if names["__properties_version1.0"] { + ids = append(ids, "msg") + } + if len(ids) != 1 { + return CandidateFormat{}, errors.New("compound document has missing or ambiguous family markers") + } + format, _ := CandidateFormatByID(ids[0]) + return format, nil +} + +func compoundDirectoryNames(reader io.ReaderAt, size int64) (map[string]bool, error) { + const ( + freeSector = uint32(0xffffffff) + endOfChain = uint32(0xfffffffe) + fatSector = uint32(0xfffffffd) + difatSector = uint32(0xfffffffc) + maxDIFATSectors = 1_024 + maxDirectoryBytes = int64(8 << 20) + ) + if size < 512 { + return nil, errors.New("compound document header is truncated") + } + header := make([]byte, 512) + if _, err := reader.ReadAt(header, 0); err != nil { + return nil, fmt.Errorf("read compound document header: %w", err) + } + if !bytes.Equal(header[:8], compoundFileMagic) || binary.LittleEndian.Uint16(header[28:30]) != 0xfffe { + return nil, errors.New("compound document header is invalid") + } + sectorShift := binary.LittleEndian.Uint16(header[30:32]) + majorVersion := binary.LittleEndian.Uint16(header[26:28]) + if (majorVersion != 3 || sectorShift != 9) && (majorVersion != 4 || sectorShift != 12) { + return nil, errors.New("compound document sector size is unsupported") + } + sectorSize := int64(1 << sectorShift) + sectorCount := size/sectorSize - 1 + if sectorCount <= 0 || size%sectorSize != 0 { + return nil, errors.New("compound document size is invalid") + } + fatEntriesPerSector := sectorSize / 4 + maxFATSectors := int((sectorCount + fatEntriesPerSector - 1) / fatEntriesPerSector) + numFAT := int(binary.LittleEndian.Uint32(header[44:48])) + firstDirectory := binary.LittleEndian.Uint32(header[48:52]) + firstDIFAT := binary.LittleEndian.Uint32(header[68:72]) + numDIFAT := int(binary.LittleEndian.Uint32(header[72:76])) + if numFAT <= 0 || numFAT > maxFATSectors || numDIFAT > maxDIFATSectors { + return nil, errors.New("compound document allocation table exceeds limits") + } + fatSectors := make([]uint32, 0, numFAT) + for offset := 76; offset < 512 && len(fatSectors) < numFAT; offset += 4 { + sector := binary.LittleEndian.Uint32(header[offset : offset+4]) + if sector != freeSector { + fatSectors = append(fatSectors, sector) + } + } + seenDIFAT := map[uint32]bool{} + for i, sector := 0, firstDIFAT; i < numDIFAT; i++ { + if int64(sector) >= sectorCount || seenDIFAT[sector] { + return nil, errors.New("compound document DIFAT chain is invalid") + } + seenDIFAT[sector] = true + data, readErr := readCompoundSector(reader, sector, sectorSize, size) + if readErr != nil { + return nil, readErr + } + for offset := 0; offset < len(data)-4 && len(fatSectors) < numFAT; offset += 4 { + fatID := binary.LittleEndian.Uint32(data[offset : offset+4]) + if fatID != freeSector { + fatSectors = append(fatSectors, fatID) + } + } + sector = binary.LittleEndian.Uint32(data[len(data)-4:]) + if i == numDIFAT-1 && sector != endOfChain { + return nil, errors.New("compound document DIFAT chain does not terminate") + } + } + if len(fatSectors) != numFAT { + return nil, errors.New("compound document FAT sector count is invalid") + } + fat := make([]uint32, 0, int(sectorCount)) + seenFAT := map[uint32]bool{} + for _, sector := range fatSectors { + if int64(sector) >= sectorCount || seenFAT[sector] { + return nil, errors.New("compound document FAT sector list is invalid") + } + seenFAT[sector] = true + data, readErr := readCompoundSector(reader, sector, sectorSize, size) + if readErr != nil { + return nil, readErr + } + for offset := 0; offset < len(data); offset += 4 { + fat = append(fat, binary.LittleEndian.Uint32(data[offset:offset+4])) + } + } + if len(fat) < int(sectorCount) { + return nil, errors.New("compound document FAT is truncated") + } + for _, sector := range fatSectors { + if fat[sector] != fatSector { + return nil, errors.New("compound document FAT sector is not self-marked") + } + } + for sector := range seenDIFAT { + if fat[sector] != difatSector { + return nil, errors.New("compound document DIFAT sector is not self-marked") + } + } + + var directory bytes.Buffer + seenDirectory := map[uint32]bool{} + for sector := firstDirectory; sector != endOfChain; { + if int64(sector) >= sectorCount || seenDirectory[sector] || int64(directory.Len())+sectorSize > maxDirectoryBytes { + return nil, errors.New("compound document directory chain exceeds limits") + } + seenDirectory[sector] = true + data, readErr := readCompoundSector(reader, sector, sectorSize, size) + if readErr != nil { + return nil, readErr + } + _, _ = directory.Write(data) + next := fat[sector] + if next == freeSector || next == fatSector || next == difatSector { + return nil, errors.New("compound document directory chain is invalid") + } + sector = next + } + entries := make([]compoundDirectoryEntry, 0, directory.Len()/128) + rootIndex := -1 + data := directory.Bytes() + for offset := 0; offset+128 <= len(data); offset += 128 { + entry := data[offset : offset+128] + entryType := entry[66] + if entryType == 0 { + entries = append(entries, compoundDirectoryEntry{}) + continue + } + if entryType != 1 && entryType != 2 && entryType != 5 { + return nil, errors.New("compound document directory entry type is invalid") + } + nameLength := int(binary.LittleEndian.Uint16(entry[64:66])) + if nameLength < 2 || nameLength > 64 || nameLength%2 != 0 { + return nil, errors.New("compound document directory name is invalid") + } + name, decodeErr := decodeUTF16LE(entry[:nameLength-2]) + if decodeErr != nil || name == "" { + return nil, errors.New("compound document directory name is invalid") + } + entries = append(entries, compoundDirectoryEntry{ + name: name, entryType: entryType, + left: binary.LittleEndian.Uint32(entry[68:72]), right: binary.LittleEndian.Uint32(entry[72:76]), + child: binary.LittleEndian.Uint32(entry[76:80]), + }) + if entryType == 5 { + if rootIndex != -1 { + return nil, errors.New("compound document has multiple root entries") + } + rootIndex = len(entries) - 1 + } + } + if rootIndex < 0 { + return nil, errors.New("compound document has no root entry") + } + names := map[string]bool{} + seen := map[uint32]bool{} + stack := []uint32{entries[rootIndex].child} + for len(stack) > 0 { + index := stack[len(stack)-1] + stack = stack[:len(stack)-1] + if index == compoundNoStream { + continue + } + if uint64(index) >= uint64(len(entries)) || seen[index] || entries[index].entryType == 0 || entries[index].entryType == 5 { + return nil, errors.New("compound document root directory tree is invalid") + } + seen[index] = true + entry := entries[index] + if entry.entryType == 2 { + names[entry.name] = true + } + stack = append(stack, entry.left, entry.right) + } + return names, nil +} + +func readCompoundSector(reader io.ReaderAt, sector uint32, sectorSize, size int64) ([]byte, error) { + offset := (int64(sector) + 1) * sectorSize + if offset < sectorSize || offset > size-sectorSize { + return nil, errors.New("compound document sector is out of range") + } + data := make([]byte, sectorSize) + if _, err := reader.ReadAt(data, offset); err != nil { + return nil, fmt.Errorf("read compound document sector: %w", err) + } + return data, nil +} + +func decodeUTF16LE(data []byte) (string, error) { + if len(data)%2 != 0 { + return "", errors.New("odd UTF-16 length") + } + runes := make([]rune, 0, len(data)/2) + for i := 0; i < len(data); i += 2 { + value := binary.LittleEndian.Uint16(data[i : i+2]) + if value == 0 || value >= 0xd800 && value <= 0xdfff { + return "", errors.New("unsupported UTF-16 directory name") + } + runes = append(runes, rune(value)) + } + return string(runes), nil +} + +func detectZIPFormat(reader io.ReaderAt, size int64) (CandidateFormat, error) { + if err := validateZIPEndRecord(reader, size); err != nil { + return CandidateFormat{}, err + } + archive, err := zip.NewReader(reader, size) + if err != nil { + return CandidateFormat{}, fmt.Errorf("open document ZIP container: %w", err) + } + if len(archive.File) > maxZIPEntries { + return CandidateFormat{}, errors.New("document ZIP container has too many entries") + } + names := make(map[string]bool, len(archive.File)) + var expanded uint64 + var mimeValue string + var contentTypes []byte + for _, entry := range archive.File { + if err := validateZIPName(entry.Name); err != nil { + return CandidateFormat{}, err + } + if entry.Mode()&os.ModeSymlink != 0 { + return CandidateFormat{}, errors.New("document ZIP container contains a symlink") + } + if entry.Flags&1 != 0 || (entry.Method != zip.Store && entry.Method != zip.Deflate) { + return CandidateFormat{}, errors.New("document ZIP container uses unsupported encryption or compression") + } + if entry.UncompressedSize64 > maxZIPSingleExpandedByte || expanded > maxZIPExpandedBytes-entry.UncompressedSize64 { + return CandidateFormat{}, errors.New("document ZIP container exceeds expanded-byte limits") + } + expanded += entry.UncompressedSize64 + if names[entry.Name] { + return CandidateFormat{}, errors.New("document ZIP container has duplicate entry names") + } + names[entry.Name] = true + if err := verifyZIPEntry(entry); err != nil { + return CandidateFormat{}, err + } + if entry.Name == "mimetype" { + value, readErr := readZIPEntry(entry, 256) + if readErr != nil { + return CandidateFormat{}, readErr + } + mimeValue = string(value) + } + if entry.Name == ooxmlContentTypesName { + value, readErr := readZIPEntry(entry, 2<<20) + if readErr != nil { + return CandidateFormat{}, readErr + } + contentTypes = value + } + } + + var id string + switch { + case names["word/document.xml"] && hasOOXMLMainType(contentTypes, "/word/document.xml", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"): + id = "docx" + case names["ppt/presentation.xml"] && hasOOXMLMainType(contentTypes, "/ppt/presentation.xml", + "application/vnd.openxmlformats-officedocument.presentationml.presentation.main+xml"): + id = "pptx" + case names["xl/workbook.xml"] && hasOOXMLMainType(contentTypes, "/xl/workbook.xml", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"): + id = "xlsx" + case mimeValue == "application/vnd.oasis.opendocument.text" && names["META-INF/manifest.xml"]: + id = "odt" + case mimeValue == "application/vnd.oasis.opendocument.spreadsheet" && names["META-INF/manifest.xml"]: + id = "ods" + case mimeValue == "application/epub+zip" && names["META-INF/container.xml"]: + id = "epub" + case hasNumbersMarker(names): + id = "numbers" + default: + return CandidateFormat{}, errors.New("zip container is not a supported document format") + } + format, _ := CandidateFormatByID(id) + return format, nil +} + +func hasOOXMLMainType(content []byte, partName, contentType string) bool { + if len(content) == 0 || !validXMLDocument(content) { + return false + } + var document struct { + XMLName xml.Name `xml:"Types"` + Overrides []struct { + PartName string `xml:"PartName,attr"` + ContentType string `xml:"ContentType,attr"` + } `xml:"Override"` + } + if err := xml.Unmarshal(content, &document); err != nil || + document.XMLName.Space != "http://schemas.openxmlformats.org/package/2006/content-types" { + return false + } + found := false + for _, override := range document.Overrides { + if override.PartName != partName { + continue + } + if found || override.ContentType != contentType { + return false + } + found = true + } + return found +} + +func validateZIPEndRecord(reader io.ReaderAt, size int64) error { + const maxTail = int64(65_557) + tailSize := min(size, maxTail) + tail := make([]byte, tailSize) + if _, err := reader.ReadAt(tail, size-tailSize); err != nil && !errors.Is(err, io.EOF) { + return fmt.Errorf("read document ZIP end record: %w", err) + } + signature := []byte{'P', 'K', 0x05, 0x06} + offset := bytes.LastIndex(tail, signature) + if offset < 0 || len(tail)-offset < 22 { + return errors.New("document ZIP container has no bounded end record") + } + record := tail[offset:] + entries := binary.LittleEndian.Uint16(record[10:12]) + entriesOnDisk := binary.LittleEndian.Uint16(record[8:10]) + centralSize := binary.LittleEndian.Uint32(record[12:16]) + centralOffset := binary.LittleEndian.Uint32(record[16:20]) + commentSize := int(binary.LittleEndian.Uint16(record[20:22])) + if binary.LittleEndian.Uint16(record[4:6]) != 0 || binary.LittleEndian.Uint16(record[6:8]) != 0 || entriesOnDisk != entries { + return errors.New("multi-disk document ZIP containers are unsupported") + } + if entries == 0xffff || centralSize == 0xffffffff || centralOffset == 0xffffffff || + int(entries) > maxZIPEntries || centralSize > maxZIPCentralDirectory { + return errors.New("document ZIP central directory exceeds limits") + } + if int64(centralOffset)+int64(centralSize) > size { + return errors.New("document ZIP central directory is out of range") + } + if len(record) != 22+commentSize { + return errors.New("document ZIP end record has invalid comment length") + } + return nil +} + +func validateZIPName(name string) error { + if name == "" || strings.ContainsRune(name, 0) || strings.ContainsAny(name, "\\:") || strings.HasPrefix(name, "/") { + return errors.New("document ZIP container has an unsafe entry name") + } + clean := path.Clean(name) + if clean == ".." || strings.HasPrefix(clean, "../") || clean != strings.TrimSuffix(name, "/") { + return errors.New("document ZIP container has a traversing entry name") + } + return nil +} + +func readZIPEntry(entry *zip.File, limit int64) ([]byte, error) { + if limit < 0 || entry.UncompressedSize64 > maxZIPSingleExpandedByte || int64(entry.UncompressedSize64) > limit { + return nil, errors.New("document ZIP marker entry exceeds limit") + } + reader, err := entry.Open() + if err != nil { + return nil, fmt.Errorf("open document ZIP marker: %w", err) + } + defer func() { _ = reader.Close() }() + value, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if err != nil { + return nil, fmt.Errorf("read document ZIP marker: %w", err) + } + if int64(len(value)) > limit { + return nil, errors.New("document ZIP marker entry exceeds limit") + } + return value, nil +} + +func verifyZIPEntry(entry *zip.File) error { + if entry.UncompressedSize64 > maxZIPSingleExpandedByte { + return errors.New("document ZIP entry exceeds verification limit") + } + reader, err := entry.Open() + if err != nil { + return fmt.Errorf("open document ZIP entry: %w", err) + } + expectedSize := int64(entry.UncompressedSize64) + written, readErr := io.Copy(io.Discard, io.LimitReader(reader, expectedSize+1)) + closeErr := reader.Close() + if readErr != nil || closeErr != nil || written != expectedSize { + return errors.New("document ZIP entry failed bounded verification") + } + return nil +} + +func hasNumbersMarker(names map[string]bool) bool { + for name := range names { + if strings.HasPrefix(name, "Index/Tables/") && strings.HasSuffix(name, ".iwa") { + return true + } + } + return false +} + +func detectTextFormat(content []byte, mediaType string) (CandidateFormat, error) { + if !utf8.Valid(content) || bytes.IndexByte(content, 0) >= 0 { + return CandidateFormat{}, errors.New("document text is not safe UTF-8") + } + candidate, ok := candidateByMediaType(mediaType) + if !ok || !isTextCandidate(candidate) { + return CandidateFormat{}, errors.New("document bytes have no supported signature") + } + trimmed := bytes.TrimSpace(content) + switch candidate.ID { + case "json": + if len(trimmed) == 0 || !jsontext.Value(trimmed).IsValid() { + return CandidateFormat{}, errors.New("declared JSON document is invalid") + } + case "jsonl": + for line := range bytes.SplitSeq(trimmed, []byte{'\n'}) { + if len(bytes.TrimSpace(line)) > 0 && !jsontext.Value(bytes.TrimSpace(line)).IsValid() { + return CandidateFormat{}, errors.New("declared JSONL document is invalid") + } + } + case "xml": + if !validXMLDocument(content) { + return CandidateFormat{}, errors.New("declared XML document is invalid") + } + case "csv": + csvReader := csv.NewReader(bytes.NewReader(content)) + csvReader.FieldsPerRecord = -1 + csvReader.ReuseRecord = true + records := 0 + for { + _, readErr := csvReader.Read() + if errors.Is(readErr, io.EOF) { + break + } + if readErr != nil { + return CandidateFormat{}, errors.New("declared CSV document is invalid") + } + records++ + } + if records == 0 { + return CandidateFormat{}, errors.New("declared CSV document is invalid") + } + case "latex": + if !bytes.Contains(content, []byte(`\documentclass`)) && !bytes.Contains(content, []byte(`\begin{document}`)) { + return CandidateFormat{}, errors.New("declared LaTeX document has no document marker") + } + case "eml": + message, err := mail.ReadMessage(bytes.NewReader(content)) + if err != nil || message.Header.Get("From") == "" || message.Header.Get("Date") == "" { + return CandidateFormat{}, errors.New("declared EML document lacks required message headers") + } + case "yaml": + if len(trimmed) == 0 || (!bytes.HasPrefix(trimmed, []byte("---")) && !bytes.Contains(trimmed, []byte(": "))) { + return CandidateFormat{}, errors.New("declared YAML document has no structural marker") + } + } + return candidate, nil +} + +func validXMLDocument(content []byte) bool { + decoder := xml.NewDecoder(bytes.NewReader(content)) + depth := 0 + roots := 0 + for { + token, err := decoder.Token() + if errors.Is(err, io.EOF) { + return roots == 1 && depth == 0 + } + if err != nil { + return false + } + switch value := token.(type) { + case xml.StartElement: + if depth == 0 { + roots++ + if roots > 1 { + return false + } + } + depth++ + case xml.EndElement: + depth-- + if depth < 0 { + return false + } + case xml.CharData: + if depth == 0 && len(bytes.TrimSpace(value)) != 0 { + return false + } + } + } +} + +func isTextCandidate(candidate CandidateFormat) bool { + switch candidate.Family { + case "text", "structured", "source", "mail", "spreadsheet": + return candidate.ID != "msg" && candidate.ID != "xls" && candidate.ID != "xlsx" && candidate.ID != "ods" && candidate.ID != "numbers" + default: + return false + } +} + +func candidateByMediaType(mediaType string) (CandidateFormat, bool) { + for _, candidate := range candidateFormats { + if candidate.MediaType == mediaType { + return candidate, true + } + } + return CandidateFormat{}, false +} + +func readPrefix(reader io.ReaderAt, size, limit int64) ([]byte, error) { + length := min(size, limit) + buffer := make([]byte, length) + read, err := reader.ReadAt(buffer, 0) + if err != nil && !errors.Is(err, io.EOF) { + return nil, fmt.Errorf("read document signature: %w", err) + } + if int64(read) != length { + return nil, errors.New("document bytes changed during signature read") + } + return buffer, nil +} diff --git a/document/internal/formatdetect/formats.go b/document/internal/formatdetect/formats.go new file mode 100644 index 00000000..60deac03 --- /dev/null +++ b/document/internal/formatdetect/formats.go @@ -0,0 +1,78 @@ +package formatdetect + +import ( + "fmt" + "slices" +) + +const ( + mediaTypeJSON = "application/json" + mediaTypePDF = "application/pdf" + formatIDPDF = "pdf" +) + +// CandidateFormat describes one locally detectable document format. A +// candidate does not authorize an upload. +type CandidateFormat struct { + ID string `json:"id"` + Family string `json:"family"` + MediaType string `json:"media_type"` + UnitKind string `json:"unit_kind"` +} + +var candidateFormats = []CandidateFormat{ + {ID: formatIDPDF, Family: formatIDPDF, MediaType: mediaTypePDF, UnitKind: "page"}, + {ID: "docx", Family: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", UnitKind: "page"}, + {ID: "doc", Family: "word", MediaType: "application/msword", UnitKind: "page"}, + {ID: "odt", Family: "word", MediaType: "application/vnd.oasis.opendocument.text", UnitKind: "page"}, + {ID: "rtf", Family: "word", MediaType: "application/rtf", UnitKind: "page"}, + {ID: "pptx", Family: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", UnitKind: "slide"}, + {ID: "ppt", Family: "presentation", MediaType: "application/vnd.ms-powerpoint", UnitKind: "slide"}, + {ID: "xlsx", Family: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", UnitKind: "sheet"}, + {ID: "xls", Family: "spreadsheet", MediaType: "application/vnd.ms-excel", UnitKind: "sheet"}, + {ID: "ods", Family: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", UnitKind: "sheet"}, + {ID: "numbers", Family: "spreadsheet", MediaType: "application/vnd.apple.numbers", UnitKind: "sheet"}, + {ID: "csv", Family: "spreadsheet", MediaType: "text/csv", UnitKind: "record"}, + {ID: "epub", Family: "ebook", MediaType: "application/epub+zip", UnitKind: "spine"}, + {ID: "txt", Family: "text", MediaType: "text/plain", UnitKind: "section"}, + {ID: "markdown", Family: "text", MediaType: "text/markdown", UnitKind: "section"}, + {ID: "rst", Family: "text", MediaType: "text/x-rst", UnitKind: "section"}, + {ID: "latex", Family: "text", MediaType: "application/x-tex", UnitKind: "section"}, + {ID: "json", Family: "structured", MediaType: mediaTypeJSON, UnitKind: "record"}, + {ID: "jsonl", Family: "structured", MediaType: "application/x-ndjson", UnitKind: "record"}, + {ID: "xml", Family: "structured", MediaType: "application/xml", UnitKind: "record"}, + {ID: "yaml", Family: "structured", MediaType: "application/yaml", UnitKind: "record"}, + {ID: "go", Family: "source", MediaType: "text/x-go", UnitKind: "section"}, + {ID: "python", Family: "source", MediaType: "text/x-python", UnitKind: "section"}, + {ID: "javascript", Family: "source", MediaType: "text/javascript", UnitKind: "section"}, + {ID: "eml", Family: "mail", MediaType: "message/rfc822", UnitKind: "message"}, + {ID: "msg", Family: "mail", MediaType: "application/vnd.ms-outlook", UnitKind: "message"}, +} + +// CandidateFormats returns a defensive copy in stable probe order. +func CandidateFormats() []CandidateFormat { + return slices.Clone(candidateFormats) +} + +// CandidateFormatByID returns the candidate with the given stable identifier. +func CandidateFormatByID(id string) (CandidateFormat, bool) { + for _, candidate := range candidateFormats { + if candidate.ID == id { + return candidate, true + } + } + return CandidateFormat{}, false +} + +// CandidateFormatByMediaType returns the one exact canonical declared type. +func CandidateFormatByMediaType(mediaType string) (CandidateFormat, bool) { + return candidateByMediaType(mediaType) +} + +// ProbeFixtureSentinel returns the synthetic phrase required in one fixture. +func ProbeFixtureSentinel(formatID string) (string, error) { + if _, ok := CandidateFormatByID(formatID); !ok { + return "", fmt.Errorf("unknown Mistral probe format %q", formatID) + } + return "docbank probe " + formatID + " cedar 7319", nil +} diff --git a/document/llamaparse/client.go b/document/llamaparse/client.go new file mode 100644 index 00000000..315b56ac --- /dev/null +++ b/document/llamaparse/client.go @@ -0,0 +1,1224 @@ +package llamaparse + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" +) + +const ( + apiHost = "api.cloud.llamaindex.ai" + apiOrigin = "https://" + apiHost + uploadPath = "/api/v1/parsing/upload" + providerID = "llamaparse.parse-v1" + timeForm = "2006-01-02T15:04:05.000000000Z" + markdownPageSeparator = "\n\n---\n\n" + maxConfiguredBytes = 512 << 20 + maxConfiguredPolls = 100_000 + maxConfiguredDuration = 24 * time.Hour + defaultUploadBytes = 50 << 20 + defaultRequestOverhead = 1 << 20 + defaultControlBytes = 64 << 10 + defaultResultBytes = 64 << 20 + defaultArtifactBytes = 16 << 20 +) + +var errExecutionIdentityChanged = errors.New("reported execution identity changed") + +var ( + _ document.RenditionProvider = (*Client)(nil) + _ document.ResumableRenditionProvider = (*Client)(nil) +) + +// SecretResolver resolves the one credential named by a frozen profile. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile freezes the hosted model, parse preset, and every network/output +// bound. API origin and routes are deliberately not configurable. +type Profile struct { + Model string + Preset string + SecretBinding string + MaxUploadBytes int64 + MaxRequestBytes int64 + MaxControlBytes int64 + MaxPolls int + PollInterval time.Duration + RequestTimeout time.Duration + MaxResultBytes int64 + MaxArtifactBytes int64 + MaxArtifacts int + MaxWallTime time.Duration + RetainImages bool +} + +// Client is a fixed-origin hosted LlamaParse rendition provider. +type Client struct { + descriptor document.RenditionDescriptor + profile Profile + secrets SecretResolver + transport http.RoundTripper +} + +// NewProvider constructs a provider around an injected hardened transport. +// It never resolves credentials or performs network access during setup. +func NewProvider(profile Profile, secrets SecretResolver, transport http.RoundTripper) (*Client, error) { + profile = defaultProfile(profile) + if err := validateProfile(profile); err != nil { + return nil, err + } + if nilValue(secrets) { + return nil, errors.New("LlamaParse named credential resolver is required") + } + if nilValue(transport) { + return nil, errors.New("LlamaParse hardened transport is required") + } + policyDigest := sha256.Sum256([]byte(fmt.Sprintf( + "llamaparse-profile/v1\x00%s\x00%s\x00%s\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%t", + profile.Model, profile.Preset, profile.SecretBinding, profile.MaxUploadBytes, + profile.MaxRequestBytes, profile.MaxControlBytes, profile.MaxPolls, + profile.PollInterval, profile.RequestTimeout, profile.MaxResultBytes, + profile.MaxArtifactBytes, profile.MaxWallTime, profile.RetainImages, + ))) + roles := []document.EvidenceArtifactRole(nil) + if profile.RetainImages { + roles = []document.EvidenceArtifactRole{document.EvidenceArtifactImage} + } + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: hex.EncodeToString(policyDigest[:]), + TrustBoundary: document.RenditionTrustHostedProvider, + SupportedFormats: []document.RenditionFormatCapability{{ + MediaFamily: "pdf", MediaType: "application/pdf", + InputKind: document.RenditionInputOriginalFile, + }}, + ReturnsMarkdown: true, ReturnsStructured: true, ArtifactRoles: roles, + }) + if err != nil { + return nil, fmt.Errorf("LlamaParse descriptor: %w", err) + } + return &Client{ + descriptor: cloneDescriptor(descriptor), profile: profile, + secrets: secrets, transport: &fixedOriginTransport{base: transport}, + }, nil +} + +// Descriptor returns a defensive copy of the immutable provider identity. +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Render starts and completes one hosted parse operation. +func (client *Client) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + return client.RenderResumable(ctx, upload, authorization, nil, nil) +} + +// RenderResumable submits exact authorized bytes or resumes one known UUID +// handle. A new handle is checkpointed before the first status request. +func (client *Client) RenderResumable( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, resume *document.RenditionResumeHandle, + checkpoint document.RenditionResumeCheckpoint, +) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("LlamaParse client is required") + } + now := time.Now().UTC() + resumeState, expiresAt, enforceExpiry, err := client.validateInvocation(now, upload, authorization, resume) + if err != nil { + return document.RenditionResult{}, err + } + operationCtx, cancel := boundedOperationContext(ctx, expiresAt, client.profile.MaxWallTime, enforceExpiry) + defer cancel() + state := operationState{startedAt: resumeState.submittedAt, completedAt: resumeState.checkpointedAt, + enforceExpiry: enforceExpiry} + var jobID string + if resume == nil { + jobID, err = client.submit(operationCtx, ctx, expiresAt, upload, &state) + if err != nil { + return document.RenditionResult{}, err + } + checkpointedAt := time.Now().UTC() + if err := client.lifecycleErrorAt( + ctx, operationCtx, expiresAt, enforceExpiry, checkpointedAt, nil, + ); err != nil { + return document.RenditionResult{}, err + } + handle := encodeResumeHandle(resumeStateV1{ + jobID: jobID, submittedAt: state.startedAt, checkpointedAt: checkpointedAt, + }) + if checkpoint != nil { + if err := checkpoint(document.RenditionResumeHandle{Value: handle}); err != nil { + return document.RenditionResult{}, err + } + } + } else { + jobID = resumeState.jobID + } + if err := client.poll(operationCtx, ctx, expiresAt, jobID, &state); err != nil { + return document.RenditionResult{}, err + } + result, err := client.result(operationCtx, ctx, expiresAt, jobID, authorization, &state) + if err != nil { + return document.RenditionResult{}, err + } + acceptedAt := time.Now().UTC() + if resume == nil { + state.completedAt = acceptedAt + } + if err := client.lifecycleErrorAt( + ctx, operationCtx, expiresAt, enforceExpiry, acceptedAt, nil, + ); err != nil { + return document.RenditionResult{}, err + } + result.Receipt = document.RenditionReceipt{ + ProviderID: client.descriptor.ID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: authorization.SourceSHA256, + OperationID: "llamaparse-" + jobID[:8], + StartedAt: state.startedAt.Format(timeForm), CompletedAt: state.completedAt.Format(timeForm), + Warnings: state.warnings, + Usage: document.RenditionUsage{ + Requests: state.requests, Retries: state.retries, InputBytes: authorization.SourceBytes, + OutputBytes: state.outputBytes, Units: int64(len(result.Evidence.Units)), + }, + RetryDelayMillis: state.pollDelay.Milliseconds(), + } + return result, nil +} + +func (client *Client) validateInvocation( + now time.Time, upload document.AuthorizedUpload, authorization document.RenditionAuthorization, + resume *document.RenditionResumeHandle, +) (resumeStateV1, time.Time, bool, error) { + expiresAt, err := time.Parse(timeForm, authorization.ExpiresAt) + if err != nil { + return resumeStateV1{}, time.Time{}, false, renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse authorization expiry is invalid", err) + } + if resume == nil { + if nilValue(upload) { + return resumeStateV1{}, time.Time{}, false, errors.New("LlamaParse authorized upload is required for submission") + } + if _, err := document.ValidateRenditionProviderRequestAt(now, client, upload, authorization); err != nil { + return resumeStateV1{}, time.Time{}, false, err + } + } else { + if !nilValue(upload) { + return resumeStateV1{}, time.Time{}, false, errors.New("LlamaParse resume must not receive source bytes") + } + parsed, parseErr := parseResumeHandle(resume.Value) + if parseErr != nil { + return resumeStateV1{}, time.Time{}, false, renditionError(document.RenditionErrorUnknownJob, + "LlamaParse resume handle is invalid", parseErr) + } + if authorization.ProviderID != client.descriptor.ID || + authorization.DescriptorFingerprint != client.descriptor.Fingerprint || + authorization.PolicyFingerprint != client.descriptor.PolicyFingerprint { + return resumeStateV1{}, time.Time{}, false, renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse resume authority changed", nil) + } + authorizedAt, authErr := time.Parse(timeForm, authorization.AuthorizedAt) + if authErr != nil || parsed.submittedAt.Before(authorizedAt) || + parsed.checkpointedAt.Before(parsed.submittedAt) || parsed.checkpointedAt.After(expiresAt) { + return resumeStateV1{}, time.Time{}, false, renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse resume receipt facts are outside the sealed authorization", authErr) + } + return parsed, expiresAt, false, nil + } + return resumeStateV1{}, expiresAt, true, nil +} + +func (client *Client) submit( + operationCtx, callerCtx context.Context, expiresAt time.Time, upload document.AuthorizedUpload, + state *operationState, +) (string, error) { + metadata := upload.Metadata() + if metadata.ByteLength > client.profile.MaxUploadBytes { + return "", renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse upload exceeds profile limit", nil) + } + source, err := readExactUpload(operationCtx, upload, metadata, client.profile.MaxUploadBytes) + if err != nil { + if operationCtx.Err() != nil { + return "", client.contextOrTransient(callerCtx, operationCtx, expiresAt, state.enforceExpiry, + "LlamaParse upload read timed out", err) + } + return "", err + } + defer clear(source) + body := new(bytes.Buffer) + writer := multipart.NewWriter(body) + for _, field := range [][2]string{ + {"model", client.profile.Model}, {"preset", client.profile.Preset}, + {"page_error_tolerance", "0"}, {"save_images", strconv.FormatBool(client.profile.RetainImages)}, + {"disable_image_extraction", strconv.FormatBool(!client.profile.RetainImages)}, + {"take_screenshot", "false"}, + } { + if err := writer.WriteField(field[0], field[1]); err != nil { + return "", renditionError(document.RenditionErrorTransient, + "LlamaParse request could not be built", err) + } + } + header := make(textproto.MIMEHeader) + header.Set("Content-Disposition", mime.FormatMediaType("form-data", map[string]string{ + "name": "file", "filename": metadata.Filename, + })) + header.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(header) + if err == nil { + _, err = part.Write(source) + } + if closeErr := writer.Close(); err == nil { + err = closeErr + } + if err != nil { + return "", renditionError(document.RenditionErrorTransient, + "LlamaParse request could not be built", err) + } + if int64(body.Len()) > client.profile.MaxRequestBytes { + return "", renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse request exceeds profile limit", nil) + } + state.startedAt = time.Now().UTC() + raw, status, err := client.do(operationCtx, http.MethodPost, uploadPath, + writer.FormDataContentType(), body.Bytes(), client.profile.MaxControlBytes, state) + if err != nil { + if providerError, ok := errors.AsType[*document.RenditionProviderError](err); ok && + providerError.Code() == document.RenditionErrorAuthentication { + return "", err + } + if callerCtx.Err() != nil { + return "", renditionError(document.RenditionErrorAmbiguousSubmission, + "LlamaParse submission outcome is ambiguous", callerCtx.Err()) + } + return "", renditionError(document.RenditionErrorAmbiguousSubmission, + "LlamaParse submission outcome is ambiguous", err) + } + if status < 200 || status >= 300 { + return "", client.httpError(status, "submission", raw) + } + var job jobResponse + if err := strictJSON(raw, &job); err != nil || validateJobID(job.ID) != nil { + return "", renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse submission schema changed", err) + } + if err := validateInitialStatus(job.Status); err != nil { + return "", err + } + return job.ID, nil +} + +func (client *Client) poll( + operationCtx, callerCtx context.Context, expiresAt time.Time, jobID string, + state *operationState, +) error { + for attempt := range client.profile.MaxPolls { + raw, status, err := client.do(operationCtx, http.MethodGet, statusPath(jobID), "", nil, + client.profile.MaxControlBytes, state) + if err != nil { + return client.contextOrTransient(callerCtx, operationCtx, expiresAt, state.enforceExpiry, + "LlamaParse status request failed", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return err + } + if status < 200 || status >= 300 { + return client.httpError(status, "status", raw) + } + var job jobResponse + if err := strictJSON(raw, &job); err != nil || job.ID != jobID { + return renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse status schema changed", err) + } + switch job.Status { + case "SUCCESS": + return nil + case "PENDING": + if attempt+1 == client.profile.MaxPolls { + break + } + state.retries++ + if err := waitContext(operationCtx, client.profile.PollInterval); err != nil { + return client.contextOrTransient(callerCtx, operationCtx, expiresAt, state.enforceExpiry, + "LlamaParse polling timed out", err) + } + state.pollDelay += client.profile.PollInterval + case "ERROR": + return renditionError(document.RenditionErrorUnsupportedInput, + "LlamaParse could not parse the input", nil) + case "PARTIAL_SUCCESS": + return renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse returned partial output", nil) + case "CANCELLED": + return renditionError(document.RenditionErrorCanceled, + "LlamaParse job was canceled", nil) + default: + return renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse status schema changed", nil) + } + } + return renditionError(document.RenditionErrorTransient, + "LlamaParse polling limit was reached", nil) +} + +func (client *Client) result( + operationCtx, callerCtx context.Context, expiresAt time.Time, jobID string, + authorization document.RenditionAuthorization, state *operationState, +) (document.RenditionResult, error) { + budget := min64(client.profile.MaxResultBytes, int64(authorization.MaxTotalResultBytes)) + raw, status, err := client.do(operationCtx, http.MethodGet, jsonResultPath(jobID), "", nil, + budget, state) + if err != nil { + return document.RenditionResult{}, client.contextOrMalformed( + callerCtx, operationCtx, expiresAt, state.enforceExpiry, "LlamaParse result request failed", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionResult{}, err + } + if status < 200 || status >= 300 { + return document.RenditionResult{}, client.httpError(status, "result", raw) + } + budget -= int64(len(raw)) + var envelope jsonResult + parseErr := strictJSON(raw, &envelope) + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, parseErr); err != nil { + return document.RenditionResult{}, err + } + if parseErr != nil || len(envelope.Pages) == 0 || !jsonObject(envelope.JobMetadata) { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse result schema changed", parseErr) + } + jobPages, err := client.validateJobMetadata(envelope.JobMetadata) + if lifecycleErr := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, err); lifecycleErr != nil { + return document.RenditionResult{}, lifecycleErr + } + if err != nil { + if errors.Is(err, errExecutionIdentityChanged) { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse model or preset changed", err) + } + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse result metadata is malformed", err) + } + var pages []resultPage + pageParseErr := strictJSON(envelope.Pages, &pages) + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, pageParseErr); err != nil { + return document.RenditionResult{}, err + } + if pageParseErr != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse page schema changed", pageParseErr) + } + if len(pages) == 0 { + if jobPages != 0 { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse page output is incomplete", nil) + } + return client.markdownFallback( + operationCtx, callerCtx, expiresAt, jobID, authorization, jobPages, budget, state) + } + if jobPages <= 0 || len(pages) != jobPages { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse page output is incomplete", nil) + } + evidence, markdown, imageName, err := client.pageEvidence(pages, authorization) + if lifecycleErr := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, err); lifecycleErr != nil { + return document.RenditionResult{}, lifecycleErr + } + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse page output is incomplete", err) + } + if len(markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse Markdown exceeds authorization", nil) + } + result := document.RenditionResult{Evidence: evidence, ProviderMarkdown: markdown} + if imageName != "" { + artifact, sourceArtifact, artifactErr := client.fetchImage( + operationCtx, callerCtx, expiresAt, jobID, imageName, authorization, budget, state) + if artifactErr != nil { + return document.RenditionResult{}, artifactErr + } + result.Artifacts = []document.RenditionArtifact{artifact} + result.Evidence.Artifacts = []document.SourceEvidenceArtifactV1{sourceArtifact} + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionResult{}, err + } + return result, nil +} + +func (client *Client) markdownFallback( + operationCtx, callerCtx context.Context, expiresAt time.Time, jobID string, + authorization document.RenditionAuthorization, expectedJobPages int, budget int64, state *operationState, +) (document.RenditionResult, error) { + raw, status, err := client.do(operationCtx, http.MethodGet, markdownResultPath(jobID), "", nil, + budget, state) + if err != nil { + return document.RenditionResult{}, client.contextOrMalformed( + callerCtx, operationCtx, expiresAt, state.enforceExpiry, "LlamaParse Markdown request failed", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionResult{}, err + } + if status < 200 || status >= 300 { + return document.RenditionResult{}, client.httpError(status, "Markdown result", raw) + } + var envelope markdownResult + parseErr := strictJSON(raw, &envelope) + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, parseErr); err != nil { + return document.RenditionResult{}, err + } + if parseErr != nil || !jsonObject(envelope.JobMetadata) || + envelope.Markdown == "" || !utf8.ValidString(envelope.Markdown) { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse Markdown result is malformed", parseErr) + } + jobPages, metadataErr := client.validateJobMetadata(envelope.JobMetadata) + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, metadataErr); err != nil { + return document.RenditionResult{}, err + } + if metadataErr != nil { + if errors.Is(metadataErr, errExecutionIdentityChanged) { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse model or preset changed", metadataErr) + } + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse result metadata is malformed", metadataErr) + } + if jobPages != expectedJobPages { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse result page counts disagree", nil) + } + markdown := []byte(envelope.Markdown) + if len(markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse Markdown exceeds authorization", nil) + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceDegradedProvenance, + Family: authorization.MediaFamily, UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", + Reason: "LlamaParse returned Markdown without page provenance", + }}, + Units: []document.SourceEvidenceUnitV1{{ + Order: 0, ProviderID: "llamaparse-markdown", Text: envelope.Markdown, + Locator: document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone, + Name: "llamaparse-markdown", + }, + }}, + } + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "LlamaParse Markdown cannot be represented", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionResult{}, err + } + state.warnings = []string{"degraded_provenance"} + return document.RenditionResult{Evidence: evidence, ProviderMarkdown: markdown}, nil +} + +func (client *Client) pageEvidence( + pages []resultPage, authorization document.RenditionAuthorization, +) (document.SourceEvidenceV1, []byte, string, error) { + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: authorization.MediaFamily, UnitKind: document.EvidenceUnitPage, + Units: make([]document.SourceEvidenceUnitV1, 0, len(pages)), + } + markdown := make([]string, 0, len(pages)) + imageName := "" + for order, page := range pages { + if page.Page == nil || *page.Page != order { + return document.SourceEvidenceV1{}, nil, "", errors.New("page sequence is not complete") + } + if page.Status != nil && *page.Status != "" && *page.Status != "SUCCESS" { + return document.SourceEvidenceV1{}, nil, "", errors.New("page has a non-success status") + } + text := "" + if page.Markdown != nil { + text = *page.Markdown + } else if page.Text != nil { + text = *page.Text + } + if !utf8.ValidString(text) || text == "" && (page.NoTextContent == nil || !*page.NoTextContent) { + return document.SourceEvidenceV1{}, nil, "", errors.New("page has no representable text") + } + if page.Confidence != nil && (math.IsNaN(*page.Confidence) || math.IsInf(*page.Confidence, 0) || + *page.Confidence < 0 || *page.Confidence > 1) { + return document.SourceEvidenceV1{}, nil, "", errors.New("page confidence is invalid") + } + unit := document.SourceEvidenceUnitV1{ + Order: order, ProviderID: fmt.Sprintf("llamaparse-page-%d", order), Text: text, + Locator: document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginZero, + Start: int64(order), End: int64(order), + }, + } + if page.Confidence != nil { + unit.Confidence = &document.SourceEvidenceConfidenceV1{ + Interpretation: document.EvidenceConfidenceProbability, + Minimum: 0, Maximum: 1, Value: *page.Confidence, + } + } + evidence.Units = append(evidence.Units, unit) + markdown = append(markdown, text) + for _, image := range page.Images { + if !client.profile.RetainImages || imageName != "" || validateArtifactName(image.Name) != nil { + return document.SourceEvidenceV1{}, nil, "", errors.New("provider images are unrepresentable") + } + imageName = image.Name + } + } + if imageName != "" && (client.profile.MaxArtifacts != 1 || authorization.MaxArtifacts < 1 || + !slices.Contains(authorization.AllowedArtifactRoles, document.EvidenceArtifactImage)) { + return document.SourceEvidenceV1{}, nil, "", errors.New("provider image is not authorized") + } + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.SourceEvidenceV1{}, nil, "", err + } + return evidence, []byte(strings.Join(markdown, markdownPageSeparator)), imageName, nil +} + +func (client *Client) validateJobMetadata(raw json.RawMessage) (int, error) { + var metadata map[string]json.RawMessage + if err := json.Unmarshal(raw, &metadata); err != nil { + return 0, err + } + encodedPages, present := metadata["job_pages"] + if !present { + return 0, errors.New("reported page count is missing") + } + var pages int + if err := strictJSON(encodedPages, &pages); err != nil || pages < 0 { + return 0, errors.New("reported page count is invalid") + } + for field, expected := range map[string]string{ + "model": client.profile.Model, "preset": client.profile.Preset, + } { + encoded, present := metadata[field] + if !present { + continue + } + var actual string + if err := strictJSON(encoded, &actual); err != nil || actual != expected { + return 0, errExecutionIdentityChanged + } + } + return pages, nil +} + +func (client *Client) fetchImage( + operationCtx, callerCtx context.Context, expiresAt time.Time, jobID, name string, + authorization document.RenditionAuthorization, budget int64, state *operationState, +) (document.RenditionArtifact, document.SourceEvidenceArtifactV1, error) { + limit := min64(client.profile.MaxArtifactBytes, int64(authorization.MaxArtifactBytes)) + limit = min64(limit, budget) + raw, status, mediaType, err := client.doWithMedia(operationCtx, http.MethodGet, + imageResultPath(jobID, name), "", nil, limit, state) + if err != nil { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, + client.contextOrMalformed(callerCtx, operationCtx, expiresAt, state.enforceExpiry, + "LlamaParse image request failed", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, err + } + if status < 200 || status >= 300 { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, + client.httpError(status, "image result", raw) + } + canonical, _, err := mime.ParseMediaType(mediaType) + if err != nil || (canonical != "image/png" && canonical != "image/jpeg" && canonical != "image/webp") { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, renditionError( + document.RenditionErrorMalformedEvidence, "LlamaParse image media type changed", err) + } + detected, err := media.DetectBytes(raw, canonical) + if err != nil || detected.Kind != media.KindImage || detected.MediaType != canonical || + detected.FrameCount != 1 || detected.Animated { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, renditionError( + document.RenditionErrorMalformedEvidence, "LlamaParse image result is malformed", err) + } + if err := client.lifecycleError(callerCtx, operationCtx, expiresAt, state.enforceExpiry, nil); err != nil { + return document.RenditionArtifact{}, document.SourceEvidenceArtifactV1{}, err + } + digest := sha256.Sum256(raw) + checksum := hex.EncodeToString(digest[:]) + pointer := "images/" + name + return document.RenditionArtifact{ + Role: document.EvidenceArtifactImage, MediaType: canonical, + Payload: raw, SHA256: checksum, + }, document.SourceEvidenceArtifactV1{ + Pointer: pointer, ProviderID: "llamaparse-image-0", + Role: document.EvidenceArtifactImage, SHA256: checksum, + }, nil +} + +func (client *Client) do( + ctx context.Context, method, path, contentType string, body []byte, limit int64, + state *operationState, +) ([]byte, int, error) { + raw, status, _, err := client.doWithMedia(ctx, method, path, contentType, body, limit, state) + return raw, status, err +} + +func (client *Client) doWithMedia( + ctx context.Context, method, path, contentType string, body []byte, limit int64, + state *operationState, +) ([]byte, int, string, error) { + if err := ctx.Err(); err != nil { + return nil, 0, "", err + } + credential, err := client.secrets.ResolveSecret(ctx, client.profile.SecretBinding) + if err != nil || !validCredential(credential) { + return nil, 0, "", renditionError(document.RenditionErrorAuthentication, + "LlamaParse credential is unavailable", err) + } + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + var reader io.Reader + if body != nil { + reader = bytes.NewReader(body) + } + request, err := http.NewRequestWithContext(requestCtx, method, apiOrigin+path, reader) + if err != nil { + return nil, 0, "", err + } + request.Header.Set("Authorization", "Bearer "+credential) + request.Header.Set("Accept", "application/json") + if contentType != "" { + request.Header.Set("Content-Type", contentType) + } + state.requests++ + response, err := client.transport.RoundTrip(request) + if err != nil { + return nil, 0, "", err + } + if response == nil || response.Body == nil { + return nil, 0, "", errors.New("provider returned an empty HTTP response") + } + raw, readErr := readBounded(response.Body, limit) + closeErr := response.Body.Close() + if readErr != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, response.StatusCode, response.Header.Get("Content-Type"), ctxErr + } + return nil, response.StatusCode, response.Header.Get("Content-Type"), readErr + } + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, response.StatusCode, response.Header.Get("Content-Type"), ctxErr + } + if closeErr != nil { + return nil, response.StatusCode, response.Header.Get("Content-Type"), + fmt.Errorf("close provider response: %w", closeErr) + } + state.outputBytes += int64(len(raw)) + return raw, response.StatusCode, response.Header.Get("Content-Type"), nil +} + +func (client *Client) httpError(status int, stage string, raw []byte) error { + _ = raw // Provider bodies stay private and are never interpolated. + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return renditionError(document.RenditionErrorAuthentication, + "LlamaParse authentication was rejected", nil) + case http.StatusTooManyRequests: + return renditionError(document.RenditionErrorRateLimited, + "LlamaParse rate limit was exhausted", nil) + case http.StatusServiceUnavailable, http.StatusInsufficientStorage: + return renditionError(document.RenditionErrorCapacity, + "LlamaParse capacity is unavailable", nil) + case http.StatusNotFound, http.StatusGone: + if stage != "submission" { + return renditionError(document.RenditionErrorUnknownJob, + "LlamaParse job is unknown or expired", nil) + } + case http.StatusBadRequest, http.StatusUnprocessableEntity, http.StatusUnsupportedMediaType, + http.StatusRequestEntityTooLarge: + return renditionError(document.RenditionErrorUnsupportedInput, + "LlamaParse rejected the input", nil) + } + if status >= 300 && status < 400 { + return renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse redirects are refused", nil) + } + if status >= 500 { + return renditionError(document.RenditionErrorTransient, + "LlamaParse service request failed", nil) + } + return renditionError(document.RenditionErrorPolicyRejected, + "LlamaParse rejected the request", nil) +} + +func (client *Client) contextOrTransient( + callerCtx, operationCtx context.Context, expiresAt time.Time, enforceExpiry bool, + message string, cause error, +) error { + if callerCtx.Err() != nil { + return renditionError(document.RenditionErrorCanceled, "LlamaParse rendering was canceled", callerCtx.Err()) + } + if enforceExpiry && !time.Now().UTC().Before(expiresAt) { + return expiredError(cause) + } + if operationCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, "LlamaParse operation timed out", cause) + } + if hasRenditionErrorCode(cause, document.RenditionErrorAuthentication) { + return cause + } + return renditionError(document.RenditionErrorTransient, message, cause) +} + +func (client *Client) contextOrMalformed( + callerCtx, operationCtx context.Context, expiresAt time.Time, enforceExpiry bool, + message string, cause error, +) error { + if callerCtx.Err() != nil || operationCtx.Err() != nil { + return client.contextOrTransient(callerCtx, operationCtx, expiresAt, enforceExpiry, message, cause) + } + if hasRenditionErrorCode(cause, document.RenditionErrorAuthentication) { + return cause + } + return renditionError(document.RenditionErrorMalformedEvidence, message, cause) +} + +func (client *Client) lifecycleError( + callerCtx, operationCtx context.Context, expiresAt time.Time, enforceExpiry bool, cause error, +) error { + return client.lifecycleErrorAt( + callerCtx, operationCtx, expiresAt, enforceExpiry, time.Now().UTC(), cause) +} + +func (*Client) lifecycleErrorAt( + callerCtx, operationCtx context.Context, expiresAt time.Time, enforceExpiry bool, + observedAt time.Time, cause error, +) error { + if callerCtx.Err() != nil { + return renditionError(document.RenditionErrorCanceled, + "LlamaParse rendering was canceled", callerCtx.Err()) + } + if enforceExpiry && !observedAt.Before(expiresAt) { + return expiredError(cause) + } + if operationCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, + "LlamaParse operation timed out", operationCtx.Err()) + } + return nil +} + +type operationState struct { + startedAt time.Time + completedAt time.Time + enforceExpiry bool + requests int64 + retries int64 + outputBytes int64 + pollDelay time.Duration + warnings []string +} + +type resumeStateV1 struct { + jobID string + submittedAt time.Time + checkpointedAt time.Time +} + +type jobResponse struct { + ID string `json:"id"` + Status string `json:"status"` + ErrorCode *string `json:"error_code,omitempty"` + ErrorMessage *string `json:"error_message,omitempty"` +} + +type jsonResult struct { + Pages json.RawMessage `json:"pages"` + JobMetadata json.RawMessage `json:"job_metadata"` +} + +type markdownResult struct { + Markdown string `json:"markdown"` + JobMetadata json.RawMessage `json:"job_metadata"` +} + +type resultPage struct { + Page *int `json:"page"` + Text *string `json:"text,omitempty"` + Markdown *string `json:"md,omitempty"` + Images []resultImage `json:"images"` + Charts json.RawMessage `json:"charts"` + Tables json.RawMessage `json:"tables"` + Layout json.RawMessage `json:"layout"` + Items json.RawMessage `json:"items"` + Status *string `json:"status,omitempty"` + Links json.RawMessage `json:"links"` + Width *float64 `json:"width,omitempty"` + Height *float64 `json:"height,omitempty"` + TriggeredAutoMode *bool `json:"triggeredAutoMode,omitempty"` + ParsingMode string `json:"parsingMode"` + StructuredData json.RawMessage `json:"structuredData,omitempty"` + NoStructuredContent bool `json:"noStructuredContent"` + NoTextContent *bool `json:"noTextContent"` + IsAudioTranscript bool `json:"isAudioTranscript,omitempty"` + DurationInSeconds *float64 `json:"durationInSeconds,omitempty"` + SlideSpeakerNotes *string `json:"slideSpeakerNotes,omitempty"` + Confidence *float64 `json:"confidence,omitempty"` + PrintedPageNumber *string `json:"printedPageNumber,omitempty"` + PageHeaderMarkdown *string `json:"pageHeaderMarkdown,omitempty"` + PageFooterMarkdown *string `json:"pageFooterMarkdown,omitempty"` +} + +type resultImage struct { + Name string `json:"name"` + Height *float64 `json:"height,omitempty"` + Width *float64 `json:"width,omitempty"` + X *float64 `json:"x,omitempty"` + Y *float64 `json:"y,omitempty"` + OriginalWidth *int `json:"original_width,omitempty"` + OriginalHeight *int `json:"original_height,omitempty"` + Type *string `json:"type,omitempty"` +} + +type fixedOriginTransport struct{ base http.RoundTripper } + +func (transport *fixedOriginTransport) RoundTrip(request *http.Request) (*http.Response, error) { + if request == nil || request.URL == nil || request.URL.Scheme != "https" || + request.URL.Host != apiHost || request.URL.User != nil || request.URL.RawQuery != "" || + request.URL.Fragment != "" || request.Host != "" && request.Host != apiHost { + return nil, errors.New("LlamaParse request destination is not fixed") + } + return transport.base.RoundTrip(request) +} + +func statusPath(jobID string) string { return "/api/v1/parsing/job/" + jobID } +func jsonResultPath(jobID string) string { return statusPath(jobID) + "/result/json" } +func markdownResultPath(jobID string) string { return statusPath(jobID) + "/result/markdown" } +func imageResultPath(jobID, name string) string { return statusPath(jobID) + "/result/image/" + name } + +func defaultProfile(profile Profile) Profile { + if profile.MaxUploadBytes == 0 { + profile.MaxUploadBytes = defaultUploadBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = profile.MaxUploadBytes + defaultRequestOverhead + } + if profile.MaxControlBytes == 0 { + profile.MaxControlBytes = defaultControlBytes + } + if profile.MaxPolls == 0 { + profile.MaxPolls = 600 + } + if profile.PollInterval == 0 { + profile.PollInterval = time.Second + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = 30 * time.Second + } + if profile.MaxResultBytes == 0 { + profile.MaxResultBytes = defaultResultBytes + } + if profile.MaxArtifactBytes == 0 { + profile.MaxArtifactBytes = defaultArtifactBytes + } + if profile.MaxWallTime == 0 { + profile.MaxWallTime = 30 * time.Minute + } + if profile.RetainImages && profile.MaxArtifacts == 0 { + profile.MaxArtifacts = 1 + } + return profile +} + +func validateProfile(profile Profile) error { + for subject, value := range map[string]string{ + "model": profile.Model, "preset": profile.Preset, "credential binding": profile.SecretBinding, + } { + if err := validateToken(value); err != nil { + return fmt.Errorf("LlamaParse %s: %w", subject, err) + } + } + if profile.MaxUploadBytes <= 0 || profile.MaxUploadBytes > maxConfiguredBytes || + profile.MaxRequestBytes <= profile.MaxUploadBytes || profile.MaxRequestBytes > maxConfiguredBytes+defaultRequestOverhead || + profile.MaxControlBytes <= 0 || profile.MaxControlBytes > 1<<20 || + profile.MaxResultBytes <= 0 || profile.MaxResultBytes > maxConfiguredBytes || + profile.MaxArtifactBytes <= 0 || profile.MaxArtifactBytes > maxConfiguredBytes || + profile.MaxPolls <= 0 || profile.MaxPolls > maxConfiguredPolls || + profile.PollInterval <= 0 || profile.PollInterval > time.Hour || + profile.RequestTimeout <= 0 || profile.RequestTimeout > maxConfiguredDuration || + profile.MaxWallTime <= 0 || profile.MaxWallTime > maxConfiguredDuration { + return errors.New("LlamaParse profile bounds are invalid") + } + if profile.RetainImages && profile.MaxArtifacts != 1 || !profile.RetainImages && profile.MaxArtifacts != 0 { + return errors.New("LlamaParse image artifact count is unrepresentable") + } + return nil +} + +func validateToken(value string) error { + if value == "" || len(value) > 128 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return errors.New("value must contain 1-128 canonical characters") + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || + char >= '0' && char <= '9' || strings.ContainsRune("_.-", char) { + continue + } + return errors.New("value contains unsupported characters") + } + return nil +} + +func validateJobID(value string) error { + if len(value) != 36 { + return errors.New("job ID is not a canonical UUID") + } + for index, char := range value { + if index == 8 || index == 13 || index == 18 || index == 23 { + if char != '-' { + return errors.New("job ID is not a canonical UUID") + } + continue + } + if (char < '0' || char > '9') && (char < 'a' || char > 'f') { + return errors.New("job ID is not a canonical UUID") + } + } + return nil +} + +func encodeResumeHandle(state resumeStateV1) string { + return fmt.Sprintf("lp1.%s.%d.%d", state.jobID, + state.submittedAt.UnixNano(), state.checkpointedAt.UnixNano()) +} + +func parseResumeHandle(value string) (resumeStateV1, error) { + if len(value) > 512 { + return resumeStateV1{}, errors.New("resume handle exceeds its bound") + } + parts := strings.Split(value, ".") + if len(parts) != 4 || parts[0] != "lp1" { + return resumeStateV1{}, errors.New("resume handle version is invalid") + } + if err := validateJobID(parts[1]); err != nil { + return resumeStateV1{}, err + } + parseTime := func(encoded string) (time.Time, error) { + nanoseconds, err := strconv.ParseInt(encoded, 10, 64) + if err != nil || nanoseconds <= 0 || strconv.FormatInt(nanoseconds, 10) != encoded { + return time.Time{}, errors.New("resume handle timestamp is invalid") + } + return time.Unix(0, nanoseconds).UTC(), nil + } + submittedAt, err := parseTime(parts[2]) + if err != nil { + return resumeStateV1{}, err + } + checkpointedAt, err := parseTime(parts[3]) + if err != nil || checkpointedAt.Before(submittedAt) { + return resumeStateV1{}, errors.New("resume handle checkpoint time is invalid") + } + return resumeStateV1{ + jobID: parts[1], submittedAt: submittedAt, checkpointedAt: checkpointedAt, + }, nil +} + +func validateInitialStatus(status string) error { + switch status { + case "PENDING", "SUCCESS": + return nil + case "ERROR": + return renditionError(document.RenditionErrorUnsupportedInput, "LlamaParse could not parse the input", nil) + case "PARTIAL_SUCCESS": + return renditionError(document.RenditionErrorMalformedEvidence, "LlamaParse returned partial output", nil) + case "CANCELLED": + return renditionError(document.RenditionErrorCanceled, "LlamaParse job was canceled", nil) + default: + return renditionError(document.RenditionErrorPolicyRejected, "LlamaParse submission schema changed", nil) + } +} + +func validateArtifactName(name string) error { + if err := validateToken(name); err != nil { + return err + } + if name == "." || name == ".." { + return errors.New("artifact name is not an opaque path segment") + } + escaped := url.PathEscape(name) + if escaped != name || strings.ContainsAny(name, "/\\:%") { + return errors.New("artifact name is not an opaque path segment") + } + return nil +} + +func validCredential(value string) bool { + if value == "" || len(value) > 4096 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return false + } + for _, char := range value { + if char < 0x21 || char == 0x7f { + return false + } + } + return true +} + +func readExactUpload(ctx context.Context, upload io.Reader, metadata document.AuthorizedUploadMetadata, limit int64) ([]byte, error) { + data, err := io.ReadAll(io.LimitReader(&contextReader{ctx: ctx, reader: upload}, limit+1)) + if err != nil { + if ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, renditionError(document.RenditionErrorTransient, "LlamaParse upload could not be read", err) + } + digest := sha256.Sum256(data) + if int64(len(data)) != metadata.ByteLength || int64(len(data)) > limit || hex.EncodeToString(digest[:]) != metadata.SHA256 { + clear(data) + return nil, renditionError(document.RenditionErrorPolicyRejected, "LlamaParse upload identity changed", nil) + } + return data, nil +} + +type contextReader struct { + ctx context.Context + reader io.Reader +} + +func (reader *contextReader) Read(buffer []byte) (int, error) { + if err := reader.ctx.Err(); err != nil { + return 0, err + } + return reader.reader.Read(buffer) +} + +func readBounded(reader io.Reader, limit int64) ([]byte, error) { + if limit <= 0 { + return nil, errors.New("response byte budget is exhausted") + } + data, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if err != nil { + return nil, err + } + if int64(len(data)) > limit { + return nil, errors.New("provider response exceeds byte limit") + } + return data, nil +} + +func strictJSON(raw []byte, target any) error { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return err + } + if decoder.More() { + return errors.New("JSON response has trailing values") + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return errors.New("JSON response has trailing data") + } + return nil +} + +func jsonObject(raw json.RawMessage) bool { + trimmed := bytes.TrimSpace(raw) + return len(trimmed) >= 2 && trimmed[0] == '{' && trimmed[len(trimmed)-1] == '}' && json.Valid(trimmed) +} + +func waitContext(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func boundedOperationContext( + ctx context.Context, expiresAt time.Time, wall time.Duration, enforceExpiry bool, +) (context.Context, context.CancelFunc) { + deadline := time.Now().Add(wall) + if enforceExpiry && expiresAt.Before(deadline) { + deadline = expiresAt + } + return context.WithDeadline(ctx, deadline) +} + +func renditionError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("LlamaParse error classification failed: %w", err) + } + return providerError +} + +func hasRenditionErrorCode(err error, code document.RenditionErrorCode) bool { + providerError, ok := errors.AsType[*document.RenditionProviderError](err) + return ok && providerError.Code() == code +} + +func expiredError(cause error) error { + return renditionError(document.RenditionErrorPolicyRejected, "LlamaParse authorization expired", cause) +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +func nilValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} + +func min64(left, right int64) int64 { + if left < right { + return left + } + return right +} diff --git a/document/llamaparse/client_test.go b/document/llamaparse/client_test.go new file mode 100644 index 00000000..b09f3042 --- /dev/null +++ b/document/llamaparse/client_test.go @@ -0,0 +1,957 @@ +package llamaparse + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media/mediatest" +) + +const testJobID = "123e4567-e89b-12d3-a456-426614174000" + +var _ document.ResumableRenditionProvider = (*Client)(nil) + +func TestClientUploadsExactAuthorizedBytesAndMapsNaturalPages(t *testing.T) { + source := []byte("%PDF-1.7\nsynthetic exact bytes\n%%EOF\n") + fixture := newFixture(t, source) + var requests []*http.Request + pollCount := 0 + var transport http.RoundTripper = roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests = append(requests, request.Clone(request.Context())) + assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) + switch request.URL.Path { + case uploadPath: + assert.Equal(t, http.MethodPost, request.Method) + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + assert.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, parameters["boundary"]) + fields := map[string]string{} + for { + part, nextErr := reader.NextPart() + if errors.Is(nextErr, io.EOF) { + break + } + require.NoError(t, nextErr) + payload, readErr := io.ReadAll(part) + require.NoError(t, readErr) + if part.FormName() == "file" { + assert.Equal(t, "synthetic.pdf", part.FileName()) + assert.Equal(t, "application/pdf", part.Header.Get("Content-Type")) + assert.Equal(t, source, payload) + } else { + fields[part.FormName()] = string(payload) + } + } + assert.Equal(t, "parse-model-v1", fields["model"]) + assert.Equal(t, "document-v1", fields["preset"]) + assert.Equal(t, "0", fields["page_error_tolerance"]) + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"PENDING"}`), nil + case statusPath(testJobID): + pollCount++ + if pollCount == 1 { + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"PENDING"}`), nil + } + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + return response(request, http.StatusOK, `{ + "pages":[ + {"page":0,"text":"First page","md":"# First page","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}, + {"page":1,"text":"Second page","md":"Second page","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false} + ], + "job_metadata":{"job_pages":2} + }`), nil + default: + t.Fatalf("unexpected route %s", request.URL.String()) + return nil, errors.New("unexpected route") + } + }) + client := fixture.client(t, transport) + var checkpoint document.RenditionResumeHandle + result, err := document.RenderRenditionWithResume( + t.Context(), client, fixture.upload(), fixture.authorization, nil, + func(handle document.RenditionResumeHandle) error { + checkpoint = handle + return nil + }, + ) + require.NoError(t, err) + + assert.True(t, strings.HasPrefix(checkpoint.Value, "lp1."+testJobID+".")) + handleParts := strings.Split(checkpoint.Value, ".") + require.Len(t, handleParts, 4) + checkpointNanos, err := strconv.ParseInt(handleParts[3], 10, 64) + require.NoError(t, err) + completedAt, err := time.Parse(timeForm, result.Receipt.CompletedAt) + require.NoError(t, err) + assert.True(t, completedAt.After(time.Unix(0, checkpointNanos)), + "fresh completion must follow the durable checkpoint") + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, int64(0), result.Evidence.Units[0].Locator.Start) + assert.Equal(t, "# First page", result.Evidence.Units[0].Text) + assert.Equal(t, "# First page\n\n---\n\nSecond page", string(result.ProviderMarkdown)) + assert.Equal(t, int64(1), result.Receipt.Usage.Retries) + assert.Equal(t, testJobID[:8], result.Receipt.OperationID[len("llamaparse-"):]) + require.Len(t, requests, 4) + for _, request := range requests { + assert.Equal(t, "https", request.URL.Scheme) + assert.Equal(t, apiHost, request.URL.Host) + assert.Empty(t, request.URL.RawQuery) + } + assert.Equal(t, 4, fixture.secrets.calls()) +} + +func TestClientResumesWithoutUploadingSource(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nresume\n%%EOF\n")) + var paths []string + var transport http.RoundTripper = roundTripFunc(func(request *http.Request) (*http.Response, error) { + paths = append(paths, request.URL.Path) + switch request.URL.Path { + case statusPath(testJobID): + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + return response(request, http.StatusOK, + `{"pages":[{"page":0,"text":"Resumed","md":"Resumed","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`), nil + default: + t.Fatalf("unexpected route %s", request.URL.Path) + return nil, errors.New("unexpected route") + } + }) + client := fixture.client(t, transport) + now := time.Now().UTC() + result, err := client.RenderResumable(t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: testResumeHandle(now.Add(-time.Second), now)}, nil) + require.NoError(t, err) + assert.Equal(t, []string{statusPath(testJobID), jsonResultPath(testJobID)}, paths) + assert.Equal(t, "Resumed", result.Evidence.Units[0].Text) +} + +func TestClientRejectsInvalidResumeFactsWithoutEgress(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\ninvalid resume\n%%EOF\n")) + authorizedAt, err := time.Parse(timeForm, fixture.authorization.AuthorizedAt) + require.NoError(t, err) + expiresAt, err := time.Parse(timeForm, fixture.authorization.ExpiresAt) + require.NoError(t, err) + validSubmittedAt := authorizedAt.Add(time.Nanosecond) + validCheckpointedAt := validSubmittedAt.Add(time.Nanosecond) + tests := []struct { + name string + handle string + want document.RenditionErrorCode + }{ + { + name: "submitted before authorization", + handle: testResumeHandle(authorizedAt.Add(-time.Nanosecond), validCheckpointedAt), + want: document.RenditionErrorPolicyRejected, + }, + { + name: "checkpoint after expiry", + handle: testResumeHandle(validSubmittedAt, expiresAt.Add(time.Nanosecond)), + want: document.RenditionErrorPolicyRejected, + }, + { + name: "noncanonical timestamp", + handle: fmt.Sprintf("lp1.%s.0%d.%d", testJobID, validSubmittedAt.UnixNano(), validCheckpointedAt.UnixNano()), + want: document.RenditionErrorUnknownJob, + }, + { + name: "unparseable timestamp", + handle: fmt.Sprintf("lp1.%s.not-a-time.%d", testJobID, validCheckpointedAt.UnixNano()), + want: document.RenditionErrorUnknownJob, + }, + { + name: "wrong version", + handle: strings.Replace(testResumeHandle(validSubmittedAt, validCheckpointedAt), "lp1.", "lp2.", 1), + want: document.RenditionErrorUnknownJob, + }, + {name: "legacy bare job ID", handle: testJobID, want: document.RenditionErrorUnknownJob}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("invalid resume handle reached egress") + return nil, errors.New("unexpected egress") + })) + + _, err := client.RenderResumable(t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: testCase.handle}, nil) + + assertCode(t, err, testCase.want) + assert.Zero(t, fixture.secrets.calls()) + }) + } +} + +func TestClientPreservesExplicitMiddleBlankPage(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nblank page\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: `{"pages":[ + {"page":0,"md":"First","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}, + {"page":1,"md":"","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":true,"triggeredAutoMode":false}, + {"page":2,"md":"Third","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false} + ],"job_metadata":{"job_pages":3}}`}, + }) + + result, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 3) + assert.Empty(t, result.Evidence.Units[1].Text) + assert.Equal(t, int64(1), result.Evidence.Units[1].Locator.Start) +} + +func TestClientUsesExplicitDegradedMarkdownOnlyWithoutPageProvenance(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nfallback\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: `{"pages":[],"job_metadata":{"job_pages":0}}`}, + markdownResultPath(testJobID): { + body: `{"markdown":"# Fallback markdown","job_metadata":{"job_pages":0}}`, + }, + }) + result, err := document.RenderRenditionWithResume( + t.Context(), fixture.client(t, transport), fixture.upload(), fixture.authorization, + nil, func(document.RenditionResumeHandle) error { return nil }) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + assert.Equal(t, []string{"degraded_provenance"}, result.Receipt.Warnings) +} + +func TestClientRejectsContradictoryFallbackPageCounts(t *testing.T) { + tests := []struct { + name string + jsonPages int + markdownPages int + }{ + {name: "structured result reports missing pages", jsonPages: 2, markdownPages: 2}, + {name: "fallback changes the reported count", jsonPages: 0, markdownPages: 1}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\ncontradictory fallback\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): { + body: fmt.Sprintf(`{"pages":[],"job_metadata":{"job_pages":%d}}`, testCase.jsonPages), + }, + markdownResultPath(testJobID): { + body: fmt.Sprintf(`{"markdown":"Fallback","job_metadata":{"job_pages":%d}}`, + testCase.markdownPages), + }, + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestClientRejectsPartialPagesAndSchemaDrift(t *testing.T) { + tests := []struct { + name string + body string + }{ + {name: "partial page status", body: `{"pages":[{"page":0,"md":"partial","status":"ERROR","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + {name: "page gap", body: `{"pages":[{"page":1,"md":"gap","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + {name: "unknown page field", body: `{"pages":[{"page":0,"md":"drift","provider_url":"https://evil.example/result","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + {name: "provider page count exceeds result", body: `{"pages":[{"page":0,"md":"partial","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":2}}`}, + {name: "missing explicit page index", body: `{"pages":[{"md":"invented zero","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + {name: "missing provider page count", body: `{"pages":[{"page":0,"md":"unproven","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{}}`}, + {name: "unexplained empty page", body: `{"pages":[{"page":0,"md":"","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\npartial\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: testCase.body}, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + assert.NotContains(t, err.Error(), "evil.example") + }) + } +} + +func TestClientRejectsReportedModelDrift(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nmodel drift\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: `{"pages":[{"page":0,"md":"drift","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1,"model":"parse-model-next","preset":"document-v1"}}`}, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorPolicyRejected) +} + +func TestClientClassifiesHostedFailuresWithoutLeakingProviderBodies(t *testing.T) { + tests := []struct { + name string + status int + want document.RenditionErrorCode + }{ + {name: "authentication", status: http.StatusUnauthorized, want: document.RenditionErrorAuthentication}, + {name: "rate", status: http.StatusTooManyRequests, want: document.RenditionErrorRateLimited}, + {name: "capacity", status: http.StatusServiceUnavailable, want: document.RenditionErrorCapacity}, + {name: "terminal input", status: http.StatusUnsupportedMediaType, want: document.RenditionErrorUnsupportedInput}, + {name: "redirect", status: http.StatusTemporaryRedirect, want: document.RenditionErrorPolicyRejected}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nerrors\n%%EOF\n")) + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := response(request, testCase.status, + `{"detail":"provider body secret synthetic-source"}`) + response.Header.Set("Location", "https://evil.example/stolen") + return response, nil + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, nil) + assertCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "provider body") + assert.NotContains(t, err.Error(), "evil.example") + }) + } +} + +func TestClientClassifiesAmbiguousSubmissionAndUnknownJobs(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nambiguous\n%%EOF\n")) + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("private transport failure") + })) + _, err := document.RenderRenditionWithResume(t.Context(), client, fixture.upload(), + fixture.authorization, nil, nil) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) + assert.NotContains(t, err.Error(), "private transport") + + for _, status := range []int{http.StatusNotFound, http.StatusGone} { + client = fixture.client(t, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return response(request, status, `{"detail":"private unknown detail"}`), nil + })) + now := time.Now().UTC() + _, err = client.RenderResumable(t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: testResumeHandle(now.Add(-time.Second), now)}, nil) + assertCode(t, err, document.RenditionErrorUnknownJob) + } +} + +func TestClientPreservesPreEgressCredentialFailureClassification(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nauth failure\n%%EOF\n")) + fixture.secrets.value = "" + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("credential failure reached egress") + return nil, errors.New("unexpected egress") + })) + + _, err := document.RenderRenditionWithResume(t.Context(), client, fixture.upload(), + fixture.authorization, nil, nil) + + assertCode(t, err, document.RenditionErrorAuthentication) +} + +func TestClientPreservesCredentialFailureClassificationAcrossResultStages(t *testing.T) { + t.Run("poll", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\npoll credential\n%%EOF\n")) + fixture.secrets.failAt = 2 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.URL.Path != uploadPath { + t.Fatalf("credential failure reached %s egress", request.URL.Path) + } + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorAuthentication) + }) + + t.Run("result", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nresult credential\n%%EOF\n")) + fixture.secrets.failAt = 3 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath, statusPath(testJobID): + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + default: + t.Fatalf("credential failure reached %s egress", request.URL.Path) + return nil, errors.New("unexpected egress") + } + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorAuthentication) + }) + + t.Run("image", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nimage credential\n%%EOF\n")) + fixture.profile.RetainImages = true + fixture.profile.MaxArtifacts = 1 + fixture.authorization.AllowedArtifactRoles = []document.EvidenceArtifactRole{document.EvidenceArtifactImage} + fixture.authorization.MaxArtifacts = 1 + fixture.authorization.MaxArtifactBytes = 1024 + fixture.secrets.failAt = 4 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath, statusPath(testJobID): + return response(request, http.StatusOK, + `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + return response(request, http.StatusOK, `{"pages":[{"page":0,"md":"Image page","images":[{"name":"figure-1.png"}],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`), nil + default: + t.Fatalf("credential failure reached %s egress", request.URL.Path) + return nil, errors.New("unexpected egress") + } + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorAuthentication) + }) +} + +func TestClientRefusesProviderAuthoredArtifactURLsAndUsesFixedImageRoute(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nimage\n%%EOF\n")) + fixture.profile.RetainImages = true + fixture.profile.MaxArtifacts = 1 + fixture.authorization.AllowedArtifactRoles = []document.EvidenceArtifactRole{document.EvidenceArtifactImage} + fixture.authorization.MaxArtifacts = 1 + fixture.authorization.MaxArtifactBytes = 128 + var paths []string + var transport http.RoundTripper = roundTripFunc(func(request *http.Request) (*http.Response, error) { + paths = append(paths, request.URL.Path) + switch request.URL.Path { + case uploadPath: + return response(request, http.StatusOK, `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case statusPath(testJobID): + return response(request, http.StatusOK, `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + return response(request, http.StatusOK, `{"pages":[{"page":0,"md":"Image page","images":[{"name":"figure-1.png"}],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`), nil + case imageResultPath(testJobID, "figure-1.png"): + result := bytesResponse(request, http.StatusOK, mediatest.PNG(4, 3, nil)) + result.Header.Set("Content-Type", "image/png") + return result, nil + default: + t.Fatalf("unexpected route %s", request.URL.String()) + return nil, errors.New("unexpected route") + } + }) + result, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + require.NoError(t, err) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, mediatest.PNG(4, 3, nil), result.Artifacts[0].Payload) + assert.Equal(t, imageResultPath(testJobID, "figure-1.png"), paths[len(paths)-1]) + + for _, name := range []string{"https://evil.example/a.png", ".", ".."} { + paths = nil + transport = routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: `{"pages":[{"page":0,"md":"unsafe","images":[{"name":"` + name + `"}],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`}, + }) + _, err = document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + } +} + +func TestClientRejectsMalformedOrMismatchedRetainedImages(t *testing.T) { + for _, testCase := range []struct { + name string + mediaType string + payload []byte + }{ + {name: "malformed", mediaType: "image/png", payload: []byte("not a PNG")}, + {name: "mismatched", mediaType: "image/jpeg", payload: mediatest.PNG(4, 3, nil)}, + {name: "animated", mediaType: "image/gif", payload: mediatest.GIF(2, 2, 2)}, + } { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nimage validation\n%%EOF\n")) + fixture.profile.RetainImages = true + fixture.profile.MaxArtifacts = 1 + fixture.authorization.AllowedArtifactRoles = []document.EvidenceArtifactRole{document.EvidenceArtifactImage} + fixture.authorization.MaxArtifacts = 1 + fixture.authorization.MaxArtifactBytes = 1024 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath, statusPath(testJobID): + return response(request, http.StatusOK, `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + return response(request, http.StatusOK, `{"pages":[{"page":0,"md":"Image page","images":[{"name":"figure-1.png"}],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`), nil + case imageResultPath(testJobID, "figure-1.png"): + result := bytesResponse(request, http.StatusOK, testCase.payload) + result.Header.Set("Content-Type", testCase.mediaType) + return result, nil + default: + return nil, errors.New("unexpected route") + } + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestClientEnforcesUploadPollResultAndArtifactBounds(t *testing.T) { + t.Run("upload identity", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nidentity\n%%EOF\n")) + upload := &testUpload{Reader: bytes.NewReader([]byte("different")), metadata: fixture.metadata} + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, + roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("mismatched source reached egress") + return nil, errors.New("mismatched source reached egress") + })), upload, fixture.authorization, nil, nil) + assertCode(t, err, document.RenditionErrorPolicyRejected) + }) + + t.Run("poll limit", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\npoll\n%%EOF\n")) + fixture.profile.MaxPolls = 2 + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"PENDING"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"PENDING"}`}, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorTransient) + }) + + t.Run("result bytes", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nresult\n%%EOF\n")) + fixture.profile.MaxResultBytes = 64 + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: strings.Repeat("x", 65)}, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) +} + +func TestClientClassifiesExactUploadReadLifecycle(t *testing.T) { + tests := []struct { + name string + configure func(*fixture) + duringRead func(context.CancelFunc) + want document.RenditionErrorCode + }{ + { + name: "caller cancellation", + duringRead: func(cancel context.CancelFunc) { + cancel() + }, + want: document.RenditionErrorCanceled, + }, + { + name: "authorization expiry", + configure: func(fixture *fixture) { + fixture.authorization.ExpiresAt = time.Now().UTC().Add(20 * time.Millisecond).Format(timeForm) + }, + duringRead: func(context.CancelFunc) { + time.Sleep(40 * time.Millisecond) + }, + want: document.RenditionErrorPolicyRejected, + }, + { + name: "wall timeout", + configure: func(fixture *fixture) { + fixture.profile.MaxWallTime = 20 * time.Millisecond + }, + duringRead: func(context.CancelFunc) { + time.Sleep(40 * time.Millisecond) + }, + want: document.RenditionErrorTransient, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nslow upload\n%%EOF\n")) + if testCase.configure != nil { + testCase.configure(fixture) + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + upload := &testUpload{ + Reader: &callbackReadCloser{ + reader: bytes.NewReader(fixture.source), + before: func() { testCase.duringRead(cancel) }, + }, + metadata: fixture.metadata, + } + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("expired upload reached egress") + return nil, errors.New("unexpected egress") + })) + + _, err := document.RenderRenditionWithResume(ctx, client, upload, + fixture.authorization, nil, nil) + + assertCode(t, err, testCase.want) + }) + } +} + +func TestClientClassifiesCancellationAndAuthorizationExpiry(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\ncancel\n%%EOF\n")) + ctx, cancel := context.WithCancel(t.Context()) + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + if request.URL.Path == uploadPath { + return response(request, http.StatusOK, `{"id":"`+testJobID+`","status":"PENDING"}`), nil + } + cancel() + <-request.Context().Done() + return nil, request.Context().Err() + }) + _, err := document.RenderRenditionWithResume(ctx, fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorCanceled) +} + +func TestClientResumesHistoricalSealedAuthorizationWithRecordedReceiptTimes(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nhistorical resume\n%%EOF\n")) + transport := routeTransport(t, map[string]routeResponse{ + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): { + body: `{"pages":[{"page":0,"md":"Historical","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}`, + }, + }) + client := fixture.client(t, transport) + historical := time.Now().UTC().Add(-time.Hour) + fixture.authorization.AuthorizedAt = historical.Format(timeForm) + fixture.authorization.ExpiresAt = historical.Add(time.Minute).Format(timeForm) + submittedAt := historical.Add(10 * time.Second) + checkpointedAt := historical.Add(11 * time.Second) + evidencePolicy, err := document.NewEvidencePolicy(100_000) + require.NoError(t, err) + normalizePolicy, err := document.NewNormalizePolicy(100_000) + require.NoError(t, err) + renditionPolicy, err := document.NewRenditionPolicy(normalizePolicy, 1_000) + require.NoError(t, err) + snapshot, err := document.SealRenditionExecutionAt(historical.Add(time.Second), client, + fixture.upload(), fixture.authorization, evidencePolicy, renditionPolicy) + require.NoError(t, err) + + result, err := document.ResumeRendition(t.Context(), client, snapshot, + document.RenditionResumeHandle{Value: testResumeHandle(submittedAt, checkpointedAt)}, nil) + + require.NoError(t, err) + assert.Equal(t, submittedAt.Format(timeForm), result.Receipt.StartedAt) + assert.Equal(t, checkpointedAt.Format(timeForm), result.Receipt.CompletedAt) +} + +func TestClientRechecksLifecycleAfterResponseBodiesAndBeforeReturn(t *testing.T) { + const resultBody = `{"pages":[{"page":0,"md":"Complete","images":[],"charts":[],"tables":[],"layout":[],"items":[],"links":[],"parsingMode":"parse_page","noStructuredContent":true,"noTextContent":false,"triggeredAutoMode":false}],"job_metadata":{"job_pages":1}}` + + t.Run("caller cancellation after complete body", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nbody cancel\n%%EOF\n")) + ctx, cancel := context.WithCancel(t.Context()) + transport := resultBodyTransport(t, resultBody, func(request *http.Request) io.ReadCloser { + return &callbackReadCloser{reader: strings.NewReader(resultBody), before: cancel} + }) + + _, err := document.RenderRenditionWithResume(ctx, fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorCanceled) + }) + + t.Run("wall timeout after complete body", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nbody timeout\n%%EOF\n")) + fixture.profile.MaxWallTime = 20 * time.Millisecond + transport := resultBodyTransport(t, resultBody, func(request *http.Request) io.ReadCloser { + return &callbackReadCloser{reader: strings.NewReader(resultBody), before: func() { time.Sleep(40 * time.Millisecond) }} + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorTransient) + }) + + t.Run("authorization expiry after complete body", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nbody expiry\n%%EOF\n")) + expiresAt := time.Now().UTC().Add(20 * time.Millisecond) + fixture.authorization.ExpiresAt = expiresAt.Format(timeForm) + transport := resultBodyTransport(t, resultBody, func(request *http.Request) io.ReadCloser { + return &callbackReadCloser{reader: strings.NewReader(resultBody), before: func() { + time.Sleep(time.Until(expiresAt) + 10*time.Millisecond) + }} + }) + + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorPolicyRejected) + }) + + t.Run("caller cancellation at final acceptance", func(t *testing.T) { + fixture := newFixture(t, []byte("%PDF-1.7\nfinal cancel\n%%EOF\n")) + ctx := newCancelWhenCheckedContext(t.Context()) + defer ctx.cancel() + transport := routeTransport(t, map[string]routeResponse{ + uploadPath: {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + statusPath(testJobID): {body: `{"id":"` + testJobID + `","status":"SUCCESS"}`}, + jsonResultPath(testJobID): {body: resultBody}, + }) + + _, err := fixture.client(t, transport).RenderResumable(ctx, fixture.upload(), + fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + + assertCode(t, err, document.RenditionErrorCanceled) + }) +} + +type fixture struct { + profile Profile + secrets *testSecrets + metadata document.AuthorizedUploadMetadata + source []byte + authorization document.RenditionAuthorization +} + +func newFixture(t *testing.T, source []byte) *fixture { + t.Helper() + profile := Profile{ + Model: "parse-model-v1", Preset: "document-v1", SecretBinding: "llamaparse-production", + MaxUploadBytes: 1024, MaxRequestBytes: 4096, MaxControlBytes: 1024, + MaxPolls: 3, PollInterval: time.Millisecond, RequestTimeout: time.Second, + MaxResultBytes: 32 << 10, MaxArtifactBytes: 1024, MaxArtifacts: 0, + MaxWallTime: time.Second, + } + secrets := &testSecrets{value: "synthetic-secret"} + client, err := NewProvider(profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("unused fixture transport") + })) + require.NoError(t, err) + descriptor := client.Descriptor() + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: "synthetic.pdf", MediaFamily: "pdf", MediaType: "application/pdf", + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("1", 64), + ProviderMetadataChecksum: strings.Repeat("2", 64), + InputKind: document.RenditionInputOriginalFile, + } + now := time.Now().UTC() + return &fixture{ + profile: profile, secrets: secrets, metadata: metadata, source: bytes.Clone(source), + authorization: document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("3", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, + InputKind: metadata.InputKind, MaxProviderMarkdownBytes: 16 << 10, + MaxTotalResultBytes: 64 << 10, + AuthorizedAt: now.Add(-time.Second).Format(timeForm), + ExpiresAt: now.Add(time.Minute).Format(timeForm), + }, + } +} + +func (fixture *fixture) client(t *testing.T, transport http.RoundTripper) *Client { + t.Helper() + client, err := NewProvider(fixture.profile, fixture.secrets, transport) + require.NoError(t, err) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + return client +} + +func (fixture *fixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +type testUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } +func (upload *testUpload) Close() error { return nil } + +type testSecrets struct { + mu sync.Mutex + value string + count int + failAt int +} + +type callbackReadCloser struct { + reader io.Reader + once sync.Once + before func() +} + +func (body *callbackReadCloser) Read(buffer []byte) (int, error) { + body.once.Do(body.before) + return body.reader.Read(buffer) +} + +func (*callbackReadCloser) Close() error { return nil } + +type cancelWhenCheckedContext struct { + context.Context + + done chan struct{} + once sync.Once +} + +func newCancelWhenCheckedContext(parent context.Context) *cancelWhenCheckedContext { + return &cancelWhenCheckedContext{Context: parent, done: make(chan struct{})} +} + +func (ctx *cancelWhenCheckedContext) Done() <-chan struct{} { return ctx.done } + +func (ctx *cancelWhenCheckedContext) Err() error { + ctx.cancel() + return context.Canceled +} + +func (ctx *cancelWhenCheckedContext) cancel() { + ctx.once.Do(func() { close(ctx.done) }) +} + +func (secrets *testSecrets) ResolveSecret(context.Context, string) (string, error) { + secrets.mu.Lock() + defer secrets.mu.Unlock() + secrets.count++ + if secrets.count == secrets.failAt { + return "", errors.New("synthetic credential failure") + } + return secrets.value, nil +} + +func (secrets *testSecrets) calls() int { + secrets.mu.Lock() + defer secrets.mu.Unlock() + return secrets.count +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type routeResponse struct { + status int + body string +} + +func routeTransport(t *testing.T, routes map[string]routeResponse) http.RoundTripper { + t.Helper() + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + route, ok := routes[request.URL.Path] + if !ok { + t.Fatalf("unexpected route %s", request.URL.String()) + } + status := route.status + if status == 0 { + status = http.StatusOK + } + return response(request, status, route.body), nil + }) +} + +func resultBodyTransport( + t *testing.T, resultBody string, body func(*http.Request) io.ReadCloser, +) http.RoundTripper { + t.Helper() + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath, statusPath(testJobID): + return response(request, http.StatusOK, `{"id":"`+testJobID+`","status":"SUCCESS"}`), nil + case jsonResultPath(testJobID): + result := response(request, http.StatusOK, resultBody) + result.Body = body(request) + return result, nil + default: + t.Fatalf("unexpected route %s", request.URL.String()) + return nil, errors.New("unexpected route") + } + }) +} + +func response(request *http.Request, status int, body string) *http.Response { + return bytesResponse(request, status, []byte(body)) +} + +func bytesResponse(request *http.Request, status int, body []byte) *http.Response { + return &http.Response{ + StatusCode: status, Status: http.StatusText(status), Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), Request: request, + } +} + +func testResumeHandle(submittedAt, checkpointedAt time.Time) string { + return fmt.Sprintf("lp1.%s.%d.%d", testJobID, submittedAt.UnixNano(), checkpointedAt.UnixNano()) +} + +func assertCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerError, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok, "expected classified provider error, got %T: %v", err, err) + assert.Equal(t, want, providerError.Code()) +} diff --git a/document/llamaparse/doc.go b/document/llamaparse/doc.go new file mode 100644 index 00000000..ad4ec10b --- /dev/null +++ b/document/llamaparse/doc.go @@ -0,0 +1,3 @@ +// Package llamaparse adapts the fixed hosted LlamaParse v1 API to Docbank's +// storage-neutral, resumable rendition provider contract. +package llamaparse diff --git a/document/marker/client.go b/document/marker/client.go new file mode 100644 index 00000000..7e830c56 --- /dev/null +++ b/document/marker/client.go @@ -0,0 +1,755 @@ +// Package marker implements the fixed uploaded-file self-hosted Marker flow. +package marker + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/jsontext" + "encoding/json/v2" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "net/url" + "path/filepath" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/formatdetect" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + providerID = "marker.self-hosted-v1" + uploadPath = "/marker/upload" + adapterContract = "marker-self-hosted-adapter/v1" + timestampForm = "2006-01-02T15:04:05.000000000Z" + defaultRequestTimeout = 10 * time.Minute + defaultMaxDocumentBytes = int64(200 << 20) + defaultMaxRequestBytes = int64(201 << 20) + defaultMaxResponseBytes = int64(64 << 20) + defaultMaxMetadataBytes = int64(4 << 20) + defaultMaxImages = 64 + defaultMaxImageBytes = int64(16 << 20) + defaultMaxUnits = 10_000 + maxTimeout = 24 * time.Hour + maxDocumentBytes = int64(1 << 30) + maxRequestBytes = int64(1<<30) + (2 << 20) + maxResponseBytes = int64(512 << 20) + maxMetadataBytes = int64(64 << 20) + maxImages = 64 + maxImageBytes = int64(256 << 20) + maxUnits = 1_000_000 + maxSecretBytes = 64 << 10 + pageSeparator = "------------------------------------------------" +) + +var _ document.RenditionProvider = (*Client)(nil) + +// SecretResolver resolves an optional operator-fronted Marker credential. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile pins the operator deployment, runtime, credential name, conversion +// mode, fixed wire contract, and every request/result bound. +type Profile struct { + Origin string + Descriptor document.RenditionDescriptor + DeploymentFingerprint string + RuntimeFingerprint string + SecretBinding string + Mode string + RequestTimeout time.Duration + MaxDocumentBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + MaxMetadataBytes int64 + MaxImages int + MaxImageBytes int64 + MaxUnits int +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + DeploymentFingerprint string `json:"deployment_fingerprint"` + RuntimeFingerprint string `json:"runtime_fingerprint"` + CredentialBinding string `json:"credential_binding"` + Mode string `json:"mode"` + OutputFormat string `json:"output_format"` + ForceOCR bool `json:"force_ocr"` + PaginateOutput bool `json:"paginate_output"` + RequestTimeoutNanos int64 `json:"request_timeout_nanos"` + MaxDocumentBytes int64 `json:"max_document_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + MaxMetadataBytes int64 `json:"max_metadata_bytes"` + MaxImages int `json:"max_images"` + MaxImageBytes int64 `json:"max_image_bytes"` + MaxUnits int `json:"max_units"` + SupportedFormats []document.RenditionFormatCapability `json:"supported_formats"` +} + +// SupportedFormats returns the exact Marker families for which Docbank has +// bounded original-file upload authority. +func SupportedFormats() []document.RenditionFormatCapability { + return []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/webp", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/gif", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/html", InputKind: document.RenditionInputOriginalFile}, + } +} + +// PolicyFingerprint returns the canonical profile identity expected in the +// rendition descriptor. Credential values and transports are never included. +func PolicyFingerprint(profile Profile) (string, error) { + normalized, err := normalizeProfile(profile) + if err != nil { + return "", err + } + identity := policyIdentity{AdapterContract: adapterContract, Origin: normalized.Origin, + Route: uploadPath, DeploymentFingerprint: normalized.DeploymentFingerprint, + RuntimeFingerprint: normalized.RuntimeFingerprint, CredentialBinding: normalized.SecretBinding, + Mode: normalized.Mode, OutputFormat: "markdown", ForceOCR: false, PaginateOutput: true, + RequestTimeoutNanos: int64(normalized.RequestTimeout), MaxDocumentBytes: normalized.MaxDocumentBytes, + MaxRequestBytes: normalized.MaxRequestBytes, MaxResponseBytes: normalized.MaxResponseBytes, + MaxMetadataBytes: normalized.MaxMetadataBytes, MaxImages: normalized.MaxImages, + MaxImageBytes: normalized.MaxImageBytes, MaxUnits: normalized.MaxUnits, + SupportedFormats: SupportedFormats()} + encoded, err := json.Marshal(identity, json.Deterministic(true)) + if err != nil { + return "", fmt.Errorf("marker: encode policy identity: %w", err) + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +// Client renders exact authorized uploads through one fixed Marker route. +type Client struct { + profile Profile + descriptor document.RenditionDescriptor + secrets SecretResolver + http *http.Client +} + +// New validates a pinned self-hosted profile. The injected transport is the +// operator's hardened destination policy; Client removes ambient cookies and +// always refuses redirects. +func New(profile Profile, secrets SecretResolver, transport http.RoundTripper) (*Client, error) { + normalized, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewRenditionDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("marker: invalid descriptor: %w", err) + } + if descriptor.ID != providerID || descriptor.TrustBoundary != document.RenditionTrustOperatorNetwork || + !descriptor.ReturnsMarkdown || descriptor.ReturnsStructured || len(descriptor.ArtifactRoles) != 0 { + return nil, errors.New("marker: descriptor result or trust contract is invalid") + } + wantFormats := SupportedFormats() + slices.SortFunc(wantFormats, compareFormats) + if !slices.Equal(descriptor.SupportedFormats, wantFormats) { + return nil, errors.New("marker: descriptor must advertise the exact supported format set") + } + fingerprint, err := PolicyFingerprint(normalized) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("marker: descriptor policy fingerprint does not match profile") + } + if normalized.SecretBinding == "" { + if !nilValue(secrets) { + return nil, errors.New("marker: secret resolver requires a named binding") + } + } else if nilValue(secrets) { + return nil, errors.New("marker: named secret binding requires a resolver") + } + if transport == nil { + return nil, errors.New("marker: hardened transport is required") + } + normalized.Descriptor = descriptor + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, + http: &http.Client{Transport: transport, CheckRedirect: providerhttp.RefuseRedirects}}, nil +} + +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +func (client *Client) Render(ctx context.Context, upload document.AuthorizedUpload, authorization document.RenditionAuthorization) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("marker: client is required") + } + if _, err := document.ValidateRenditionProviderRequest(client, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > client.profile.MaxDocumentBytes { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, "Marker input exceeds the document byte limit", nil) + } + if !filenameMatches(metadata) { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, "Marker input filename does not match the authorized format", nil) + } + expiresAt, err := time.Parse(timestampForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, "Marker authorization expiry is invalid", nil) + } + operationCtx, cancel := operationContext(ctx, expiresAt, client.profile.RequestTimeout) + defer cancel() + if err := checkOperation(operationCtx, expiresAt); err != nil { + return document.RenditionResult{}, err + } + source, err := readExact(operationCtx, upload, metadata) + if err != nil { + return document.RenditionResult{}, err + } + defer clear(source) + expectedNaturalUnits := 0 + switch metadata.MediaFamily { + case "pdf": + pages, countErr := formatdetect.CountPDFPages(source) + if countErr != nil || pages <= 0 { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, "Marker PDF page authority is invalid", countErr) + } + if pages > int64(client.profile.MaxUnits) { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, "Marker PDF exceeds the unit limit", nil) + } + expectedNaturalUnits = int(pages) + case "image": + detected, detectErr := media.DetectBytes(source, metadata.MediaType) + if detectErr != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, "Marker image authority is invalid", detectErr) + } + if detected.Kind != media.KindImage || detected.MediaType != metadata.MediaType { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, "Marker image identity does not match the authorization", nil) + } + if detected.FrameCount != 1 || detected.Animated { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, "Marker requires a single-frame image", nil) + } + expectedNaturalUnits = 1 + } + body, contentType, err := buildMultipart(metadata, source, client.profile.Mode, client.profile.MaxRequestBytes) + if err != nil { + return document.RenditionResult{}, err + } + started := time.Now().UTC() + request, err := http.NewRequestWithContext(operationCtx, http.MethodPost, client.profile.Origin+uploadPath, bytes.NewReader(body)) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorTransient, "could not prepare Marker request", err) + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", contentType) + if err := client.authorize(request); err != nil { + return document.RenditionResult{}, err + } + if err := checkOperation(operationCtx, expiresAt); err != nil { + return document.RenditionResult{}, err + } + response, err := client.http.Do(request) + if err != nil { + if operationErr := checkOperation(operationCtx, expiresAt); operationErr != nil { + return document.RenditionResult{}, operationErr + } + return document.RenditionResult{}, providerError(document.RenditionErrorAmbiguousSubmission, "Marker submission outcome is unknown", err) + } + defer func() { _ = response.Body.Close() }() + responseBody, err := readBounded(operationCtx, expiresAt, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return document.RenditionResult{}, err + } + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return document.RenditionResult{}, statusError(response.StatusCode) + } + mediaType, _, mediaErr := mime.ParseMediaType(response.Header.Get("Content-Type")) + if mediaErr != nil || mediaType != "application/json" { + return document.RenditionResult{}, malformedError("Marker response content type is invalid", mediaErr) + } + if len(responseBody) > authorization.MaxTotalResultBytes { + return document.RenditionResult{}, malformedError("Marker response exceeds authorization", nil) + } + result, warnings, err := client.parseResult(responseBody, metadata.MediaFamily, expectedNaturalUnits) + if err != nil { + return document.RenditionResult{}, err + } + if len(result.markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, malformedError("Marker Markdown exceeds authorization", nil) + } + completed := time.Now().UTC() + return document.RenditionResult{Evidence: result.evidence, ProviderMarkdown: result.markdown, + Receipt: document.RenditionReceipt{ProviderID: client.descriptor.ID, + DescriptorFingerprint: client.descriptor.Fingerprint, PolicyFingerprint: authorization.PolicyFingerprint, + SourceSHA256: metadata.SHA256, OperationID: "marker-" + authorization.RenditionRequestFingerprint[:24], + StartedAt: started.Format(timestampForm), CompletedAt: completed.Format(timestampForm), Warnings: warnings, + Usage: document.RenditionUsage{Requests: 1, InputBytes: int64(len(source)), OutputBytes: int64(len(responseBody)), Units: int64(len(result.evidence.Units))}}}, nil +} + +type parsedResult struct { + markdown []byte + evidence document.SourceEvidenceV1 +} + +func (client *Client) parseResult(body []byte, family string, expectedNaturalUnits int) (parsedResult, []string, error) { + var wire struct { + Format *string `json:"format"` + Output *string `json:"output"` + Images map[string]string `json:"images"` + Metadata jsontext.Value `json:"metadata"` + Success *bool `json:"success"` + Error string `json:"error,omitempty"` + } + if err := json.Unmarshal(body, &wire, json.RejectUnknownMembers(true)); err != nil { + return parsedResult{}, nil, malformedError("Marker result JSON is invalid", err) + } + if wire.Success == nil || !*wire.Success || wire.Format == nil || *wire.Format != "markdown" || + wire.Output == nil || *wire.Output == "" || wire.Images == nil || len(wire.Metadata) == 0 || wire.Error != "" { + return parsedResult{}, nil, malformedError("Marker result is incomplete or unsuccessful", nil) + } + if len(wire.Metadata) > int(client.profile.MaxMetadataBytes) { + return parsedResult{}, nil, malformedError("Marker metadata exceeds byte limit", nil) + } + if err := validateImages(wire.Images, client.profile.MaxImages, client.profile.MaxImageBytes); err != nil { + return parsedResult{}, nil, err + } + stats, err := parseMetadata(wire.Metadata, client.profile.MaxUnits) + if err != nil { + return parsedResult{}, nil, err + } + if expectedNaturalUnits > 0 && !statsProveUnits(stats, expectedNaturalUnits) { + return parsedResult{}, nil, malformedError("Marker result does not prove complete source units", nil) + } + markdown := []byte(*wire.Output) + evidence, natural := naturalEvidence(family, *wire.Output, stats) + if !natural { + evidence = degradedEvidence(family, *wire.Output) + } + warnings := []string(nil) + if len(wire.Images) != 0 { + warnings = []string{"provider_images_not_retained"} + } + return parsedResult{markdown: markdown, evidence: evidence}, warnings, nil +} + +type pageStat struct{ PageID int } + +func statsProveUnits(stats []pageStat, expected int) bool { + if len(stats) != expected { + return false + } + for index, stat := range stats { + if stat.PageID != index { + return false + } + } + return true +} + +func parseMetadata(raw jsontext.Value, maximum int) ([]pageStat, error) { + var metadata struct { + TableOfContents jsontext.Value `json:"table_of_contents"` + PageStats []struct { + PageID *int `json:"page_id"` + TextExtractionMethod string `json:"text_extraction_method"` + BlockCounts jsontext.Value `json:"block_counts"` + BlockMetadata jsontext.Value `json:"block_metadata"` + } `json:"page_stats"` + } + if err := json.Unmarshal(raw, &metadata, json.RejectUnknownMembers(true)); err != nil { + return nil, malformedError("Marker metadata schema changed", err) + } + if len(metadata.TableOfContents) == 0 || metadata.PageStats == nil || len(metadata.PageStats) > maximum { + return nil, malformedError("Marker metadata is incomplete or exceeds the unit limit", nil) + } + stats := make([]pageStat, len(metadata.PageStats)) + for index, stat := range metadata.PageStats { + if stat.PageID == nil || *stat.PageID < 0 || len(stat.TextExtractionMethod) > 128 || + len(stat.BlockCounts) == 0 || len(stat.BlockMetadata) == 0 { + return nil, malformedError("Marker page metadata is incomplete", nil) + } + stats[index] = pageStat{PageID: *stat.PageID} + } + return stats, nil +} + +func naturalEvidence(family, markdown string, stats []pageStat) (document.SourceEvidenceV1, bool) { + if family != "pdf" && family != "image" { + return document.SourceEvidenceV1{}, false + } + parts, ok := splitPages(markdown, stats) + if !ok { + return document.SourceEvidenceV1{}, false + } + evidence := document.SourceEvidenceV1{ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceComplete, Family: family, UnitKind: document.EvidenceUnitPage, + Units: make([]document.SourceEvidenceUnitV1, len(parts))} + for index, text := range parts { + evidence.Units[index] = document.SourceEvidenceUnitV1{Order: index, + ProviderID: "marker-page-" + strconv.Itoa(index), Text: text, + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorPage, + IndexOrigin: document.EvidenceIndexOriginZero, Start: int64(index), End: int64(index)}} + } + return evidence, true +} + +func splitPages(markdown string, stats []pageStat) ([]string, bool) { + if len(stats) == 0 { + return nil, false + } + lines := strings.Split(markdown, "\n") + parts := make([]string, 0, len(stats)) + current := make([]string, 0) + seen := false + for _, line := range lines { + index, separator := pageLine(line) + if separator { + if seen { + parts = append(parts, strings.TrimSpace(strings.Join(current, "\n"))) + current = current[:0] + } + if index != len(parts) || index >= len(stats) || stats[index].PageID != index { + return nil, false + } + seen = true + continue + } + if !seen && strings.TrimSpace(line) != "" { + return nil, false + } + if seen { + current = append(current, line) + } + } + if seen { + parts = append(parts, strings.TrimSpace(strings.Join(current, "\n"))) + } + return parts, len(parts) == len(stats) +} + +func pageLine(line string) (int, bool) { + open := strings.IndexByte(line, '{') + closeIndex := strings.IndexByte(line, '}') + if open != 0 || closeIndex <= 1 || line[closeIndex+1:] != pageSeparator { + return 0, false + } + value, err := strconv.Atoi(line[1:closeIndex]) + return value, err == nil && value >= 0 +} + +func degradedEvidence(family, markdown string) document.SourceEvidenceV1 { + return document.SourceEvidenceV1{ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceDegradedProvenance, Family: family, UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{Kind: document.EvidenceOmissionField, + Field: "natural_provenance", Reason: "Marker returned no source-native unit mapping"}}, + Units: []document.SourceEvidenceUnitV1{{Order: 0, Text: markdown, + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone}}}} +} + +func validateImages(images map[string]string, maximum int, maxBytes int64) error { + if len(images) > maximum { + return malformedError("Marker returned too many images", nil) + } + for name, encoded := range images { + if name == "" || len(name) > 1024 || !utf8.ValidString(name) || strings.ContainsRune(name, 0) || + int64(base64.StdEncoding.DecodedLen(len(encoded))) > maxBytes { + return malformedError("Marker image output is invalid or exceeds limits", nil) + } + decoder := base64.NewDecoder(base64.StdEncoding.Strict(), strings.NewReader(encoded)) + count, err := io.Copy(io.Discard, io.LimitReader(decoder, maxBytes+1)) + if err != nil || count > maxBytes { + return malformedError("Marker image output is invalid or exceeds limits", err) + } + } + return nil +} + +func buildMultipart(metadata document.AuthorizedUploadMetadata, source []byte, mode string, maximum int64) ([]byte, string, error) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + header := make(textproto.MIMEHeader) + header.Set("Content-Disposition", multipart.FileContentDisposition("file", metadata.Filename)) + header.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(header) + if err == nil { + _, err = part.Write(source) + } + for _, field := range [][2]string{{"mode", mode}, {"force_ocr", "false"}, {"paginate_output", "true"}, {"output_format", "markdown"}} { + if err == nil { + err = writer.WriteField(field[0], field[1]) + } + } + if err == nil { + err = writer.Close() + } + if err != nil { + return nil, "", providerError(document.RenditionErrorTransient, "could not prepare Marker upload", err) + } + if int64(body.Len()) > maximum { + return nil, "", providerError(document.RenditionErrorPolicyRejected, "Marker multipart request exceeds byte limit", nil) + } + return body.Bytes(), writer.FormDataContentType(), nil +} + +func (client *Client) authorize(request *http.Request) error { + if client.profile.SecretBinding == "" { + return nil + } + secret, err := client.secrets.ResolveSecret(request.Context(), client.profile.SecretBinding) + if err != nil || secret == "" || len(secret) > maxSecretBytes || strings.ContainsAny(secret, "\r\n\x00") { + return providerError(document.RenditionErrorAuthentication, "Marker credential is unavailable or invalid", err) + } + request.Header.Set("Authorization", "Bearer "+secret) + return nil +} + +func normalizeProfile(profile Profile) (Profile, error) { + origin, err := validateOrigin(profile.Origin) + if err != nil { + return Profile{}, err + } + profile.Origin = origin + if !validFingerprint(profile.DeploymentFingerprint) || !validFingerprint(profile.RuntimeFingerprint) { + return Profile{}, errors.New("marker: deployment and runtime fingerprints must be lowercase SHA-256") + } + if profile.SecretBinding != "" && !validToken(profile.SecretBinding) { + return Profile{}, errors.New("marker: secret binding is invalid") + } + if profile.Mode != "fast" && profile.Mode != "balanced" { + return Profile{}, errors.New("marker: mode must be fast or balanced") + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.MaxDocumentBytes == 0 { + profile.MaxDocumentBytes = defaultMaxDocumentBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultMaxRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.MaxMetadataBytes == 0 { + profile.MaxMetadataBytes = defaultMaxMetadataBytes + } + if profile.MaxImages == 0 { + profile.MaxImages = defaultMaxImages + } + if profile.MaxImageBytes == 0 { + profile.MaxImageBytes = defaultMaxImageBytes + } + if profile.MaxUnits == 0 { + profile.MaxUnits = defaultMaxUnits + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxTimeout || profile.MaxDocumentBytes <= 0 || profile.MaxDocumentBytes > maxDocumentBytes || + profile.MaxRequestBytes <= profile.MaxDocumentBytes || profile.MaxRequestBytes > maxRequestBytes || + profile.MaxResponseBytes <= 0 || profile.MaxResponseBytes > maxResponseBytes || profile.MaxMetadataBytes <= 0 || profile.MaxMetadataBytes > maxMetadataBytes || + profile.MaxMetadataBytes > profile.MaxResponseBytes || profile.MaxImages < 1 || profile.MaxImages > maxImages || + profile.MaxImageBytes <= 0 || profile.MaxImageBytes > maxImageBytes || profile.MaxImageBytes > profile.MaxResponseBytes || + profile.MaxUnits < 1 || profile.MaxUnits > maxUnits { + return Profile{}, errors.New("marker: execution bounds are invalid") + } + return profile, nil +} + +func validateOrigin(raw string) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Opaque != "" || parsed.ForceQuery || parsed.Fragment != "" || + (parsed.Path != "" && parsed.Path != "/") || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return "", errors.New("marker: origin must be one HTTP(S) origin without path, credentials, query, or fragment") + } + return parsed.Scheme + "://" + parsed.Host, nil +} + +func filenameMatches(metadata document.AuthorizedUploadMetadata) bool { + ext := strings.ToLower(filepath.Ext(metadata.Filename)) + switch metadata.MediaType { + case "application/pdf": + return ext == ".pdf" + case "image/jpeg": + return ext == ".jpg" || ext == ".jpeg" + case "image/png": + return ext == ".png" + case "image/webp": + return ext == ".webp" + case "image/gif": + return ext == ".gif" + case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": + return ext == ".docx" + case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": + return ext == ".xlsx" + case "application/vnd.openxmlformats-officedocument.presentationml.presentation": + return ext == ".pptx" + case "application/epub+zip": + return ext == ".epub" + case "text/html": + return ext == ".html" || ext == ".htm" + default: + return false + } +} + +func readExact(ctx context.Context, upload io.Reader, metadata document.AuthorizedUploadMetadata) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, providerError(document.RenditionErrorCanceled, "Marker rendering canceled", err) + } + value, err := io.ReadAll(io.LimitReader(upload, metadata.ByteLength+1)) + if err != nil { + return nil, providerError(document.RenditionErrorTransient, "could not read authorized Marker upload", err) + } + digest := sha256.Sum256(value) + if int64(len(value)) != metadata.ByteLength || hex.EncodeToString(digest[:]) != metadata.SHA256 { + return nil, providerError(document.RenditionErrorPolicyRejected, "authorized Marker upload identity mismatch", nil) + } + return value, nil +} + +func operationContext(ctx context.Context, expiresAt time.Time, timeout time.Duration) (context.Context, context.CancelFunc) { + deadline := time.Now().Add(timeout) + if caller, ok := ctx.Deadline(); ok && caller.Before(deadline) { + deadline = caller + } + if expiresAt.Before(deadline) { + deadline = expiresAt + } + return context.WithDeadline(ctx, deadline) +} + +func checkOperation(ctx context.Context, expiresAt time.Time) error { + if errors.Is(ctx.Err(), context.Canceled) { + return providerError(document.RenditionErrorCanceled, "Marker rendering canceled", ctx.Err()) + } + if !time.Now().Before(expiresAt) { + return expiredError() + } + if err := ctx.Err(); err != nil { + return providerError(document.RenditionErrorCanceled, "Marker rendering canceled", err) + } + return nil +} + +func readBounded(ctx context.Context, expiresAt time.Time, reader io.Reader, maximum int64) ([]byte, error) { + value, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + if operationErr := checkOperation(ctx, expiresAt); operationErr != nil { + return nil, operationErr + } + return nil, providerError(document.RenditionErrorAmbiguousSubmission, "could not read Marker result", err) + } + if operationErr := checkOperation(ctx, expiresAt); operationErr != nil { + return nil, operationErr + } + if int64(len(value)) > maximum { + return nil, malformedError("Marker response exceeds byte limit", nil) + } + return value, nil +} + +func statusError(status int) error { + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return providerError(document.RenditionErrorAuthentication, "Marker authentication failed", nil) + case http.StatusTooManyRequests: + return providerError(document.RenditionErrorRateLimited, "Marker rate limit is exhausted", nil) + case http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout, http.StatusInsufficientStorage: + return providerError(document.RenditionErrorCapacity, "Marker capacity is unavailable", nil) + case http.StatusInternalServerError: + return providerError(document.RenditionErrorTransient, "Marker is temporarily unavailable", nil) + case http.StatusUnsupportedMediaType, http.StatusUnprocessableEntity: + return providerError(document.RenditionErrorUnsupportedInput, "Marker rejected the input format", nil) + case http.StatusBadRequest, http.StatusRequestEntityTooLarge: + return providerError(document.RenditionErrorPolicyRejected, "Marker rejected the upload", nil) + default: + return malformedError("Marker returned an unexpected HTTP status", nil) + } +} + +func expiredError() error { + return providerError(document.RenditionErrorPolicyRejected, "Marker authorization expired", nil) +} +func malformedError(message string, cause error) error { + return providerError(document.RenditionErrorMalformedEvidence, message, cause) +} + +func providerError(code document.RenditionErrorCode, message string, cause error) error { + value, err := document.NewRenditionProviderError(code, message, 0, cause) + if err == nil { + return value + } + return fmt.Errorf("marker error classification: %w", err) +} + +func validFingerprint(value string) bool { + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == sha256.Size && value == strings.ToLower(value) +} + +func validToken(value string) bool { + if value == "" || len(value) > 128 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return false + } + for _, character := range value { + if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' || character >= '0' && character <= '9' || strings.ContainsRune("_.-", character) { + continue + } + return false + } + return true +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +func compareFormats(left, right document.RenditionFormatCapability) int { + if value := strings.Compare(left.MediaFamily, right.MediaFamily); value != 0 { + return value + } + if value := strings.Compare(left.MediaType, right.MediaType); value != 0 { + return value + } + return strings.Compare(string(left.InputKind), string(right.InputKind)) +} + +func nilValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/marker/client_test.go b/document/marker/client_test.go new file mode 100644 index 00000000..3b04f94b --- /dev/null +++ b/document/marker/client_test.go @@ -0,0 +1,530 @@ +package marker + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media/mediatest" +) + +type testUpload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (*testUpload) Close() error { return nil } +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("missing secret") + } + return value, nil +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (fn roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return fn(request) } + +type callbackReadCloser struct { + reader io.Reader + once sync.Once + before func() +} + +func (body *callbackReadCloser) Read(value []byte) (int, error) { + body.once.Do(body.before) + return body.reader.Read(value) +} + +func (*callbackReadCloser) Close() error { return nil } + +type errorReadCloser struct { + once sync.Once + before func() + err error +} + +func (body *errorReadCloser) Read([]byte) (int, error) { + body.once.Do(body.before) + return 0, body.err +} + +func (*errorReadCloser) Close() error { return nil } + +func TestClientUploadsExactBytesToFixedRouteAndMapsProvenPages(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(2)) + var calls atomic.Int64 + client := newClient(t, fixture.profile, testSecrets{"marker-front": "synthetic-secret"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + calls.Add(1) + assert.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "https://marker.internal/marker/upload", request.URL.String()) + assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) + assertMultipart(t, request, fixture.metadata, fixture.source) + return jsonResponse(request, http.StatusOK, `{"format":"markdown","output":"{0}------------------------------------------------\n\n# First\n\n{1}------------------------------------------------\n\nSecond","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}},{"page_id":1,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}}]},"success":true}`), nil + })) + + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, int64(1), calls.Load()) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, "# First", result.Evidence.Units[0].Text) + assert.Equal(t, int64(1), result.Evidence.Units[1].Locator.Start) + assert.Equal(t, fixture.metadata.SHA256, result.Receipt.SourceSHA256) +} + +func TestDescriptorCoversMarkerConverterFamilies(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + want := []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/webp", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/gif", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/html", InputKind: document.RenditionInputOriginalFile}, + } + assert.ElementsMatch(t, want, fixture.descriptor.SupportedFormats) + + changed := fixture.profile + changed.RuntimeFingerprint = strings.Repeat("9", 64) + fingerprint, err := PolicyFingerprint(changed) + require.NoError(t, err) + assert.NotEqual(t, fixture.descriptor.PolicyFingerprint, fingerprint) + changed = fixture.profile + changed.SecretBinding = "other" + fingerprint, err = PolicyFingerprint(changed) + require.NoError(t, err) + assert.NotEqual(t, fixture.descriptor.PolicyFingerprint, fingerprint) +} + +func TestClientRequiresPinnedProfileAndOptionalCredentialPair(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + changed := fixture.profile + changed.RuntimeFingerprint = strings.Repeat("9", 64) + _, err := New(changed, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, `{}`)) + require.ErrorContains(t, err, "policy fingerprint") + + withoutCredential := fixture.profile + withoutCredential.SecretBinding = "" + withoutCredential.Descriptor = descriptorFor(t, withoutCredential) + _, err = New(withoutCredential, nil, staticTransport(http.StatusOK, `{}`)) + require.NoError(t, err) + _, err = New(withoutCredential, testSecrets{}, staticTransport(http.StatusOK, `{}`)) + require.ErrorContains(t, err, "named binding") +} + +func TestClientDegradesTransformedFamiliesWithoutInventingUnits(t *testing.T) { + for _, input := range []struct{ family, mediaType, filename string }{ + {"word", "application/vnd.openxmlformats-officedocument.wordprocessingml.document", "notes.docx"}, + {"spreadsheet", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "book.xlsx"}, + {"presentation", "application/vnd.openxmlformats-officedocument.presentationml.presentation", "deck.pptx"}, + {"ebook", "application/epub+zip", "book.epub"}, + {"text", "text/html", "page.html"}, + {"text", "text/html", "page.htm"}, + } { + t.Run(input.family, func(t *testing.T) { + fixture := newFixture(t, input.family, input.mediaType, input.filename, []byte("synthetic source")) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, + `{"format":"markdown","output":"{0}------------------------------------------------\n\nReadable","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}}]},"success":true}`)) + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + }) + } +} + +func TestClientRejectsUnsupportedFilenameIdentityAndBoundsBeforeEgress(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + var calls atomic.Int64 + transport := roundTripFunc(func(*http.Request) (*http.Response, error) { calls.Add(1); return nil, errors.New("unexpected") }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + + wrong := fixture + wrong.metadata.Filename = "report.txt" + wrong.authorization.ProviderMetadataChecksum = wrong.metadata.ProviderMetadataChecksum + _, err := client.Render(t.Context(), wrong.upload(), wrong.authorization) + assertProviderCode(t, err, document.RenditionErrorUnsupportedInput) + + wrong = fixture + wrong.metadata.SHA256 = strings.Repeat("0", 64) + wrong.authorization.SourceSHA256 = wrong.metadata.SHA256 + _, err = client.Render(t.Context(), wrong.upload(), wrong.authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + + tooLarge := fixture.profile + tooLarge.MaxDocumentBytes = 2 + tooLarge.MaxRequestBytes = 1024 + tooLarge.Descriptor = descriptorFor(t, tooLarge) + client = newClient(t, tooLarge, testSecrets{"marker-front": "secret"}, transport) + tooLargeFixture := fixture.withDescriptor(tooLarge.Descriptor) + _, err = client.Render(t.Context(), tooLargeFixture.upload(), tooLargeFixture.authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.Zero(t, calls.Load()) +} + +func TestClientRejectsRedirectMalformedPartialAndOversizedResults(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + for _, testCase := range []struct { + name string + status int + body string + want document.RenditionErrorCode + }{ + {"redirect", http.StatusFound, "", document.RenditionErrorMalformedEvidence}, + {"malformed", http.StatusOK, `{`, document.RenditionErrorMalformedEvidence}, + {"partial", http.StatusOK, `{"format":"markdown","output":"some","images":{},"metadata":{"table_of_contents":[],"page_stats":[]},"success":false,"error":"private"}`, document.RenditionErrorMalformedEvidence}, + {"schema drift", http.StatusOK, `{"format":"markdown","output":"some","images":{},"metadata":{"table_of_contents":[],"page_stats":[]},"success":true,"new_field":1}`, document.RenditionErrorMalformedEvidence}, + {"wrong format", http.StatusOK, `{"format":"html","output":"some","images":{},"metadata":{"table_of_contents":[],"page_stats":[]},"success":true}`, document.RenditionErrorMalformedEvidence}, + } { + t.Run(testCase.name, func(t *testing.T) { + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(testCase.status, testCase.body)) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "private") + }) + } + + bounded := fixture.profile + bounded.MaxResponseBytes = 64 + bounded.MaxMetadataBytes = 32 + bounded.Descriptor = descriptorFor(t, bounded) + client := newClient(t, bounded, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, strings.Repeat("x", 65))) + boundedFixture := fixture.withDescriptor(bounded.Descriptor) + _, err := client.Render(t.Context(), boundedFixture.upload(), boundedFixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) +} + +func TestClientNeverFollowsRedirects(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + var calls atomic.Int64 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + calls.Add(1) + response := jsonResponse(request, http.StatusFound, "") + response.Header.Set("Location", "https://attacker.invalid/result") + return response, nil + }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Equal(t, int64(1), calls.Load()) +} + +func TestClientRequiresLocalProofOfCompletePDFUnitsAndTotalResultBound(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(2)) + partial := `{"format":"markdown","output":"{0}------------------------------------------------\n\nOnly first","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}}]},"success":true}` + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, partial)) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + + fixture = newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + fixture.authorization.MaxTotalResultBytes = 128 + complete := `{"format":"markdown","output":"{0}------------------------------------------------\n\nComplete","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}}]},"success":true}` + client = newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, complete)) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) +} + +func TestClientProvesStillImageIdentityBeforeEgress(t *testing.T) { + const complete = `{"format":"markdown","output":"{0}------------------------------------------------\n\nImage text","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"surya","block_counts":[],"block_metadata":{}}]},"success":true}` + still := newFixture(t, "image", "image/png", "scan.png", mediatest.PNG(4, 3, nil)) + client := newClient(t, still.profile, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, complete)) + result, err := document.RenderRendition(t.Context(), client, still.upload(), still.authorization) + require.NoError(t, err) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + require.Len(t, result.Evidence.Units, 1) + + for _, testCase := range []struct { + name string + mediaType string + filename string + source []byte + }{ + {name: "malformed", mediaType: "image/png", filename: "scan.png", source: []byte("not a png")}, + {name: "mismatched", mediaType: "image/jpeg", filename: "scan.jpg", source: mediatest.PNG(4, 3, nil)}, + {name: "animated", mediaType: "image/gif", filename: "scan.gif", source: mediatest.GIF(2, 2, 2)}, + } { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, "image", testCase.mediaType, testCase.filename, testCase.source) + var calls atomic.Int64 + transport := roundTripFunc(func(*http.Request) (*http.Response, error) { + calls.Add(1) + return nil, errors.New("unexpected egress") + }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorUnsupportedInput) + assert.Zero(t, calls.Load()) + }) + } +} + +func TestClientRechecksExpiryAndCancellationWhileReadingResponse(t *testing.T) { + const complete = `{"format":"markdown","output":"{0}------------------------------------------------\n\nComplete","images":{},"metadata":{"table_of_contents":[],"page_stats":[{"page_id":0,"text_extraction_method":"pdftext","block_counts":[],"block_metadata":{}}]},"success":true}` + + t.Run("expiry after complete body", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + expiresAt := time.Now().UTC().Add(30 * time.Millisecond) + fixture.authorization.ExpiresAt = expiresAt.Format(timestampForm) + body := &callbackReadCloser{reader: strings.NewReader(complete), before: func() { + time.Sleep(time.Until(expiresAt) + 10*time.Millisecond) + }} + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: body, Request: request}, nil + }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + + _, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) + + t.Run("cancellation after complete body", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + ctx, cancel := context.WithCancel(t.Context()) + body := &callbackReadCloser{reader: strings.NewReader(complete), before: cancel} + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: body, Request: request}, nil + }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + + _, err := document.RenderRendition(ctx, client, fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorCanceled) + }) + + t.Run("cancellation on body read error", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + ctx, cancel := context.WithCancel(t.Context()) + body := &errorReadCloser{before: cancel, err: errors.New("private body failure")} + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: body, Request: request}, nil + }) + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, transport) + + _, err := document.RenderRendition(ctx, client, fixture.upload(), fixture.authorization) + + assertProviderCode(t, err, document.RenditionErrorCanceled) + }) +} + +func TestClientClassifiesAuthCapacityTransportExpiryAndCancellation(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + for _, testCase := range []struct { + status int + want document.RenditionErrorCode + }{ + {http.StatusUnauthorized, document.RenditionErrorAuthentication}, + {http.StatusTooManyRequests, document.RenditionErrorRateLimited}, + {http.StatusServiceUnavailable, document.RenditionErrorCapacity}, + {http.StatusUnsupportedMediaType, document.RenditionErrorUnsupportedInput}, + } { + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(testCase.status, "private")) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "private") + } + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("private transport") + })) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorAmbiguousSubmission) + + expired := fixture + expired.authorization.ExpiresAt = time.Now().UTC().Add(-time.Second).Format(timestampForm) + _, err = client.Render(t.Context(), expired.upload(), expired.authorization) + require.Error(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, err = client.Render(ctx, fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorCanceled) +} + +func TestClientBoundsImageAndMetadataPayloads(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + for _, body := range []string{ + `{"format":"markdown","output":"safe","images":{"a":"c3ludGhldGljIGltYWdl"},"metadata":{"table_of_contents":[],"page_stats":[]},"success":true}`, + `{"format":"markdown","output":"safe","images":{},"metadata":{"table_of_contents":"` + strings.Repeat("x", 257) + `","page_stats":[]},"success":true}`, + } { + client := newClient(t, fixture.profile, testSecrets{"marker-front": "secret"}, staticTransport(http.StatusOK, body)) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + } +} + +type fixture struct { + profile Profile + descriptor document.RenditionDescriptor + metadata document.AuthorizedUploadMetadata + authorization document.RenditionAuthorization + source []byte +} + +func newFixture(t *testing.T, family, mediaType, filename string, source []byte) fixture { + t.Helper() + profile := Profile{Origin: "https://marker.internal", SecretBinding: "marker-front", Mode: "fast", + DeploymentFingerprint: strings.Repeat("d", 64), RuntimeFingerprint: strings.Repeat("r", 64), + RequestTimeout: time.Second, MaxDocumentBytes: 1 << 20, MaxRequestBytes: 2 << 20, + MaxResponseBytes: 1 << 20, MaxMetadataBytes: 256, MaxImages: 1, MaxImageBytes: 8, MaxUnits: 32} + // Runtime fingerprints are SHA-256 values, so use hexadecimal input. + profile.RuntimeFingerprint = strings.Repeat("e", 64) + profile.Descriptor = descriptorFor(t, profile) + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{Filename: filename, MediaFamily: family, MediaType: mediaType, + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), InputKind: document.RenditionInputOriginalFile} + started := time.Now().UTC().Add(-time.Minute) + authorization := document.RenditionAuthorization{ProviderID: profile.Descriptor.ID, DescriptorFingerprint: profile.Descriptor.Fingerprint, + PolicyFingerprint: profile.Descriptor.PolicyFingerprint, RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, MediaFamily: family, MediaType: mediaType, + InputKind: document.RenditionInputOriginalFile, MaxProviderMarkdownBytes: 4096, MaxTotalResultBytes: 32768, + AuthorizedAt: started.Format(timestampForm), ExpiresAt: started.Add(10 * time.Minute).Format(timestampForm)} + return fixture{profile: profile, descriptor: profile.Descriptor, metadata: metadata, authorization: authorization, source: source} +} + +func descriptorFor(t *testing.T, profile Profile) document.RenditionDescriptor { + t.Helper() + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ID: providerID, + ContractVersion: document.RenditionProviderContractVersion, PolicyFingerprint: fingerprint, + TrustBoundary: document.RenditionTrustOperatorNetwork, SupportedFormats: SupportedFormats(), ReturnsMarkdown: true}) + require.NoError(t, err) + return descriptor +} + +func (fixture fixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func (fixture fixture) withDescriptor(descriptor document.RenditionDescriptor) fixture { + fixture.descriptor, fixture.profile.Descriptor = descriptor, descriptor + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + return fixture +} + +func newClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, transport) + require.NoError(t, err) + return client +} + +func staticTransport(status int, body string) http.RoundTripper { + return roundTripFunc(func(request *http.Request) (*http.Response, error) { return jsonResponse(request, status, body), nil }) +} + +func jsonResponse(request *http.Request, status int, body string) *http.Response { + return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(body)), Request: request} +} + +func assertMultipart(t *testing.T, request *http.Request, metadata document.AuthorizedUploadMetadata, source []byte) { + t.Helper() + mediaType, params, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + require.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, params["boundary"]) + fields := map[string]string{} + for { + part, partErr := reader.NextPart() + if errors.Is(partErr, io.EOF) { + break + } + require.NoError(t, partErr) + value, readErr := io.ReadAll(part) + require.NoError(t, readErr) + if part.FormName() == "file" { + assert.Equal(t, metadata.Filename, part.FileName()) + assert.Equal(t, metadata.MediaType, part.Header.Get("Content-Type")) + assert.Equal(t, source, value) + } else { + fields[part.FormName()] = string(value) + } + } + assert.Equal(t, map[string]string{"mode": "fast", "force_ocr": "false", "paginate_output": "true", "output_format": "markdown"}, fields) +} + +func assertProviderCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok, "%T: %v", err, err) + assert.Equal(t, want, providerErr.Code()) +} + +func TestPolicyFingerprintIsCanonical(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "report.pdf", testPDF(1)) + first, err := PolicyFingerprint(fixture.profile) + require.NoError(t, err) + encoded, err := json.Marshal(map[string]string{"fingerprint": first}, json.Deterministic(true)) + require.NoError(t, err) + assert.Contains(t, string(encoded), first) +} + +func testPDF(pageCount int) []byte { + var output bytes.Buffer + output.WriteString("%PDF-1.4\n%synthetic-marker-fixture\n") + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", pdfKids(pageCount), pageCount), + } + for index := range pageCount { + objects = append(objects, fmt.Sprintf("<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents %d 0 R >>", pageCount+3+index)) + } + for range pageCount { + objects = append(objects, "<< /Length 0 >>\nstream\n\nendstream") + } + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} + +func pdfKids(pageCount int) string { + children := make([]string, pageCount) + for index := range pageCount { + children[index] = fmt.Sprintf("%d 0 R", index+3) + } + return strings.Join(children, " ") +} diff --git a/document/media/detect.go b/document/media/detect.go index 56ac6290..b50986a7 100644 --- a/document/media/detect.go +++ b/document/media/detect.go @@ -10,6 +10,7 @@ import ( "io" "math" "math/bits" + "sort" _ "image/jpeg" // Register the JPEG decoder used by image.DecodeConfig. _ "image/png" // Register the PNG decoder used by image.DecodeConfig. @@ -99,8 +100,13 @@ func sniff(data []byte) (Metadata, error) { if !ok { return Metadata{}, ErrMalformedMedia } + mediaType := "video/mp4" + if info.container == "quicktime" { + mediaType = "video/quicktime" + } return Metadata{ - Format: FormatMP4, Kind: KindVideo, MediaType: "video/mp4", + Format: FormatMP4, Kind: KindVideo, MediaType: mediaType, + Container: info.container, Codec: info.codec, Width: info.width, Height: info.height, DurationMS: info.durationMS, DurationKnown: info.durationKnown, }, nil @@ -408,14 +414,21 @@ func isMP4(data []byte) bool { type mp4Info struct { width, height, durationMS int64 + frameCount int64 durationKnown bool pictureTracks int + container, codec string + sampleAuthority bool moovCount, mvhdCount int + mdatCount int timescale uint64 movieDuration uint64 trackDurations []uint64 mediaDurations []mp4Duration + mediaDataRanges []mp4ByteRange + sampleRanges []mp4SampleRange + source []byte unknownDuration bool } @@ -423,22 +436,69 @@ type mp4Duration struct { value, timescale uint64 } +type mp4ByteRange struct { + start, end uint64 +} + +type mp4SampleRange struct { + mp4ByteRange + + descriptionIndex uint32 +} + +type mp4SampleToChunk struct { + firstChunk, samplesPerChunk, descriptionIndex uint32 +} + func mp4Metadata(data []byte) (mp4Info, bool) { - var info mp4Info + mediaDataRanges, ok := mp4MediaDataRanges(data) + if !ok { + return mp4Info{}, false + } + info := mp4Info{mediaDataRanges: mediaDataRanges, source: data} if !scanMP4Boxes(data, &info, nil, "", 0) || info.moovCount != 1 || info.mvhdCount > 1 || - info.pictureTracks == 0 || info.width <= 0 || info.height <= 0 { + info.pictureTracks == 0 || info.width <= 0 || info.height <= 0 || info.codec == "" { return mp4Info{}, false } + info.container = canonicalMP4Container(data) + info.sampleAuthority = info.sampleAuthority && info.mdatCount > 0 && + mp4SampleRangesDoNotOverlap(info.sampleRanges) info.resolveDuration() return info, true } +func mp4MediaDataRanges(data []byte) ([]mp4ByteRange, bool) { + var ranges []mp4ByteRange + for offset := 0; offset < len(data); { + headerLen, size, ok := mp4BoxHeader(data, offset) + if !ok { + return nil, false + } + if string(data[offset+4:offset+8]) == "mdat" { + ranges = append(ranges, mp4ByteRange{ + start: uint64(offset + headerLen), //nolint:gosec // offsets are non-negative and bounded by len(data) + end: uint64(offset + size), //nolint:gosec // offsets are non-negative and bounded by len(data) + }) + } + offset += size + } + return ranges, true +} + const maxMP4Depth = 6 type mp4TrackInfo struct { tkhdCount, edtsCount, elstCount int hdlrCount, mdhdCount int stsdCount, sttsCount, cttsCount, sampleSizeCount int + stscCount int + chunkOffsetCount int + chunkOffsets []uint64 + sampleToChunks []mp4SampleToChunk + sampleDescriptionCount uint32 + sampleDescriptions []mp4CodecConfiguration + defaultSampleSize uint32 + sampleSizes []uint32 handlerType string presentationWidth, presentationHeight int64 codedWidth, codedHeight int64 @@ -448,6 +508,7 @@ type mp4TrackInfo struct { timedSamples, compositionSamples uint64 sampleCount uint64 hasVisualSamples bool + codec string unknownSampleEntries int } @@ -501,6 +562,8 @@ func scanMP4Boxes(data []byte, info *mp4Info, track *mp4TrackInfo, parent string if info.moovCount > 1 || !scanMP4Boxes(payload, info, nil, kind, depth+1) { return false } + case kind == "mdat" && parent == "": + info.mdatCount++ case kind == "trak" && parent == "moov": trackInfo := &mp4TrackInfo{} if !scanMP4Boxes(payload, info, trackInfo, kind, depth+1) || !trackInfo.finish(info) { @@ -560,16 +623,30 @@ func scanMP4Boxes(data []byte, info *mp4Info, track *mp4TrackInfo, parent string if track == nil || !parseSampleSize(kind, payload, track) { return false } + case kind == "stsc" && parent == "stbl": + if track == nil || !parseSTSC(payload, track) { + return false + } + case (kind == "stco" || kind == "co64") && parent == "stbl": + if track == nil || !parseChunkOffsets(kind, payload, track) { + return false + } case kind == "hdlr" && parent == "mdia": if track == nil || !parseHDLR(payload, track) { return false } + case kind == "hdlr" && parent == "minf": + // QuickTime may carry a separate data-handler declaration here. + // It is not the media handler that establishes track modality. + if len(payload) < 12 || payload[0] != 0 { + return false + } case kind == "mdhd" && parent == "mdia": if track == nil || !parseMDHD(payload, info, track) { return false } case kind == "moov" || kind == "trak" || kind == "edts" || kind == "elst" || kind == "mdia" || kind == "minf" || - kind == "stbl" || kind == "stsd" || kind == "stts" || kind == "ctts" || kind == "stsz" || kind == "stz2" || kind == "mvhd" || + kind == "stbl" || kind == "stsd" || kind == "stts" || kind == "ctts" || kind == "stsz" || kind == "stz2" || kind == "stsc" || kind == "stco" || kind == "co64" || kind == "mvhd" || kind == "tkhd" || kind == "hdlr" || kind == "mdhd": // A structural header outside its authoritative position is not a // file this package can bound. @@ -680,18 +757,25 @@ func parseSTSD(payload []byte, track *mp4TrackInfo) bool { } width := int64(binary.BigEndian.Uint16(entry[24:26])) height := int64(binary.BigEndian.Uint16(entry[26:28])) - codecWidth, codecHeight, ok := visualCodecDimensions(kind, entry) + configuration, codecWidth, codecHeight, ok := visualCodecDimensions(kind, entry) if width <= 0 || height <= 0 || !ok { return false } + if track.codec != "" && track.codec != configuration.codec { + return false + } track.hasVisualSamples = true + track.codec = configuration.codec track.codedWidth = max(track.codedWidth, width, codecWidth) track.codedHeight = max(track.codedHeight, height, codecHeight) + track.sampleDescriptions = append(track.sampleDescriptions, configuration) } else { track.unknownSampleEntries++ + track.sampleDescriptions = append(track.sampleDescriptions, mp4CodecConfiguration{}) } offset += size } + track.sampleDescriptionCount = seen return seen == want } @@ -829,7 +913,7 @@ func parseELST(payload []byte, track *mp4TrackInfo) bool { } func parseSampleSize(kind string, payload []byte, track *mp4TrackInfo) bool { - if len(payload) < 12 || payload[0] != 0 { + if len(payload) < 12 || payload[0] != 0 || payload[1] != 0 || payload[2] != 0 || payload[3] != 0 { return false } track.sampleSizeCount++ @@ -843,10 +927,15 @@ func parseSampleSize(kind string, payload []byte, track *mp4TrackInfo) bool { samples := int(count) // #nosec G115 -- capped at MaxInt32 above. switch kind { case "stsz": - if binary.BigEndian.Uint32(payload[4:8]) == 0 { + track.defaultSampleSize = binary.BigEndian.Uint32(payload[4:8]) + if track.defaultSampleSize == 0 { if len(payload) != 12+samples*4 { return false } + track.sampleSizes = make([]uint32, 0, samples) + for offset := 12; offset < len(payload); offset += 4 { + track.sampleSizes = append(track.sampleSizes, binary.BigEndian.Uint32(payload[offset:offset+4])) + } } else if len(payload) != 12 { return false } @@ -865,6 +954,26 @@ func parseSampleSize(kind string, payload []byte, track *mp4TrackInfo) bool { if len(payload) != 12+tableBytes { return false } + track.sampleSizes = make([]uint32, 0, samples) + switch payload[7] { + case 4: + for index := range samples { + value := payload[12+index/2] + if index%2 == 0 { + track.sampleSizes = append(track.sampleSizes, uint32(value>>4)) + } else { + track.sampleSizes = append(track.sampleSizes, uint32(value&0x0f)) + } + } + case 8: + for _, value := range payload[12:] { + track.sampleSizes = append(track.sampleSizes, uint32(value)) + } + case 16: + for offset := 12; offset < len(payload); offset += 2 { + track.sampleSizes = append(track.sampleSizes, uint32(binary.BigEndian.Uint16(payload[offset:offset+2]))) + } + } default: return false } @@ -872,6 +981,63 @@ func parseSampleSize(kind string, payload []byte, track *mp4TrackInfo) bool { return true } +func parseSTSC(payload []byte, track *mp4TrackInfo) bool { + if len(payload) < 8 || payload[0] != 0 || payload[1] != 0 || payload[2] != 0 || payload[3] != 0 { + return false + } + track.stscCount++ + if track.stscCount != 1 { + return false + } + count := binary.BigEndian.Uint32(payload[4:8]) + if count > math.MaxInt32 || len(payload) != 8+int(count)*12 { + return false + } + track.sampleToChunks = make([]mp4SampleToChunk, 0, count) + var previousFirst uint32 + for offset := 8; offset < len(payload); offset += 12 { + entry := mp4SampleToChunk{ + firstChunk: binary.BigEndian.Uint32(payload[offset : offset+4]), + samplesPerChunk: binary.BigEndian.Uint32(payload[offset+4 : offset+8]), + descriptionIndex: binary.BigEndian.Uint32(payload[offset+8 : offset+12]), + } + if entry.firstChunk == 0 || entry.samplesPerChunk == 0 || entry.descriptionIndex == 0 || + len(track.sampleToChunks) == 0 && entry.firstChunk != 1 || previousFirst >= entry.firstChunk { + return false + } + track.sampleToChunks = append(track.sampleToChunks, entry) + previousFirst = entry.firstChunk + } + return true +} + +func parseChunkOffsets(kind string, payload []byte, track *mp4TrackInfo) bool { + if len(payload) < 8 || payload[0] != 0 || payload[1] != 0 || payload[2] != 0 || payload[3] != 0 { + return false + } + track.chunkOffsetCount++ + if track.chunkOffsetCount != 1 { + return false + } + count := binary.BigEndian.Uint32(payload[4:8]) + entryBytes := 4 + if kind == "co64" { + entryBytes = 8 + } + if count > math.MaxInt32 || len(payload) != 8+int(count)*entryBytes { + return false + } + track.chunkOffsets = make([]uint64, 0, count) + for offset := 8; offset < len(payload); offset += entryBytes { + if kind == "co64" { + track.chunkOffsets = append(track.chunkOffsets, binary.BigEndian.Uint64(payload[offset:offset+8])) + } else { + track.chunkOffsets = append(track.chunkOffsets, uint64(binary.BigEndian.Uint32(payload[offset:offset+4]))) + } + } + return true +} + func parseHDLR(payload []byte, track *mp4TrackInfo) bool { if len(payload) < 12 || payload[0] != 0 { return false @@ -886,7 +1052,7 @@ func parseHDLR(payload []byte, track *mp4TrackInfo) bool { func isVisualSampleEntry(kind string) bool { switch kind { - case "avc1", "avc2", "avc3", "avc4", "hvc1", "hev1", "vp08", "vp09", "av01", "mp4v": + case "avc1", "hvc1", "vp09", "av01": return true default: return false @@ -915,15 +1081,134 @@ func (track *mp4TrackInfo) finish(info *mp4Info) bool { return isNonVisualHandler(track.handlerType) && !presentation && !track.hasVisualSamples } if !presentation || !track.hasVisualSamples || track.stsdCount != 1 || track.unknownSampleEntries != 0 || - track.codedWidth <= 0 || track.codedHeight <= 0 { + track.codedWidth <= 0 || track.codedHeight <= 0 || track.codec == "" { return false } + if track.sampleCount > math.MaxInt64 { + return false + } + trackSampleAuthority := track.proveSampleRanges(info) + if info.pictureTracks == 0 { + info.sampleAuthority = trackSampleAuthority + } else { + info.sampleAuthority = info.sampleAuthority && trackSampleAuthority + } + sampleCount := int64(track.sampleCount) // #nosec G115 -- checked against MaxInt64 above. + if sampleCount > math.MaxInt64-info.frameCount { + return false + } + info.frameCount += sampleCount info.pictureTracks++ + if info.codec != "" && info.codec != track.codec { + return false + } + info.codec = track.codec info.width = max(info.width, track.presentationWidth, track.codedWidth) info.height = max(info.height, track.presentationHeight, track.codedHeight) return true } +const maxMP4MappedSamples = 10_000 + +func (track *mp4TrackInfo) proveSampleRanges(info *mp4Info) bool { + if track.stscCount != 1 || track.chunkOffsetCount != 1 || len(track.sampleToChunks) == 0 || + len(track.chunkOffsets) == 0 || track.sampleCount == 0 || track.sampleCount > maxMP4MappedSamples || + uint64(track.sampleToChunks[len(track.sampleToChunks)-1].firstChunk) > uint64(len(track.chunkOffsets)) { + return false + } + sampleIndex := uint64(0) + toChunkIndex := 0 + for chunkIndex, chunkOffset := range track.chunkOffsets { + chunkNumber := uint32(chunkIndex + 1) + for toChunkIndex+1 < len(track.sampleToChunks) && + track.sampleToChunks[toChunkIndex+1].firstChunk <= chunkNumber { + toChunkIndex++ + } + mapping := track.sampleToChunks[toChunkIndex] + if mapping.firstChunk > chunkNumber || mapping.descriptionIndex > track.sampleDescriptionCount || + uint64(mapping.samplesPerChunk) > track.sampleCount-sampleIndex { + return false + } + configuration := track.sampleDescriptions[mapping.descriptionIndex-1] + cursor := chunkOffset + for range mapping.samplesPerChunk { + size, ok := track.sampleSize(sampleIndex) + if !ok || size == 0 || math.MaxUint64-cursor < size { + return false + } + sampleRange := mp4SampleRange{ + start: cursor, end: cursor + size, + descriptionIndex: mapping.descriptionIndex, + } + if !mp4RangeInsideMediaData(sampleRange.mp4ByteRange, info.mediaDataRanges) { + return false + } + if sampleRange.end > uint64(len(info.source)) { + return false + } + start := int(sampleRange.start) //nolint:gosec // proven no greater than the source length above + end := int(sampleRange.end) //nolint:gosec // proven no greater than the source length above + if !validMP4CodecSample(configuration, info.source[start:end]) { + return false + } + info.sampleRanges = append(info.sampleRanges, sampleRange) + cursor += size + sampleIndex++ + } + } + return sampleIndex == track.sampleCount +} + +func (track *mp4TrackInfo) sampleSize(index uint64) (uint64, bool) { + if index >= track.sampleCount { + return 0, false + } + if track.defaultSampleSize != 0 { + return uint64(track.defaultSampleSize), true + } + if index >= uint64(len(track.sampleSizes)) { + return 0, false + } + return uint64(track.sampleSizes[index]), true +} + +func mp4RangeInsideMediaData(sample mp4ByteRange, mediaData []mp4ByteRange) bool { + if sample.start >= sample.end { + return false + } + for _, bounds := range mediaData { + if sample.start >= bounds.start && sample.end <= bounds.end { + return true + } + } + return false +} + +func mp4SampleRangesDoNotOverlap(samples []mp4SampleRange) bool { + if len(samples) == 0 { + return false + } + sort.Slice(samples, func(left, right int) bool { + if samples[left].start == samples[right].start { + return samples[left].end < samples[right].end + } + return samples[left].start < samples[right].start + }) + for index := 1; index < len(samples); index++ { + if samples[index].start < samples[index-1].end { + return false + } + } + return true +} + +func canonicalMP4Container(data []byte) string { + if len(data) >= 12 && string(data[8:12]) == "qt " { + return "quicktime" + } + return "mp4" +} + func isNonVisualHandler(handler string) bool { switch handler { case "soun", "hint", "subt", "text", "sbtl", "clcp", "meta", "tmcd": diff --git a/document/media/detect_test.go b/document/media/detect_test.go index e90213c1..e660d06c 100644 --- a/document/media/detect_test.go +++ b/document/media/detect_test.go @@ -85,6 +85,75 @@ func TestDetectBytesRecognizesSupportedContainers(t *testing.T) { } } +// TestDetectReportsVerifiedMP4CodecAndQuickTimeContainer catches a detector +// that treats every ISO base-media file as generic MP4 or omits the verified +// H.264 sample-entry identity needed for Gemini video eligibility. +func TestDetectReportsVerifiedMP4CodecAndQuickTimeContainer(t *testing.T) { + t.Parallel() + for _, tt := range []struct { + name, codec, container string + data []byte + }{ + {name: "QuickTime H.264", codec: "h264", container: "quicktime", data: mediatest.H264MOV()}, + {name: "MP4 H.265", codec: "h265", container: "mp4", data: mediatest.H265MP4()}, + {name: "MP4 VP9", codec: "vp9", container: "mp4", data: mediatest.VP9MP4()}, + {name: "MP4 AV1", codec: "av1", container: "mp4", data: mediatest.AV1MP4()}, + } { + t.Run(tt.name, func(t *testing.T) { + metadata, err := media.DetectBytes(tt.data, "video/mp4") + require.NoError(t, err) + assert.Equal(t, tt.container, metadata.Container) + assert.Equal(t, tt.codec, metadata.Codec) + assert.Equal(t, int64(1_000), metadata.DurationMS) + assert.True(t, metadata.DurationKnown) + }) + } + + // Two authoritative visual sample entries must not collapse into an + // arbitrary codec choice. + conflicting := mp4TwoPictureTracks(16, 16, 16, 16) + secondEntry := bytes.LastIndex(conflicting, []byte("avc1")) + secondConfig := bytes.LastIndex(conflicting, []byte("avcC")) + require.NotEqual(t, -1, secondEntry) + require.NotEqual(t, -1, secondConfig) + copy(conflicting[secondEntry:secondEntry+4], "av01") + copy(conflicting[secondConfig:secondConfig+4], "av1C") + conflicting[secondConfig+4] = 0x81 + _, err := media.DetectBytes(conflicting, "video/mp4") + require.ErrorIs(t, err, media.ErrMalformedMedia) +} + +func TestDetectRejectsRenamedAVCConfigurationAsAnotherCodec(t *testing.T) { + t.Parallel() + for _, testCase := range []struct { + name, sampleEntry, configBox string + }{ + {name: "VP9", sampleEntry: "vp09", configBox: "vpcC"}, + {name: "AV1", sampleEntry: "av01", configBox: "av1C"}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + data := decodableAVCMP4(t) + entry := bytes.Index(data, []byte("avc1")) + config := bytes.Index(data, []byte("avcC")) + require.NotEqual(t, -1, entry) + require.NotEqual(t, -1, config) + copy(data[entry:entry+4], testCase.sampleEntry) + copy(data[config:config+4], testCase.configBox) + if testCase.configBox == "av1C" { + copy(data[config+4:config+8], []byte{0x81, 0, 0, 0}) + } + + _, err := media.DetectBytes(data, "video/mp4") + require.ErrorIs(t, err, media.ErrMalformedMedia) + }) + } +} + +func quickTimeH264MP4() []byte { + return mediatest.H264MOV() +} + func TestDetectBytesRejectsUnsupportedAndMalformedInput(t *testing.T) { tests := []struct { name string diff --git a/document/media/doc.go b/document/media/doc.go index 34d2a964..b3d0288a 100644 --- a/document/media/doc.go +++ b/document/media/doc.go @@ -1,12 +1,12 @@ -// Package media detects still images, animated images, and video from bytes -// and evaluates them against a bounded eligibility policy. +// Package media inspects provider-bound documents and media against finite +// local limits, and retains the lower-level visual detector used by callers. // -// Detection sniffs container signatures and reads only the metadata needed -// to bound provider input: dimensions, frame count, and duration. It never -// decodes pixels or samples. The declared media type is recorded for callers -// but is never trusted for detection. +// Capability inspection binds exact bytes to their declared identity, policy, +// provider descriptor, processing profile, and disclosure authority. It +// rejects formats whose expansion, semantic units, external references, or +// decode work cannot be bounded locally. Visual detection reads container +// metadata without decoding pixels or samples. // // The package does not perform filesystem, network, storage, database, -// queue, daemon, vault, or application work, and it has no notion of -// attachment ownership, roles, or hashes. +// queue, daemon, vault, or application work. package media diff --git a/document/media/inspect.go b/document/media/inspect.go new file mode 100644 index 00000000..eed58392 --- /dev/null +++ b/document/media/inspect.go @@ -0,0 +1,938 @@ +package media + +import ( + "archive/zip" + "bytes" + "crypto/sha256" + "encoding/binary" + "encoding/csv" + "encoding/hex" + "encoding/json/v2" + "encoding/xml" + "errors" + "fmt" + "io" + "math" + "mime" + "net/url" + "path/filepath" + "slices" + "strings" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/formatdetect" +) + +const ( + capabilityRecordVersion = 1 + maxInspectionSourceBytes = int64(1 << 30) + maxID3v2TagBytes = 1 << 20 +) + +// CapabilityReason is a stable capability inspection outcome. +type CapabilityReason string + +const ( + CapabilityReasonEligible CapabilityReason = "eligible" + CapabilityReasonUnsupported CapabilityReason = "unsupported_media_type" + CapabilityReasonUnboundedFamily CapabilityReason = "unbounded_media_family" + CapabilityReasonMalformed CapabilityReason = "malformed_input" + CapabilityReasonSourceBytes CapabilityReason = "source_bytes_exceeded" + CapabilityReasonExpandedBytes CapabilityReason = "expanded_bytes_exceeded" + CapabilityReasonEntryBytes CapabilityReason = "entry_bytes_exceeded" + CapabilityReasonEntryCount CapabilityReason = "entry_count_exceeded" + CapabilityReasonNestedContainer CapabilityReason = "nested_container" + CapabilityReasonEncryptedContainer CapabilityReason = "encrypted_container" + CapabilityReasonExternalReference CapabilityReason = "external_reference" + CapabilityReasonSemanticUnits CapabilityReason = "semantic_units_exceeded" + CapabilityReasonVisualBounds CapabilityReason = "visual_bounds_exceeded" +) + +// InspectionPolicy binds finite preflight limits and downstream authority to +// one exact source. Every supported family must have an enforceable bound. +type InspectionPolicy struct { + Filename string `json:"filename"` + DeclaredMediaType string `json:"declared_media_type"` + ExpectedBytes int64 `json:"expected_bytes"` + ExpectedSHA256 string `json:"expected_sha256"` + + DescriptorFingerprint string `json:"descriptor_fingerprint"` + ProfileFingerprint string `json:"profile_fingerprint"` + DisclosureFingerprint string `json:"disclosure_fingerprint"` + InputKind document.RenditionInputKind `json:"input_kind"` + + MaxSourceBytes int64 `json:"max_source_bytes"` + MaxExpandedBytes int64 `json:"max_expanded_bytes"` + MaxEntryBytes int64 `json:"max_entry_bytes"` + MaxEntries int64 `json:"max_entries"` + MaxNestingDepth int64 `json:"max_nesting_depth"` + MaxTextLines int64 `json:"max_text_lines"` + MaxCharacters int64 `json:"max_characters"` + MaxRecords int64 `json:"max_records"` + MaxPages int64 `json:"max_pages"` + MaxSlides int64 `json:"max_slides"` + MaxSheets int64 `json:"max_sheets"` + MaxCells int64 `json:"max_cells"` + MaxSpineItems int64 `json:"max_spine_items"` + MaxResources int64 `json:"max_resources"` + MaxPixels int64 `json:"max_pixels"` + MaxFrames int64 `json:"max_frames"` + MaxDurationMS int64 `json:"max_duration_ms"` +} + +// CapabilityMeasurements records only finite, locally verified measurements. +type CapabilityMeasurements struct { + CompressedBytes int64 `json:"compressed_bytes"` + ExpandedBytes int64 `json:"expanded_bytes"` + Entries int64 `json:"entries"` + MaxEntryBytes int64 `json:"max_entry_bytes"` + NestingDepth int64 `json:"nesting_depth"` + TextLines int64 `json:"text_lines"` + Characters int64 `json:"characters"` + Records int64 `json:"records"` + Pages int64 `json:"pages"` + Slides int64 `json:"slides"` + Sheets int64 `json:"sheets"` + Cells int64 `json:"cells"` + SpineItems int64 `json:"spine_items"` + Resources int64 `json:"resources"` + Pixels int64 `json:"pixels"` + Frames int64 `json:"frames"` + DurationMS int64 `json:"duration_ms"` +} + +// CapabilityRecord is a checksum-sealed decision over exact source bytes. +// Mutating any exported field invalidates Checksum. +type CapabilityRecord struct { + Version int `json:"version"` + Eligible bool `json:"eligible"` + Reason CapabilityReason `json:"reason"` + MediaFamily string `json:"media_family"` + MediaType string `json:"media_type"` + Format string `json:"format"` + SourceBytes int64 `json:"source_bytes"` + SourceSHA256 string `json:"source_sha256"` + Measurements CapabilityMeasurements `json:"measurements"` + + PolicyFingerprint string `json:"policy_fingerprint"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + ProfileFingerprint string `json:"profile_fingerprint"` + DisclosureFingerprint string `json:"disclosure_fingerprint"` + InputKind document.RenditionInputKind `json:"input_kind"` + Policy InspectionPolicy `json:"policy"` + Checksum string `json:"checksum"` + + localAuthority bool +} + +type capabilityRecordIdentity CapabilityRecord + +// InspectCapability reads one bounded source, verifies its declared identity, and +// returns a checksum-sealed finite capability decision. +func InspectCapability(reader io.Reader, policy InspectionPolicy) (CapabilityRecord, error) { + if reader == nil { + return CapabilityRecord{}, errors.New("media: inspection source is required") + } + if err := validateInspectionPolicy(policy); err != nil { + return CapabilityRecord{}, err + } + data, err := io.ReadAll(io.LimitReader(reader, policy.MaxSourceBytes+1)) + if err != nil { + return CapabilityRecord{}, fmt.Errorf("media: read inspection source: %w", err) + } + if int64(len(data)) > policy.MaxSourceBytes { + return sealCapabilityRecord(policy, data, CapabilityRecord{ + Eligible: false, Reason: CapabilityReasonSourceBytes, + }) + } + digest := sha256.Sum256(data) + digestHex := hex.EncodeToString(digest[:]) + if int64(len(data)) != policy.ExpectedBytes { + return CapabilityRecord{}, fmt.Errorf("media: source byte length %d does not match declared %d", + len(data), policy.ExpectedBytes) + } + if digestHex != policy.ExpectedSHA256 { + return CapabilityRecord{}, errors.New("media: source SHA-256 does not match declaration") + } + + record := CapabilityRecord{Eligible: false, Reason: CapabilityReasonUnsupported} + baseType, _, _ := mime.ParseMediaType(policy.DeclaredMediaType) + ext := strings.ToLower(filepath.Ext(policy.Filename)) + switch { + case isTextFamily(ext, baseType): + record = inspectText(data, ext, baseType, policy) + case isZIPFamily(ext, baseType, data): + record = inspectZIP(data, ext, baseType, policy) + case strings.HasPrefix(baseType, "image/") || strings.HasPrefix(baseType, "video/"): + record = inspectVisualCapability(data, baseType, policy) + case baseType == "application/pdf" || ext == ".pdf": + record = inspectPDF(data, policy) + case baseType == "audio/wav" || baseType == "audio/x-wav" || ext == ".wav": + record = inspectWAV(data, policy) + case baseType == "audio/mpeg" || ext == ".mp3": + record = inspectMP3(data, policy) + case strings.HasPrefix(baseType, "audio/"): + record = CapabilityRecord{Eligible: false, Reason: CapabilityReasonUnboundedFamily, + MediaFamily: "audio", MediaType: baseType, Format: strings.TrimPrefix(ext, ".")} + case baseType == "application/msword" || baseType == "application/vnd.ms-powerpoint" || + baseType == "application/vnd.ms-excel": + record = CapabilityRecord{Eligible: false, Reason: CapabilityReasonUnboundedFamily, + MediaFamily: "document", MediaType: baseType, Format: strings.TrimPrefix(ext, ".")} + } + if record.Eligible && requiresDocumentFormatDetection(record.MediaFamily) { + detected, detectErr := formatdetect.DetectFormat( + bytes.NewReader(data), int64(len(data)), baseType) + if detectErr != nil || !filenameAllowsFormat(ext, detected.ID) { + record.Eligible = false + record.Reason = CapabilityReasonMalformed + } else { + record.MediaFamily = detected.Family + record.MediaType = detected.MediaType + record.Format = detected.ID + } + } + if record.Eligible && (record.MediaFamily == string(KindImage) || + record.MediaFamily == string(KindVideo)) { + quickTimeIdentity := record.MediaFamily == string(KindVideo) && + baseType == "video/quicktime" && ext == ".mov" + if !quickTimeIdentity && (baseType != record.MediaType || !filenameAllowsVisualFormat(ext, record.Format)) { + record.Eligible = false + record.Reason = CapabilityReasonMalformed + } + } + if record.Eligible && record.MediaFamily == "audio" { + validIdentity := record.Format == "wav" && ext == ".wav" && + (baseType == "audio/wav" || baseType == "audio/x-wav") || + record.Format == "mp3" && ext == ".mp3" && baseType == "audio/mpeg" + if !validIdentity { + record.Eligible = false + record.Reason = CapabilityReasonMalformed + } + } + return sealCapabilityRecord(policy, data, record) +} + +// ValidateCapabilityRecord verifies canonical authority fields and checksum. +func ValidateCapabilityRecord(record CapabilityRecord) error { + if record.Version != capabilityRecordVersion { + return errors.New("media: capability record version is invalid") + } + for subject, value := range map[string]string{ + "source SHA-256": record.SourceSHA256, "policy fingerprint": record.PolicyFingerprint, + "descriptor fingerprint": record.DescriptorFingerprint, + "profile fingerprint": record.ProfileFingerprint, + "disclosure fingerprint": record.DisclosureFingerprint, "checksum": record.Checksum, + } { + if !validSHA256(value) { + return fmt.Errorf("media: capability record %s is invalid", subject) + } + } + if record.SourceBytes <= 0 { + return errors.New("media: capability record source bytes are invalid") + } + if err := validateInspectionPolicy(record.Policy); err != nil { + return fmt.Errorf("media: capability record policy is invalid: %w", err) + } + policyEncoded, err := json.Marshal(record.Policy, json.Deterministic(true)) + if err != nil { + return fmt.Errorf("media: encode capability record policy: %w", err) + } + if sha256Hex(policyEncoded) != record.PolicyFingerprint { + return errors.New("media: capability record policy fingerprint does not match policy") + } + if record.DescriptorFingerprint != record.Policy.DescriptorFingerprint || + record.ProfileFingerprint != record.Policy.ProfileFingerprint || + record.DisclosureFingerprint != record.Policy.DisclosureFingerprint || + record.InputKind != record.Policy.InputKind { + return errors.New("media: capability record authority does not match policy") + } + identity := record + identity.Checksum = "" + identity.localAuthority = false + encoded, err := json.Marshal(capabilityRecordIdentity(identity), json.Deterministic(true)) + if err != nil { + return fmt.Errorf("media: encode capability record: %w", err) + } + if sha256Hex(encoded) != record.Checksum { + return errors.New("media: capability record checksum is invalid") + } + if record.Eligible != (record.Reason == CapabilityReasonEligible) { + return errors.New("media: capability eligibility and reason disagree") + } + return nil +} + +// InspectionPolicy returns the policy sealed into a locally produced record. +// Records decoded from an external representation deliberately lack this +// authority and cannot authorize an upload. +func (record CapabilityRecord) InspectionPolicy() (InspectionPolicy, bool) { + return record.Policy, record.localAuthority +} + +func sealCapabilityRecord( + policy InspectionPolicy, data []byte, record CapabilityRecord, +) (CapabilityRecord, error) { + digest := sha256.Sum256(data) + policyEncoded, err := json.Marshal(policy, json.Deterministic(true)) + if err != nil { + return CapabilityRecord{}, fmt.Errorf("media: encode inspection policy: %w", err) + } + record.Version = capabilityRecordVersion + record.SourceBytes = int64(len(data)) + record.SourceSHA256 = hex.EncodeToString(digest[:]) + record.PolicyFingerprint = sha256Hex(policyEncoded) + record.DescriptorFingerprint = policy.DescriptorFingerprint + record.ProfileFingerprint = policy.ProfileFingerprint + record.DisclosureFingerprint = policy.DisclosureFingerprint + record.InputKind = policy.InputKind + record.Policy = policy + record.Checksum = "" + record.localAuthority = false + encoded, err := json.Marshal(capabilityRecordIdentity(record), json.Deterministic(true)) + if err != nil { + return CapabilityRecord{}, fmt.Errorf("media: encode capability record: %w", err) + } + record.Checksum = sha256Hex(encoded) + record.localAuthority = true + return record, nil +} + +func validateInspectionPolicy(policy InspectionPolicy) error { + if policy.Filename == "" || policy.Filename == "." || policy.Filename == ".." || + strings.ContainsAny(policy.Filename, "/\\\x00") || + filepath.Base(policy.Filename) != policy.Filename { + return errors.New("media: inspection filename must be a base name") + } + baseType, _, err := mime.ParseMediaType(policy.DeclaredMediaType) + if err != nil || baseType == "" { + return errors.New("media: declared media type is invalid") + } + if policy.ExpectedBytes <= 0 || policy.ExpectedBytes > maxInspectionSourceBytes || + policy.MaxSourceBytes <= 0 || policy.MaxSourceBytes > maxInspectionSourceBytes || + policy.ExpectedBytes > policy.MaxSourceBytes { + return errors.New("media: source byte bounds are invalid") + } + for subject, value := range map[string]string{ + "expected SHA-256": policy.ExpectedSHA256, + "descriptor fingerprint": policy.DescriptorFingerprint, + "profile fingerprint": policy.ProfileFingerprint, + "disclosure fingerprint": policy.DisclosureFingerprint, + } { + if !validSHA256(value) { + return fmt.Errorf("media: %s is invalid", subject) + } + } + if policy.InputKind != document.RenditionInputOriginalFile && + policy.InputKind != document.RenditionInputDerivedUpload { + return errors.New("media: inspection input kind is invalid") + } + limits := []int64{policy.MaxExpandedBytes, policy.MaxEntryBytes, policy.MaxEntries, + policy.MaxNestingDepth, policy.MaxTextLines, policy.MaxCharacters, policy.MaxPages, policy.MaxSlides, + policy.MaxSheets, policy.MaxCells, policy.MaxSpineItems, policy.MaxResources} + for _, limit := range limits { + if limit <= 0 { + return errors.New("media: every finite inspection limit must be positive") + } + } + if policy.MaxNestingDepth != 1 { + return errors.New("media: the registered container inspector requires max nesting depth one") + } + if policy.MaxDurationMS < 0 || policy.MaxFrames < 0 || policy.MaxPixels < 0 || policy.MaxRecords < 0 { + return errors.New("media: optional inspection limits must not be negative") + } + return nil +} + +func inspectText(data []byte, ext, mediaType string, policy InspectionPolicy) CapabilityRecord { + record := CapabilityRecord{MediaFamily: "text", MediaType: mediaType, + Format: strings.TrimPrefix(ext, ".")} + if !utf8.Valid(data) || bytes.IndexByte(data, 0) >= 0 { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.Characters = int64(utf8.RuneCount(data)) + record.Measurements.TextLines = int64(bytes.Count(data, []byte{'\n'})) + if len(data) != 0 && data[len(data)-1] != '\n' { + record.Measurements.TextLines++ + } + switch mediaType { + case "text/csv": + csvReader := csv.NewReader(bytes.NewReader(data)) + csvReader.FieldsPerRecord = -1 + for { + _, err := csvReader.Read() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.Records++ + } + case "application/x-ndjson", "application/jsonl": + for line := range bytes.SplitSeq(data, []byte{'\n'}) { + if len(bytes.TrimSpace(line)) != 0 { + record.Measurements.Records++ + } + } + case "application/json", "application/xml", "message/rfc822": + record.Measurements.Records = 1 + } + if mediaType == "application/xml" { + external, err := inspectXML(data, &record.Measurements, xmlMeasureNone) + if err != nil { + record.Reason = CapabilityReasonMalformed + return record + } + if external { + record.Reason = CapabilityReasonExternalReference + return record + } + } + if record.Measurements.TextLines > policy.MaxTextLines || + record.Measurements.Characters > policy.MaxCharacters || + (policy.MaxRecords > 0 && record.Measurements.Records > policy.MaxRecords) { + record.Reason = CapabilityReasonSemanticUnits + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +func inspectZIP(data []byte, ext, mediaType string, policy InspectionPolicy) CapabilityRecord { + record := CapabilityRecord{MediaType: mediaType, Format: strings.TrimPrefix(ext, ".")} + switch ext { + case ".pptx", ".pptm", ".odp": + record.MediaFamily = "presentation" + case ".xlsx", ".xlsm", ".ods": + record.MediaFamily = "spreadsheet" + case ".epub": + record.MediaFamily = "ebook" + default: + record.Reason = CapabilityReasonUnboundedFamily + return record + } + zr, err := zip.NewReader(bytes.NewReader(data), int64(len(data))) + if err != nil { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.CompressedBytes = int64(len(data)) + record.Measurements.Entries = int64(len(zr.File)) + if record.Measurements.Entries > policy.MaxEntries { + record.Reason = CapabilityReasonEntryCount + return record + } + for _, file := range zr.File { + if file.Flags&1 != 0 { + record.Reason = CapabilityReasonEncryptedContainer + return record + } + if file.UncompressedSize64 > math.MaxInt64 { + record.Reason = CapabilityReasonEntryBytes + return record + } + size := int64(file.UncompressedSize64) // #nosec G115 -- bounded above by MaxInt64 + if size > policy.MaxEntryBytes { + record.Reason = CapabilityReasonEntryBytes + return record + } + if size > record.Measurements.MaxEntryBytes { + record.Measurements.MaxEntryBytes = size + } + if size > policy.MaxExpandedBytes-record.Measurements.ExpandedBytes { + record.Reason = CapabilityReasonExpandedBytes + return record + } + record.Measurements.ExpandedBytes += size + if looksNestedContainer(file.Name) { + record.Measurements.NestingDepth = 2 + record.Reason = CapabilityReasonNestedContainer + return record + } + body, readErr := readZIPEntry(file, policy.MaxEntryBytes) + if readErr != nil { + record.Reason = CapabilityReasonMalformed + return record + } + if len(body) >= 4 && bytes.Equal(body[:4], []byte("PK\x03\x04")) { + record.Measurements.NestingDepth = 2 + record.Reason = CapabilityReasonNestedContainer + return record + } + name := strings.ToLower(file.Name) + xmlContent := strings.HasSuffix(name, ".xml") || strings.HasSuffix(name, ".rels") || + strings.HasSuffix(name, ".opf") || ext == ".epub" && + (strings.HasSuffix(name, ".xhtml") || strings.HasSuffix(name, ".html") || + strings.HasSuffix(name, ".htm") || strings.HasSuffix(name, ".svg")) + if xmlContent { + mode := xmlMeasureNone + switch { + case strings.HasPrefix(name, "xl/worksheets/"): + mode = xmlMeasureOOXMLSheet + case name == "content.xml" && ext == ".ods": + mode = xmlMeasureODSSheet + case strings.HasSuffix(name, ".opf") && ext == ".epub": + mode = xmlMeasureEPUBPackage + } + external, xmlErr := inspectXML(body, &record.Measurements, mode) + if xmlErr != nil { + record.Reason = CapabilityReasonMalformed + return record + } + if external { + record.Reason = CapabilityReasonExternalReference + return record + } + } + if ext == ".epub" && strings.HasSuffix(name, ".css") && hasCSSReference(body) { + record.Reason = CapabilityReasonExternalReference + return record + } + if strings.HasPrefix(name, "ppt/slides/slide") && strings.HasSuffix(name, ".xml") { + record.Measurements.Slides++ + } + if strings.HasPrefix(name, "xl/worksheets/") && strings.HasSuffix(name, ".xml") { + record.Measurements.Sheets++ + } + } + if record.Measurements.Slides > policy.MaxSlides || + record.Measurements.Sheets > policy.MaxSheets || record.Measurements.Cells > policy.MaxCells || + record.Measurements.SpineItems > policy.MaxSpineItems || + record.Measurements.Resources > policy.MaxResources { + record.Reason = CapabilityReasonSemanticUnits + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +type xmlMeasurement uint8 + +const ( + xmlMeasureNone xmlMeasurement = iota + xmlMeasureOOXMLSheet + xmlMeasureODSSheet + xmlMeasureEPUBPackage +) + +func inspectXML( + data []byte, measurements *CapabilityMeasurements, mode xmlMeasurement, +) (bool, error) { + decoder := xml.NewDecoder(bytes.NewReader(data)) + for { + token, err := decoder.Token() + if err != nil { + if errors.Is(err, io.EOF) { + return false, nil + } + return false, fmt.Errorf("inspect XML token: %w", err) + } + if directive, ok := token.(xml.Directive); ok && + strings.Contains(strings.ToUpper(string(directive)), "DOCTYPE") { + return true, nil + } + start, ok := token.(xml.StartElement) + if !ok { + continue + } + switch name := strings.ToLower(start.Name.Local); { + case mode == xmlMeasureOOXMLSheet && name == "c": + measurements.Cells++ + case mode == xmlMeasureODSSheet && name == "table": + measurements.Sheets++ + case mode == xmlMeasureODSSheet && name == "table-cell": + measurements.Cells++ + case mode == xmlMeasureEPUBPackage && name == "itemref": + measurements.SpineItems++ + case mode == xmlMeasureEPUBPackage && name == "item": + measurements.Resources++ + } + for _, attribute := range start.Attr { + value := strings.TrimSpace(attribute.Value) + name := strings.ToLower(attribute.Name.Local) + if name == "targetmode" && strings.EqualFold(value, "External") { + return true, nil + } + if name != "target" && name != "href" && name != "src" && name != "schemalocation" { + continue + } + parsed, parseErr := url.Parse(value) + if strings.HasPrefix(value, "//") || parseErr == nil && parsed.IsAbs() { + return true, nil + } + } + } +} + +func inspectVisualCapability(data []byte, declaredType string, policy InspectionPolicy) CapabilityRecord { + metadata, err := DetectBytes(data, declaredType) + if err != nil { + return CapabilityRecord{Reason: CapabilityReasonMalformed, MediaType: declaredType} + } + if declaredType == "video/quicktime" && + (metadata.Kind != KindVideo || metadata.Container != "quicktime") { + return CapabilityRecord{Reason: CapabilityReasonMalformed, MediaType: declaredType} + } + record := CapabilityRecord{MediaFamily: string(metadata.Kind), MediaType: metadata.MediaType, + Format: string(metadata.Format), Measurements: CapabilityMeasurements{ + Pixels: metadata.Pixels(), Frames: int64(metadata.FrameCount), DurationMS: metadata.DurationMS, + }} + if metadata.Kind == KindVideo { + info, ok := mp4Metadata(data) + if !ok || !info.sampleAuthority { + return CapabilityRecord{Reason: CapabilityReasonMalformed, MediaType: declaredType} + } + record.Measurements.Frames = info.frameCount + } + if policy.MaxPixels <= 0 || record.Measurements.Pixels > policy.MaxPixels || + policy.MaxFrames <= 0 || record.Measurements.Frames <= 0 || + record.Measurements.Frames > policy.MaxFrames || + metadata.Kind == KindVideo && (policy.MaxDurationMS <= 0 || !metadata.DurationKnown || + metadata.DurationMS > policy.MaxDurationMS) { + record.Reason = CapabilityReasonVisualBounds + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +func inspectPDF(data []byte, policy InspectionPolicy) CapabilityRecord { + record := CapabilityRecord{MediaFamily: "pdf", MediaType: "application/pdf", Format: "pdf"} + pages, err := formatdetect.CountPDFPages(data) + if err != nil { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.Pages = pages + if record.Measurements.Pages <= 0 { + record.Reason = CapabilityReasonMalformed + return record + } + if record.Measurements.Pages > policy.MaxPages { + record.Reason = CapabilityReasonSemanticUnits + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +// CountPDFPages resolves the catalog page tree without treating page-like +// bytes in streams, strings, comments, or unrelated objects as authority. +func CountPDFPages(data []byte) (int64, error) { + return formatdetect.CountPDFPages(data) +} + +// PDFInfoFields resolves string fields only from the PDF's final +// cross-reference authority, excluding unrelated objects and stream content. +func PDFInfoFields(data []byte) (map[string]string, error) { + return formatdetect.PDFInfoFields(data) +} + +func inspectWAV(data []byte, policy InspectionPolicy) CapabilityRecord { + record := CapabilityRecord{MediaFamily: "audio", MediaType: "audio/wav", Format: "wav"} + if len(data) < 44 || !bytes.Equal(data[:4], []byte("RIFF")) || + !bytes.Equal(data[8:12], []byte("WAVE")) || + uint64(binary.LittleEndian.Uint32(data[4:8]))+8 != uint64(len(data)) { + record.Reason = CapabilityReasonMalformed + return record + } + var byteRate, sampleRate, audioBytes uint32 + var audioFormat, channels, blockAlign, bitsPerSample uint16 + var seenFormat, seenAudio bool + offset := 12 + for offset+8 <= len(data) { + chunkSize := binary.LittleEndian.Uint32(data[offset+4 : offset+8]) + if uint64(chunkSize) > uint64(len(data)) { + record.Reason = CapabilityReasonMalformed + return record + } + size := int(chunkSize) // #nosec G115 -- bounded by len(data), which is an int + body := offset + 8 + if size < 0 || body+size > len(data) { + record.Reason = CapabilityReasonMalformed + return record + } + switch string(data[offset : offset+4]) { + case "fmt ": + if seenFormat || size < 16 { + record.Reason = CapabilityReasonMalformed + return record + } + seenFormat = true + audioFormat = binary.LittleEndian.Uint16(data[body : body+2]) + channels = binary.LittleEndian.Uint16(data[body+2 : body+4]) + sampleRate = binary.LittleEndian.Uint32(data[body+4 : body+8]) + byteRate = binary.LittleEndian.Uint32(data[body+8 : body+12]) + blockAlign = binary.LittleEndian.Uint16(data[body+12 : body+14]) + bitsPerSample = binary.LittleEndian.Uint16(data[body+14 : body+16]) + case "data": + if seenAudio { + record.Reason = CapabilityReasonMalformed + return record + } + seenAudio = true + audioBytes = chunkSize + } + offset = body + size + size%2 + } + bytesPerSample := uint64(bitsPerSample) / 8 + wantBlockAlign := uint64(channels) * bytesPerSample + wantByteRate := uint64(sampleRate) * wantBlockAlign + if offset != len(data) || !seenFormat || !seenAudio || + (audioFormat != 1 && audioFormat != 3) || channels == 0 || sampleRate == 0 || + bitsPerSample == 0 || bitsPerSample%8 != 0 || wantBlockAlign == 0 || + wantBlockAlign > math.MaxUint16 || uint64(blockAlign) != wantBlockAlign || + wantByteRate > math.MaxUint32 || uint64(byteRate) != wantByteRate || + audioBytes == 0 || uint64(audioBytes)%wantBlockAlign != 0 || policy.MaxDurationMS <= 0 { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.DurationMS = (int64(audioBytes)*1000 + int64(byteRate) - 1) / int64(byteRate) + if record.Measurements.DurationMS > policy.MaxDurationMS { + record.Reason = CapabilityReasonVisualBounds + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +type mp3FrameHeader struct { + version byte + sampleRate uint64 + samples, length uint64 +} + +func inspectMP3(data []byte, policy InspectionPolicy) CapabilityRecord { + record := CapabilityRecord{MediaFamily: "audio", MediaType: "audio/mpeg", Format: "mp3"} + if policy.MaxDurationMS <= 0 { + record.Reason = CapabilityReasonMalformed + return record + } + audioStart, audioEnd, ok := mp3AudioBounds(data) + if !ok { + record.Reason = CapabilityReasonMalformed + return record + } + var reference mp3FrameHeader + var seenFrame bool + var totalSamples uint64 + for offset := audioStart; offset < audioEnd; { + frame, ok := parseMP3FrameHeader(data[offset:audioEnd]) + if !ok || seenFrame && (frame.version != reference.version || frame.sampleRate != reference.sampleRate) { + record.Reason = CapabilityReasonMalformed + return record + } + if !seenFrame { + reference = frame + seenFrame = true + } + remaining := uint64(audioEnd - offset) //nolint:gosec // offset is below audioEnd by the loop condition + if frame.length > remaining || math.MaxUint64-totalSamples < frame.samples { + record.Reason = CapabilityReasonMalformed + return record + } + totalSamples += frame.samples + offset += int(frame.length) // #nosec G115 -- bounded by the remaining input bytes above. + } + if totalSamples == 0 { + record.Reason = CapabilityReasonMalformed + return record + } + durationMS, ok := mp3DurationMilliseconds(totalSamples, reference.sampleRate) + if !ok { + record.Reason = CapabilityReasonMalformed + return record + } + record.Measurements.DurationMS = durationMS + if durationMS > policy.MaxDurationMS { + record.Reason = CapabilityReasonVisualBounds + return record + } + record.Eligible, record.Reason = true, CapabilityReasonEligible + return record +} + +func mp3AudioBounds(data []byte) (int, int, bool) { + start, end := 0, len(data) + if bytes.HasPrefix(data, []byte("ID3")) { + if len(data) < 10 { + return 0, 0, false + } + version, revision, flags := data[3], data[4], data[5] + var allowedFlags byte + switch version { + case 2: + allowedFlags = 0xc0 + case 3: + allowedFlags = 0xe0 + case 4: + allowedFlags = 0xf0 + default: + return 0, 0, false + } + if revision == 0xff || flags & ^allowedFlags != 0 { + return 0, 0, false + } + tagSize := 0 + for _, value := range data[6:10] { + if value&0x80 != 0 { + return 0, 0, false + } + tagSize = tagSize<<7 | int(value) + } + if tagSize > maxID3v2TagBytes { + return 0, 0, false + } + footerBytes := 0 + if version == 4 && flags&0x10 != 0 { + footerBytes = 10 + } + start = 10 + tagSize + footerBytes + if start > end { + return 0, 0, false + } + if footerBytes != 0 { + footer := data[start-footerBytes : start] + if !bytes.Equal(footer[:3], []byte("3DI")) || footer[3] != version || footer[4] != revision || + footer[5] != flags || !bytes.Equal(footer[6:10], data[6:10]) { + return 0, 0, false + } + } + } + if end-start >= 128 && bytes.Equal(data[end-128:end-125], []byte("TAG")) { + end -= 128 + } + return start, end, start < end +} + +func parseMP3FrameHeader(data []byte) (mp3FrameHeader, bool) { + if len(data) < 4 || data[0] != 0xff || data[1]&0xe0 != 0xe0 { + return mp3FrameHeader{}, false + } + versionID := (data[1] >> 3) & 0x03 + if versionID == 1 || (data[1]>>1)&0x03 != 1 { + return mp3FrameHeader{}, false + } + if data[3]&0x03 == 2 { // emphasis=2 is reserved by MPEG audio. + return mp3FrameHeader{}, false + } + bitrateIndex := (data[2] >> 4) & 0x0f + sampleRateIndex := (data[2] >> 2) & 0x03 + if bitrateIndex == 0 || bitrateIndex == 15 || sampleRateIndex == 3 { + return mp3FrameHeader{}, false + } + var sampleRate, bitrate, samples, coefficient uint64 + switch versionID { + case 3: // MPEG-1 Layer III. + sampleRate = []uint64{44_100, 48_000, 32_000}[sampleRateIndex] + bitrate = []uint64{32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320}[bitrateIndex-1] + samples, coefficient = 1_152, 144_000 + case 2: // MPEG-2 Layer III. + sampleRate = []uint64{22_050, 24_000, 16_000}[sampleRateIndex] + bitrate = []uint64{8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160}[bitrateIndex-1] + samples, coefficient = 576, 72_000 + case 0: // MPEG-2.5 Layer III. + sampleRate = []uint64{11_025, 12_000, 8_000}[sampleRateIndex] + bitrate = []uint64{8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160}[bitrateIndex-1] + samples, coefficient = 576, 72_000 + default: + return mp3FrameHeader{}, false + } + length := coefficient*bitrate/sampleRate + uint64((data[2]>>1)&1) + return mp3FrameHeader{version: versionID, sampleRate: sampleRate, samples: samples, length: length}, length >= 4 +} + +func mp3DurationMilliseconds(samples, sampleRate uint64) (int64, bool) { + if samples == 0 || sampleRate == 0 || samples > math.MaxUint64/1_000 { + return 0, false + } + milliseconds := (samples*1_000 + sampleRate - 1) / sampleRate + if milliseconds > math.MaxInt64 { + return 0, false + } + return int64(milliseconds), true +} + +func hasCSSReference(data []byte) bool { + lower := strings.ToLower(string(data)) + return strings.Contains(lower, "url") || strings.Contains(lower, "@import") +} + +func readZIPEntry(file *zip.File, limit int64) ([]byte, error) { + reader, err := file.Open() + if err != nil { + return nil, fmt.Errorf("open ZIP entry: %w", err) + } + defer func() { _ = reader.Close() }() + data, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if err != nil || int64(len(data)) > limit || uint64(len(data)) != file.UncompressedSize64 { + return nil, errors.New("ZIP entry exceeds bound") + } + return data, nil +} + +func isTextFamily(ext, mediaType string) bool { + candidate, ok := formatdetect.CandidateFormatByMediaType(mediaType) + if ok && slicesContains([]string{ + "txt", "markdown", "rst", "latex", "json", "jsonl", "xml", "yaml", + "go", "python", "javascript", "eml", "csv", + }, candidate.ID) { + return true + } + return slicesContains([]string{".txt", ".md", ".csv", ".json", ".jsonl", ".xml", ".yaml", ".yml", ".eml"}, ext) +} + +func isZIPFamily(ext, mediaType string, data []byte) bool { + return len(data) >= 4 && bytes.Equal(data[:4], []byte("PK\x03\x04")) || + strings.Contains(mediaType, "officedocument") || mediaType == "application/epub+zip" || + slicesContains([]string{".pptx", ".pptm", ".xlsx", ".xlsm", ".odp", ".ods", ".epub"}, ext) +} + +func looksNestedContainer(name string) bool { + ext := strings.ToLower(filepath.Ext(name)) + return slicesContains([]string{".zip", ".epub", ".pptx", ".xlsx", ".docx", ".ods", ".odp"}, ext) +} + +func requiresDocumentFormatDetection(family string) bool { + return slicesContains([]string{"text", "structured", "source", "mail", "pdf", "presentation", "spreadsheet", "ebook"}, family) +} + +func filenameAllowsFormat(extension, format string) bool { + allowed := map[string][]string{ + "pdf": {".pdf"}, "pptx": {".pptx"}, "xlsx": {".xlsx"}, "ods": {".ods"}, + "epub": {".epub"}, "csv": {".csv"}, "txt": {".txt"}, + "markdown": {".md", ".markdown"}, "rst": {".rst"}, "latex": {".tex", ".latex"}, + "json": {".json"}, "jsonl": {".jsonl", ".ndjson"}, "xml": {".xml"}, + "yaml": {".yaml", ".yml"}, "go": {".go"}, "python": {".py"}, + "javascript": {".js", ".mjs", ".cjs"}, "eml": {".eml"}, + } + return slices.Contains(allowed[format], extension) +} + +func filenameAllowsVisualFormat(extension, format string) bool { + allowed := map[string][]string{ + string(FormatJPEG): {".jpg", ".jpeg"}, string(FormatPNG): {".png"}, + string(FormatWebP): {".webp"}, string(FormatGIF): {".gif"}, string(FormatMP4): {".mp4"}, + } + return slices.Contains(allowed[format], extension) +} + +func slicesContains(values []string, candidate string) bool { + return slices.Contains(values, candidate) +} + +func validSHA256(value string) bool { + if len(value) != sha256.Size*2 || strings.ToLower(value) != value { + return false + } + _, err := hex.DecodeString(value) + return err == nil +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/document/media/inspect_test.go b/document/media/inspect_test.go new file mode 100644 index 00000000..33203843 --- /dev/null +++ b/document/media/inspect_test.go @@ -0,0 +1,1166 @@ +package media_test + +import ( + "archive/zip" + "bytes" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json/v2" + "fmt" + "math" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/media/mediatest" +) + +func TestInspectBindsFiniteTextToPolicyAndSource(t *testing.T) { + t.Parallel() + data := []byte("alpha\nbeta\n") + policy := inspectionPolicy(data, "notes.txt", "text/plain") + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + require.True(t, record.Eligible) + assert.Equal(t, "text", record.MediaFamily) + assert.Equal(t, "text/plain", record.MediaType) + assert.Equal(t, int64(2), record.Measurements.TextLines) + assert.Equal(t, int64(11), record.Measurements.Characters) + assert.Equal(t, sha256Hex(data), record.SourceSHA256) + assert.NotEmpty(t, record.PolicyFingerprint) + assert.NotEmpty(t, record.Checksum) + require.NoError(t, media.ValidateCapabilityRecord(record)) + + mutated := record + mutated.Measurements.TextLines++ + require.ErrorContains(t, media.ValidateCapabilityRecord(mutated), "checksum") + + encoded, err := json.Marshal(record, json.Deterministic(true)) + require.NoError(t, err) + assert.Contains(t, string(encoded), `"max_text_lines":1000`) + var decoded media.CapabilityRecord + require.NoError(t, json.Unmarshal(encoded, &decoded)) + assert.Equal(t, policy, decoded.Policy) + require.NoError(t, media.ValidateCapabilityRecord(decoded)) + _, local := decoded.InspectionPolicy() + assert.False(t, local, "a portable record must not recreate local upload authority") +} + +func TestInspectRejectsDeclaredSourceMismatchAndExcessiveUnits(t *testing.T) { + t.Parallel() + data := []byte("alpha\nbeta\n") + policy := inspectionPolicy(data, "notes.txt", "text/plain") + policy.ExpectedSHA256 = strings.Repeat("0", 64) + _, err := media.InspectCapability(bytes.NewReader(data), policy) + require.ErrorContains(t, err, "SHA-256") + + policy = inspectionPolicy(data, "notes.txt", "text/plain") + policy.MaxTextLines = 1 + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonSemanticUnits, record.Reason) + + jsonl := []byte("{\"a\":1}\n{\"b\":2}\n") + policy = inspectionPolicy(jsonl, "records.jsonl", "application/x-ndjson") + policy.MaxRecords = 1 + record, err = media.InspectCapability(bytes.NewReader(jsonl), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, int64(2), record.Measurements.Records) + assert.Equal(t, media.CapabilityReasonSemanticUnits, record.Reason) +} + +func TestInspectBoundsZIPContainers(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + data []byte + policy func([]byte) media.InspectionPolicy + reason media.CapabilityReason + }{ + { + name: "aggregate expansion", + data: zipBytes(t, validPPTXEntries(zipEntry{name: "slides/one.xml", body: strings.Repeat("x", 32)})), + policy: func(data []byte) media.InspectionPolicy { + p := inspectionPolicy(data, "deck.pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation") + p.MaxExpandedBytes = 16 + return p + }, + reason: media.CapabilityReasonExpandedBytes, + }, + { + name: "nested archive", + data: zipBytes(t, validPPTXEntries(zipEntry{name: "nested.zip", body: "PK\x03\x04nested"})), + policy: func(data []byte) media.InspectionPolicy { + return inspectionPolicy(data, "deck.pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation") + }, + reason: media.CapabilityReasonNestedContainer, + }, + { + name: "external relationship", + data: zipBytes(t, validPPTXEntries(zipEntry{name: "_rels/.rels", body: ``})), + policy: func(data []byte) media.InspectionPolicy { + return inspectionPolicy(data, "deck.pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation") + }, + reason: media.CapabilityReasonExternalReference, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + record, err := media.InspectCapability(bytes.NewReader(tt.data), tt.policy(tt.data)) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, tt.reason, record.Reason) + if tt.reason == media.CapabilityReasonNestedContainer { + assert.Equal(t, int64(2), record.Measurements.NestingDepth) + } + }) + } +} + +func TestInspectRejectsEncryptedZIPEntry(t *testing.T) { + t.Parallel() + data := zipBytes(t, validPPTXEntries(zipEntry{name: "slides/one.xml", body: "safe"})) + // Both the local and central directory general-purpose flags carry bit 0. + data[6] |= 1 + central := bytes.Index(data, []byte("PK\x01\x02")) + require.NotEqual(t, -1, central) + data[central+8] |= 1 + + record, err := media.InspectCapability(bytes.NewReader(data), inspectionPolicy(data, "deck.pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation")) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonEncryptedContainer, record.Reason) +} + +func TestInspectRejectsMalformedAndExternallyReferentialContainerXML(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + reason media.CapabilityReason + }{ + {name: "malformed XML", body: ``, reason: media.CapabilityReasonMalformed}, + {name: "external DTD", body: ``, reason: media.CapabilityReasonExternalReference}, + {name: "public DTD", body: ``, reason: media.CapabilityReasonExternalReference}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + data := zipBytes(t, []zipEntry{{name: "xl/worksheets/sheet1.xml", body: tt.body}}) + record, err := media.InspectCapability(bytes.NewReader(data), inspectionPolicy(data, "book.xlsx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, tt.reason, record.Reason) + }) + } +} + +func TestInspectRejectsExternalReferenceInEPUBContent(t *testing.T) { + t.Parallel() + data := zipBytes(t, validEPUBEntries( + zipEntry{name: "OPS/content.opf", body: ``}, + zipEntry{name: "OPS/chapter.xhtml", body: ``}, + )) + record, err := media.InspectCapability(bytes.NewReader(data), + inspectionPolicy(data, "book.epub", "application/epub+zip")) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonExternalReference, record.Reason) +} + +func TestInspectCountsFinitePresentationSpreadsheetAndEPUBUnits(t *testing.T) { + t.Parallel() + tests := []struct { + name string + filename string + mediaType string + entries []zipEntry + family string + assertions func(*testing.T, media.CapabilityRecord) + }{ + { + name: "slides", filename: "deck.pptx", mediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", + entries: validPPTXEntries( + zipEntry{name: "ppt/slides/slide1.xml", body: ""}, + zipEntry{name: "ppt/slides/slide2.xml", body: ""}, + ), family: "presentation", + assertions: func(t *testing.T, r media.CapabilityRecord) { + t.Helper() + assert.Equal(t, int64(2), r.Measurements.Slides) + }, + }, + { + name: "sheets and cells", filename: "book.xlsx", mediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + entries: validXLSXEntries(zipEntry{name: "xl/worksheets/sheet1.xml", body: ``}), family: "spreadsheet", + assertions: func(t *testing.T, r media.CapabilityRecord) { + t.Helper() + assert.Equal(t, int64(1), r.Measurements.Sheets) + assert.Equal(t, int64(2), r.Measurements.Cells) + }, + }, + { + name: "spine and resources", filename: "book.epub", mediaType: "application/epub+zip", + entries: validEPUBEntries(zipEntry{name: "OPS/content.opf", body: ``}), family: "ebook", + assertions: func(t *testing.T, r media.CapabilityRecord) { + t.Helper() + assert.Equal(t, int64(1), r.Measurements.SpineItems) + assert.Equal(t, int64(2), r.Measurements.Resources) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + data := zipBytes(t, tt.entries) + record, err := media.InspectCapability(bytes.NewReader(data), inspectionPolicy(data, tt.filename, tt.mediaType)) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, tt.family, record.MediaFamily) + tt.assertions(t, record) + }) + } +} + +func TestInspectRejectsDeceptiveDeclaredFormatAndFilename(t *testing.T) { + t.Parallel() + pdf := syntheticPDF("deceptive") + tests := []struct { + name, filename, mediaType string + data []byte + }{ + {name: "PDF with text filename", filename: "report.txt", mediaType: "application/pdf", data: pdf}, + {name: "text declared as PDF", filename: "report.pdf", mediaType: "application/pdf", data: []byte("not a PDF")}, + {name: "OOXML declared generic ZIP", filename: "deck.pptx", mediaType: "application/zip", data: zipBytes(t, validPPTXEntries())}, + {name: "PPTX without main markers", filename: "deck.pptx", mediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", data: zipBytes(t, []zipEntry{{name: "ppt/slides/slide1.xml", body: ""}})}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + record, err := media.InspectCapability(bytes.NewReader(tt.data), inspectionPolicy(tt.data, tt.filename, tt.mediaType)) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + } +} + +func TestInspectBindsVisualAndStandaloneXMLIdentity(t *testing.T) { + t.Parallel() + png := mediatest.PNG(4, 3, nil) + for _, filename := range []string{"image.txt", "image.jpeg"} { + policy := inspectionPolicy(png, filename, "image/png") + policy.MaxPixels, policy.MaxFrames = 100, 2 + record, err := media.InspectCapability(bytes.NewReader(png), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + } + policy := inspectionPolicy(png, "image.png", "image/png") + policy.MaxPixels, policy.MaxFrames = 100, 2 + record, err := media.InspectCapability(bytes.NewReader(png), policy) + require.NoError(t, err) + assert.True(t, record.Eligible) + + xml := []byte(``) + record, err = media.InspectCapability(bytes.NewReader(xml), inspectionPolicy(xml, "record.xml", "application/xml")) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonExternalReference, record.Reason) +} + +func TestInspectRejectsUnregisteredFamily(t *testing.T) { + t.Parallel() + data := []byte("legacy binary document") + record, err := media.InspectCapability(bytes.NewReader(data), inspectionPolicy(data, "report.doc", "application/msword")) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonUnboundedFamily, record.Reason) +} + +func TestInspectBoundsPDFPagesAndWAVDuration(t *testing.T) { + t.Parallel() + pdf := syntheticPDFPages(2, "bounds") + pdfPolicy := inspectionPolicy(pdf, "report.pdf", "application/pdf") + pdfPolicy.MaxPages = 1 + record, err := media.InspectCapability(bytes.NewReader(pdf), pdfPolicy) + require.NoError(t, err) + assert.Equal(t, int64(2), record.Measurements.Pages) + assert.Equal(t, media.CapabilityReasonSemanticUnits, record.Reason) + + wav := wavBytes(8_000, 8_000) + wavPolicy := inspectionPolicy(wav, "sample.wav", "audio/wav") + wavPolicy.MaxDurationMS = 500 + record, err = media.InspectCapability(bytes.NewReader(wav), wavPolicy) + require.NoError(t, err) + assert.Equal(t, int64(1_000), record.Measurements.DurationMS) + assert.Equal(t, media.CapabilityReasonVisualBounds, record.Reason) +} + +func TestInspectRejectsForgedWAVRateAndTrailingBytes(t *testing.T) { + t.Parallel() + tests := map[string]func([]byte) []byte{ + "forged byte rate": func(data []byte) []byte { + binary.LittleEndian.PutUint32(data[28:32], 1<<30) + return data + }, + "trailing bytes outside RIFF": func(data []byte) []byte { + return append(data, []byte("trailing")...) + }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + wav := mutate(wavBytes(8_000, 8_000)) + policy := inspectionPolicy(wav, "sample.wav", "audio/wav") + policy.MaxDurationMS = 2_000 + record, err := media.InspectCapability(bytes.NewReader(wav), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + } +} + +func TestInspectRejectsPathLikePortableFilename(t *testing.T) { + t.Parallel() + data := []byte("alpha\n") + for _, filename := range []string{".", "..", `folder\notes.txt`, "folder/notes.txt"} { + policy := inspectionPolicy(data, filename, "text/plain") + _, err := media.InspectCapability(bytes.NewReader(data), policy) + require.ErrorContains(t, err, "filename", filename) + } +} + +func TestInspectCountsOnlyPDFPageTreeObjects(t *testing.T) { + t.Parallel() + pdf := syntheticPDF("/Type /Page in a comment") + policy := inspectionPolicy(pdf, "report.pdf", "application/pdf") + policy.MaxPages = 1 + record, err := media.InspectCapability(bytes.NewReader(pdf), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, int64(1), record.Measurements.Pages) +} + +func TestInspectCountsXRefStreamPDFPageTree(t *testing.T) { + t.Parallel() + for _, testCase := range []struct { + name string + index string + }{ + {name: "default Index"}, + {name: "explicit Index sections", index: "/Index [0 3 3 3]"}, + } { + t.Run(testCase.name, func(t *testing.T) { + pdf := syntheticXRefStreamPDFWith(xrefStreamFixture{index: testCase.index}) + + pages, err := media.CountPDFPages(pdf) + require.NoError(t, err) + assert.Equal(t, int64(1), pages) + info, err := media.PDFInfoFields(pdf) + require.NoError(t, err) + assert.Equal(t, map[string]string{"Title": "Synthetic xref stream"}, info) + + policy := inspectionPolicy(pdf, "report.pdf", "application/pdf") + policy.MaxPages = 1 + record, err := media.InspectCapability(bytes.NewReader(pdf), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, int64(1), record.Measurements.Pages) + }) + } +} + +func TestInspectRejectsUnprovenXRefStreamAuthority(t *testing.T) { + t.Parallel() + for _, testCase := range []struct { + name string + fixture xrefStreamFixture + }{ + { + name: "short stream does not cover Size", + fixture: xrefStreamFixture{mutate: func(entries []byte, _ []int) []byte { + return entries[:7] + }}, + }, + { + name: "root entry has forged offset", + fixture: xrefStreamFixture{mutate: func(entries []byte, offsets []int) []byte { + putSyntheticXRefEntry(entries, 1, 1, uint32(offsets[3]), 0) + return entries + }}, + }, + { + name: "info entry has forged offset", + fixture: xrefStreamFixture{mutate: func(entries []byte, offsets []int) []byte { + putSyntheticXRefEntry(entries, 4, 1, uint32(offsets[0]), 0) + return entries + }}, + }, + { + name: "root entry uses a compressed object", + fixture: xrefStreamFixture{mutate: func(entries []byte, _ []int) []byte { + putSyntheticXRefEntry(entries, 1, 2, 5, 0) + return entries + }}, + }, + { + name: "filtered stream is unsupported", + fixture: xrefStreamFixture{dictionaryExtra: "/Filter /FlateDecode"}, + }, + { + name: "duplicate Index sections", + fixture: xrefStreamFixture{index: "/Index [0 4 3 3]", mutate: appendSyntheticXRefEntry(3)}, + }, + { + name: "out of range Index section", + fixture: xrefStreamFixture{index: "/Index [0 6 6 1]", mutate: appendSyntheticXRefEntry(0)}, + }, + { + name: "in-use entry points outside the document", + fixture: xrefStreamFixture{mutate: func(entries []byte, _ []int) []byte { + putSyntheticXRefEntry(entries, 2, 1, math.MaxUint32, 0) + return entries + }}, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + pdf := syntheticXRefStreamPDFWith(testCase.fixture) + _, err := media.CountPDFPages(pdf) + require.Error(t, err) + _, err = media.PDFInfoFields(pdf) + require.Error(t, err) + + policy := inspectionPolicy(pdf, "report.pdf", "application/pdf") + policy.MaxPages = 1 + record, err := media.InspectCapability(bytes.NewReader(pdf), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + } +} + +func TestInspectRejectsExcessivelyDeepPDFPageTree(t *testing.T) { + t.Parallel() + pdf := syntheticDeepPDF(300) + policy := inspectionPolicy(pdf, "report.pdf", "application/pdf") + policy.MaxPages = 1 + record, err := media.InspectCapability(bytes.NewReader(pdf), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +func TestInspectBoundsVideoFramesFromValidatedSampleTables(t *testing.T) { + t.Parallel() + video := mediatest.H265MP4() + decoy := make([]byte, 20) + binary.BigEndian.PutUint32(decoy[:4], 20) + copy(decoy[4:8], "stsz") + binary.BigEndian.PutUint32(decoy[16:20], 100) + video = append(video, mediatest.Box("free", decoy)...) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 1 + policy.MaxDurationMS = 1_000 + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, int64(1), record.Measurements.Frames) + + policy.MaxFrames = 0 + record, err = media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonVisualBounds, record.Reason) +} + +// TestInspectCapabilityRejectsVideoWithoutSampleToChunkAuthority catches a +// capability gate that accepts codec headers and sample sizes without proving +// how those samples are assigned to media chunks. +func TestInspectCapabilityRejectsVideoWithoutSampleToChunkAuthority(t *testing.T) { + t.Parallel() + video := decodableAVCMP4(t) + stsc := bytes.Index(video, []byte("stsc")) + require.NotEqual(t, -1, stsc) + copy(video[stsc:stsc+4], "free") + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsVideoWithoutChunkOffsets catches a capability +// gate that cannot establish where declared chunks reside in the source. +func TestInspectCapabilityRejectsVideoWithoutChunkOffsets(t *testing.T) { + t.Parallel() + video := decodableAVCMP4(t) + stco := bytes.Index(video, []byte("stco")) + require.NotEqual(t, -1, stco) + copy(video[stco:stco+4], "free") + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsVideoWithoutMediaData catches a capability gate +// that accepts chunk offsets without proving they point into an mdat payload. +func TestInspectCapabilityRejectsVideoWithoutMediaData(t *testing.T) { + t.Parallel() + video := decodableAVCMP4(t) + mdat := bytes.Index(video, []byte("mdat")) + require.NotEqual(t, -1, mdat) + copy(video[mdat:mdat+4], "free") + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsVideoChunkOutsideMediaData catches a capability +// gate that parses an offset table but never resolves its absolute positions. +func TestInspectCapabilityRejectsVideoChunkOutsideMediaData(t *testing.T) { + t.Parallel() + video := decodableAVCMP4(t) + stco := bytes.Index(video, []byte("stco")) + require.NotEqual(t, -1, stco) + binary.BigEndian.PutUint32(video[stco+12:stco+16], uint32(len(video)+1)) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsVideoSampleOutsideMediaData catches a capability +// gate that checks only each chunk's starting offset, not the declared sample +// sizes assigned to that chunk. +func TestInspectCapabilityRejectsVideoSampleOutsideMediaData(t *testing.T) { + t.Parallel() + video := decodableAVCMP4(t) + stsz := bytes.Index(video, []byte("stsz")) + require.NotEqual(t, -1, stsz) + binary.BigEndian.PutUint32(video[stsz+16:stsz+20], uint32(len(video))) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsOverlappingVideoSamples catches a capability +// gate that independently bounds samples but permits two chunk authorities to +// claim the same media bytes. +func TestInspectCapabilityRejectsOverlappingVideoSamples(t *testing.T) { + t.Parallel() + video := mp4WithOverlappingVideoChunks(t) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsSampleToChunkRunWithoutChunk catches a mapper +// that ignores a well-formed stsc run whose first chunk is absent from the +// authoritative offset table. +func TestInspectCapabilityRejectsSampleToChunkRunWithoutChunk(t *testing.T) { + t.Parallel() + video := mp4WithUnusedSampleToChunkRun(t) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +func mp4WithUnusedSampleToChunkRun(t *testing.T) []byte { + t.Helper() + video := decodableAVCMP4(t) + stsc := bytes.Index(video, []byte("stsc")) + require.NotEqual(t, -1, stsc) + boxStart := stsc - 4 + oldSize := int(binary.BigEndian.Uint32(video[boxStart:stsc])) + require.Equal(t, 28, oldSize) + replacement := make([]byte, 40) + binary.BigEndian.PutUint32(replacement[:4], uint32(len(replacement))) + copy(replacement[4:8], "stsc") + binary.BigEndian.PutUint32(replacement[12:16], 2) + copy(replacement[16:28], video[stsc+12:stsc+24]) + binary.BigEndian.PutUint32(replacement[28:32], 2) + binary.BigEndian.PutUint32(replacement[32:36], 1) + binary.BigEndian.PutUint32(replacement[36:40], 1) + video = append(append(append([]byte(nil), video[:boxStart]...), replacement...), video[boxStart+oldSize:]...) + for _, kind := range []string{"stbl", "minf", "mdia", "trak", "moov"} { + index := bytes.Index(video, []byte(kind)) + require.GreaterOrEqual(t, index, 4) + size := binary.BigEndian.Uint32(video[index-4 : index]) + binary.BigEndian.PutUint32(video[index-4:index], size+12) + } + stco := bytes.Index(video, []byte("stco")) + require.NotEqual(t, -1, stco) + offset := binary.BigEndian.Uint32(video[stco+12 : stco+16]) + binary.BigEndian.PutUint32(video[stco+12:stco+16], offset+12) + return video +} + +func mp4WithOverlappingVideoChunks(t *testing.T) []byte { + t.Helper() + video := decodableAVCMP4(t) + stsc := bytes.Index(video, []byte("stsc")) + stco := bytes.Index(video, []byte("stco")) + require.NotEqual(t, -1, stsc) + require.NotEqual(t, -1, stco) + binary.BigEndian.PutUint32(video[stsc+16:stsc+20], 1) + oldOffset := binary.BigEndian.Uint32(video[stco+12 : stco+16]) + replacement := make([]byte, 24) + binary.BigEndian.PutUint32(replacement[:4], uint32(len(replacement))) + copy(replacement[4:8], "stco") + binary.BigEndian.PutUint32(replacement[12:16], 2) + binary.BigEndian.PutUint32(replacement[16:20], oldOffset+4) + binary.BigEndian.PutUint32(replacement[20:24], oldOffset+4) + boxStart := stco - 4 + video = append(append(append([]byte(nil), video[:boxStart]...), replacement...), video[boxStart+20:]...) + for _, kind := range []string{"stbl", "minf", "mdia", "trak", "moov"} { + index := bytes.Index(video, []byte(kind)) + require.GreaterOrEqual(t, index, 4) + size := binary.BigEndian.Uint32(video[index-4 : index]) + binary.BigEndian.PutUint32(video[index-4:index], size+4) + } + return video +} + +// TestInspectCapabilityAcceptsVideoWith64BitChunkOffsets catches a sample +// authority implementation that narrows valid ISO BMFF layouts to stco even +// when an equivalent bounded co64 table is present. +func TestInspectCapabilityAcceptsVideoWith64BitChunkOffsets(t *testing.T) { + t.Parallel() + video := mp4With64BitChunkOffsets(t) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 2 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, int64(2), record.Measurements.Frames) +} + +// TestInspectCapabilityRejectsSamplesThatDoNotMatchVideoCodec catches a +// capability gate that proves container tables and codec configuration but +// never verifies that the mapped mdat sample is media for that codec. +func TestInspectCapabilityRejectsSamplesThatDoNotMatchVideoCodec(t *testing.T) { + t.Parallel() + for _, testCase := range []struct { + name, filename, mediaType string + data []byte + corrupt func([]byte, int) + }{ + { + name: "H264", filename: "clip.mov", mediaType: "video/quicktime", data: mediatest.H264MOV(), + corrupt: func(data []byte, sample int) { clear(data[sample : sample+4]) }, + }, + { + name: "H265", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.H265MP4(), + corrupt: func(data []byte, sample int) { clear(data[sample : sample+4]) }, + }, + { + name: "VP9", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.VP9MP4(), + corrupt: func(data []byte, sample int) { data[sample] = 0 }, + }, + { + name: "AV1", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.AV1MP4(), + corrupt: func(data []byte, sample int) { data[sample] |= 0x80 }, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + sampleOffset := firstMP4SampleOffset(t, testCase.data) + testCase.corrupt(testCase.data, sampleOffset) + policy := inspectionPolicy(testCase.data, testCase.filename, testCase.mediaType) + policy.MaxPixels = 64 * 64 + policy.MaxFrames = 1 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(testCase.data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + } +} + +func TestInspectCapabilityAcceptsMappedDecodableSupportedVideoCodecs(t *testing.T) { + t.Parallel() + for _, testCase := range []struct { + name, filename, mediaType string + data []byte + pixels int64 + }{ + {name: "H264 MOV", filename: "clip.mov", mediaType: "video/quicktime", data: mediatest.H264MOV(), pixels: 16 * 16}, + {name: "H265 MP4", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.H265MP4(), pixels: 16 * 16}, + {name: "VP9 MP4", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.VP9MP4(), pixels: 16 * 16}, + {name: "AV1 MP4", filename: "clip.mp4", mediaType: "video/mp4", data: mediatest.AV1MP4(), pixels: 64 * 64}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + policy := inspectionPolicy(testCase.data, testCase.filename, testCase.mediaType) + policy.MaxPixels = testCase.pixels + policy.MaxFrames = 1 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(testCase.data), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, "video", record.MediaFamily) + assert.Equal(t, int64(1), record.Measurements.Frames) + assert.Equal(t, int64(1_000), record.Measurements.DurationMS) + }) + } +} + +func TestInspectCapabilityRejectsHeaderOnlyVideoDeclaration(t *testing.T) { + t.Parallel() + video := mediatest.MP4(16, 16, 1_000) + policy := inspectionPolicy(video, "clip.mp4", "video/mp4") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 1 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(video), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +func firstMP4SampleOffset(t *testing.T, data []byte) int { + t.Helper() + stco := bytes.Index(data, []byte("stco")) + require.NotEqual(t, -1, stco) + offset := uint64(binary.BigEndian.Uint32(data[stco+12 : stco+16])) + require.Less(t, offset, uint64(len(data))) + return int(offset) +} + +func mp4With64BitChunkOffsets(t *testing.T) []byte { + t.Helper() + video := decodableAVCMP4(t) + stco := bytes.Index(video, []byte("stco")) + require.NotEqual(t, -1, stco) + oldOffset := binary.BigEndian.Uint32(video[stco+12 : stco+16]) + replacement := make([]byte, 24) + binary.BigEndian.PutUint32(replacement[:4], uint32(len(replacement))) + copy(replacement[4:8], "co64") + binary.BigEndian.PutUint32(replacement[12:16], 1) + binary.BigEndian.PutUint64(replacement[16:24], uint64(oldOffset)+4) + boxStart := stco - 4 + video = append(append(append([]byte(nil), video[:boxStart]...), replacement...), video[boxStart+20:]...) + for _, kind := range []string{"stbl", "minf", "mdia", "trak", "moov"} { + index := bytes.Index(video, []byte(kind)) + require.GreaterOrEqual(t, index, 4) + size := binary.BigEndian.Uint32(video[index-4 : index]) + binary.BigEndian.PutUint32(video[index-4:index], size+4) + } + return video +} + +// TestInspectCapabilityProvesGeminiMP3Duration catches a regression where +// valid MPEG Layer III frames are left in the unbounded audio family instead +// of contributing their literal, finite duration to capability proof. +func TestInspectCapabilityProvesGeminiMP3Duration(t *testing.T) { + t.Parallel() + data := syntheticMP3Frames(10) + policy := inspectionPolicy(data, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 262 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, "audio", record.MediaFamily) + assert.Equal(t, "audio/mpeg", record.MediaType) + assert.Equal(t, "mp3", record.Format) + assert.Equal(t, int64(262), record.Measurements.DurationMS) +} + +func TestInspectCapabilityProvesRealMP3WithBoundedID3Tags(t *testing.T) { + t.Parallel() + audio := mediatest.MP3() + id3v2 := append([]byte{'I', 'D', '3', 4, 0, 0, 0, 0, 0, 4}, []byte("TEST")...) + id3v1 := make([]byte, 128) + copy(id3v1, "TAG") + tagged := slices.Concat(id3v2, audio, id3v1) + + barePolicy := inspectionPolicy(audio, "sample.mp3", "audio/mpeg") + barePolicy.MaxDurationMS = 1_000 + bare, err := media.InspectCapability(bytes.NewReader(audio), barePolicy) + require.NoError(t, err) + require.True(t, bare.Eligible, bare.Reason) + + taggedPolicy := inspectionPolicy(tagged, "sample.mp3", "audio/mpeg") + taggedPolicy.MaxDurationMS = 1_000 + withTags, err := media.InspectCapability(bytes.NewReader(tagged), taggedPolicy) + require.NoError(t, err) + require.True(t, withTags.Eligible, withTags.Reason) + assert.Equal(t, bare.Measurements.DurationMS, withTags.Measurements.DurationMS) + + for _, testCase := range []struct { + name string + tag []byte + }{ + {name: "truncated header", tag: []byte("ID3\x04\x00")}, + {name: "unknown version", tag: []byte{'I', 'D', '3', 5, 0, 0, 0, 0, 0, 0}}, + {name: "unknown flags", tag: []byte{'I', 'D', '3', 4, 0, 1, 0, 0, 0, 0}}, + {name: "non synchsafe size", tag: []byte{'I', 'D', '3', 4, 0, 0, 0, 0, 0x80, 0}}, + {name: "declared body exceeds input", tag: []byte{'I', 'D', '3', 4, 0, 0, 0, 0, 4, 0}}, + {name: "declared body exceeds tag bound", tag: []byte{'I', 'D', '3', 4, 0, 0, 0, 0x40, 0, 1}}, + } { + t.Run(testCase.name, func(t *testing.T) { + candidate := append(append([]byte(nil), testCase.tag...), audio...) + policy := inspectionPolicy(candidate, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 1_000 + record, inspectErr := media.InspectCapability(bytes.NewReader(candidate), policy) + require.NoError(t, inspectErr) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + } +} + +// TestInspectCapabilityProvesGeminiQuickTimeVideo catches an identity check +// that rejects a locally verified MOV merely because detection normalizes it +// to the shared video/mp4 media type and mp4 format. +func TestInspectCapabilityProvesGeminiQuickTimeVideo(t *testing.T) { + t.Parallel() + data := quickTimeH264MP4() + policy := inspectionPolicy(data, "clip.mov", "video/quicktime") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 1 + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + require.True(t, record.Eligible, record.Reason) + assert.Equal(t, "video", record.MediaFamily) + assert.Equal(t, "video/quicktime", record.MediaType) + assert.Equal(t, "mp4", record.Format) +} + +func TestInspectCapabilityRejectsMP4ClaimingQuickTimeIdentity(t *testing.T) { + t.Parallel() + data := mediatest.MP4(16, 16, 500) + policy := inspectionPolicy(data, "clip.mov", "video/quicktime") + policy.MaxPixels = 16 * 16 + policy.MaxFrames = 1 + policy.MaxDurationMS = 500 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) +} + +// TestInspectCapabilityRejectsUnboundedOrOverlongGeminiMP3 catches a parser +// that accepts mixed frame streams or fails to enforce their derived duration. +func TestInspectCapabilityRejectsUnboundedOrOverlongGeminiMP3(t *testing.T) { + t.Parallel() + t.Run("mixed MPEG versions are malformed", func(t *testing.T) { + data := syntheticMP3Frames(2) + data[417+1] = 0xf3 // MPEG-2 Layer III, unlike the preceding MPEG-1 frame. + policy := inspectionPolicy(data, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + + t.Run("MPEG-2.5 cannot reset the stream identity", func(t *testing.T) { + data := append(syntheticMPEG25Frame(), syntheticMP3Frames(1)...) + policy := inspectionPolicy(data, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + + t.Run("reserved emphasis is malformed", func(t *testing.T) { + data := syntheticMP3Frames(1) + data[3] = 0x02 + policy := inspectionPolicy(data, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 1_000 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonMalformed, record.Reason) + }) + + t.Run("duration above the policy is rejected", func(t *testing.T) { + data := syntheticMP3Frames(4) + policy := inspectionPolicy(data, "sample.mp3", "audio/mpeg") + policy.MaxDurationMS = 100 + + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + assert.False(t, record.Eligible) + assert.Equal(t, media.CapabilityReasonVisualBounds, record.Reason) + assert.Equal(t, int64(105), record.Measurements.DurationMS) + }) +} + +func inspectionPolicy(data []byte, filename, mediaType string) media.InspectionPolicy { + return media.InspectionPolicy{ + Filename: filename, DeclaredMediaType: mediaType, + ExpectedBytes: int64(len(data)), ExpectedSHA256: sha256Hex(data), + DescriptorFingerprint: strings.Repeat("a", 64), ProfileFingerprint: strings.Repeat("b", 64), + DisclosureFingerprint: strings.Repeat("c", 64), InputKind: document.RenditionInputOriginalFile, + MaxSourceBytes: 1 << 20, MaxExpandedBytes: 1 << 20, MaxEntryBytes: 1 << 20, + MaxEntries: 100, MaxNestingDepth: 1, MaxTextLines: 1_000, MaxCharacters: 1 << 20, + MaxPages: 100, MaxSlides: 100, MaxSheets: 100, MaxCells: 10_000, MaxSpineItems: 1_000, MaxResources: 10_000, + } +} + +func syntheticMP3Frames(count int) []byte { + const frameBytes = 417 // MPEG-1 Layer III, 128 kbps, 44.1 kHz, no padding. + frames := make([]byte, count*frameBytes) + for offset := 0; offset < len(frames); offset += frameBytes { + frames[offset], frames[offset+1], frames[offset+2], frames[offset+3] = 0xff, 0xfb, 0x90, 0 + } + return frames +} + +func syntheticMPEG25Frame() []byte { + const frameBytes = 522 // MPEG-2.5 Layer III, 80 kbps, 11.025 kHz, no padding. + frame := make([]byte, frameBytes) + frame[0], frame[1], frame[2], frame[3] = 0xff, 0xe3, 0x90, 0 + return frame +} + +type zipEntry struct{ name, body string } + +func validPPTXEntries(extra ...zipEntry) []zipEntry { + return append([]zipEntry{ + {name: "[Content_Types].xml", body: ``}, + {name: "ppt/presentation.xml", body: ``}, + }, extra...) +} + +func validXLSXEntries(extra ...zipEntry) []zipEntry { + return append([]zipEntry{ + {name: "[Content_Types].xml", body: ``}, + {name: "xl/workbook.xml", body: ``}, + }, extra...) +} + +func validEPUBEntries(extra ...zipEntry) []zipEntry { + return append([]zipEntry{ + {name: "mimetype", body: "application/epub+zip"}, + {name: "META-INF/container.xml", body: ``}, + }, extra...) +} + +func zipBytes(t *testing.T, entries []zipEntry) []byte { + t.Helper() + var out bytes.Buffer + zw := zip.NewWriter(&out) + for _, entry := range entries { + w, err := zw.Create(entry.name) + require.NoError(t, err) + _, err = w.Write([]byte(entry.body)) + require.NoError(t, err) + } + require.NoError(t, zw.Close()) + return out.Bytes() +} + +func sha256Hex(data []byte) string { + digest := sha256.Sum256(data) + return hex.EncodeToString(digest[:]) +} + +func wavBytes(byteRate, dataBytes uint32) []byte { + data := make([]byte, 44+dataBytes) + copy(data[0:4], "RIFF") + binary.LittleEndian.PutUint32(data[4:8], uint32(len(data)-8)) // #nosec G115 -- synthetic fixture is bounded + copy(data[8:12], "WAVE") + copy(data[12:16], "fmt ") + binary.LittleEndian.PutUint32(data[16:20], 16) + binary.LittleEndian.PutUint16(data[20:22], 1) + binary.LittleEndian.PutUint16(data[22:24], 1) + binary.LittleEndian.PutUint32(data[24:28], byteRate) + binary.LittleEndian.PutUint32(data[28:32], byteRate) + binary.LittleEndian.PutUint16(data[32:34], 1) + binary.LittleEndian.PutUint16(data[34:36], 8) + copy(data[36:40], "data") + binary.LittleEndian.PutUint32(data[40:44], dataBytes) + return data +} + +func syntheticPDF(label string) []byte { + return syntheticPDFPages(1, label) +} + +func syntheticPDFPages(pageCount int, label string) []byte { + kids := make([]string, 0, pageCount) + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R /Note (/Type /Page in a string) >>", + "", + } + for index := range pageCount { + kids = append(kids, fmt.Sprintf("%d 0 R", index+3)) + objects = append(objects, "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>") + } + objects[1] = fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", strings.Join(kids, " "), pageCount) + objects = append(objects, + "<< /Type /Page /Parent 2 0 R /Note (orphan page object) >>", + "<< /Length 11 >>\nstream\n/Type /Page\nendstream", + ) + var output bytes.Buffer + _, _ = fmt.Fprintf(&output, "%%PDF-1.4\n%%%x\n", label) + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} + +type xrefStreamFixture struct { + index, dictionaryExtra string + mutate func([]byte, []int) []byte +} + +func syntheticXRefStreamPDFWith(fixture xrefStreamFixture) []byte { + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>", + "<< /Title (Synthetic xref stream) >>", + } + var output bytes.Buffer + _, _ = output.WriteString("%PDF-1.7\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + entries := make([]byte, 6*7) + putSyntheticXRefEntry(entries, 0, 0, 0, 65_535) + for index, offset := range offsets { + putSyntheticXRefEntry(entries, index+1, 1, uint32(offset), 0) // #nosec G115 -- bounded fixture. + } + putSyntheticXRefEntry(entries, 5, 1, uint32(xref), 0) // #nosec G115 -- bounded fixture. + if fixture.mutate != nil { + entries = fixture.mutate(entries, offsets) + } + _, _ = fmt.Fprintf(&output, + "5 0 obj\n<< /Type /XRef /Size 6 /Root 1 0 R /Info 4 0 R /W [1 4 2] /Length %d %s %s >>\nstream\n", + len(entries), fixture.index, fixture.dictionaryExtra) + _, _ = output.Write(entries) + _, _ = fmt.Fprintf(&output, "\nendstream\nendobj\nstartxref\n%d\n%%%%EOF\n", xref) + return output.Bytes() +} + +func appendSyntheticXRefEntry(index int) func([]byte, []int) []byte { + return func(entries []byte, _ []int) []byte { + return append(entries, entries[index*7:(index+1)*7]...) + } +} + +func putSyntheticXRefEntry(entries []byte, index int, kind byte, offset uint32, generation uint16) { + entry := entries[index*7 : (index+1)*7] + entry[0] = kind + binary.BigEndian.PutUint32(entry[1:5], offset) + binary.BigEndian.PutUint16(entry[5:7], generation) +} + +func syntheticDeepPDF(depth int) []byte { + objects := make([]string, 0, depth+2) + objects = append(objects, "<< /Type /Catalog /Pages 2 0 R >>") + for index := range depth { + number := index + 2 + child := number + 1 + parent := "" + if index != 0 { + parent = fmt.Sprintf(" /Parent %d 0 R", number-1) + } + objects = append(objects, fmt.Sprintf( + "<< /Type /Pages /Kids [%d 0 R] /Count 1%s >>", child, parent)) + } + objects = append(objects, fmt.Sprintf( + "<< /Type /Page /Parent %d 0 R /MediaBox [0 0 612 792] >>", depth+1)) + var output bytes.Buffer + _, _ = output.WriteString("%PDF-1.4\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, + "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} diff --git a/document/media/mediatest/codec.go b/document/media/mediatest/codec.go new file mode 100644 index 00000000..545cde2c --- /dev/null +++ b/document/media/mediatest/codec.go @@ -0,0 +1,77 @@ +package mediatest + +import "encoding/base64" + +// The fixtures below are synthetic solid-color, single-frame videos generated +// from FFmpeg 9.0.1's public lavfi color source. Generation used no external +// input, and every result was decoded back to a null sink before its base64 was +// recorded. Tests decode these constants directly and require no codec binary. + +// H264MOV returns a 16x16 red H.264 QuickTime fixture. SHA-256: +// d6d619c2ff67a1216d81ec675cc61c14d70e88bd11c6514d04a1f0414960b1ee +// +// ffmpeg -f lavfi -i color=c=red:s=16x16:r=1:d=1 -frames:v 1 -an +// -c:v libx264 -pix_fmt yuv420p -preset veryslow -tune stillimage +// -bsf:v filter_units=remove_types=6 -map_metadata -1 -movflags +faststart +// -f mov h264.mov. +func H264MOV() []byte { + return decodeFixture(h264MOVBase64) +} + +// H265MP4 returns a 16x16 green H.265 MP4 fixture. SHA-256: +// 5d7367612c8bf7f9688722fde62718eb404986c6503d01bd9727a3e7fb30ee4b +// +// ffmpeg -f lavfi -i color=c=green:s=16x16:r=1:d=1 -frames:v 1 -an +// -c:v libx265 -tag:v hvc1 -pix_fmt yuv420p -preset ultrafast +// -x265-params log-level=error:info=0 -map_metadata -1 -movflags +faststart +// h265.mp4. +func H265MP4() []byte { + return decodeFixture(h265MP4Base64) +} + +// VP9MP4 returns a 16x16 blue VP9 MP4 fixture. SHA-256: +// af5eeff8d624378f43cea29b8ef33649fbcfa4fd9e8600674a206d9c12620fda +// +// ffmpeg -f lavfi -i color=c=blue:s=16x16:r=1:d=1 -frames:v 1 -an +// -c:v libvpx-vp9 -lossless 1 -deadline best -cpu-used 0 -row-mt 0 +// -map_metadata -1 -movflags +faststart vp9.mp4. +func VP9MP4() []byte { + return decodeFixture(vp9MP4Base64) +} + +// AV1MP4 returns a 64x64 yellow AV1 MP4 fixture. SHA-256: +// 70360fe0090fcc26bed290c1d75725ee0483478b56f674282032fcc7d4fbbce8 +// +// ffmpeg -f lavfi -i color=c=yellow:s=64x64:r=1:d=1 -frames:v 1 -an +// -c:v libsvtav1 -preset 13 -crf 63 -pix_fmt yuv420p -map_metadata -1 +// -movflags +faststart av1.mp4. +func AV1MP4() []byte { + return decodeFixture(av1MP4Base64) +} + +// MP3 returns a 50 ms silent MPEG-1 Layer III fixture without metadata tags. +// SHA-256: 17944e919ea30c51a87c9c849e80ef9dcfda166cd275d0d448ab63b3ae9b95dc +// +// ffmpeg -f lavfi -i anullsrc=r=44100:cl=mono -t 0.05 -c:a libmp3lame +// -b:a 32k -write_xing 0 -id3v2_version 0 -map_metadata -1 -f mp3 sample.mp3 +func MP3() []byte { + return decodeFixture(mp3Base64) +} + +func decodeFixture(encoded string) []byte { + data, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + panic(err) + } + return data +} + +const h264MOVBase64 = "AAAAFGZ0eXBxdCAgAAACAHF0ICAAAALubW9vdgAAAGxtdmhkAAAAAAAAAAAAAAAAAAAD6AAAA+gAAQAAAQAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAlp0cmFrAAAAXHRraGQAAAADAAAAAAAAAAAAAAABAAAAAAAAA+gAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAABAAAAAQAAAAAAAkZWR0cwAAABxlbHN0AAAAAAAAAAEAAAPoAAAAAAABAAAAAAHSbWRpYQAAACBtZGhkAAAAAAAAAAAAAAAAAABAAAAAQAB//wAAAAAALWhkbHIAAAAAbWhscnZpZGUAAAAAAAAAAAAAAAAMVmlkZW9IYW5kbGVyAAABfW1pbmYAAAAUdm1oZAAAAAEAAAAAAAAAAAAAACxoZGxyAAAAAGRobHJ1cmwgAAAAAAAAAAAAAAAAC0RhdGFIYW5kbGVyAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAAABEXN0YmwAAACtc3RzZAAAAAAAAAABAAAAnWF2YzEAAAAAAAAAAQAAAABGRk1QAAACAAAAAgAAEAAQAEgAAABIAAAAAAAAAAEUTGF2YzYzLjEuMTAxIGxpYngyNjQAAAAAAAAAAAAAAAAY//8AAAA3YXZjQwFkAAr/4QAZZ2QACqxyBF7ARAAAAwAEAAADAAg8SJYRgAEAB2joQ48TITD9+PgAAAAAEHBhc3AAAAABAAAAAQAAABhzdHRzAAAAAAAAAAEAAAABAABAAAAAABxzdHNjAAAAAAAAAAEAAAABAAAAAQAAAAEAAAAUc3RzegAAAAAAAAAbAAAAAQAAABRzdGNvAAAAAAAAAAEAAAMSAAAAIHVkdGEAAAAYqXN3cgAMVcRMYXZmNjMuMS4xMDEAAAAId2lkZQAAACNtZGF0AAAAF2WIgQAC+c/+7oK+BTbNED8HLCrLyfnB" + +const h265MP4Base64 = "AAAAHGZ0eXBpc29tAAACAGlzb21pc28ybXA0MQAAA15tb292AAAAbG12aGQAAAAAAAAAAAAAAAAAAAPoAAAD6AABAAABAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAAACiXRyYWsAAABcdGtoZAAAAAMAAAAAAAAAAAAAAAEAAAAAAAAD6AAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAEAAAABAAAAAAACRlZHRzAAAAHGVsc3QAAAAAAAAAAQAAA+gAAAAAAAEAAAAAAgFtZGlhAAAAIG1kaGQAAAAAAAAAAAAAAAAAAEAAAABAAFXEAAAAAAAtaGRscgAAAAAAAAAAdmlkZQAAAAAAAAAAAAAAAFZpZGVvSGFuZGxlcgAAAAGsbWluZgAAABR2bWhkAAAAAQAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAAABbHN0YmwAAAEIc3RzZAAAAAAAAAABAAAA+Gh2YzEAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAEAAQAEgAAABIAAAAAAAAAAEUTGF2YzYzLjEuMTAxIGxpYngyNjUAAAAAAAAAAAAAAAAY//8AAAB0aHZjQwEBYAAAAJAAAAAAAB7wAPz9+PgAAA8DoAABABhAAQwB//8BYAAAAwCQAAADAAADAB6VlAmhAAEAKEIBAQFgAAADAJAAAAMAAAMAHqCIRZZWVbwvAWgIAAADAAgAAAMACECiAAEABkQBwHPAiQAAAApmaWVsAQAAAAAQcGFzcAAAAAEAAAABAAAAFGJ0cnQAAAAAAAAAmAAAAAAAAAAYc3R0cwAAAAAAAAABAAAAAQAAQAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAEAAAABAAAAFHN0c3oAAAAAAAAAEwAAAAEAAAAUc3RjbwAAAAAAAAABAAADigAAAGF1ZHRhAAAAWW1ldGEAAAAAAAAAIWhkbHIAAAAAAAAAAG1kaXJhcHBsAAAAAAAAAAAAAAAALGlsc3QAAAAkqXRvbwAAABxkYXRhAAAAAQAAAABMYXZmNjMuMS4xMDEAAAAIZnJlZQAAABttZGF0AAAADygBrE7UF/8MSf6GYflf4A==" + +const vp9MP4Base64 = "AAAAHGZ0eXBpc29tAAACAGlzb21pc28ybXA0MQAAAv5tb292AAAAbG12aGQAAAAAAAAAAAAAAAAAAAPoAAAD6AABAAABAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAAACKXRyYWsAAABcdGtoZAAAAAMAAAAAAAAAAAAAAAEAAAAAAAAD6AAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAEAAAAAAEAAAABAAAAAAACRlZHRzAAAAHGVsc3QAAAAAAAAAAQAAA+gAAAAAAAEAAAAAAaFtZGlhAAAAIG1kaGQAAAAAAAAAAAAAAAAAAEAAAABAAFXEAAAAAAAtaGRscgAAAAAAAAAAdmlkZQAAAAAAAAAAAAAAAFZpZGVvSGFuZGxlcgAAAAFMbWluZgAAABR2bWhkAAAAAQAAAAAAAAAAAAAAJGRpbmYAAAAcZHJlZgAAAAAAAAABAAAADHVybCAAAAABAAABDHN0YmwAAACoc3RzZAAAAAAAAAABAAAAmHZwMDkAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAEAAQAEgAAABIAAAAAAAAAAEXTGF2YzYzLjEuMTAxIGxpYnZweC12cDkAAAAAAAAAAAAY//8AAAAUdnBjQwEAAAAACoICAgIAAAAAAApmaWVsAQAAAAAQcGFzcAAAAAEAAAABAAAAFGJ0cnQAAAAAAAABUAAAAAAAAAAYc3R0cwAAAAAAAAABAAAAAQAAQAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAEAAAABAAAAFHN0c3oAAAAAAAAAKgAAAAEAAAAUc3RjbwAAAAAAAAABAAADKgAAAGF1ZHRhAAAAWW1ldGEAAAAAAAAAIWhkbHIAAAAAAAAAAG1kaXJhcHBsAAAAAAAAAAAAAAAALGlsc3QAAAAkqXRvbwAAABxkYXRhAAAAAQAAAABMYXZmNjMuMS4xMDEAAAAIZnJlZQAAADJtZGF0gkmDQgAA8AD2ADgkHBgAAAAgAAAQv///5zlAGP////2yMBw////jFQAA" + +const av1MP4Base64 = "AAAAIGZ0eXBpc29tAAACAGlzb21hdjAxaXNvMm1wNDEAAAMDbW9vdgAAAGxtdmhkAAAAAAAAAAAAAAAAAAAD6AAAA+gAAQAAAQAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAi50cmFrAAAAXHRraGQAAAADAAAAAAAAAAAAAAABAAAAAAAAA+gAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAEAAAABAAAAAAAAkZWR0cwAAABxlbHN0AAAAAAAAAAEAAAPoAAAAAAABAAAAAAGmbWRpYQAAACBtZGhkAAAAAAAAAAAAAAAAAABAAAAAQABVxAAAAAAALWhkbHIAAAAAAAAAAHZpZGUAAAAAAAAAAAAAAABWaWRlb0hhbmRsZXIAAAABUW1pbmYAAAAUdm1oZAAAAAEAAAAAAAAAAAAAACRkaW5mAAAAHGRyZWYAAAAAAAAAAQAAAAx1cmwgAAAAAQAAARFzdGJsAAAArXN0c2QAAAAAAAAAAQAAAJ1hdjAxAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAEAAQABIAAAASAAAAAAAAAABFkxhdmM2My4xLjEwMSBsaWJzdnRhdjEAAAAAAAAAAAAAGP//AAAAGWF2MUOBAAwACgsCAAAFFX/8AvgAQAAAAApmaWVsAQAAAAAQcGFzcAAAAAEAAAABAAAAFGJ0cnQAAAAAAAABCAAAAQgAAAAYc3R0cwAAAAAAAAABAAAAAQAAQAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAEAAAABAAAAFHN0c3oAAAAAAAAAIQAAAAEAAAAUc3RjbwAAAAAAAAABAAADMwAAAGF1ZHRhAAAAWW1ldGEAAAAAAAAAIWhkbHIAAAAAAAAAAG1kaXJhcHBsAAAAAAAAAAAAAAAALGlsc3QAAAAkqXRvbwAAABxkYXRhAAAAAQAAAABMYXZmNjMuMS4xMDEAAAAIZnJlZQAAACltZGF0CgsCAAAFFX/8SvkAQDISEADzgj/+aYMAAAiQzzTwUe1S" + +const mp3Base64 = "//sQxAADwAABpAAAACAAADSAAAAETEFNRTQuMFVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVX/+xLEKYPAAAGkAAAAIAAANIAAAARVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVX/+xDEU4PAAAGkAAAAIAAANIAAAARVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVQ==" diff --git a/document/media/mediatest/mediatest.go b/document/media/mediatest/mediatest.go index ced04f1a..b0ba5f13 100644 --- a/document/media/mediatest/mediatest.go +++ b/document/media/mediatest/mediatest.go @@ -7,6 +7,7 @@ package mediatest import ( "bytes" "encoding/binary" + "fmt" "image" "image/color" "image/gif" @@ -34,6 +35,54 @@ func PNG(width, height int, fill color.Color) []byte { return out.Bytes() } +// WAV returns 10 ms of synthetic unsigned 8-bit mono PCM at 8 kHz. +func WAV() []byte { + const sampleRate = 8_000 + const audioBytes = 80 + data := make([]byte, 44+audioBytes) + copy(data[0:4], "RIFF") + binary.LittleEndian.PutUint32(data[4:8], uint32(len(data)-8)) //nolint:gosec // fixed tiny fixture + copy(data[8:12], "WAVE") + copy(data[12:16], "fmt ") + binary.LittleEndian.PutUint32(data[16:20], 16) + binary.LittleEndian.PutUint16(data[20:22], 1) + binary.LittleEndian.PutUint16(data[22:24], 1) + binary.LittleEndian.PutUint32(data[24:28], sampleRate) + binary.LittleEndian.PutUint32(data[28:32], sampleRate) + binary.LittleEndian.PutUint16(data[32:34], 1) + binary.LittleEndian.PutUint16(data[34:36], 8) + copy(data[36:40], "data") + binary.LittleEndian.PutUint32(data[40:44], audioBytes) + for index := 44; index < len(data); index++ { + data[index] = 128 + } + return data +} + +// PDF returns a synthetic one-page PDF 1.4 file with a literal cross-reference +// table. It is assembled entirely in Go from public format syntax. +func PDF() []byte { + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 16 16] >>", + } + var output bytes.Buffer + _, _ = output.WriteString("%PDF-1.4\n%synthetic\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} + // GIF returns a GIF with the given number of frames. Frames alternate between // two palette colors so an animated result differs from its first frame. func GIF(width, height, frames int) []byte { diff --git a/document/media/mp4_codec.go b/document/media/mp4_codec.go index 0a664421..41a19cf9 100644 --- a/document/media/mp4_codec.go +++ b/document/media/mp4_codec.go @@ -5,23 +5,33 @@ import ( "math" ) +type mp4CodecConfiguration struct { + codec string + nalLengthBytes int +} + // visualCodecDimensions reads dimensions from an out-of-band codec // configuration. avc3/hev1 and the other recognized visual sample-entry kinds // can carry configuration changes in media samples, which this metadata-only // detector cannot bound, so they fail closed. -func visualCodecDimensions(kind string, entry []byte) (int64, int64, bool) { +func visualCodecDimensions(kind string, entry []byte) (mp4CodecConfiguration, int64, int64, bool) { if len(entry) < 78 { - return 0, 0, false + return mp4CodecConfiguration{}, 0, 0, false } + var configuration mp4CodecConfiguration var configKind string var parse func([]byte) (int64, int64, bool) switch kind { case "avc1": - configKind, parse = "avcC", avcConfigDimensions + configuration.codec, configKind, parse = "h264", "avcC", avcConfigDimensions case "hvc1": - configKind, parse = "hvcC", hevcConfigDimensions + configuration.codec, configKind, parse = "h265", "hvcC", hevcConfigDimensions + case "vp09": + configuration.codec, configKind, parse = "vp9", "vpcC", visualEntryDimensions + case "av01": + configuration.codec, configKind, parse = "av1", "av1C", visualEntryDimensions default: - return 0, 0, false + return mp4CodecConfiguration{}, 0, 0, false } children := entry[78:] var width, height int64 @@ -29,22 +39,154 @@ func visualCodecDimensions(kind string, entry []byte) (int64, int64, bool) { for offset := 0; offset < len(children); { headerLen, size, ok := mp4BoxHeader(children, offset) if !ok { - return 0, 0, false + return mp4CodecConfiguration{}, 0, 0, false } if string(children[offset+4:offset+8]) == configKind { configs++ if configs > 1 { - return 0, 0, false + return mp4CodecConfiguration{}, 0, 0, false } - codecWidth, codecHeight, ok := parse(children[offset+headerLen : offset+size]) + config := children[offset+headerLen : offset+size] + if (kind == "vp09" && !validVP9Config(config)) || (kind == "av01" && !validAV1Config(config)) { + return mp4CodecConfiguration{}, 0, 0, false + } + codecWidth, codecHeight, ok := parse(config) if !ok { - return 0, 0, false + return mp4CodecConfiguration{}, 0, 0, false + } + switch kind { + case "avc1": + configuration.nalLengthBytes = int(config[4]&0x03) + 1 + case "hvc1": + configuration.nalLengthBytes = int(config[21]&0x03) + 1 + } + if kind == "vp09" || kind == "av01" { + codecWidth = int64(binary.BigEndian.Uint16(entry[24:26])) + codecHeight = int64(binary.BigEndian.Uint16(entry[26:28])) } width, height = codecWidth, codecHeight } offset += size } - return width, height, configs == 1 && width > 0 && height > 0 + return configuration, width, height, configs == 1 && width > 0 && height > 0 +} + +func visualEntryDimensions(config []byte) (int64, int64, bool) { + return 1, 1, len(config) > 0 +} + +func validVP9Config(config []byte) bool { + if len(config) < 12 || config[0] != 1 || config[1] != 0 || config[2] != 0 || config[3] != 0 || + config[4] > 3 { + return false + } + bitDepth := config[6] >> 4 + if bitDepth != 8 && bitDepth != 10 && bitDepth != 12 || config[6]>>1&0x07 > 3 { + return false + } + initializationDataSize := int(binary.BigEndian.Uint16(config[10:12])) + return len(config) == 12+initializationDataSize +} + +func validAV1Config(config []byte) bool { + if len(config) < 4 || config[0] != 0x81 || config[3]&0xe0 != 0 || + config[3]&0x10 == 0 && config[3]&0x0f != 0 { + return false + } + return len(config) == 4 || validAV1OBUs(config[4:], false, true) +} + +func validMP4CodecSample(configuration mp4CodecConfiguration, sample []byte) bool { + switch configuration.codec { + case "h264": + return validLengthPrefixedNALSample(sample, configuration.nalLengthBytes, false) + case "h265": + return validLengthPrefixedNALSample(sample, configuration.nalLengthBytes, true) + case "vp9": + return len(sample) > 0 && sample[0]&0x03 == 0x02 + case "av1": + return validAV1OBUs(sample, true, false) + default: + return false + } +} + +func validLengthPrefixedNALSample(sample []byte, lengthBytes int, hevc bool) bool { + if lengthBytes < 1 || lengthBytes > 4 { + return false + } + hasPicture := false + for offset := 0; offset < len(sample); { + if offset+lengthBytes > len(sample) { + return false + } + size := uint64(0) + for _, value := range sample[offset : offset+lengthBytes] { + size = size<<8 | uint64(value) + } + offset += lengthBytes + remaining := uint64(len(sample) - offset) //nolint:gosec // offset is bounded by the sample length + if size == 0 || size > remaining { + return false + } + nal := sample[offset : offset+int(size)] + if hevc { + if len(nal) < 2 || nal[0]&0x80 != 0 || nal[1]&0x07 == 0 { + return false + } + hasPicture = hasPicture || (nal[0]>>1)&0x3f <= 31 + } else { + if nal[0]&0x80 != 0 || nal[0]&0x1f == 0 || nal[0]&0x1f > 23 { + return false + } + nalType := nal[0] & 0x1f + hasPicture = hasPicture || nalType >= 1 && nalType <= 5 + } + offset += int(size) + } + return len(sample) > 0 && hasPicture +} + +func validAV1OBUs(data []byte, requireFrame, requireSequence bool) bool { + hasFrame, hasSequence := false, false + for offset := 0; offset < len(data); { + header := data[offset] + offset++ + obuType := header >> 3 & 0x0f + if header&0x81 != 0 || header&0x02 == 0 || obuType == 0 || obuType >= 9 && obuType <= 14 { + return false + } + if header&0x04 != 0 { + if offset >= len(data) || data[offset]&0x07 != 0 { + return false + } + offset++ + } + size, sizeBytes, ok := readAV1LEB128(data[offset:]) + remaining := uint64(len(data) - offset - sizeBytes) //nolint:gosec // offsets are bounded by data above + if !ok || size > remaining { + return false + } + offset += sizeBytes + if size == 0 && obuType != 2 && obuType != 15 { + return false + } + hasSequence = hasSequence || obuType == 1 + hasFrame = hasFrame || obuType == 3 || obuType == 6 || obuType == 7 + offset += int(size) //nolint:gosec // size is bounded by the remaining data above + } + return len(data) > 0 && (!requireFrame || hasFrame) && (!requireSequence || hasSequence) +} + +func readAV1LEB128(data []byte) (uint64, int, bool) { + var value uint64 + for index := 0; index < len(data) && index < 8; index++ { + value |= uint64(data[index]&0x7f) << (index * 7) + if data[index]&0x80 == 0 { + return value, index + 1, true + } + } + return 0, 0, false } func avcConfigDimensions(config []byte) (int64, int64, bool) { diff --git a/document/media/policy.go b/document/media/policy.go index a3695884..11840367 100644 --- a/document/media/policy.go +++ b/document/media/policy.go @@ -136,7 +136,7 @@ func withinPixels(width, height, limit int64) bool { return width <= limit && height <= limit && width <= limit/height } -// Inspect detects and evaluates one input. Detection failures become the +// Inspect detects and evaluates one visual input. Detection failures become the // matching Reason rather than an error so callers can record stable outcomes; // only read failures are returned as errors. Inputs longer than the policy // limit are refused before they are read. diff --git a/document/media/types.go b/document/media/types.go index b145381f..3d2e1106 100644 --- a/document/media/types.go +++ b/document/media/types.go @@ -34,6 +34,10 @@ type Metadata struct { Kind Kind `json:"kind"` // MediaType is the canonical media type for Format. MediaType string `json:"media_type"` + // Container is the canonical video container, when Kind is video. + Container string `json:"container,omitzero"` + // Codec is the locally verified visual codec, when Kind is video. + Codec string `json:"codec,omitzero"` // DeclaredMediaType is the caller-declared media type, recorded verbatim. DeclaredMediaType string `json:"declared_media_type,omitzero"` // Size is the input length in bytes. diff --git a/document/mistral/client.go b/document/mistral/client.go index 081db086..c7e85f2a 100644 --- a/document/mistral/client.go +++ b/document/mistral/client.go @@ -66,6 +66,7 @@ type Result struct { ReturnedModel string UnitsProcessed int ProviderBytes *int64 + ResponseBytes int64 Metrics RequestMetrics } @@ -122,7 +123,7 @@ type wireUsage struct { // Client calls the single endpoint derived from its policy. type Client struct { policy Policy - apiKey string + credential func(context.Context) (string, error) maxRetries int maxRetryDelay time.Duration http *http.Client @@ -136,6 +137,21 @@ func NewClient(policy Policy, config ClientConfig) (*Client, error) { if config.APIKey == "" || config.APIKey != strings.TrimSpace(config.APIKey) { return nil, errors.New("mistral OCR API key is required and must not contain surrounding whitespace") } + apiKey := config.APIKey + return newClientWithCredential(policy, config, func(context.Context) (string, error) { + return apiKey, nil + }) +} + +func newClientWithCredential( + policy Policy, config ClientConfig, credential func(context.Context) (string, error), +) (*Client, error) { + if policy.digest == "" { + return nil, errors.New("mistral policy is invalid; use NewPolicy") + } + if credential == nil { + return nil, errors.New("mistral OCR credential resolver is required") + } if config.Timeout == 0 { config.Timeout = DefaultTimeout } @@ -169,7 +185,7 @@ func NewClient(policy Policy, config ClientConfig) (*Client, error) { return http.ErrUseLastResponse } return &Client{ - policy: policy, apiKey: config.APIKey, maxRetries: config.MaxRetries, + policy: policy, credential: credential, maxRetries: config.MaxRetries, maxRetryDelay: config.MaxRetryDelay, http: httpClient, }, nil } @@ -220,6 +236,18 @@ func (c *Client) process( options requestOptions, method UnitBoundMethod, maxUnits int, +) (Result, error) { + return c.processWith(ctx, snapshotForAttempt, readVerifiedDocument, nil, options, method, maxUnits) +} + +func (c *Client) processWith( + ctx context.Context, + snapshotForAttempt func() (preparedSnapshot, error), + readDocument func(context.Context, preparedSnapshot) ([]byte, error), + beforeEgress func() error, + options requestOptions, + method UnitBoundMethod, + maxUnits int, ) (Result, error) { requests := 0 var providerLatency time.Duration @@ -238,7 +266,7 @@ func (c *Client) process( return Result{}, newProcessError(err, requests, providerLatency) } result, retryHeader, requested, latency, processErr := c.processOnce( - ctx, snapshot, prefix, suffix, encodedLength, method, maxUnits, + ctx, snapshot, readDocument, beforeEgress, prefix, suffix, encodedLength, method, maxUnits, ) if requested { requests++ @@ -284,12 +312,14 @@ func (c *Client) validatePreparedSnapshot( func (c *Client) processOnce( ctx context.Context, snapshot preparedSnapshot, + readDocument func(context.Context, preparedSnapshot) ([]byte, error), + beforeEgress func() error, prefix, suffix []byte, encodedLength int64, method UnitBoundMethod, maxUnits int, ) (Result, string, bool, time.Duration, error) { - documentBytes, err := readVerifiedDocument(ctx, snapshot) + documentBytes, err := readDocument(ctx, snapshot) if err != nil { return Result{}, "", false, 0, err } @@ -315,7 +345,25 @@ func (c *Client) processOnce( } request.ContentLength = encodedLength request.Header.Set("Content-Type", mediaTypeJSON) - request.Header.Set("Authorization", "Bearer "+c.apiKey) + if beforeEgress != nil { + if err := beforeEgress(); err != nil { + return Result{}, "", false, 0, err + } + } + apiKey, err := c.credential(ctx) + if err != nil { + return Result{}, "", false, 0, &credentialError{cause: err} + } + if apiKey == "" || len(apiKey) > 64<<10 || apiKey != strings.TrimSpace(apiKey) || + strings.ContainsAny(apiKey, "\r\n\x00") { + return Result{}, "", false, 0, &credentialError{} + } + request.Header.Set("Authorization", "Bearer "+apiKey) + if beforeEgress != nil { + if err := beforeEgress(); err != nil { + return Result{}, "", false, 0, err + } + } started := time.Now() response, err := c.http.Do(request) @@ -332,7 +380,7 @@ func (c *Client) processOnce( return Result{}, response.Header.Get("Retry-After"), true, latency, &transientError{status: response.StatusCode} } if response.StatusCode >= http.StatusBadRequest { - return Result{}, "", true, latency, fmt.Errorf("mistral OCR HTTP %d: %w", response.StatusCode, ErrPermanentResponse) + return Result{}, "", true, latency, &permanentResponseError{status: response.StatusCode} } if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { return Result{}, "", true, latency, fmt.Errorf("mistral OCR unexpected HTTP %d", response.StatusCode) @@ -409,7 +457,9 @@ func (c *Client) processOnce( if err := validateWireResult(wire, c.policy.values.Model, snapshot, method, maxUnits); err != nil { return Result{}, "", true, latency, err } - return providerNeutralResult(wire, snapshot.format), "", true, latency, nil + result := providerNeutralResult(wire, snapshot.format) + result.ResponseBytes = int64(len(body)) + return result, "", true, latency, nil } func newProcessError(err error, requests int, providerLatency time.Duration) error { @@ -428,6 +478,19 @@ type transientError struct { cause error } +type credentialError struct{ cause error } + +func (e *credentialError) Error() string { return "Mistral OCR credential is unavailable" } +func (e *credentialError) Unwrap() error { return e.cause } + +type permanentResponseError struct{ status int } + +func (e *permanentResponseError) Error() string { + return fmt.Sprintf("mistral OCR HTTP %d: %s", e.status, ErrPermanentResponse) +} + +func (e *permanentResponseError) Unwrap() error { return ErrPermanentResponse } + func (e *transientError) Error() string { if e.cause != nil { return e.cause.Error() @@ -655,7 +718,8 @@ func validateWireResult( return errors.New("mistral OCR response omitted model") } if result.Model != expectedModel { - return fmt.Errorf("mistral OCR response model %q does not match requested model", result.Model) + return fmt.Errorf("mistral OCR response model %q does not match requested model: %w", + result.Model, ErrCapabilityContract) } if result.Pages == nil { return errors.New("mistral OCR response omitted pages") diff --git a/document/mistral/embedding.go b/document/mistral/embedding.go new file mode 100644 index 00000000..b738e12d --- /dev/null +++ b/document/mistral/embedding.go @@ -0,0 +1,515 @@ +package mistral + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + json "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "net/http" + "net/netip" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/manifestjson" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + EmbeddingProviderID = "mistral.embeddings-v1" + EmbeddingDocumentFormatterV1 = "mistral/document/v1" + EmbeddingQueryFormatterV1 = "mistral/query/v1" + EmbeddingScalarFloat32 = "float32" + EmbeddingModel = "mistral-embed" + EmbeddingHostedAliasRevision = "mutable-alias-export-only" + + embeddingOrigin = "https://api.mistral.ai" + embeddingPath = "/v1/embeddings" + embeddingAdapterContract = "docbank-mistral-embeddings/v1" + + defaultEmbeddingTimeout = 30 * time.Second + defaultEmbeddingMaxBatchItems = 128 + defaultEmbeddingMaxInputBytes = int64(1 << 20) + defaultEmbeddingMaxRequestBytes = int64(2 << 20) + defaultEmbeddingMaxResponseBytes = int64(32 << 20) + maxEmbeddingSecretBytes = 64 << 10 + embeddingUnitLengthTolerance = 1e-4 +) + +type EmbeddingProfile struct { + Endpoint string + EgressPolicy providerhttp.EgressPolicy + Descriptor document.EmbeddingDescriptor + ModelInput document.ModelInputContract + SecretBinding string + RequestTimeout time.Duration + MaxRetries int + MaxRetryDelay time.Duration + MaxBatchItems int + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 +} + +type embeddingPolicyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + ModelInput document.ModelInputContract `json:"model_input"` + CredentialBinding string `json:"credential_binding"` + Egress embeddingEgressIdentity `json:"egress"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxRetries int `json:"max_retries"` + MaxRetryDelay int64 `json:"max_retry_delay_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` +} + +type embeddingEgressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +var ErrEmbeddingCapacity = errors.New("mistral embedding capacity exceeded") + +type EmbeddingClient struct { + profile EmbeddingProfile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +var _ document.EmbeddingProvider = (*EmbeddingClient)(nil) + +type embeddingWireRequest struct { + Input []string `json:"input"` + Model string `json:"model"` + EncodingFormat string `json:"encoding_format"` +} + +type embeddingWireResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Data []embeddingWireItem `json:"data"` + Model string `json:"model"` + Usage embeddingWireUsage `json:"usage"` +} + +type embeddingWireItem struct { + Object string `json:"object"` + Embedding []float32 `json:"embedding"` + Index *int `json:"index"` +} + +type embeddingWireUsage struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + TotalTokens int64 `json:"total_tokens"` + PromptAudioSeconds *float64 `json:"prompt_audio_seconds"` + PromptTokensDetails *embeddingTokenDetails `json:"prompt_tokens_details,omitempty"` + PromptTokenDetails *embeddingTokenDetails `json:"prompt_token_details,omitempty"` + NumCachedTokens *int64 `json:"num_cached_tokens,omitempty"` + ServiceTier string `json:"service_tier,omitempty"` +} + +type embeddingTokenDetails struct { + CachedTokens int64 `json:"cached_tokens"` +} + +func EmbeddingPolicyFingerprint(profile EmbeddingProfile) (string, error) { + normalized, descriptorIdentity, err := normalizeEmbeddingProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(embeddingPolicyIdentity{ + AdapterContract: embeddingAdapterContract, Origin: normalized.Endpoint, Route: embeddingPath, + Descriptor: descriptorIdentity, ModelInput: normalized.ModelInput, + CredentialBinding: normalized.SecretBinding, Egress: embeddingEgressPolicyIdentity(normalized.EgressPolicy), + RequestTimeout: int64(normalized.RequestTimeout), MaxRetries: normalized.MaxRetries, + MaxRetryDelay: int64(normalized.MaxRetryDelay), MaxBatchItems: normalized.MaxBatchItems, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, + }, json.Deterministic(true)) + if err != nil { + return "", fmt.Errorf("mistral embedding: encode policy identity: %w", err) + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +func NewEmbeddingProvider(profile EmbeddingProfile, secrets SecretResolver, resolver providerhttp.Resolver) (*EmbeddingClient, error) { + normalized, _, err := normalizeEmbeddingProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("mistral embedding: invalid descriptor: %w", err) + } + fingerprint, err := EmbeddingPolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("mistral embedding: descriptor policy fingerprint does not match profile") + } + if descriptor.SupportsTextQuery { + return nil, errors.New("mistral embedding: hosted mutable alias is export-only") + } + if normalized.SecretBinding == "" || nilValue(secrets) { + return nil, errors.New("mistral embedding: named secret binding and resolver are required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, fmt.Errorf("mistral embedding: invalid sealed egress policy: %w", err) + } + normalized.Descriptor = cloneEmbeddingDescriptor(descriptor) + return &EmbeddingClient{profile: normalized, descriptor: cloneEmbeddingDescriptor(descriptor), secrets: secrets, http: &http.Client{Transport: transport, CheckRedirect: providerhttp.RefuseRedirects}}, nil +} + +func (client *EmbeddingClient) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneEmbeddingDescriptor(client.descriptor) +} + +func (client *EmbeddingClient) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if client == nil { + return document.EmbeddingResult{}, errors.New("mistral embedding: client is required") + } + if ctx == nil { + return document.EmbeddingResult{}, errors.New("mistral embedding: context is required") + } + if err := document.ValidateEmbeddingProviderRequest(client, inputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || authorization.MaxInputBytes > client.profile.MaxInputBytes || authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, errors.New("mistral embedding: authorization exceeds profile capacity") + } + rendered := make([]string, len(inputs)) + for index, input := range inputs { + if input.Role == document.EmbeddingRoleDocument { + rendered[index] = client.profile.ModelInput.EncodeDocument(input.Text) + } else { + rendered[index] = client.profile.ModelInput.EncodeQuery(input.Text) + } + } + payload, err := json.Marshal(embeddingWireRequest{Input: rendered, Model: client.descriptor.Model, EncodingFormat: "float"}) + if err != nil { + return document.EmbeddingResult{}, errors.New("mistral embedding: could not encode request") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + return document.EmbeddingResult{}, fmt.Errorf("%w: request byte limit", ErrEmbeddingCapacity) + } + started := time.Now() + metrics := RequestMetrics{} + for attempt := 1; ; attempt++ { + result, retryHeader, retry, requested, attemptErr := client.embeddingAttempt(ctx, payload, inputs) + if requested { + metrics.Requests++ + } + if attemptErr == nil { + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, err + } + return result, nil + } + if !retry || attempt >= client.profile.MaxRetries { + metrics.Latency = time.Since(started) + return document.EmbeddingResult{}, &processError{err: attemptErr, metrics: metrics} + } + metrics.Retries++ + if waitErr := waitContext(ctx, retryAfter(retryHeader, attempt, client.profile.MaxRetryDelay)); waitErr != nil { + metrics.Latency = time.Since(started) + return document.EmbeddingResult{}, &processError{err: fmt.Errorf("mistral embedding: request canceled: %w", waitErr), metrics: metrics} + } + } +} + +func (client *EmbeddingClient) embeddingAttempt(ctx context.Context, payload []byte, inputs []document.EmbeddingInput) (document.EmbeddingResult, string, bool, bool, error) { + attemptCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + secret, err := client.secrets.ResolveSecret(attemptCtx, client.profile.SecretBinding) + if err != nil || !validEmbeddingSecret(secret) { + if contextErr := attemptCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, "", false, false, contextErr + } + return document.EmbeddingResult{}, "", false, false, fmt.Errorf("%w: credential unavailable", ErrPermanentResponse) + } + request, err := http.NewRequestWithContext(attemptCtx, http.MethodPost, client.profile.Endpoint+embeddingPath, bytes.NewReader(payload)) + if err != nil { + return document.EmbeddingResult{}, "", false, false, fmt.Errorf("%w: request construction", ErrPermanentResponse) + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+secret) + response, err := client.http.Do(request) + if err != nil { + if contextErr := attemptCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, "", false, true, contextErr + } + return document.EmbeddingResult{}, "", true, true, fmt.Errorf("%w: transport", ErrTransientResponse) + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode >= 300 && response.StatusCode < 400 { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: redirect HTTP %d", ErrPermanentResponse, response.StatusCode) + } + if response.StatusCode == http.StatusRequestEntityTooLarge { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: HTTP %d", ErrEmbeddingCapacity, response.StatusCode) + } + if response.StatusCode == http.StatusTooManyRequests || response.StatusCode >= 500 { + return document.EmbeddingResult{}, response.Header.Get("Retry-After"), true, true, fmt.Errorf("%w: HTTP %d", ErrTransientResponse, response.StatusCode) + } + if response.StatusCode != http.StatusOK { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: HTTP %d", ErrPermanentResponse, response.StatusCode) + } + if err := validateEmbeddingContentType(response.Header.Get("Content-Type")); err != nil { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: response content type", ErrPermanentResponse) + } + body, err := readEmbeddingBody(attemptCtx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + if contextErr := attemptCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, "", false, true, contextErr + } + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: bounded response", ErrPermanentResponse) + } + if err := manifestjson.RejectDuplicateKeys(body, "mistral embedding response"); err != nil { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: duplicate response member", ErrPermanentResponse) + } + var decoded embeddingWireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return document.EmbeddingResult{}, "", false, true, fmt.Errorf("%w: malformed response", ErrPermanentResponse) + } + result, err := client.validateAndOrder(decoded, inputs) + if err != nil { + err = fmt.Errorf("%w: invalid response", ErrPermanentResponse) + } + return result, "", false, true, err +} + +func (client *EmbeddingClient) validateAndOrder(response embeddingWireResponse, inputs []document.EmbeddingInput) (document.EmbeddingResult, error) { + if response.ID == "" || response.Object != "list" || response.Model != client.descriptor.Model { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider model or response contract drifted") + } + if response.Usage.PromptTokens < 0 || response.Usage.CompletionTokens < 0 || response.Usage.TotalTokens < response.Usage.PromptTokens+response.Usage.CompletionTokens { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider usage is invalid") + } + for _, details := range []*embeddingTokenDetails{response.Usage.PromptTokensDetails, response.Usage.PromptTokenDetails} { + if details != nil && (details.CachedTokens < 0 || details.CachedTokens > response.Usage.PromptTokens) { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider usage is invalid") + } + } + if response.Usage.NumCachedTokens != nil && (*response.Usage.NumCachedTokens < 0 || *response.Usage.NumCachedTokens > response.Usage.PromptTokens) { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider usage is invalid") + } + if response.Usage.ServiceTier != "" && response.Usage.ServiceTier != "standard" && response.Usage.ServiceTier != "priority" { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider usage is invalid") + } + if len(response.Data) != len(inputs) { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider response has a missing vector") + } + vectors := make([]document.EmbeddingVector, len(inputs)) + seen := make([]bool, len(inputs)) + for _, item := range response.Data { + if item.Object != "embedding" || item.Index == nil || *item.Index < 0 || *item.Index >= len(inputs) || seen[*item.Index] { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider response index contract drifted") + } + if err := client.validateVector(item.Embedding); err != nil { + return document.EmbeddingResult{}, err + } + seen[*item.Index] = true + vectors[*item.Index] = document.EmbeddingVector{Key: inputs[*item.Index].Key, Values: slices.Clone(item.Embedding)} + } + if slices.Contains(seen, false) { + return document.EmbeddingResult{}, errors.New("mistral embedding: provider response has a missing vector index") + } + return document.EmbeddingResult{Vectors: vectors}, nil +} + +func (client *EmbeddingClient) validateVector(vector []float32) error { + if len(vector) != client.descriptor.Dimension { + return errors.New("mistral embedding: provider vector dimension does not match profile") + } + var norm float64 + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("mistral embedding: provider vector contains a non-finite value") + } + norm += float64(value) * float64(value) + } + if norm == 0 { + return errors.New("mistral embedding: provider returned a zero vector") + } + if client.descriptor.Normalization == document.VectorNormalizationUnitLength && math.Abs(norm-1) > embeddingUnitLengthTolerance { + return errors.New("mistral embedding: provider vector normalization does not match profile") + } + return nil +} + +func normalizeEmbeddingProfile(profile EmbeddingProfile) (EmbeddingProfile, document.EmbeddingDescriptor, error) { + if profile.Endpoint == "" { + profile.Endpoint = embeddingOrigin + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultEmbeddingTimeout + } + if profile.MaxRetries == 0 { + profile.MaxRetries = DefaultMaxRetries + } + if profile.MaxRetryDelay == 0 { + profile.MaxRetryDelay = DefaultMaxRetryDelay + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultEmbeddingMaxBatchItems + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultEmbeddingMaxInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultEmbeddingMaxRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultEmbeddingMaxResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > MaxTimeout || profile.MaxRetries < 1 || profile.MaxRetries > MaxRetries || profile.MaxRetryDelay <= 0 || profile.MaxRetryDelay > MaxRetryDelay || profile.MaxBatchItems < 1 || profile.MaxBatchItems > 1000 || profile.MaxInputBytes < 1 || profile.MaxRequestBytes < 1 || profile.MaxResponseBytes < 1 { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, errors.New("mistral embedding: execution bounds are invalid") + } + if !validEmbeddingBinding(profile.SecretBinding) { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, errors.New("mistral embedding: secret binding is invalid") + } + if err := normalizeAndValidateEmbeddingEgress(&profile); err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, err + } + descriptorIdentity := profile.Descriptor + descriptorIdentity.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptorIdentity.Fingerprint = "" + var err error + descriptorIdentity, err = document.NewEmbeddingDescriptor(descriptorIdentity) + if err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, fmt.Errorf("mistral embedding: invalid descriptor identity: %w", err) + } + descriptorIdentity.PolicyFingerprint = "" + descriptorIdentity.Fingerprint = "" + if descriptorIdentity.SupportsTextQuery { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, errors.New("mistral embedding: hosted mutable alias is export-only") + } + if descriptorIdentity.ID != EmbeddingProviderID || descriptorIdentity.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptorIdentity.Model != EmbeddingModel || descriptorIdentity.ModelRevision != EmbeddingHostedAliasRevision || descriptorIdentity.Dimension != 1024 || descriptorIdentity.ScalarEncoding != EmbeddingScalarFloat32 || descriptorIdentity.DocumentFormatter != EmbeddingDocumentFormatterV1 || descriptorIdentity.QueryFormatter != EmbeddingQueryFormatterV1 || !slices.Equal(descriptorIdentity.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || !slices.Equal(descriptorIdentity.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeText}) || descriptorIdentity.ModelInput != profile.ModelInput || descriptorIdentity.CompatibilityID != profile.ModelInput.CompatibilityID || profile.ModelInput.Document.Mode != document.ModelInputModeText || profile.ModelInput.Query.Mode != document.ModelInputModeText { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, errors.New("mistral embedding: descriptor does not match text adapter contract") + } + return profile, descriptorIdentity, nil +} + +func normalizeAndValidateEmbeddingEgress(profile *EmbeddingProfile) error { + if profile.EgressPolicy.ConnectTimeout == 0 { + profile.EgressPolicy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if profile.EgressPolicy.KeepAlive == 0 { + profile.EgressPolicy.KeepAlive = providerhttp.DefaultKeepAlive + } + if profile.EgressPolicy.TLSHandshakeTimeout == 0 { + profile.EgressPolicy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if profile.EgressPolicy.ProxyMode == "" { + profile.EgressPolicy.ProxyMode = providerhttp.ProxyDisabled + } + if profile.EgressPolicy.TLS.RootCAs != nil { + return errors.New("mistral embedding: custom egress roots cannot enter canonical identity") + } + parsed, err := url.Parse(profile.Endpoint) + if err != nil || parsed.User != nil || parsed.Path != "" || parsed.RawQuery != "" || parsed.Fragment != "" { + return errors.New("mistral embedding: endpoint must be an exact provider origin") + } + port := parsed.Port() + if port == "" { + switch parsed.Scheme { + case "https": + port = "443" + case "http": + port = "80" + } + } + if parsed.Scheme != profile.EgressPolicy.Scheme || !strings.EqualFold(parsed.Hostname(), profile.EgressPolicy.Host) || port != strconv.FormatUint(uint64(profile.EgressPolicy.Port), 10) { + return errors.New("mistral embedding: endpoint and egress authority differ") + } + profile.Endpoint = strings.TrimSuffix(profile.Endpoint, "/") + slices.SortFunc(profile.EgressPolicy.AllowedCIDRs, func(a, b netip.Prefix) int { return strings.Compare(a.Masked().String(), b.Masked().String()) }) + slices.Sort(profile.EgressPolicy.TLS.SPKISHA256) + return nil +} + +func embeddingEgressPolicyIdentity(policy providerhttp.EgressPolicy) embeddingEgressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.Masked().String() + } + return embeddingEgressIdentity{Scheme: policy.Scheme, Host: strings.ToLower(policy.Host), Port: policy.Port, AllowedCIDRs: cidrs, ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func readEmbeddingBody(ctx context.Context, reader io.Reader, maximum int64) ([]byte, error) { + body, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("mistral embedding: response read canceled: %w", contextErr) + } + return nil, errors.New("mistral embedding: could not read provider response") + } + if int64(len(body)) > maximum { + return nil, errors.New("mistral embedding: provider response byte limit exceeded") + } + return body, nil +} + +func validateEmbeddingContentType(value string) error { + mediaType, parameters, err := mime.ParseMediaType(value) + if err != nil || mediaType != "application/json" || (len(parameters) != 0 && (len(parameters) != 1 || !strings.EqualFold(parameters["charset"], "utf-8"))) { + return errors.New("mistral embedding: provider response content type is invalid") + } + return nil +} + +func validEmbeddingSecret(secret string) bool { + if secret == "" || len(secret) > maxEmbeddingSecretBytes { + return false + } + for _, character := range secret { + if unicode.IsControl(character) || unicode.IsSpace(character) { + return false + } + } + return true +} + +func validEmbeddingBinding(value string) bool { + return value != "" && len(value) <= 128 && value == strings.TrimSpace(value) && strings.IndexFunc(value, unicode.IsControl) < 0 +} + +func cloneEmbeddingDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} diff --git a/document/mistral/embedding_review_test.go b/document/mistral/embedding_review_test.go new file mode 100644 index 00000000..51db8f66 --- /dev/null +++ b/document/mistral/embedding_review_test.go @@ -0,0 +1,47 @@ +package mistral + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/netip" + "net/url" + "strconv" + "testing" + + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document/providerhttp" +) + +type mistralEmbeddingResolver struct{ address netip.Addr } + +func (resolver mistralEmbeddingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return []netip.Addr{resolver.address}, nil +} + +func mistralFixture(t *testing.T, handler http.Handler) (string, providerhttp.EgressPolicy, providerhttp.Resolver) { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + parsed, err := url.Parse(server.URL) + require.NoError(t, err) + _, portText, err := net.SplitHostPort(parsed.Host) + require.NoError(t, err) + port, err := strconv.ParseUint(portText, 10, 16) + require.NoError(t, err) + return "http://mistral.invalid:" + portText, providerhttp.EgressPolicy{ + Scheme: "http", Host: "mistral.invalid", Port: uint16(port), + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, ProxyMode: providerhttp.ProxyDisabled, + }, mistralEmbeddingResolver{address: netip.MustParseAddr("127.0.0.1")} +} + +func TestHostedMistralAliasIsAlwaysExportOnlyAndSealed(t *testing.T) { + endpoint, egress, resolver := mistralFixture(t, http.NotFoundHandler()) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile.Descriptor.SupportsTextQuery = true + _, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "synthetic-secret"}, resolver) + require.ErrorContains(t, err, "export-only") +} diff --git a/document/mistral/embedding_test.go b/document/mistral/embedding_test.go new file mode 100644 index 00000000..e5add100 --- /dev/null +++ b/document/mistral/embedding_test.go @@ -0,0 +1,335 @@ +package mistral + +import ( + "context" + json "encoding/json/v2" + "errors" + "net/http" + "net/netip" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +type embeddingSecretMap map[string]string + +func (secrets embeddingSecretMap) ResolveSecret(_ context.Context, name string) (string, error) { + return secrets[name], nil +} + +func TestEmbeddingProviderAppliesDocumentEnvelopeAndRestoresIndices(t *testing.T) { + var request struct { + Input []string `json:"input"` + Model string `json:"model"` + EncodingFormat string `json:"encoding_format"` + } + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, incoming *http.Request) { + assert.Equal(t, "/v1/embeddings", incoming.URL.Path) + assert.Equal(t, "Bearer synthetic-secret", incoming.Header.Get("Authorization")) + assert.NoError(t, json.UnmarshalRead(incoming.Body, &request, json.RejectUnknownMembers(true))) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(mistralEmbeddingBody(t, []int{1, 0})) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "synthetic-secret"}, resolver) + require.NoError(t, err) + inputs := []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage"}, + {Key: "second", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "second"}, + } + result, err := document.ExecuteEmbedding(t.Context(), provider, inputs, embeddingAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.Equal(t, []string{"document envelope: passage", "document envelope: second"}, request.Input) + assert.Equal(t, []string{"first", "second"}, []string{result.Vectors[0].Key, result.Vectors[1].Key}) +} + +func TestMistralEmbeddingPolicyFingerprintCoversCompleteContractAndEgress(t *testing.T) { + profile := testEmbeddingProfile(t) + baseline, err := EmbeddingPolicyFingerprint(profile) + require.NoError(t, err) + changed := profile + changed.ModelInput = testMistralModelInput(t, "changed: {{content}}", "query envelope: {{content}}") + changed.Descriptor.ModelInput, changed.Descriptor.CompatibilityID = changed.ModelInput, changed.ModelInput.CompatibilityID + recomputeMistralEmbeddingProfile(t, &changed) + assert.NotEqual(t, baseline, changed.Descriptor.PolicyFingerprint) + changed = profile + changed.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24")} + recomputeMistralEmbeddingProfile(t, &changed) + assert.NotEqual(t, baseline, changed.Descriptor.PolicyFingerprint) +} + +func TestMistralEmbeddingRejectsMalformedAndPartialResponsesPrivately(t *testing.T) { + one, two := mistralUnitEmbedding(0), mistralUnitEmbedding(1) + valid := func(indices []int, vectors [][]float32) []byte { + items := make([]map[string]any, len(indices)) + for index := range indices { + items[index] = map[string]any{"object": "embedding", "embedding": vectors[index], "index": indices[index]} + } + body, err := json.Marshal(map[string]any{"id": "synthetic", "object": "list", "data": items, "model": EmbeddingModel, "usage": map[string]any{"prompt_tokens": 2, "completion_tokens": 0, "total_tokens": 2, "prompt_audio_seconds": nil}}) + require.NoError(t, err) + return body + } + tests := []struct { + name string + body []byte + }{ + {"malformed", []byte(`{"object":`)}, {"unknown", []byte(`{"unknown":"PRIVATE_RAW_BODY"}`)}, {"duplicate", []byte(`{"model":"mistral-embed","model":"PRIVATE_RAW_BODY"}`)}, + {"partial indices", valid([]int{0}, [][]float32{one})}, {"duplicate index", valid([]int{0, 0}, [][]float32{one, two})}, {"out of range", valid([]int{0, 2}, [][]float32{one, two})}, {"wrong dimension", valid([]int{0, 1}, [][]float32{one[:1023], two})}, + {"zero vector", valid([]int{0, 1}, [][]float32{make([]float32, 1024), two})}, {"non finite", []byte(`{"id":"x","object":"list","data":[{"object":"embedding","embedding":[1e999],"index":0}],"model":"mistral-embed","usage":{"prompt_tokens":1,"completion_tokens":0,"total_tokens":1,"prompt_audio_seconds":null}}`)}, + } + inputs := []document.EmbeddingInput{{Key: "a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "PRIVATE_INPUT"}, {Key: "b", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "other"}} + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(test.body) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "PRIVATE_SECRET"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), inputs, embeddingAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.NotContains(t, err.Error(), "PRIVATE") + }) + } +} + +func TestMistralEmbeddingClassifiesCapacityPermanentAndTransient(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + for _, test := range []struct { + name string + status int + kind error + }{{"capacity", http.StatusRequestEntityTooLarge, ErrEmbeddingCapacity}, {"permanent", http.StatusBadRequest, ErrPermanentResponse}, {"rate limit", http.StatusTooManyRequests, ErrTransientResponse}, {"transient", http.StatusServiceUnavailable, ErrTransientResponse}} { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int32 + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + calls.Add(1) + writer.WriteHeader(test.status) + _, _ = writer.Write([]byte("PRIVATE_BODY")) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy, profile.MaxRetries, profile.MaxRetryDelay = endpoint, egress, 2, time.Millisecond + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, embeddingAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, test.kind) + assert.NotContains(t, err.Error(), "PRIVATE_BODY") + metrics := MetricsFromError(err) + assert.Equal(t, int(calls.Load()), metrics.Requests) + if errors.Is(test.kind, ErrTransientResponse) { + assert.Equal(t, 2, metrics.Requests) + assert.Equal(t, 1, metrics.Retries) + } + }) + } +} + +func TestMistralEmbeddingRejectsIdentityDrift(t *testing.T) { + for _, test := range []struct { + name string + mutate func(*EmbeddingProfile) + }{ + {"model", func(profile *EmbeddingProfile) { profile.Descriptor.Model = "mistral-embed-drift" }}, + {"revision", func(profile *EmbeddingProfile) { profile.Descriptor.ModelRevision = "invented-revision" }}, + {"compatibility", func(profile *EmbeddingProfile) { profile.Descriptor.CompatibilityID = "mistral/other-space/v1" }}, + } { + t.Run(test.name, func(t *testing.T) { + profile := testEmbeddingProfile(t) + test.mutate(&profile) + _, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, nil) + require.Error(t, err) + }) + } +} + +func TestMistralEmbeddingResponseLimitAndTransportExhaustion(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + t.Run("response limit", func(t *testing.T) { + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(strings.Repeat("x", 8193))) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy, profile.MaxResponseBytes = endpoint, egress, 8192 + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + authorization := embeddingAuthorization(profile.Descriptor) + authorization.MaxResponseBytes = 8192 + _, err = provider.Embed(t.Context(), input, authorization) + require.ErrorIs(t, err, ErrPermanentResponse) + }) + t.Run("transport exhaustion", func(t *testing.T) { + profile := testEmbeddingProfile(t) + profile.MaxRetries, profile.MaxRetryDelay = 2, time.Millisecond + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, failingMistralEmbeddingResolver{}) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, embeddingAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, ErrTransientResponse) + assert.Equal(t, 2, MetricsFromError(err).Requests) + }) +} + +type failingMistralEmbeddingResolver struct{} + +func (failingMistralEmbeddingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return nil, errors.New("synthetic DNS failure") +} + +func TestMistralEmbeddingRequestBoundRedirectAndCancellation(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Location", "http://elsewhere.invalid/private") + writer.WriteHeader(http.StatusTemporaryRedirect) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, embeddingAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, ErrPermanentResponse) + bounded := profile + bounded.MaxRequestBytes = 1 + recomputeMistralEmbeddingProfile(t, &bounded) + provider, err = NewEmbeddingProvider(bounded, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, embeddingAuthorization(bounded.Descriptor)) + require.ErrorIs(t, err, ErrEmbeddingCapacity) +} + +func TestMistralEmbeddingAcceptsDocumentedOptionalUsageFields(t *testing.T) { + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + body := mistralEmbeddingBody(t, []int{0}) + body = []byte(strings.Replace(string(body), `"prompt_tokens":2`, `"prompt_tokens":2,"prompt_tokens_details":{"cached_tokens":1},"prompt_token_details":{"cached_tokens":1},"num_cached_tokens":1,"service_tier":"standard"`, 1)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(body) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}}, embeddingAuthorization(profile.Descriptor)) + require.NoError(t, err) +} + +func TestMistralEmbeddingCancellationPreservesIdentityAndMetrics(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + t.Run("retry wait", func(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond) + defer cancel() + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Retry-After", "60") + writer.WriteHeader(http.StatusServiceUnavailable) + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(ctx, input, embeddingAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Equal(t, RequestMetrics{Requests: 1, Retries: 1}, withoutLatency(MetricsFromError(err))) + }) + + t.Run("response read", func(t *testing.T) { + started := make(chan struct{}) + endpoint, egress, resolver := mistralFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusOK) + flusher, ok := writer.(http.Flusher) + if !assert.True(t, ok) { + return + } + flusher.Flush() + close(started) + <-request.Context().Done() + })) + profile := testEmbeddingProfile(t) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + recomputeMistralEmbeddingProfile(t, &profile) + provider, err := NewEmbeddingProvider(profile, embeddingSecretMap{"credential:mistral-embed": "secret"}, resolver) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + _, embedErr := provider.Embed(ctx, input, embeddingAuthorization(profile.Descriptor)) + done <- embedErr + }() + <-started + cancel() + err = <-done + require.ErrorIs(t, err, context.Canceled) + assert.Equal(t, RequestMetrics{Requests: 1}, withoutLatency(MetricsFromError(err))) + }) +} + +func withoutLatency(metrics RequestMetrics) RequestMetrics { metrics.Latency = 0; return metrics } + +func recomputeMistralEmbeddingProfile(t *testing.T, profile *EmbeddingProfile) { + t.Helper() + profile.Descriptor.PolicyFingerprint = strings.Repeat("0", 64) + profile.Descriptor.Fingerprint = "" + profile.Descriptor, _ = document.NewEmbeddingDescriptor(profile.Descriptor) + fingerprint, err := EmbeddingPolicyFingerprint(*profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor.Fingerprint = "" + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) +} + +func mistralUnitEmbedding(hot int) []float32 { + vector := make([]float32, 1024) + vector[hot] = 1 + return vector +} + +func testEmbeddingProfile(t *testing.T) EmbeddingProfile { + t.Helper() + modelInput := testMistralModelInput(t, "document envelope: {{content}}", "query envelope: {{content}}") + profile := EmbeddingProfile{Endpoint: embeddingOrigin, EgressPolicy: productionMistralEgress(), ModelInput: modelInput, SecretBinding: "credential:mistral-embed", MaxBatchItems: 8, MaxInputBytes: 4096, MaxRequestBytes: 8192, MaxResponseBytes: 1 << 20, Descriptor: document.EmbeddingDescriptor{ID: EmbeddingProviderID, ContractVersion: document.EmbeddingProviderContractVersion, PolicyFingerprint: strings.Repeat("0", 64), TrustBoundary: document.EmbeddingTrustHostedProvider, Model: EmbeddingModel, ModelRevision: EmbeddingHostedAliasRevision, Dimension: 1024, Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: EmbeddingScalarFloat32, DocumentFormatter: EmbeddingDocumentFormatterV1, QueryFormatter: EmbeddingQueryFormatterV1, InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: modelInput.CompatibilityID, ModelInput: modelInput, SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}}} + recomputeMistralEmbeddingProfile(t, &profile) + return profile +} + +func productionMistralEgress() providerhttp.EgressPolicy { + return providerhttp.EgressPolicy{Scheme: "https", Host: "api.mistral.ai", Port: 443, AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")}, ProxyMode: providerhttp.ProxyDisabled} +} + +func testMistralModelInput(t *testing.T, documentTemplate, queryTemplate string) document.ModelInputContract { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: "mistral/test-space/v1", Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: documentTemplate}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: queryTemplate}}) + require.NoError(t, err) + return contract +} + +func embeddingAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 8, MaxInputBytes: 4096, MaxResponseBytes: 1 << 20} +} + +func mistralEmbeddingBody(t *testing.T, indices []int) []byte { + t.Helper() + items := make([]map[string]any, len(indices)) + for position, index := range indices { + items[position] = map[string]any{"object": "embedding", "embedding": mistralUnitEmbedding(index + 1), "index": index} + } + body, err := json.Marshal(map[string]any{"id": "synthetic-response", "object": "list", "data": items, "model": EmbeddingModel, "usage": map[string]any{"prompt_tokens": 2, "completion_tokens": 0, "total_tokens": 2, "prompt_audio_seconds": nil}}) + require.NoError(t, err) + return body +} diff --git a/document/mistral/formats.go b/document/mistral/formats.go index b2c5bda0..721278ca 100644 --- a/document/mistral/formats.go +++ b/document/mistral/formats.go @@ -2,7 +2,8 @@ package mistral import ( "fmt" - "slices" + + "go.kenn.io/docbank/document/internal/formatdetect" ) const ( @@ -13,55 +14,18 @@ const ( // CandidateFormat describes one locally detectable document format. A // candidate does not authorize an upload. -type CandidateFormat struct { - ID string `json:"id"` - Family string `json:"family"` - MediaType string `json:"media_type"` - UnitKind string `json:"unit_kind"` -} +type CandidateFormat = formatdetect.CandidateFormat -var candidateFormats = []CandidateFormat{ - {ID: formatIDPDF, Family: formatIDPDF, MediaType: mediaTypePDF, UnitKind: "page"}, - {ID: "docx", Family: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", UnitKind: "page"}, - {ID: "doc", Family: "word", MediaType: "application/msword", UnitKind: "page"}, - {ID: "odt", Family: "word", MediaType: "application/vnd.oasis.opendocument.text", UnitKind: "page"}, - {ID: "rtf", Family: "word", MediaType: "application/rtf", UnitKind: "page"}, - {ID: "pptx", Family: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", UnitKind: "slide"}, - {ID: "ppt", Family: "presentation", MediaType: "application/vnd.ms-powerpoint", UnitKind: "slide"}, - {ID: "xlsx", Family: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", UnitKind: "sheet"}, - {ID: "xls", Family: "spreadsheet", MediaType: "application/vnd.ms-excel", UnitKind: "sheet"}, - {ID: "ods", Family: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", UnitKind: "sheet"}, - {ID: "numbers", Family: "spreadsheet", MediaType: "application/vnd.apple.numbers", UnitKind: "sheet"}, - {ID: "csv", Family: "spreadsheet", MediaType: "text/csv", UnitKind: "record"}, - {ID: "epub", Family: "ebook", MediaType: "application/epub+zip", UnitKind: "spine"}, - {ID: "txt", Family: "text", MediaType: "text/plain", UnitKind: "section"}, - {ID: "markdown", Family: "text", MediaType: "text/markdown", UnitKind: "section"}, - {ID: "rst", Family: "text", MediaType: "text/x-rst", UnitKind: "section"}, - {ID: "latex", Family: "text", MediaType: "application/x-tex", UnitKind: "section"}, - {ID: "json", Family: "structured", MediaType: mediaTypeJSON, UnitKind: "record"}, - {ID: "jsonl", Family: "structured", MediaType: "application/x-ndjson", UnitKind: "record"}, - {ID: "xml", Family: "structured", MediaType: "application/xml", UnitKind: "record"}, - {ID: "yaml", Family: "structured", MediaType: "application/yaml", UnitKind: "record"}, - {ID: "go", Family: "source", MediaType: "text/x-go", UnitKind: "section"}, - {ID: "python", Family: "source", MediaType: "text/x-python", UnitKind: "section"}, - {ID: "javascript", Family: "source", MediaType: "text/javascript", UnitKind: "section"}, - {ID: "eml", Family: "mail", MediaType: "message/rfc822", UnitKind: "message"}, - {ID: "msg", Family: "mail", MediaType: "application/vnd.ms-outlook", UnitKind: "message"}, -} +var candidateFormats = formatdetect.CandidateFormats() // CandidateFormats returns a defensive copy in stable probe order. func CandidateFormats() []CandidateFormat { - return slices.Clone(candidateFormats) + return formatdetect.CandidateFormats() } // CandidateFormatByID returns the candidate with the given stable identifier. func CandidateFormatByID(id string) (CandidateFormat, bool) { - for _, candidate := range candidateFormats { - if candidate.ID == id { - return candidate, true - } - } - return CandidateFormat{}, false + return formatdetect.CandidateFormatByID(id) } // ProbeFixtureSentinel returns the synthetic phrase required in one fixture. diff --git a/document/mistral/rendition.go b/document/mistral/rendition.go new file mode 100644 index 00000000..553e5e1f --- /dev/null +++ b/document/mistral/rendition.go @@ -0,0 +1,482 @@ +package mistral + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "net/http" + "reflect" + "slices" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/formatdetect" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + renditionProviderID = "mistral.ocr-v1" + renditionTimeForm = "2006-01-02T15:04:05.000000000Z" +) + +var _ document.RenditionProvider = (*RenditionClient)(nil) + +// SecretResolver resolves only the profile-bound Mistral OCR credential. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile binds one rendition adapter to the existing Mistral OCR policy and +// its complete capability evidence. +type Profile struct { + Policy Policy + CapabilityManifest CapabilityManifest + Descriptor document.RenditionDescriptor + SecretBinding string + Timeout time.Duration + MaxRetries int + MaxRetryDelay time.Duration +} + +// RenditionClient adapts the bounded Mistral OCR operation to Docbank source +// evidence. OCR remains separate from Mistral embedding operations. +type RenditionClient struct { + descriptor document.RenditionDescriptor + policy Policy + manifest CapabilityManifest + ocr *Client +} + +// NewRenditionProvider validates a fixed hosted OCR profile without resolving +// its named credential or performing network access. +func NewRenditionProvider( + profile Profile, secrets SecretResolver, httpClient *http.Client, +) (*RenditionClient, error) { + if profile.Policy.digest == "" { + return nil, errors.New("mistral rendition policy is invalid; use NewPolicy") + } + manifest := cloneCapabilityManifest(profile.CapabilityManifest) + if err := manifest.ValidateComplete(); err != nil { + return nil, fmt.Errorf("mistral rendition capability manifest: %w", err) + } + descriptor, err := document.NewRenditionDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("mistral rendition descriptor: %w", err) + } + if descriptor.ID != renditionProviderID { + return nil, fmt.Errorf("mistral rendition descriptor ID must be %s", renditionProviderID) + } + if descriptor.TrustBoundary != document.RenditionTrustHostedProvider || + !descriptor.ReturnsMarkdown || !descriptor.ReturnsStructured || len(descriptor.ArtifactRoles) != 0 { + return nil, errors.New("mistral rendition descriptor result or trust contract is invalid") + } + fingerprint, err := profile.Policy.Fingerprint(manifest) + if err != nil { + return nil, fmt.Errorf("mistral rendition policy fingerprint: %w", err) + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("mistral rendition descriptor policy fingerprint does not match capability evidence") + } + for _, format := range descriptor.SupportedFormats { + candidate, ok := renditionCandidate(format) + if !ok || format.InputKind != document.RenditionInputOriginalFile { + return nil, errors.New("mistral rendition descriptor format is not locally detectable") + } + authority, authorizeErr := profile.Policy.Authorize(manifest, candidate.ID) + if authorizeErr != nil || authority.method == UnitBoundNone { + return nil, fmt.Errorf("mistral rendition format %q has no enforceable upload authority", candidate.ID) + } + } + if profile.SecretBinding == "" || nilValue(secrets) { + return nil, errors.New("mistral rendition requires a named secret binding and resolver") + } + if err := validateRenditionToken(profile.SecretBinding, "secret binding"); err != nil { + return nil, err + } + if httpClient == nil { + return nil, errors.New("mistral rendition HTTP client is required") + } + if profile.Timeout == 0 { + profile.Timeout = DefaultTimeout + } + if profile.MaxRetries == 0 { + profile.MaxRetries = DefaultMaxRetries + } + if profile.MaxRetryDelay == 0 { + profile.MaxRetryDelay = DefaultMaxRetryDelay + } + if profile.Timeout < 0 || profile.Timeout > MaxTimeout || + profile.MaxRetries < 0 || profile.MaxRetries > MaxRetries || + profile.MaxRetryDelay < 0 || profile.MaxRetryDelay > MaxRetryDelay { + return nil, errors.New("mistral rendition execution bounds are invalid") + } + isolate := *httpClient + isolate.Jar = nil + isolate.CheckRedirect = providerhttp.RefuseRedirects + if isolate.Timeout <= 0 || isolate.Timeout > profile.Timeout { + isolate.Timeout = profile.Timeout + } + ocr, err := newClientWithCredential(profile.Policy, ClientConfig{ + Timeout: profile.Timeout, MaxRetries: profile.MaxRetries, + MaxRetryDelay: profile.MaxRetryDelay, HTTPClient: &isolate, + }, func(ctx context.Context) (string, error) { + return secrets.ResolveSecret(ctx, profile.SecretBinding) + }) + if err != nil { + return nil, fmt.Errorf("mistral rendition client: %w", err) + } + return &RenditionClient{ + descriptor: cloneRenditionDescriptor(descriptor), policy: profile.Policy, manifest: manifest, + ocr: ocr, + }, nil +} + +// Descriptor returns an immutable copy of the configured provider identity. +func (client *RenditionClient) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneRenditionDescriptor(client.descriptor) +} + +// Render verifies the one-shot upload again, re-detects its exact media type, +// and maps bounded OCR output into provider-neutral source evidence. +func (client *RenditionClient) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("mistral rendition client is required") + } + if _, err := document.ValidateRenditionProviderRequest(client, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + expiresAt, err := time.Parse(renditionTimeForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Mistral authorization expiry is invalid", err) + } + startedAt := time.Now().UTC() + if !startedAt.Before(expiresAt) { + return document.RenditionResult{}, expiredRenditionError() + } + operationCtx, cancel := context.WithDeadline(ctx, expiresAt) + defer cancel() + source, err := readRenditionUpload(operationCtx, upload, metadata, client.policy.values.MaxDocumentBytes) + if err != nil { + if authorizationDeadlineExpired(ctx, operationCtx) { + return document.RenditionResult{}, expiredRenditionError() + } + return document.RenditionResult{}, err + } + defer clear(source) + candidate, err := formatdetect.DetectFormat(bytes.NewReader(source), int64(len(source)), metadata.MediaType) + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorUnsupportedInput, + "Mistral input format could not be verified", err) + } + if candidate.Family != metadata.MediaFamily || candidate.MediaType != metadata.MediaType { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Mistral input identity does not match authorization", nil) + } + localUnits := int64(0) + if candidate.ID == formatIDPDF { + localUnits, err = formatdetect.CountPDFPages(source) + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorUnsupportedInput, + "Mistral PDF page count could not be verified", err) + } + if localUnits <= 0 || localUnits > int64(client.policy.values.MaxUnits) { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Mistral PDF exceeds the complete unit limit", nil) + } + } + formatAuthorization, err := client.policy.Authorize(client.manifest, candidate.ID) + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorUnsupportedInput, + "Mistral input has no enforceable upload authority", err) + } + snapshot := preparedSnapshot{ + size: int64(len(source)), sha256: metadata.SHA256, format: candidate, + mediaType: metadata.MediaType, + } + snapshotForAttempt := func() (preparedSnapshot, error) { + digest := sha256.Sum256(source) + if int64(len(source)) != metadata.ByteLength || hex.EncodeToString(digest[:]) != metadata.SHA256 { + return preparedSnapshot{}, errors.New("mistral rendition source identity changed") + } + return snapshot, nil + } + readDocument := func(attemptCtx context.Context, attempt preparedSnapshot) ([]byte, error) { + if err := attemptCtx.Err(); err != nil { + return nil, err + } + if attempt.size != int64(len(source)) || attempt.sha256 != metadata.SHA256 { + return nil, errors.New("mistral rendition source identity changed") + } + return bytes.Clone(source), nil + } + checkExpiry := func() error { + if !time.Now().UTC().Before(expiresAt) { + return expiredRenditionError() + } + return nil + } + options := probeRequestOptions(candidate, client.policy.values.MaxUnits, + client.policy.values.ExtractHeader, client.policy.values.ExtractFooter) + providerResult, err := client.ocr.processWith( + operationCtx, snapshotForAttempt, readDocument, checkExpiry, options, + formatAuthorization.method, client.policy.values.MaxUnits, + ) + if err != nil { + if authorizationDeadlineExpired(ctx, operationCtx) { + return document.RenditionResult{}, expiredRenditionError() + } + return document.RenditionResult{}, classifyRenditionError(ctx, err) + } + if candidate.ID == formatIDPDF && int64(providerResult.UnitsProcessed) != localUnits { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Mistral OCR page count changed", ErrCapabilityContract) + } + completedAt := time.Now().UTC() + if !completedAt.Before(expiresAt) { + return document.RenditionResult{}, expiredRenditionError() + } + evidence, markdown, err := mistralEvidence(providerResult.Document) + if err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "Mistral OCR output is malformed", err) + } + if len(markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, renditionError(document.RenditionErrorMalformedEvidence, + "Mistral OCR Markdown exceeds authorization", nil) + } + return document.RenditionResult{ + Evidence: evidence, ProviderMarkdown: markdown, + Receipt: document.RenditionReceipt{ + ProviderID: client.descriptor.ID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "mistral-" + authorization.RenditionRequestFingerprint[:24], + StartedAt: startedAt.Format(renditionTimeForm), CompletedAt: completedAt.Format(renditionTimeForm), + Usage: document.RenditionUsage{ + Requests: int64(providerResult.Metrics.Requests), Retries: int64(providerResult.Metrics.Retries), + InputBytes: metadata.ByteLength, OutputBytes: providerResult.ResponseBytes, + Units: int64(providerResult.UnitsProcessed), + }, + }, + }, nil +} + +func readRenditionUpload( + ctx context.Context, upload io.Reader, metadata document.AuthorizedUploadMetadata, maximum int64, +) ([]byte, error) { + if metadata.ByteLength > maximum { + return nil, renditionError(document.RenditionErrorPolicyRejected, + "Mistral input exceeds the policy byte limit", nil) + } + data, err := io.ReadAll(io.LimitReader(&contextReader{ctx: ctx, reader: upload}, maximum+1)) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, renditionError(document.RenditionErrorCanceled, "Mistral rendering canceled", ctxErr) + } + return nil, renditionError(document.RenditionErrorTransient, + "Mistral input could not be read", err) + } + digest := sha256.Sum256(data) + if int64(len(data)) != metadata.ByteLength || int64(len(data)) > maximum || + hex.EncodeToString(digest[:]) != metadata.SHA256 { + clear(data) + return nil, renditionError(document.RenditionErrorPolicyRejected, + "Mistral input identity does not match authorization", nil) + } + return data, nil +} + +func mistralEvidence(source document.SourceDocument) (document.SourceEvidenceV1, []byte, error) { + unitKind, locatorKind, indexed, ok := renditionUnitKinds(source.UnitKind) + if !ok || len(source.Units) == 0 { + return document.SourceEvidenceV1{}, nil, errors.New("unsupported or empty Mistral OCR unit sequence") + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: source.Family, UnitKind: unitKind, + Units: make([]document.SourceEvidenceUnitV1, 0, len(source.Units)), + } + markdownPages := make([]string, 0, len(source.Units)) + for order, sourceUnit := range source.Units { + if sourceUnit.Index != order || !utf8.ValidString(sourceUnit.Markdown) || + !utf8.ValidString(sourceUnit.Header) || !utf8.ValidString(sourceUnit.Footer) { + return document.SourceEvidenceV1{}, nil, errors.New("invalid Mistral OCR unit") + } + text := joinMistralUnit(sourceUnit) + locator := document.SourceEvidenceLocatorV1{ + Kind: locatorKind, IndexOrigin: document.EvidenceIndexOriginNone, + } + if indexed { + locator.IndexOrigin = document.EvidenceIndexOriginZero + locator.Start, locator.End = int64(order), int64(order) + if locatorKind == document.EvidenceLocatorSheet { + locator.Name = fmt.Sprintf("mistral-unit-%d", order) + } + } else { + locator.Name = fmt.Sprintf("mistral-unit-%d", order) + } + evidence.Units = append(evidence.Units, document.SourceEvidenceUnitV1{ + Order: order, ProviderID: fmt.Sprintf("mistral-unit-%d", order), Text: text, Locator: locator, + }) + markdownPages = append(markdownPages, sourceUnit.Markdown) + } + markdown := []byte(strings.Join(markdownPages, "\n\n---\n\n")) + if len(markdown) == 0 { + return document.SourceEvidenceV1{}, nil, errors.New("mistral OCR returned empty Markdown") + } + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.SourceEvidenceV1{}, nil, err + } + return evidence, markdown, nil +} + +func joinMistralUnit(unit document.SourceUnit) string { + parts := make([]string, 0, 3) + for _, part := range []string{unit.Header, unit.Markdown, unit.Footer} { + if part != "" { + parts = append(parts, part) + } + } + return strings.Join(parts, "\n\n") +} + +func renditionUnitKinds( + unitKind string, +) (document.EvidenceUnitKind, document.EvidenceLocatorKind, bool, bool) { + switch unitKind { + case "page": + return document.EvidenceUnitPage, document.EvidenceLocatorPage, true, true + case "slide": + return document.EvidenceUnitSlide, document.EvidenceLocatorSlide, true, true + case "sheet": + return document.EvidenceUnitSheet, document.EvidenceLocatorSheet, true, true + case "record": + return document.EvidenceUnitRecord, document.EvidenceLocatorRecord, true, true + case "spine": + return document.EvidenceUnitSpine, document.EvidenceLocatorSpine, true, true + case "line": + return document.EvidenceUnitLine, document.EvidenceLocatorLine, true, true + case "message": + return document.EvidenceUnitMessage, document.EvidenceLocatorMessage, false, true + case "section": + return document.EvidenceUnitSection, document.EvidenceLocatorSection, false, true + default: + return "", "", false, false + } +} + +func classifyRenditionError(ctx context.Context, cause error) error { + if providerError, ok := errors.AsType[*document.RenditionProviderError](cause); ok { + return providerError + } + if ctxErr := ctx.Err(); ctxErr != nil { + return renditionError(document.RenditionErrorCanceled, "Mistral rendering canceled", ctxErr) + } + if _, ok := errors.AsType[*credentialError](cause); ok { + return renditionError(document.RenditionErrorAuthentication, + "Mistral credential is unavailable", cause) + } + if transient, ok := errors.AsType[*transientError](cause); ok && transient.status == http.StatusTooManyRequests { + return renditionError(document.RenditionErrorRateLimited, "Mistral rate limit was exhausted", cause) + } + if transient, ok := errors.AsType[*transientError](cause); ok && + (transient.status == http.StatusServiceUnavailable || transient.status == http.StatusInsufficientStorage) { + return renditionError(document.RenditionErrorCapacity, "Mistral capacity is unavailable", cause) + } + if permanent, ok := errors.AsType[*permanentResponseError](cause); ok { + switch permanent.status { + case http.StatusUnauthorized, http.StatusForbidden: + return renditionError(document.RenditionErrorAuthentication, "Mistral authentication was rejected", cause) + case http.StatusUnsupportedMediaType: + return renditionError(document.RenditionErrorUnsupportedInput, "Mistral input format was rejected", cause) + default: + return renditionError(document.RenditionErrorPolicyRejected, "Mistral rejected the OCR request", cause) + } + } + switch { + case errors.Is(cause, ErrTransientResponse): + return renditionError(document.RenditionErrorTransient, "Mistral request retries were exhausted", cause) + case errors.Is(cause, ErrCapabilityContract): + return renditionError(document.RenditionErrorPolicyRejected, "Mistral OCR capability changed", cause) + case errors.Is(cause, ErrPermanentResponse): + return renditionError(document.RenditionErrorPolicyRejected, "Mistral rejected the OCR request", cause) + case errors.Is(cause, ErrResponseTooLarge): + return renditionError(document.RenditionErrorMalformedEvidence, "Mistral OCR response exceeds policy", cause) + default: + return renditionError(document.RenditionErrorMalformedEvidence, "Mistral OCR response is malformed", cause) + } +} + +func expiredRenditionError() error { + return renditionError(document.RenditionErrorPolicyRejected, "Mistral authorization expired", nil) +} + +func authorizationDeadlineExpired(callerCtx, operationCtx context.Context) bool { + return callerCtx.Err() == nil && errors.Is(operationCtx.Err(), context.DeadlineExceeded) +} + +func renditionError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("mistral rendition error classification: %w", err) + } + return providerError +} + +func renditionCandidate(format document.RenditionFormatCapability) (CandidateFormat, bool) { + for _, candidate := range candidateFormats { + if candidate.Family == format.MediaFamily && candidate.MediaType == format.MediaType { + return candidate, true + } + } + return CandidateFormat{}, false +} + +func cloneRenditionDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +func validateRenditionToken(value, subject string) error { + if value == "" || len(value) > 128 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return fmt.Errorf("mistral rendition %s must contain 1-128 characters", subject) + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || + char >= '0' && char <= '9' || strings.ContainsRune("_.-", char) { + continue + } + return fmt.Errorf("mistral rendition %s contains unsupported characters", subject) + } + return nil +} + +func nilValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/mistral/rendition_test.go b/document/mistral/rendition_test.go new file mode 100644 index 00000000..4d7c56ab --- /dev/null +++ b/document/mistral/rendition_test.go @@ -0,0 +1,570 @@ +package mistral + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +var _ document.RenditionProvider = (*RenditionClient)(nil) + +type renditionSecrets map[string]string + +func (secrets renditionSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + return secrets[name], nil +} + +func TestNewRenditionProviderAdvertisesOnlyManifestAuthorizedFormats(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + + client, err := NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: descriptor, + SecretBinding: "mistral-ocr", + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, http.DefaultClient) + require.NoError(t, err) + assert.Equal(t, descriptor, client.Descriptor()) + + broader := renditionDescriptor(t, policy, manifest, "pdf", "doc") + _, err = NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: broader, + SecretBinding: "mistral-ocr", + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, http.DefaultClient) + require.ErrorContains(t, err, "no enforceable upload authority") + + incomplete := manifest + incomplete.Results = incomplete.Results[:1] + _, err = NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: incomplete, Descriptor: descriptor, + SecretBinding: "mistral-ocr", + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, http.DefaultClient) + require.ErrorContains(t, err, "capability manifest") +} + +func TestRenditionClientMapsExactMistralOCRResponse(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + source := testPDF("rendition") + fixture := renditionFixture(t, descriptor, source) + var uploaded []byte + httpClient := &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + assert.Equal(t, "Bearer synthetic-key", request.Header.Get("Authorization")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + var wire struct { + Document struct { + URL string `json:"document_url"` + } `json:"document"` + } + require.NoError(t, json.Unmarshal(body, &wire)) + uploaded, err = base64.StdEncoding.DecodeString(strings.TrimPrefix( + wire.Document.URL, "data:application/pdf;base64,", + )) + require.NoError(t, err) + response := fmt.Sprintf(`{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"# Synthetic","header":"Header","footer":"Footer","dimensions":{"dpi":144,"height":792,"width":612}}],"usage_info":{"pages_processed":1,"doc_size_bytes":%d}}`, len(source)) + return &http.Response{ + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(response)), Request: request, + }, nil + })} + client, err := NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: descriptor, + SecretBinding: "mistral-ocr", Timeout: time.Second, MaxRetries: 1, + MaxRetryDelay: time.Millisecond, + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, httpClient) + require.NoError(t, err) + + result, err := document.RenderRendition(t.Context(), client, fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, source, uploaded) + assert.Equal(t, document.SourceEvidenceContractV1, result.Evidence.ContractVersion) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, "pdf", result.Evidence.Family) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 1) + assert.Equal(t, "Header\n\n# Synthetic\n\nFooter", result.Evidence.Units[0].Text) + assert.Equal(t, document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginZero, + Start: 0, End: 0, + }, result.Evidence.Units[0].Locator) + assert.Equal(t, "# Synthetic", string(result.ProviderMarkdown)) + assert.Equal(t, fixture.metadata.SHA256, result.Receipt.SourceSHA256) + assert.Equal(t, int64(1), result.Receipt.Usage.Requests) + assert.Equal(t, int64(len(source)), result.Receipt.Usage.InputBytes) + assert.Equal(t, int64(1), result.Receipt.Usage.Units) + assert.NotContains(t, fmt.Sprintf("%+v", result.Receipt), "synthetic-key") +} + +func TestRenditionClientClassifiesHTTPAndModelFailures(t *testing.T) { + for _, testCase := range []struct { + name string + status int + body string + want document.RenditionErrorCode + }{ + {name: "authentication", status: http.StatusUnauthorized, want: document.RenditionErrorAuthentication}, + {name: "unsupported", status: http.StatusUnsupportedMediaType, want: document.RenditionErrorUnsupportedInput}, + {name: "too large", status: http.StatusRequestEntityTooLarge, want: document.RenditionErrorPolicyRejected}, + {name: "capacity", status: http.StatusServiceUnavailable, want: document.RenditionErrorCapacity}, + {name: "rate limited", status: http.StatusTooManyRequests, want: document.RenditionErrorRateLimited}, + {name: "model drift", status: http.StatusOK, body: `{"model":"mistral-ocr-next","pages":[{"index":0,"markdown":"safe"}],"usage_info":{"pages_processed":1}}`, want: document.RenditionErrorPolicyRejected}, + {name: "malformed", status: http.StatusOK, body: `{`, want: document.RenditionErrorMalformedEvidence}, + } { + t.Run(testCase.name, func(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("classified")) + body := testCase.body + if body == "" { + body = `{"private":"provider detail"}` + } + client := newRenditionTestClient(t, policy, manifest, descriptor, + renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: testCase.status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), Request: request, + }, nil + })) + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "provider detail") + }) + } +} + +func TestRenditionClientRechecksIdentityAndDetectsMediaBeforeEgress(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("identity")) + var requests atomic.Int64 + client := newRenditionTestClient(t, policy, manifest, descriptor, + renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("unexpected request") + })) + + short, ok := fixture.upload().(*renditionUpload) + require.True(t, ok) + short.Reader = bytes.NewReader(fixture.source[:len(fixture.source)-1]) + _, err := client.Render(t.Context(), short, fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + + text := []byte("not a PDF") + textFixture := renditionFixture(t, descriptor, text) + _, err = client.Render(t.Context(), textFixture.upload(), textFixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorUnsupportedInput) + assert.Zero(t, requests.Load()) +} + +type countingSecrets struct { + calls atomic.Int64 + err error +} + +func (secrets *countingSecrets) ResolveSecret(context.Context, string) (string, error) { + secrets.calls.Add(1) + if secrets.err != nil { + return "", secrets.err + } + return "synthetic-key", nil +} + +func TestRenditionClientResolvesSecretPerAttemptAndStopsAtExpiry(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("expiry")) + fixture.authorization.ExpiresAt = time.Now().UTC().Add(20 * time.Millisecond).Format("2006-01-02T15:04:05.000000000Z") + secrets := &countingSecrets{} + var requests atomic.Int64 + client := newRenditionTestClient(t, policy, manifest, descriptor, secrets, + roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"1"}, "Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{}`)), Request: request, + }, nil + })) + + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + assert.Equal(t, int64(1), requests.Load()) + assert.Equal(t, int64(1), secrets.calls.Load(), "expired retry must not resolve another credential") +} + +func TestRenditionClientExpiryCancelsInFlightUpload(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("slow-upload")) + expiresAt := time.Now().UTC().Add(150 * time.Millisecond) + fixture.authorization.ExpiresAt = expiresAt.Format("2006-01-02T15:04:05.000000000Z") + requestStarted := make(chan struct{}) + client, err := NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: descriptor, + SecretBinding: "mistral-ocr", Timeout: time.Second, MaxRetries: 1, + MaxRetryDelay: time.Millisecond, + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + close(requestStarted) + buffer := make([]byte, 1) + for { + if _, readErr := request.Body.Read(buffer); readErr != nil { + return nil, readErr + } + timer := time.NewTimer(5 * time.Millisecond) + select { + case <-request.Context().Done(): + timer.Stop() + return nil, request.Context().Err() + case <-timer.C: + } + } + })}) + require.NoError(t, err) + callerCtx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + started := time.Now() + _, err = client.Render(callerCtx, fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + require.ErrorContains(t, err, "authorization expired") + assert.Less(t, time.Since(started), 600*time.Millisecond) + select { + case <-requestStarted: + default: + t.Fatal("slow upload never reached the transport") + } +} + +func TestRenditionClientRejectsPDFAboveCompleteUnitLimitBeforeEgress(t *testing.T) { + policy := testPolicy(t, 1<<20, MinUnits) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testMultipagePDF(MinUnits+1)) + var requests atomic.Int64 + client := newRenditionTestClient(t, policy, manifest, descriptor, + renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, errors.New("unexpected provider request") + })) + + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + require.ErrorContains(t, err, "unit limit") + assert.Zero(t, requests.Load()) +} + +func TestRenditionClientRejectsIncompletePDFResponse(t *testing.T) { + policy := testPolicy(t, 1<<20, MinUnits) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testMultipagePDF(2)) + client := newRenditionTestClient(t, policy, manifest, descriptor, + renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"first"}],"usage_info":{"pages_processed":1}}`, + )), Request: request, + }, nil + })) + + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + require.ErrorContains(t, err, "page count changed") +} + +func TestRenditionClientBoundsResponseAndKeepsCredentialErrorsSafe(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + policy.values.MaxResponseBytes = 64 + var err error + policy.digest, err = policyValuesDigest(policy.values) + require.NoError(t, err) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("bounded")) + client := newRenditionTestClient(t, policy, manifest, descriptor, + renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(strings.Repeat("x", 65))), Request: request, + }, nil + })) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorMalformedEvidence) + + privateCause := errors.New("secret provider token=private-value") + secrets := &countingSecrets{err: privateCause} + client = newRenditionTestClient(t, policy, manifest, descriptor, secrets, + roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("credential failure must precede egress") + return nil, errors.New("unexpected egress") + })) + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorAuthentication) + assert.NotContains(t, err.Error(), "private-value") +} + +func TestRenditionClientResolvesNamedSecretForEveryEgress(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("retry-secret")) + secrets := &countingSecrets{} + var requests atomic.Int64 + client := newRenditionTestClient(t, policy, manifest, descriptor, secrets, + roundTripFunc(func(request *http.Request) (*http.Response, error) { + attempt := requests.Add(1) + if attempt == 1 { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"0"}}, Body: io.NopCloser(strings.NewReader(`{}`)), Request: request, + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"safe"}],"usage_info":{"pages_processed":1}}`, + )), Request: request, + }, nil + })) + result, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + require.NoError(t, err) + assert.Equal(t, int64(2), requests.Load()) + assert.Equal(t, int64(2), secrets.calls.Load()) + assert.Equal(t, int64(1), result.Receipt.Usage.Retries) +} + +func TestRenditionClientRefusesRedirects(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + descriptor := renditionDescriptor(t, policy, manifest, "pdf") + fixture := renditionFixture(t, descriptor, testPDF("redirect")) + var targetRequests atomic.Int64 + target := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { targetRequests.Add(1) }) + targetServer := newTLSServer(t, target) + redirectServer := newTLSServer(t, http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + http.Redirect(response, request, targetServer.URL, http.StatusTemporaryRedirect) + })) + client := renditionServerClient(t, policy, manifest, descriptor, redirectServer) + + _, err := client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertRenditionCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Zero(t, targetRequests.Load()) +} + +func TestMistralEvidenceMapsNaturalFamilies(t *testing.T) { + for _, testCase := range []struct { + name, family, sourceKind string + wantKind document.EvidenceUnitKind + wantLocator document.EvidenceLocatorKind + }{ + {name: "pdf", family: "pdf", sourceKind: "page", wantKind: document.EvidenceUnitPage, wantLocator: document.EvidenceLocatorPage}, + {name: "word", family: "word", sourceKind: "page", wantKind: document.EvidenceUnitPage, wantLocator: document.EvidenceLocatorPage}, + {name: "presentation", family: "presentation", sourceKind: "slide", wantKind: document.EvidenceUnitSlide, wantLocator: document.EvidenceLocatorSlide}, + {name: "spreadsheet sheet", family: "spreadsheet", sourceKind: "sheet", wantKind: document.EvidenceUnitSheet, wantLocator: document.EvidenceLocatorSheet}, + {name: "spreadsheet record", family: "spreadsheet", sourceKind: "record", wantKind: document.EvidenceUnitRecord, wantLocator: document.EvidenceLocatorRecord}, + {name: "ebook", family: "ebook", sourceKind: "spine", wantKind: document.EvidenceUnitSpine, wantLocator: document.EvidenceLocatorSpine}, + {name: "structured", family: "structured", sourceKind: "record", wantKind: document.EvidenceUnitRecord, wantLocator: document.EvidenceLocatorRecord}, + {name: "text", family: "text", sourceKind: "section", wantKind: document.EvidenceUnitSection, wantLocator: document.EvidenceLocatorSection}, + {name: "source", family: "source", sourceKind: "section", wantKind: document.EvidenceUnitSection, wantLocator: document.EvidenceLocatorSection}, + {name: "mail", family: "mail", sourceKind: "message", wantKind: document.EvidenceUnitMessage, wantLocator: document.EvidenceLocatorMessage}, + } { + t.Run(testCase.name, func(t *testing.T) { + evidence, markdown, err := mistralEvidence(document.SourceDocument{ + Family: testCase.family, UnitKind: testCase.sourceKind, + Units: []document.SourceUnit{{Index: 0, Markdown: "synthetic"}}, + }) + require.NoError(t, err) + assert.Equal(t, testCase.wantKind, evidence.UnitKind) + assert.Equal(t, testCase.wantLocator, evidence.Units[0].Locator.Kind) + if testCase.wantLocator == document.EvidenceLocatorSheet { + assert.Equal(t, "mistral-unit-0", evidence.Units[0].Locator.Name) + } + assert.Equal(t, "synthetic", string(markdown)) + }) + } +} + +func newRenditionTestClient( + t *testing.T, policy Policy, manifest CapabilityManifest, descriptor document.RenditionDescriptor, + secrets SecretResolver, transport http.RoundTripper, +) *RenditionClient { + t.Helper() + drainingTransport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + _, err := io.Copy(io.Discard, request.Body) + require.NoError(t, err) + return transport.RoundTrip(request) + }) + client, err := NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: descriptor, + SecretBinding: "mistral-ocr", Timeout: time.Second, MaxRetries: 1, + MaxRetryDelay: 50 * time.Millisecond, + }, secrets, &http.Client{Transport: drainingTransport}) + require.NoError(t, err) + return client +} + +func newTLSServer(t *testing.T, handler http.Handler) *httptest.Server { + t.Helper() + server := httptest.NewTLSServer(handler) + t.Cleanup(server.Close) + return server +} + +func renditionServerClient( + t *testing.T, policy Policy, manifest CapabilityManifest, + descriptor document.RenditionDescriptor, server *httptest.Server, +) *RenditionClient { + t.Helper() + target, err := url.Parse(server.URL) + require.NoError(t, err) + base := server.Client() + transport := base.Transport + client, err := NewRenditionProvider(Profile{ + Policy: policy, CapabilityManifest: manifest, Descriptor: descriptor, + SecretBinding: "mistral-ocr", Timeout: time.Second, MaxRetries: 1, + MaxRetryDelay: 50 * time.Millisecond, + }, renditionSecrets{"mistral-ocr": "synthetic-key"}, &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + clone := request.Clone(request.Context()) + clone.URL.Scheme, clone.URL.Host = target.Scheme, target.Host + return transport.RoundTrip(clone) + })}) + require.NoError(t, err) + return client +} + +func assertRenditionCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerError, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok, "%T: %v", err, err) + assert.Equal(t, want, providerError.Code()) +} + +type renditionUpload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (upload *renditionUpload) Close() error { return nil } + +func (upload *renditionUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +type renditionTestFixture struct { + metadata document.AuthorizedUploadMetadata + authorization document.RenditionAuthorization + source []byte +} + +func renditionFixture( + t *testing.T, descriptor document.RenditionDescriptor, source []byte, +) renditionTestFixture { + t.Helper() + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: "document.pdf", MediaFamily: "pdf", MediaType: "application/pdf", + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + } + started := time.Now().UTC().Add(-time.Minute) + return renditionTestFixture{metadata: metadata, source: source, authorization: document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), SourceSHA256: metadata.SHA256, + SourceBytes: metadata.ByteLength, CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, MediaFamily: metadata.MediaFamily, + MediaType: metadata.MediaType, InputKind: metadata.InputKind, + MaxProviderMarkdownBytes: 4096, MaxTotalResultBytes: 32768, + AuthorizedAt: started.Format("2006-01-02T15:04:05.000000000Z"), + ExpiresAt: started.Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + }} +} + +func (fixture renditionTestFixture) upload() document.AuthorizedUpload { + return &renditionUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func renditionDescriptor( + t *testing.T, policy Policy, manifest CapabilityManifest, formatIDs ...string, +) document.RenditionDescriptor { + t.Helper() + formats := make([]document.RenditionFormatCapability, 0, len(formatIDs)) + for _, formatID := range formatIDs { + candidate, ok := CandidateFormatByID(formatID) + require.True(t, ok) + formats = append(formats, document.RenditionFormatCapability{ + MediaFamily: candidate.Family, MediaType: candidate.MediaType, + InputKind: document.RenditionInputOriginalFile, + }) + } + fingerprint, err := policy.Fingerprint(manifest) + require.NoError(t, err) + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: "mistral.ocr-v1", ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: fingerprint, TrustBoundary: document.RenditionTrustHostedProvider, + SupportedFormats: formats, ReturnsMarkdown: true, ReturnsStructured: true, + }) + require.NoError(t, err) + require.NotContains(t, descriptor.Fingerprint, " ") + return descriptor +} + +func testMultipagePDF(pages int) []byte { + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + } + kids := make([]string, pages) + for index := range pages { + kids[index] = fmt.Sprintf("%d 0 R", index+3) + } + objects = append(objects, fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", strings.Join(kids, " "), pages)) + for index := range pages { + objects = append(objects, fmt.Sprintf( + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents %d 0 R >>", + pages+index+3, + )) + } + for index := range pages { + objects = append(objects, fmt.Sprintf("<< /Length 0 >>\nstream\n\nendstream %% %d", index)) + } + var output bytes.Buffer + output.WriteString("%PDF-1.4\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, + "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} diff --git a/document/mistral/sniff.go b/document/mistral/sniff.go index 5f8d5f9e..45468173 100644 --- a/document/mistral/sniff.go +++ b/document/mistral/sniff.go @@ -1,1063 +1,27 @@ package mistral import ( - "archive/zip" - "bytes" - "encoding/binary" - "encoding/csv" - "encoding/json/jsontext" - "encoding/xml" - "errors" - "fmt" "io" - "math" - "mime" - "net/mail" - "os" - "path" - "strconv" - "strings" - "unicode/utf8" -) -const ( - maxSniffBytes = int64(8 << 20) - maxTextSniffBytes = int64(50 << 20) - maxZIPEntries = 10_000 - maxZIPCentralDirectory = uint32(16 << 20) - maxZIPExpandedBytes = uint64(500 << 20) - maxZIPSingleExpandedByte = uint64(100 << 20) - maxPDFTailBytes = int64(64 << 10) - maxPDFXRefBytes = int64(4 << 10) - ooxmlContentTypesName = "[Content_Types].xml" + "go.kenn.io/docbank/document/internal/formatdetect" ) -var ( - compoundFileMagic = []byte{0xd0, 0xcf, 0x11, 0xe0, 0xa1, 0xb1, 0x1a, 0xe1} +const ( + maxZIPEntries = 10_000 + ooxmlContentTypesName = "[Content_Types].xml" + compoundNoStream = uint32(0xffffffff) ) -const compoundNoStream = uint32(0xffffffff) - -type compoundDirectoryEntry struct { - name string - entryType byte - left, right, child uint32 -} +var compoundFileMagic = []byte{0xd0, 0xcf, 0x11, 0xe0, 0xa1, 0xb1, 0x1a, 0xe1} -// DetectFormat validates a provider candidate from bounded bytes. Declared -// type is a hint only: container families must prove internal markers, while -// inherently ambiguous text formats also require syntactically safe UTF-8. +// DetectFormat validates a provider candidate from bounded bytes. Core media +// inspection and Mistral preparation share the same fail-closed detector. func DetectFormat(reader io.ReaderAt, size int64, declaredMediaType string) (CandidateFormat, error) { - if reader == nil || size <= 0 { - return CandidateFormat{}, errors.New("document format detection requires nonempty bytes") - } - if size > MaxDocumentBytes { - return CandidateFormat{}, errors.New("document exceeds the format-detection byte limit") - } - mediaType, parameters, err := mime.ParseMediaType(declaredMediaType) - if err != nil || len(parameters) != 0 || mediaType != strings.ToLower(mediaType) { - return CandidateFormat{}, errors.New("document format detection requires a canonical media type") - } - prefix, err := readPrefix(reader, size, maxSniffBytes) - if err != nil { - return CandidateFormat{}, err - } - - var detected CandidateFormat - switch { - case bytes.HasPrefix(prefix, []byte("%PDF-")): - if err = validatePDFStructure(reader, size, prefix); err == nil { - detected, _ = CandidateFormatByID(formatIDPDF) - } - case bytes.HasPrefix(prefix, []byte(`{\rtf`)): - detected, _ = CandidateFormatByID("rtf") - case bytes.HasPrefix(prefix, compoundFileMagic): - detected, err = detectCompoundFormat(reader, size) - case bytes.HasPrefix(prefix, []byte("PK\x03\x04")) || bytes.HasPrefix(prefix, []byte("PK\x05\x06")): - detected, err = detectZIPFormat(reader, size) - default: - if size > maxTextSniffBytes { - return CandidateFormat{}, errors.New("document text exceeds type-detection limit") - } - content, readErr := readPrefix(reader, size, maxTextSniffBytes) - if readErr != nil { - return CandidateFormat{}, readErr - } - detected, err = detectTextFormat(content, mediaType) - } - if err != nil { - return CandidateFormat{}, err - } - if detected.MediaType != mediaType { - return CandidateFormat{}, fmt.Errorf("document bytes are %s, not declared %s", detected.MediaType, mediaType) - } - return detected, nil -} - -func validatePDFStructure(reader io.ReaderAt, size int64, prefix []byte) error { - if len(prefix) < 9 || !validPDFVersion(prefix[5:8]) || !isPDFWhitespace(prefix[8]) { - return errors.New("PDF header is invalid") - } - tailLength := min(size, maxPDFTailBytes) - tailOffset := size - tailLength - tail := make([]byte, tailLength) - read, err := reader.ReadAt(tail, tailOffset) - if err != nil && !errors.Is(err, io.EOF) { - return fmt.Errorf("read PDF trailer: %w", err) - } - if int64(read) != tailLength { - return errors.New("document bytes changed during PDF trailer read") - } - eofIndex := bytes.LastIndex(tail, []byte("%%EOF")) - if eofIndex < 0 || len(trimPDFWhitespace(tail[eofIndex+len("%%EOF"):])) != 0 { - return errors.New("PDF end marker is missing or not final") - } - beforeEOF := tail[:eofIndex] - startXRefIndex := bytes.LastIndex(beforeEOF, []byte("startxref")) - if startXRefIndex < 0 { - return errors.New("PDF startxref is missing") - } - offsetText := trimPDFWhitespace(beforeEOF[startXRefIndex+len("startxref"):]) - digitEnd := 0 - for digitEnd < len(offsetText) && offsetText[digitEnd] >= '0' && offsetText[digitEnd] <= '9' { - digitEnd++ - } - if digitEnd == 0 || len(trimPDFWhitespace(offsetText[digitEnd:])) != 0 { - return errors.New("PDF startxref offset is invalid") - } - xrefOffset, err := strconv.ParseInt(string(offsetText[:digitEnd]), 10, 64) - if err != nil || xrefOffset <= 0 || xrefOffset >= tailOffset+int64(eofIndex) { - return errors.New("PDF startxref offset is outside the document") - } - xrefLength := min(size-xrefOffset, maxPDFXRefBytes) - xref := make([]byte, xrefLength) - read, err = reader.ReadAt(xref, xrefOffset) - if err != nil && !errors.Is(err, io.EOF) { - return fmt.Errorf("read PDF cross-reference data: %w", err) - } - if int64(read) != xrefLength { - return errors.New("document bytes changed during PDF cross-reference read") - } - if validPDFTableXRef(xref, beforeEOF[:startXRefIndex]) || validPDFStreamXRef(xref) { - return nil - } - return errors.New("PDF cross-reference data is invalid") -} - -func validPDFVersion(version []byte) bool { - return len(version) == 3 && version[1] == '.' && - ((version[0] == '1' && version[2] >= '0' && version[2] <= '7') || - (version[0] == '2' && version[2] == '0')) -} - -func validPDFTableXRef(xref, beforeStartXRef []byte) bool { - position := 0 - line, ok := nextPDFLine(xref, &position) - if !ok || !bytes.Equal(trimPDFWhitespace(line), []byte("xref")) { - return false - } - header, ok := nextNonemptyPDFLine(xref, &position) - if !ok { - return false - } - headerFields := bytes.Fields(header) - if len(headerFields) != 2 { - return false - } - first, firstErr := strconv.ParseUint(string(headerFields[0]), 10, 64) - count, countErr := strconv.ParseUint(string(headerFields[1]), 10, 64) - if firstErr != nil || countErr != nil || count == 0 || first > math.MaxUint64-count { - return false - } - validatedRecords := uint64(0) - for validatedRecords < count { - record, ok := nextPDFLine(xref, &position) - if !ok { - break - } - if !validPDFXRefRecord(record) { - if position == len(xref) { - break - } - return false - } - validatedRecords++ - } - if validatedRecords == 0 { - return false - } - - return validPDFTrailer(beforeStartXRef, first, count) -} - -func validPDFStreamXRef(xref []byte) bool { - streamIndex := firstPDFKeyword(xref, "stream") - if streamIndex < 0 { - return false - } - tokens, ok := tokenizePDF(xref[:streamIndex]) - if !ok || len(tokens) < 5 || tokens[2] != "obj" { - return false - } - objectNumber, err := strconv.ParseUint(tokens[0], 10, 64) - if err != nil || objectNumber == 0 { - return false - } - generation, err := strconv.ParseUint(tokens[1], 10, 64) - if err != nil || generation > 65_535 { - return false - } - dictionary, next, ok := parsePDFDictionaryTokens(tokens, 3, 0) - if !ok || next != len(tokens) { - return false - } - _, sizeOK := pdfPositiveInteger(dictionary["Size"]) - return sizeOK && len(dictionary["Type"]) == 1 && dictionary["Type"][0] == "/XRef" && - validPDFRootReference(dictionary["Root"]) && validPDFWidths(dictionary["W"]) && - validPDFStreamLength(dictionary["Length"]) -} - -func validPDFTrailer(data []byte, first, count uint64) bool { - for end := len(data); end > 0; { - trailerIndex := lastPDFKeyword(data[:end], "trailer") - if trailerIndex < 0 { - return false - } - trailer, ok := parsePDFDictionary(data[trailerIndex+len("trailer"):]) - if ok { - if !validPDFRootReference(trailer["Root"]) { - return false - } - size, sizeOK := pdfPositiveInteger(trailer["Size"]) - return sizeOK && first+count <= size - } - end = trailerIndex - } - return false -} - -func nextPDFLine(data []byte, position *int) ([]byte, bool) { - if *position >= len(data) { - return nil, false - } - start := *position - for *position < len(data) && data[*position] != '\n' && data[*position] != '\r' { - *position++ - } - line := data[start:*position] - if *position < len(data) && data[*position] == '\r' { - *position++ - } - if *position < len(data) && data[*position] == '\n' { - *position++ - } - return line, true -} - -func nextNonemptyPDFLine(data []byte, position *int) ([]byte, bool) { - for { - line, ok := nextPDFLine(data, position) - if !ok { - return nil, false - } - if line = trimPDFWhitespace(line); len(line) != 0 { - return line, true - } - } -} - -func validPDFXRefRecord(line []byte) bool { - fields := bytes.Fields(line) - if len(fields) != 3 || len(fields[0]) != 10 || len(fields[1]) != 5 || len(fields[2]) != 1 || - (fields[2][0] != 'n' && fields[2][0] != 'f') { - return false - } - return decimalBytes(fields[0]) && decimalBytes(fields[1]) -} - -func decimalBytes(value []byte) bool { - if len(value) == 0 { - return false - } - for _, char := range value { - if char < '0' || char > '9' { - return false - } - } - return true -} - -func lastPDFKeyword(data []byte, keyword string) int { - for end := len(data); end > 0; { - index := bytes.LastIndex(data[:end], []byte(keyword)) - if index < 0 { - return -1 - } - beforeOK := index == 0 || isPDFTokenBoundary(data[index-1]) - after := index + len(keyword) - afterOK := after == len(data) || isPDFTokenBoundary(data[after]) - if beforeOK && afterOK { - return index - } - end = index - } - return -1 -} - -func firstPDFKeyword(data []byte, keyword string) int { - for start := 0; start < len(data); { - relative := bytes.Index(data[start:], []byte(keyword)) - if relative < 0 { - return -1 - } - index := start + relative - beforeOK := index == 0 || isPDFTokenBoundary(data[index-1]) - after := index + len(keyword) - afterOK := after == len(data) || isPDFTokenBoundary(data[after]) - if beforeOK && afterOK { - return index - } - start = index + 1 - } - return -1 -} - -func isPDFTokenBoundary(char byte) bool { - return isPDFWhitespace(char) || strings.ContainsRune("()<>[]{}/%", rune(char)) -} - -func tokenizePDF(data []byte) ([]string, bool) { - tokens := make([]string, 0, 32) - for position := 0; position < len(data); { - char := data[position] - if isPDFWhitespace(char) { - position++ - continue - } - if char == '%' { - for position < len(data) && data[position] != '\r' && data[position] != '\n' { - position++ - } - continue - } - if position+1 < len(data) && (string(data[position:position+2]) == "<<" || - string(data[position:position+2]) == ">>") { - tokens = append(tokens, string(data[position:position+2])) - position += 2 - continue - } - if char == '(' { - start, depth, escaped := position, 0, false - for ; position < len(data); position++ { - current := data[position] - if escaped { - escaped = false - continue - } - if current == '\\' { - escaped = true - continue - } - if current == '(' { - depth++ - } else if current == ')' { - depth-- - if depth == 0 { - position++ - break - } - } - } - if depth != 0 { - return nil, false - } - tokens = append(tokens, string(data[start:position])) - continue - } - if char == '<' { - start := position - position++ - for position < len(data) && data[position] != '>' { - position++ - } - if position == len(data) { - return nil, false - } - position++ - tokens = append(tokens, string(data[start:position])) - continue - } - if strings.ContainsRune("[]{}", rune(char)) { - tokens = append(tokens, string(char)) - position++ - continue - } - start := position - if char == '/' { - position++ - } - for position < len(data) && !isPDFTokenBoundary(data[position]) { - position++ - } - if position == start || (position == start+1 && char == '/') { - return nil, false - } - tokens = append(tokens, string(data[start:position])) - } - return tokens, true -} - -func parsePDFDictionary(data []byte) (map[string][]string, bool) { - tokens, ok := tokenizePDF(data) - if !ok { - return nil, false - } - dictionary, next, ok := parsePDFDictionaryTokens(tokens, 0, 0) - return dictionary, ok && next == len(tokens) -} - -func parsePDFDictionaryTokens( - tokens []string, - position int, - depth int, -) (map[string][]string, int, bool) { - if depth > 32 || position >= len(tokens) || tokens[position] != "<<" { - return nil, position, false - } - position++ - dictionary := make(map[string][]string) - for position < len(tokens) && tokens[position] != ">>" { - key := tokens[position] - if len(key) < 2 || key[0] != '/' { - return nil, position, false - } - key = key[1:] - if _, exists := dictionary[key]; exists { - return nil, position, false - } - position++ - valueStart := position - var ok bool - position, ok = skipPDFObject(tokens, position, depth+1) - if !ok { - return nil, position, false - } - dictionary[key] = tokens[valueStart:position] - } - if position >= len(tokens) || tokens[position] != ">>" { - return nil, position, false - } - return dictionary, position + 1, true -} - -func skipPDFObject(tokens []string, position, depth int) (int, bool) { - if depth > 32 || position >= len(tokens) { - return position, false - } - switch tokens[position] { - case "<<": - _, next, ok := parsePDFDictionaryTokens(tokens, position, depth) - return next, ok - case "[": - position++ - for position < len(tokens) && tokens[position] != "]" { - var ok bool - position, ok = skipPDFObject(tokens, position, depth+1) - if !ok { - return position, false - } - } - return position + 1, position < len(tokens) - case ">>", "]": - return position, false - default: - if position+2 < len(tokens) && decimalString(tokens[position]) && - decimalString(tokens[position+1]) && tokens[position+2] == "R" { - return position + 3, true - } - return position + 1, true - } -} - -func pdfPositiveInteger(value []string) (uint64, bool) { - if len(value) != 1 || !decimalString(value[0]) { - return 0, false - } - parsed, err := strconv.ParseUint(value[0], 10, 64) - return parsed, err == nil && parsed > 0 -} - -func validPDFRootReference(value []string) bool { - if len(value) != 3 || value[2] != "R" { - return false - } - objectNumber, objectErr := strconv.ParseUint(value[0], 10, 64) - generation, generationErr := strconv.ParseUint(value[1], 10, 64) - return objectErr == nil && objectNumber > 0 && generationErr == nil && generation <= 65_535 -} - -func validPDFWidths(value []string) bool { - if len(value) != 5 || value[0] != "[" || value[4] != "]" { - return false - } - total := uint64(0) - for _, width := range value[1:4] { - parsed, err := strconv.ParseUint(width, 10, 64) - if err != nil || parsed > 8 { - return false - } - total += parsed - } - return total > 0 -} - -func validPDFStreamLength(value []string) bool { - if _, ok := pdfPositiveInteger(value); ok { - return true - } - return validPDFRootReference(value) -} - -func decimalString(value string) bool { - return decimalBytes([]byte(value)) -} - -func trimPDFWhitespace(value []byte) []byte { - for len(value) > 0 && isPDFWhitespace(value[0]) { - value = value[1:] - } - for len(value) > 0 && isPDFWhitespace(value[len(value)-1]) { - value = value[:len(value)-1] - } - return value -} - -func isPDFWhitespace(char byte) bool { - return char == 0 || char == '\t' || char == '\n' || char == '\f' || char == '\r' || char == ' ' -} - -func detectCompoundFormat(reader io.ReaderAt, size int64) (CandidateFormat, error) { - names, err := compoundDirectoryNames(reader, size) - if err != nil { - return CandidateFormat{}, err - } - ids := make([]string, 0, 2) - if names["WordDocument"] { - ids = append(ids, "doc") - } - if names["Workbook"] || names["Book"] { - ids = append(ids, "xls") - } - if names["PowerPoint Document"] { - ids = append(ids, "ppt") - } - if names["__properties_version1.0"] { - ids = append(ids, "msg") - } - if len(ids) != 1 { - return CandidateFormat{}, errors.New("compound document has missing or ambiguous family markers") - } - format, _ := CandidateFormatByID(ids[0]) - return format, nil + return formatdetect.DetectFormat(reader, size, declaredMediaType) } func compoundDirectoryNames(reader io.ReaderAt, size int64) (map[string]bool, error) { - const ( - freeSector = uint32(0xffffffff) - endOfChain = uint32(0xfffffffe) - fatSector = uint32(0xfffffffd) - difatSector = uint32(0xfffffffc) - maxDIFATSectors = 1_024 - maxDirectoryBytes = int64(8 << 20) - ) - if size < 512 { - return nil, errors.New("compound document header is truncated") - } - header := make([]byte, 512) - if _, err := reader.ReadAt(header, 0); err != nil { - return nil, fmt.Errorf("read compound document header: %w", err) - } - if !bytes.Equal(header[:8], compoundFileMagic) || binary.LittleEndian.Uint16(header[28:30]) != 0xfffe { - return nil, errors.New("compound document header is invalid") - } - sectorShift := binary.LittleEndian.Uint16(header[30:32]) - majorVersion := binary.LittleEndian.Uint16(header[26:28]) - if (majorVersion != 3 || sectorShift != 9) && (majorVersion != 4 || sectorShift != 12) { - return nil, errors.New("compound document sector size is unsupported") - } - sectorSize := int64(1 << sectorShift) - sectorCount := size/sectorSize - 1 - if sectorCount <= 0 || size%sectorSize != 0 { - return nil, errors.New("compound document size is invalid") - } - fatEntriesPerSector := sectorSize / 4 - maxFATSectors := int((sectorCount + fatEntriesPerSector - 1) / fatEntriesPerSector) - numFAT := int(binary.LittleEndian.Uint32(header[44:48])) - firstDirectory := binary.LittleEndian.Uint32(header[48:52]) - firstDIFAT := binary.LittleEndian.Uint32(header[68:72]) - numDIFAT := int(binary.LittleEndian.Uint32(header[72:76])) - if numFAT <= 0 || numFAT > maxFATSectors || numDIFAT > maxDIFATSectors { - return nil, errors.New("compound document allocation table exceeds limits") - } - fatSectors := make([]uint32, 0, numFAT) - for offset := 76; offset < 512 && len(fatSectors) < numFAT; offset += 4 { - sector := binary.LittleEndian.Uint32(header[offset : offset+4]) - if sector != freeSector { - fatSectors = append(fatSectors, sector) - } - } - seenDIFAT := map[uint32]bool{} - for i, sector := 0, firstDIFAT; i < numDIFAT; i++ { - if int64(sector) >= sectorCount || seenDIFAT[sector] { - return nil, errors.New("compound document DIFAT chain is invalid") - } - seenDIFAT[sector] = true - data, readErr := readCompoundSector(reader, sector, sectorSize, size) - if readErr != nil { - return nil, readErr - } - for offset := 0; offset < len(data)-4 && len(fatSectors) < numFAT; offset += 4 { - fatID := binary.LittleEndian.Uint32(data[offset : offset+4]) - if fatID != freeSector { - fatSectors = append(fatSectors, fatID) - } - } - sector = binary.LittleEndian.Uint32(data[len(data)-4:]) - if i == numDIFAT-1 && sector != endOfChain { - return nil, errors.New("compound document DIFAT chain does not terminate") - } - } - if len(fatSectors) != numFAT { - return nil, errors.New("compound document FAT sector count is invalid") - } - fat := make([]uint32, 0, int(sectorCount)) - seenFAT := map[uint32]bool{} - for _, sector := range fatSectors { - if int64(sector) >= sectorCount || seenFAT[sector] { - return nil, errors.New("compound document FAT sector list is invalid") - } - seenFAT[sector] = true - data, readErr := readCompoundSector(reader, sector, sectorSize, size) - if readErr != nil { - return nil, readErr - } - for offset := 0; offset < len(data); offset += 4 { - fat = append(fat, binary.LittleEndian.Uint32(data[offset:offset+4])) - } - } - if len(fat) < int(sectorCount) { - return nil, errors.New("compound document FAT is truncated") - } - for _, sector := range fatSectors { - if fat[sector] != fatSector { - return nil, errors.New("compound document FAT sector is not self-marked") - } - } - for sector := range seenDIFAT { - if fat[sector] != difatSector { - return nil, errors.New("compound document DIFAT sector is not self-marked") - } - } - - var directory bytes.Buffer - seenDirectory := map[uint32]bool{} - for sector := firstDirectory; sector != endOfChain; { - if int64(sector) >= sectorCount || seenDirectory[sector] || int64(directory.Len())+sectorSize > maxDirectoryBytes { - return nil, errors.New("compound document directory chain exceeds limits") - } - seenDirectory[sector] = true - data, readErr := readCompoundSector(reader, sector, sectorSize, size) - if readErr != nil { - return nil, readErr - } - _, _ = directory.Write(data) - next := fat[sector] - if next == freeSector || next == fatSector || next == difatSector { - return nil, errors.New("compound document directory chain is invalid") - } - sector = next - } - entries := make([]compoundDirectoryEntry, 0, directory.Len()/128) - rootIndex := -1 - data := directory.Bytes() - for offset := 0; offset+128 <= len(data); offset += 128 { - entry := data[offset : offset+128] - entryType := entry[66] - if entryType == 0 { - entries = append(entries, compoundDirectoryEntry{}) - continue - } - if entryType != 1 && entryType != 2 && entryType != 5 { - return nil, errors.New("compound document directory entry type is invalid") - } - nameLength := int(binary.LittleEndian.Uint16(entry[64:66])) - if nameLength < 2 || nameLength > 64 || nameLength%2 != 0 { - return nil, errors.New("compound document directory name is invalid") - } - name, decodeErr := decodeUTF16LE(entry[:nameLength-2]) - if decodeErr != nil || name == "" { - return nil, errors.New("compound document directory name is invalid") - } - entries = append(entries, compoundDirectoryEntry{ - name: name, entryType: entryType, - left: binary.LittleEndian.Uint32(entry[68:72]), right: binary.LittleEndian.Uint32(entry[72:76]), - child: binary.LittleEndian.Uint32(entry[76:80]), - }) - if entryType == 5 { - if rootIndex != -1 { - return nil, errors.New("compound document has multiple root entries") - } - rootIndex = len(entries) - 1 - } - } - if rootIndex < 0 { - return nil, errors.New("compound document has no root entry") - } - names := map[string]bool{} - seen := map[uint32]bool{} - stack := []uint32{entries[rootIndex].child} - for len(stack) > 0 { - index := stack[len(stack)-1] - stack = stack[:len(stack)-1] - if index == compoundNoStream { - continue - } - if uint64(index) >= uint64(len(entries)) || seen[index] || entries[index].entryType == 0 || entries[index].entryType == 5 { - return nil, errors.New("compound document root directory tree is invalid") - } - seen[index] = true - entry := entries[index] - if entry.entryType == 2 { - names[entry.name] = true - } - stack = append(stack, entry.left, entry.right) - } - return names, nil -} - -func readCompoundSector(reader io.ReaderAt, sector uint32, sectorSize, size int64) ([]byte, error) { - offset := (int64(sector) + 1) * sectorSize - if offset < sectorSize || offset > size-sectorSize { - return nil, errors.New("compound document sector is out of range") - } - data := make([]byte, sectorSize) - if _, err := reader.ReadAt(data, offset); err != nil { - return nil, fmt.Errorf("read compound document sector: %w", err) - } - return data, nil -} - -func decodeUTF16LE(data []byte) (string, error) { - if len(data)%2 != 0 { - return "", errors.New("odd UTF-16 length") - } - runes := make([]rune, 0, len(data)/2) - for i := 0; i < len(data); i += 2 { - value := binary.LittleEndian.Uint16(data[i : i+2]) - if value == 0 || value >= 0xd800 && value <= 0xdfff { - return "", errors.New("unsupported UTF-16 directory name") - } - runes = append(runes, rune(value)) - } - return string(runes), nil -} - -func detectZIPFormat(reader io.ReaderAt, size int64) (CandidateFormat, error) { - if err := validateZIPEndRecord(reader, size); err != nil { - return CandidateFormat{}, err - } - archive, err := zip.NewReader(reader, size) - if err != nil { - return CandidateFormat{}, fmt.Errorf("open document ZIP container: %w", err) - } - if len(archive.File) > maxZIPEntries { - return CandidateFormat{}, errors.New("document ZIP container has too many entries") - } - names := make(map[string]bool, len(archive.File)) - var expanded uint64 - var mimeValue string - var contentTypes []byte - for _, entry := range archive.File { - if err := validateZIPName(entry.Name); err != nil { - return CandidateFormat{}, err - } - if entry.Mode()&os.ModeSymlink != 0 { - return CandidateFormat{}, errors.New("document ZIP container contains a symlink") - } - if entry.Flags&1 != 0 || (entry.Method != zip.Store && entry.Method != zip.Deflate) { - return CandidateFormat{}, errors.New("document ZIP container uses unsupported encryption or compression") - } - if entry.UncompressedSize64 > maxZIPSingleExpandedByte || expanded > maxZIPExpandedBytes-entry.UncompressedSize64 { - return CandidateFormat{}, errors.New("document ZIP container exceeds expanded-byte limits") - } - expanded += entry.UncompressedSize64 - if names[entry.Name] { - return CandidateFormat{}, errors.New("document ZIP container has duplicate entry names") - } - names[entry.Name] = true - if err := verifyZIPEntry(entry); err != nil { - return CandidateFormat{}, err - } - if entry.Name == "mimetype" { - value, readErr := readZIPEntry(entry, 256) - if readErr != nil { - return CandidateFormat{}, readErr - } - mimeValue = string(value) - } - if entry.Name == ooxmlContentTypesName { - value, readErr := readZIPEntry(entry, 2<<20) - if readErr != nil { - return CandidateFormat{}, readErr - } - contentTypes = value - } - } - - var id string - switch { - case names["word/document.xml"] && hasOOXMLMainType(contentTypes, "/word/document.xml", - "application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"): - id = "docx" - case names["ppt/presentation.xml"] && hasOOXMLMainType(contentTypes, "/ppt/presentation.xml", - "application/vnd.openxmlformats-officedocument.presentationml.presentation.main+xml"): - id = "pptx" - case names["xl/workbook.xml"] && hasOOXMLMainType(contentTypes, "/xl/workbook.xml", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"): - id = "xlsx" - case mimeValue == "application/vnd.oasis.opendocument.text" && names["META-INF/manifest.xml"]: - id = "odt" - case mimeValue == "application/vnd.oasis.opendocument.spreadsheet" && names["META-INF/manifest.xml"]: - id = "ods" - case mimeValue == "application/epub+zip" && names["META-INF/container.xml"]: - id = "epub" - case hasNumbersMarker(names): - id = "numbers" - default: - return CandidateFormat{}, errors.New("zip container is not a supported document format") - } - format, _ := CandidateFormatByID(id) - return format, nil -} - -func hasOOXMLMainType(content []byte, partName, contentType string) bool { - if len(content) == 0 || !validXMLDocument(content) { - return false - } - var document struct { - XMLName xml.Name `xml:"Types"` - Overrides []struct { - PartName string `xml:"PartName,attr"` - ContentType string `xml:"ContentType,attr"` - } `xml:"Override"` - } - if err := xml.Unmarshal(content, &document); err != nil || - document.XMLName.Space != "http://schemas.openxmlformats.org/package/2006/content-types" { - return false - } - found := false - for _, override := range document.Overrides { - if override.PartName != partName { - continue - } - if found || override.ContentType != contentType { - return false - } - found = true - } - return found -} - -func validateZIPEndRecord(reader io.ReaderAt, size int64) error { - const maxTail = int64(65_557) - tailSize := min(size, maxTail) - tail := make([]byte, tailSize) - if _, err := reader.ReadAt(tail, size-tailSize); err != nil && !errors.Is(err, io.EOF) { - return fmt.Errorf("read document ZIP end record: %w", err) - } - signature := []byte{'P', 'K', 0x05, 0x06} - offset := bytes.LastIndex(tail, signature) - if offset < 0 || len(tail)-offset < 22 { - return errors.New("document ZIP container has no bounded end record") - } - record := tail[offset:] - entries := binary.LittleEndian.Uint16(record[10:12]) - entriesOnDisk := binary.LittleEndian.Uint16(record[8:10]) - centralSize := binary.LittleEndian.Uint32(record[12:16]) - centralOffset := binary.LittleEndian.Uint32(record[16:20]) - commentSize := int(binary.LittleEndian.Uint16(record[20:22])) - if binary.LittleEndian.Uint16(record[4:6]) != 0 || binary.LittleEndian.Uint16(record[6:8]) != 0 || entriesOnDisk != entries { - return errors.New("multi-disk document ZIP containers are unsupported") - } - if entries == 0xffff || centralSize == 0xffffffff || centralOffset == 0xffffffff || - int(entries) > maxZIPEntries || centralSize > maxZIPCentralDirectory { - return errors.New("document ZIP central directory exceeds limits") - } - if int64(centralOffset)+int64(centralSize) > size { - return errors.New("document ZIP central directory is out of range") - } - if len(record) != 22+commentSize { - return errors.New("document ZIP end record has invalid comment length") - } - return nil -} - -func validateZIPName(name string) error { - if name == "" || strings.ContainsRune(name, 0) || strings.ContainsAny(name, "\\:") || strings.HasPrefix(name, "/") { - return errors.New("document ZIP container has an unsafe entry name") - } - clean := path.Clean(name) - if clean == ".." || strings.HasPrefix(clean, "../") || clean != strings.TrimSuffix(name, "/") { - return errors.New("document ZIP container has a traversing entry name") - } - return nil -} - -func readZIPEntry(entry *zip.File, limit int64) ([]byte, error) { - if limit < 0 || entry.UncompressedSize64 > maxZIPSingleExpandedByte || int64(entry.UncompressedSize64) > limit { - return nil, errors.New("document ZIP marker entry exceeds limit") - } - reader, err := entry.Open() - if err != nil { - return nil, fmt.Errorf("open document ZIP marker: %w", err) - } - defer func() { _ = reader.Close() }() - value, err := io.ReadAll(io.LimitReader(reader, limit+1)) - if err != nil { - return nil, fmt.Errorf("read document ZIP marker: %w", err) - } - if int64(len(value)) > limit { - return nil, errors.New("document ZIP marker entry exceeds limit") - } - return value, nil -} - -func verifyZIPEntry(entry *zip.File) error { - if entry.UncompressedSize64 > maxZIPSingleExpandedByte { - return errors.New("document ZIP entry exceeds verification limit") - } - reader, err := entry.Open() - if err != nil { - return fmt.Errorf("open document ZIP entry: %w", err) - } - expectedSize := int64(entry.UncompressedSize64) - written, readErr := io.Copy(io.Discard, io.LimitReader(reader, expectedSize+1)) - closeErr := reader.Close() - if readErr != nil || closeErr != nil || written != expectedSize { - return errors.New("document ZIP entry failed bounded verification") - } - return nil -} - -func hasNumbersMarker(names map[string]bool) bool { - for name := range names { - if strings.HasPrefix(name, "Index/Tables/") && strings.HasSuffix(name, ".iwa") { - return true - } - } - return false -} - -func detectTextFormat(content []byte, mediaType string) (CandidateFormat, error) { - if !utf8.Valid(content) || bytes.IndexByte(content, 0) >= 0 { - return CandidateFormat{}, errors.New("document text is not safe UTF-8") - } - candidate, ok := candidateByMediaType(mediaType) - if !ok || !isTextCandidate(candidate) { - return CandidateFormat{}, errors.New("document bytes have no supported signature") - } - trimmed := bytes.TrimSpace(content) - switch candidate.ID { - case "json": - if len(trimmed) == 0 || !jsontext.Value(trimmed).IsValid() { - return CandidateFormat{}, errors.New("declared JSON document is invalid") - } - case "jsonl": - for line := range bytes.SplitSeq(trimmed, []byte{'\n'}) { - if len(bytes.TrimSpace(line)) > 0 && !jsontext.Value(bytes.TrimSpace(line)).IsValid() { - return CandidateFormat{}, errors.New("declared JSONL document is invalid") - } - } - case "xml": - if !validXMLDocument(content) { - return CandidateFormat{}, errors.New("declared XML document is invalid") - } - case "csv": - csvReader := csv.NewReader(bytes.NewReader(content)) - csvReader.FieldsPerRecord = -1 - csvReader.ReuseRecord = true - records := 0 - for { - _, readErr := csvReader.Read() - if errors.Is(readErr, io.EOF) { - break - } - if readErr != nil { - return CandidateFormat{}, errors.New("declared CSV document is invalid") - } - records++ - } - if records == 0 { - return CandidateFormat{}, errors.New("declared CSV document is invalid") - } - case "latex": - if !bytes.Contains(content, []byte(`\documentclass`)) && !bytes.Contains(content, []byte(`\begin{document}`)) { - return CandidateFormat{}, errors.New("declared LaTeX document has no document marker") - } - case "eml": - message, err := mail.ReadMessage(bytes.NewReader(content)) - if err != nil || message.Header.Get("From") == "" || message.Header.Get("Date") == "" { - return CandidateFormat{}, errors.New("declared EML document lacks required message headers") - } - case "yaml": - if len(trimmed) == 0 || (!bytes.HasPrefix(trimmed, []byte("---")) && !bytes.Contains(trimmed, []byte(": "))) { - return CandidateFormat{}, errors.New("declared YAML document has no structural marker") - } - } - return candidate, nil -} - -func validXMLDocument(content []byte) bool { - decoder := xml.NewDecoder(bytes.NewReader(content)) - depth := 0 - roots := 0 - for { - token, err := decoder.Token() - if errors.Is(err, io.EOF) { - return roots == 1 && depth == 0 - } - if err != nil { - return false - } - switch value := token.(type) { - case xml.StartElement: - if depth == 0 { - roots++ - if roots > 1 { - return false - } - } - depth++ - case xml.EndElement: - depth-- - if depth < 0 { - return false - } - case xml.CharData: - if depth == 0 && len(bytes.TrimSpace(value)) != 0 { - return false - } - } - } -} - -func isTextCandidate(candidate CandidateFormat) bool { - switch candidate.Family { - case "text", "structured", "source", "mail", "spreadsheet": - return candidate.ID != "msg" && candidate.ID != "xls" && candidate.ID != "xlsx" && candidate.ID != "ods" && candidate.ID != "numbers" - default: - return false - } + return formatdetect.CompoundDirectoryNames(reader, size) } func candidateByMediaType(mediaType string) (CandidateFormat, bool) { @@ -1068,16 +32,3 @@ func candidateByMediaType(mediaType string) (CandidateFormat, bool) { } return CandidateFormat{}, false } - -func readPrefix(reader io.ReaderAt, size, limit int64) ([]byte, error) { - length := min(size, limit) - buffer := make([]byte, length) - read, err := reader.ReadAt(buffer, 0) - if err != nil && !errors.Is(err, io.EOF) { - return nil, fmt.Errorf("read document signature: %w", err) - } - if int64(read) != length { - return nil, errors.New("document bytes changed during signature read") - } - return buffer, nil -} diff --git a/document/mistral/sniff_test.go b/document/mistral/sniff_test.go index 00ccd727..0403d107 100644 --- a/document/mistral/sniff_test.go +++ b/document/mistral/sniff_test.go @@ -76,6 +76,20 @@ func TestDetectFormatRejectsMismatchUnsafeZIPAndAmbiguousCompound(t *testing.T) _, err = DetectFormat(bytes.NewReader(prefixStreamType), int64(len(prefixStreamType)), "application/pdf") require.ErrorContains(err, "cross-reference data") + for name, malformedStream := range map[string][]byte{ + "short stream": testPDFXRefStreamWith(func(entries []byte, _ []int) []byte { + return entries[:7] + }), + "forged root entry": testPDFXRefStreamWith(func(entries []byte, offsets []int) []byte { + putTestXRefEntry(entries, 1, 1, uint32(offsets[2]), 0) + return entries + }), + } { + _, detectErr := DetectFormat( + bytes.NewReader(malformedStream), int64(len(malformedStream)), "application/pdf") + require.ErrorContains(detectErr, "cross-reference data", name) + } + polyglotPDF := append(testPDF("polyglot"), []byte("PK\x03\x04synthetic")...) _, err = DetectFormat(bytes.NewReader(polyglotPDF), int64(len(polyglotPDF)), "application/pdf") require.ErrorContains(err, "not final") diff --git a/document/mistral/test_helpers_test.go b/document/mistral/test_helpers_test.go index 8c3ede1c..0d147518 100644 --- a/document/mistral/test_helpers_test.go +++ b/document/mistral/test_helpers_test.go @@ -3,6 +3,7 @@ package mistral import ( "bytes" "crypto/sha256" + "encoding/binary" "encoding/hex" "fmt" "io" @@ -113,15 +114,45 @@ func testPDF(label string) []byte { } func testPDFXRefStream() []byte { + return testPDFXRefStreamWith(nil) +} + +func testPDFXRefStreamWith(mutate func([]byte, []int) []byte) []byte { var output bytes.Buffer output.WriteString("%PDF-1.5\n") + offsets := make([]int, 3) + for index, object := range []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + "<< /Type /Page /Parent 2 0 R >>", + } { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } xref := output.Len() - _, _ = fmt.Fprintf(&output, "1 0 obj\n<< /Type /XRef /Size 2 /Root 2 0 R /W [1 2 1] /Length 4 >>\nstream\n") - output.Write([]byte{0, 0, 0, 0}) + entries := make([]byte, 5*7) + putTestXRefEntry(entries, 0, 0, 0, 65_535) + for index, offset := range offsets { + putTestXRefEntry(entries, index+1, 1, uint32(offset), 0) // #nosec G115 -- bounded fixture. + } + putTestXRefEntry(entries, 4, 1, uint32(xref), 0) // #nosec G115 -- bounded fixture. + if mutate != nil { + entries = mutate(entries, offsets) + } + _, _ = fmt.Fprintf(&output, + "4 0 obj\n<< /Type /XRef /Size 5 /Root 1 0 R /W [1 4 2] /Length %d >>\nstream\n", len(entries)) + output.Write(entries) _, _ = fmt.Fprintf(&output, "\nendstream\nendobj\nstartxref\n%d\n%%%%EOF\n", xref) return output.Bytes() } +func putTestXRefEntry(entries []byte, index int, kind byte, offset uint32, generation uint16) { + entry := entries[index*7 : (index+1)*7] + entry[0] = kind + binary.BigEndian.PutUint32(entry[1:5], offset) + binary.BigEndian.PutUint16(entry[5:7], generation) +} + type errorReader struct { err error reads *int diff --git a/document/normalize.go b/document/normalize.go index 0fb1ec31..84be5fc4 100644 --- a/document/normalize.go +++ b/document/normalize.go @@ -17,6 +17,7 @@ import ( "github.com/yuin/goldmark" goldmarkast "github.com/yuin/goldmark/ast" "github.com/yuin/goldmark/extension" + extensionast "github.com/yuin/goldmark/extension/ast" "github.com/yuin/goldmark/renderer" goldmarkhtml "github.com/yuin/goldmark/renderer/html" "github.com/yuin/goldmark/util" @@ -294,12 +295,6 @@ func joinDocumentUnitEvidence(header, body, footer string) (string, int) { return strings.Join(parts, "\n\n"), bodyOffset } -// sanitizeMarkdown converts untrusted provider Markdown into inert canonical -// Markdown-like text. It preserves NormalizeDocument's frozen behavior. -func sanitizeMarkdown(markdown string, maxLinkChars, maxSourceBytes int) (string, []canonicalHeadingMark, bool, error) { - return sanitizeMarkdownWithActiveHTML(markdown, maxLinkChars, maxSourceBytes, false) -} - // sanitizeRenditionMarkdown applies the normalizer's frozen text rules plus // removes the body of active HTML elements from a durable rendition. func sanitizeRenditionMarkdown(markdown string, maxLinkChars, maxSourceBytes, maxRunes int) (string, bool, bool, error) { @@ -312,13 +307,15 @@ func sanitizeRenditionMarkdown(markdown string, maxLinkChars, maxSourceBytes, ma markdown, sourceTruncated := truncateUTF8Bytes(markdown, maxSourceBytes) var rendered bytes.Buffer listTightness := make(map[int]bool) + taskCheckboxes := make(map[int]bool) rawSpans := make([]renditionRawSpan, 0) parser := goldmark.New( goldmark.WithExtensions(extension.GFM), goldmark.WithRendererOptions( goldmarkhtml.WithUnsafe(), renderer.WithNodeRenderers(util.Prioritized(&renditionHTMLRenderer{ - output: &rendered, tightness: listTightness, rawSpans: &rawSpans, + output: &rendered, tightness: listTightness, taskCheckboxes: taskCheckboxes, + rawSpans: &rawSpans, }, 0)), ), ) @@ -326,7 +323,9 @@ func sanitizeRenditionMarkdown(markdown string, maxLinkChars, maxSourceBytes, ma if err := parser.Convert(source, &rendered); err != nil { return "", false, false, fmt.Errorf("parse provider Markdown: %w", err) } - writer := renditionHTMLWriter{maxLinkChars: maxLinkChars, listTightness: listTightness} + writer := renditionHTMLWriter{ + maxLinkChars: maxLinkChars, listTightness: listTightness, taskCheckboxes: taskCheckboxes, + } if err := writer.consumeFragments(rendered.Bytes(), rawSpans); err != nil { return "", false, false, err } @@ -340,9 +339,10 @@ func sanitizeRenditionMarkdown(markdown string, maxLinkChars, maxSourceBytes, ma // fragment gets an independent tokenizer so raw-text state cannot escape the // source AST node that introduced it. type renditionHTMLRenderer struct { - output *bytes.Buffer - tightness map[int]bool - rawSpans *[]renditionRawSpan + output *bytes.Buffer + tightness map[int]bool + taskCheckboxes map[int]bool + rawSpans *[]renditionRawSpan } type renditionRawSpan struct { @@ -357,10 +357,33 @@ type renditionRawSpan struct { func (r *renditionHTMLRenderer) RegisterFuncs(registerer renderer.NodeRendererFuncRegisterer) { registerer.Register(goldmarkast.KindList, r.renderList) + registerer.Register(extensionast.KindTaskCheckBox, r.renderTaskCheckBox) registerer.Register(goldmarkast.KindHTMLBlock, r.renderHTMLBlock) registerer.Register(goldmarkast.KindRawHTML, r.renderRawHTML) } +func (r *renditionHTMLRenderer) renderTaskCheckBox( + writer util.BufWriter, + _ []byte, + node goldmarkast.Node, + entering bool, +) (goldmarkast.WalkStatus, error) { + if !entering { + return goldmarkast.WalkContinue, nil + } + checkbox, ok := node.(*extensionast.TaskCheckBox) + if !ok { + return goldmarkast.WalkStop, fmt.Errorf("render task checkbox node: unexpected %T", node) + } + r.taskCheckboxes[r.offset(writer)] = checkbox.IsChecked + if checkbox.IsChecked { + _, _ = writer.WriteString(` `) + } else { + _, _ = writer.WriteString(` `) + } + return goldmarkast.WalkContinue, nil +} + func (r *renditionHTMLRenderer) offset(writer util.BufWriter) int { return r.output.Len() + writer.Buffered() } @@ -451,10 +474,10 @@ func (r *renditionHTMLRenderer) renderRawHTML( return goldmarkast.WalkSkipChildren, nil } -func sanitizeMarkdownWithActiveHTML( - markdown string, - maxLinkChars, maxSourceBytes int, - dropActiveHTML bool, +// sanitizeMarkdown converts untrusted provider Markdown into inert canonical +// Markdown-like text. It preserves NormalizeDocument's frozen behavior. +func sanitizeMarkdown( + markdown string, maxLinkChars, maxSourceBytes int, ) (string, []canonicalHeadingMark, bool, error) { if markdown == "" { return "", nil, false, nil @@ -471,9 +494,7 @@ func sanitizeMarkdownWithActiveHTML( if err := parser.Convert([]byte(markdown), &rendered); err != nil { return "", nil, false, fmt.Errorf("parse provider Markdown: %w", err) } - writer := canonicalHTMLWriter{ - maxLinkChars: maxLinkChars, dropActiveHTML: dropActiveHTML, escapeMarkdown: dropActiveHTML, - } + writer := canonicalHTMLWriter{maxLinkChars: maxLinkChars} if err := writer.consume(bytes.NewReader(rendered.Bytes())); err != nil { return "", nil, false, err } @@ -497,6 +518,7 @@ const ( renditionText renditionInlineKind = iota renditionInlineCode renditionLinkInline + renditionTaskMarkerInline ) type renditionInline struct { @@ -543,6 +565,7 @@ type renditionHTMLWriter struct { rawFragment bool blocks []renditionBlock listTightness map[int]bool + taskCheckboxes map[int]bool renderedOffset int currentKind renditionBlockKind currentLevel int @@ -889,7 +912,17 @@ func (w *renditionHTMLWriter) startTag(token html.Token, tokenOffset int, suppre } return } - if tag == "input" && !w.rawFragment && !suppressText { + if tag == "input" && !w.rawFragment { + if checked, generated := w.taskCheckboxes[tokenOffset]; generated { + w.flushPendingSpace() + marker := "[ ]" + if checked { + marker = "[x]" + } + w.appendInline(renditionInline{kind: renditionTaskMarkerInline, text: marker}) + w.pendingSpace = true + return + } if marker, ok := parserGeneratedCheckboxMarker(token); ok { w.writeText(marker) return @@ -1290,6 +1323,8 @@ func (w *renditionHTMLWriter) appendFlattenedInlines(inlines []renditionInline) switch inline.kind { case renditionText, renditionInlineCode: w.appendText(inline.text) + case renditionTaskMarkerInline: + w.appendInline(inline) case renditionLinkInline: w.appendFlattenedInlines(inline.children) } @@ -2018,6 +2053,12 @@ func appendRenditionInlinesWithFallback( } output.WriteString(value) fallback.mark(output) + case renditionTaskMarkerInline: + if available >= 0 && utf8.RuneCountInString(inline.text) > remaining { + return true + } + output.WriteString(inline.text) + fallback.mark(output) case renditionLinkInline: overhead := utf8.RuneCountInString(inline.destination) + 4 if available >= 0 && overhead > remaining { @@ -2072,6 +2113,12 @@ func appendRenditionPlainLabel( } output.WriteString(value) fallback.markEscapedText(textStart, inline.text) + case renditionTaskMarkerInline: + if available >= 0 && utf8.RuneCountInString(inline.text) > remaining { + return true + } + output.WriteString(inline.text) + fallback.mark(output) case renditionLinkInline: if appendRenditionPlainLabel(output, inline.children, remaining, fallback) { return true @@ -2167,32 +2214,15 @@ func serializeRenditionTable(rows [][][]renditionInline) string { } type canonicalHTMLWriter struct { - output strings.Builder - maxLinkChars int - inPre bool - skipTag string - skipDepth int - cellIndex int - links []string - renditionLinks []*renditionLink - preFenceOpen bool - pendingSpace bool - dropActiveHTML bool - escapeMarkdown bool - preLanguage string - preContent strings.Builder - inlineCode bool - inlineContent strings.Builder - inTable bool - inTableCell bool - tableRows [][]string - tableRow []string - tableCell strings.Builder -} - -type renditionLink struct { - destination string - text strings.Builder + output strings.Builder + maxLinkChars int + inPre bool + skipTag string + skipDepth int + cellIndex int + links []string + preFenceOpen bool + pendingSpace bool } func (w *canonicalHTMLWriter) consume(reader io.Reader) error { @@ -2233,13 +2263,6 @@ func (w *canonicalHTMLWriter) startTag(token html.Token, selfClosing bool) { w.skipDepth = 1 return } - if w.dropActiveHTML && isActiveHTML(tag) { - if !selfClosing && !isHTMLVoidElement(tag) { - w.skipTag = tag - w.skipDepth = 1 - } - return - } if tag == "svg" { if !selfClosing { w.skipTag = tag @@ -2249,12 +2272,6 @@ func (w *canonicalHTMLWriter) startTag(token html.Token, selfClosing bool) { } switch tag { case "table": - if w.escapeMarkdown { - w.block() - w.inTable = true - w.tableRows = nil - return - } w.block() case "h1", "h2", "h3", "h4", "h5", "h6": w.block() @@ -2267,31 +2284,15 @@ func (w *canonicalHTMLWriter) startTag(token html.Token, selfClosing bool) { case "br": w.line() case "tr": - if w.escapeMarkdown && w.inTable { - w.tableRow = nil - return - } w.line() w.cellIndex = 0 case "td", "th": - if w.escapeMarkdown && w.inTable { - w.inTableCell = true - w.tableCell.Reset() - return - } if w.cellIndex > 0 { w.output.WriteString(" | ") } w.cellIndex++ case "pre": w.block() - if w.escapeMarkdown { - w.inPre = true - w.preFenceOpen = true - w.preLanguage = "" - w.preContent.Reset() - return - } w.output.WriteString("```") w.inPre = true w.preFenceOpen = true @@ -2300,26 +2301,15 @@ func (w *canonicalHTMLWriter) startTag(token html.Token, selfClosing bool) { for _, attribute := range token.Attr { if attribute.Key == "class" && strings.HasPrefix(attribute.Val, "language-") { if language := safeCodeLanguage(strings.TrimPrefix(attribute.Val, "language-")); language != "" { - if w.escapeMarkdown { - w.preLanguage = language - } else { - w.output.WriteString(language) - } + w.output.WriteString(language) } } } - if !w.escapeMarkdown { - w.output.WriteByte('\n') - } + w.output.WriteByte('\n') w.preFenceOpen = false } else if !w.inPre { w.flushPendingSpace() - if w.escapeMarkdown { - w.inlineCode = true - w.inlineContent.Reset() - } else { - w.output.WriteByte('`') - } + w.output.WriteByte('`') } case "img": for _, attribute := range token.Attr { @@ -2348,26 +2338,10 @@ func (w *canonicalHTMLWriter) startTag(token html.Token, selfClosing bool) { } } case "a": - if w.escapeMarkdown { - w.flushPendingSpace() - link := "" - for _, attribute := range token.Attr { - if attribute.Key == "href" { - link = safeRenditionLink(attribute.Val, w.maxLinkChars) - break - } - } - w.renditionLinks = append(w.renditionLinks, &renditionLink{destination: link}) - return - } link := "" for _, attribute := range token.Attr { if attribute.Key == "href" { - if w.escapeMarkdown { - link = safeRenditionLink(attribute.Val, w.maxLinkChars) - } else { - link = safeStoredLink(attribute.Val, w.maxLinkChars) - } + link = safeStoredLink(attribute.Val, w.maxLinkChars) break } } @@ -2415,37 +2389,10 @@ func (w *canonicalHTMLWriter) endTag(tag string) { case "li": w.line() case "tr": - if w.escapeMarkdown && w.inTable { - w.tableRows = append(w.tableRows, append([]string(nil), w.tableRow...)) - return - } w.line() - case "td", "th": - if w.escapeMarkdown && w.inTable && w.inTableCell { - w.tableRow = append(w.tableRow, w.tableCell.String()) - w.tableCell.Reset() - w.inTableCell = false - return - } case "table": - if w.escapeMarkdown && w.inTable { - w.inTable = false - w.writeRenditionRaw(serializeGFMTable(w.tableRows)) - w.tableRows = nil - w.block() - return - } w.block() case "pre": - if w.escapeMarkdown { - w.writeSafeCodeFence(w.preLanguage, w.preContent.String()) - w.inPre = false - w.preFenceOpen = false - w.preLanguage = "" - w.preContent.Reset() - w.block() - return - } if w.preFenceOpen { w.output.WriteByte('\n') } @@ -2457,28 +2404,9 @@ func (w *canonicalHTMLWriter) endTag(tag string) { case "code": if !w.inPre { w.flushPendingSpace() - if w.escapeMarkdown { - w.writeSafeInlineCode(w.inlineContent.String()) - w.inlineCode = false - w.inlineContent.Reset() - } else { - w.output.WriteByte('`') - } + w.output.WriteByte('`') } case "a": - if w.escapeMarkdown { - if len(w.renditionLinks) == 0 { - return - } - link := w.renditionLinks[len(w.renditionLinks)-1] - w.renditionLinks = w.renditionLinks[:len(w.renditionLinks)-1] - if link.destination == "" { - w.writeRenditionRaw(link.text.String()) - } else { - w.writeRenditionRaw("[" + link.text.String() + "](" + link.destination + ")") - } - return - } w.flushPendingSpace() if len(w.links) == 0 { return @@ -2526,10 +2454,6 @@ func isHTMLBlockElement(tag string) bool { func (w *canonicalHTMLWriter) writeText(value string) { if w.inPre { - if w.escapeMarkdown { - w.preContent.WriteString(stripUnsafeControls(value)) - return - } if w.preFenceOpen { w.output.WriteByte('\n') w.preFenceOpen = false @@ -2537,10 +2461,6 @@ func (w *canonicalHTMLWriter) writeText(value string) { w.output.WriteString(stripUnsafeControls(value)) return } - if w.inlineCode { - w.inlineContent.WriteString(stripUnsafeControls(value)) - return - } value = stripUnsafeControls(value) for _, character := range value { if unicode.IsSpace(character) { @@ -2553,84 +2473,9 @@ func (w *canonicalHTMLWriter) writeText(value string) { } func (w *canonicalHTMLWriter) writeMarkdownRune(character rune) { - if w.escapeMarkdown && strings.ContainsRune("\\\\`*_{}[]<>#!|~", character) { - w.writeRenditionRaw("\\") - } - if w.escapeMarkdown { - w.writeRenditionRaw(string(character)) - return - } w.output.WriteRune(character) } -func (w *canonicalHTMLWriter) writeRenditionRaw(value string) { - if len(w.renditionLinks) > 0 { - w.renditionLinks[len(w.renditionLinks)-1].text.WriteString(value) - return - } - if w.inTableCell { - w.tableCell.WriteString(value) - return - } - w.output.WriteString(value) -} - -func (w *canonicalHTMLWriter) writeSafeInlineCode(content string) { - fence := strings.Repeat("`", maxBacktickRun(content)+1) - w.writeRenditionRaw(fence) - if strings.HasPrefix(content, "`") || strings.HasSuffix(content, "`") { - w.writeRenditionRaw(" ") - } - w.writeRenditionRaw(content) - if strings.HasPrefix(content, "`") || strings.HasSuffix(content, "`") { - w.writeRenditionRaw(" ") - } - w.writeRenditionRaw(fence) -} - -func serializeGFMTable(rows [][]string) string { - if len(rows) == 0 { - return "" - } - columns := 0 - for _, row := range rows { - columns = max(columns, len(row)) - } - if columns == 0 { - return "" - } - format := func(row []string) string { - cells := make([]string, columns) - copy(cells, row) - return "| " + strings.Join(cells, " | ") + " |" - } - delimiter := make([]string, columns) - for index := range delimiter { - delimiter[index] = "---" - } - lines := []string{format(rows[0]), "| " + strings.Join(delimiter, " | ") + " |"} - for _, row := range rows[1:] { - lines = append(lines, format(row)) - } - return strings.Join(lines, "\n") -} - -func (w *canonicalHTMLWriter) writeSafeCodeFence(language, content string) { - if w.inTableCell { - w.writeSafeInlineCode(strings.ReplaceAll(strings.Join(strings.Fields(content), " "), "|", "\\|")) - return - } - fence := strings.Repeat("`", max(3, maxBacktickRun(content)+1)) - w.writeRenditionRaw(fence) - w.writeRenditionRaw(language) - w.writeRenditionRaw("\n") - w.writeRenditionRaw(content) - if !strings.HasSuffix(content, "\n") { - w.writeRenditionRaw("\n") - } - w.writeRenditionRaw(fence) -} - func maxBacktickRun(value string) int { maximum := 0 current := 0 @@ -2646,13 +2491,6 @@ func maxBacktickRun(value string) int { } func (w *canonicalHTMLWriter) flushPendingSpace() { - if w.escapeMarkdown { - if w.pendingSpace && w.renditionTargetHasContent() { - w.writeRenditionRaw(" ") - } - w.pendingSpace = false - return - } if w.pendingSpace && w.output.Len() > 0 && !strings.HasSuffix(w.output.String(), "\n") && !strings.HasSuffix(w.output.String(), " ") { w.output.WriteByte(' ') @@ -2660,21 +2498,7 @@ func (w *canonicalHTMLWriter) flushPendingSpace() { w.pendingSpace = false } -func (w *canonicalHTMLWriter) renditionTargetHasContent() bool { - if len(w.renditionLinks) > 0 { - return w.renditionLinks[len(w.renditionLinks)-1].text.Len() > 0 - } - if w.inTableCell { - return w.tableCell.Len() > 0 - } - return w.output.Len() > 0 && !strings.HasSuffix(w.output.String(), "\n") -} - func (w *canonicalHTMLWriter) line() { - if w.escapeMarkdown && w.inTable { - w.pendingSpace = false - return - } w.pendingSpace = false if w.output.Len() > 0 && !strings.HasSuffix(w.output.String(), "\n") { w.output.WriteByte('\n') @@ -2682,10 +2506,6 @@ func (w *canonicalHTMLWriter) line() { } func (w *canonicalHTMLWriter) block() { - if w.escapeMarkdown && w.inTable { - w.pendingSpace = false - return - } w.line() if w.output.Len() > 0 && !strings.HasSuffix(w.output.String(), "\n\n") { w.output.WriteByte('\n') diff --git a/document/openaiembed/client.go b/document/openaiembed/client.go new file mode 100644 index 00000000..39f9b2c8 --- /dev/null +++ b/document/openaiembed/client.go @@ -0,0 +1,720 @@ +// Package openaiembed implements one bounded OpenAI-compatible text embedding +// endpoint for operator-controlled local deployments. +package openaiembed + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + json "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "net" + "net/http" + "net/netip" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + // ProviderID is the fixed provider-neutral adapter identity. + ProviderID = "openai-compatible.embeddings-v1" + // DocumentFormatterV1 identifies the adapter's document formatting path. + DocumentFormatterV1 = "openai-compatible/document/v1" + // QueryFormatterV1 identifies the adapter's query formatting path. + QueryFormatterV1 = "openai-compatible/query/v1" + // ScalarEncodingFloat32 is the only scalar representation returned by the adapter. + ScalarEncodingFloat32 = "float32" + + embeddingsPath = "/v1/embeddings" + adapterContract = "docbank-openai-compatible-embeddings/v1" + + defaultRequestTimeout = 30 * time.Second + defaultMaxBatchItems = 128 + defaultMaxInputBytes = int64(1 << 20) + defaultMaxRequestBytes = int64(2 << 20) + defaultMaxResponseBytes = int64(32 << 20) + + maxRequestTimeout = 5 * time.Minute + maxBatchItems = 10_000 + maxInputBytes = int64(1 << 30) + maxRequestBytes = int64(1 << 30) + maxResponseBytes = int64(1 << 30) + maxSecretBytes = 64 << 10 + maxIdentityBytes = 128 + unitLengthTolerance = 1e-4 +) + +var _ document.EmbeddingProvider = (*Client)(nil) + +// SecretResolver resolves only the optional credential binding named by a +// profile. Secret values never enter profile or vector-space identity. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile freezes one OpenAI-compatible endpoint, explicit model-input +// contract, immutable deployment identity, and all transport capacity bounds. +// Exactly one of DeploymentEpoch and ProviderRevisionHeader is required. +type Profile struct { + Origin string + Descriptor document.EmbeddingDescriptor + ModelInput document.ModelInputContract + DeploymentContract *DeploymentContract + SecretBinding string + DeploymentEpoch string + ProviderRevisionHeader string + RequestTimeout time.Duration + MaxBatchItems int + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + EgressPolicy providerhttp.EgressPolicy +} + +var ( + ErrTransientResponse = errors.New("openaiembed: transient provider response") + ErrCapacityResponse = errors.New("openaiembed: provider capacity exceeded") + ErrPermanentResponse = errors.New("openaiembed: permanent provider response") +) + +type ProviderError struct { + Kind error + StatusCode int + RetryDelay time.Duration + RetrySet bool +} + +func (err *ProviderError) Error() string { + return fmt.Sprintf("openaiembed: HTTP %d: %v", err.StatusCode, err.Kind) +} +func (err *ProviderError) Unwrap() error { return err.Kind } +func RetryAfter(err error) (time.Duration, bool) { + providerErr, ok := errors.AsType[*ProviderError](err) + if !ok || !providerErr.RetrySet { + return 0, false + } + return providerErr.RetryDelay, true +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + ModelInput document.ModelInputContract `json:"model_input"` + DeploymentContract *DeploymentContract `json:"deployment_contract,omitempty"` + CredentialBinding string `json:"credential_binding"` + DeploymentEpoch string `json:"deployment_epoch,omitempty"` + ProviderRevisionHeader string `json:"provider_revision_header,omitempty"` + RequestTimeoutNanos int64 `json:"request_timeout_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + Egress *openAIEgressIdentity `json:"egress,omitempty"` +} + +type openAIEgressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +// Client calls exactly POST /v1/embeddings on the profile origin. +type Client struct { + profile Profile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +type wireRequest struct { + Input []string `json:"input"` + Model string `json:"model"` + EncodingFormat string `json:"encoding_format"` +} + +type wireResponse struct { + Object string `json:"object"` + Data []wireEmbedding `json:"data"` + Model string `json:"model"` + Usage *wireUsage `json:"usage,omitempty"` +} + +type wireEmbedding struct { + Object string `json:"object"` + Embedding []float32 `json:"embedding"` + Index *int `json:"index"` +} + +type wireUsage struct { + PromptTokens int64 `json:"prompt_tokens"` + TotalTokens int64 `json:"total_tokens"` +} + +// PolicyFingerprint returns the canonical profile identity expected in the +// embedding descriptor. Credential values and HTTP client state are excluded. +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptorIdentity, err := normalizeProfile(profile) + if err != nil { + return "", err + } + identity := policyIdentity{ + AdapterContract: adapterContract, Origin: normalized.Origin, Route: embeddingsPath, + Descriptor: descriptorIdentity, ModelInput: normalized.ModelInput, + DeploymentContract: normalized.DeploymentContract, + CredentialBinding: normalized.SecretBinding, DeploymentEpoch: normalized.DeploymentEpoch, + ProviderRevisionHeader: normalized.ProviderRevisionHeader, + RequestTimeoutNanos: int64(normalized.RequestTimeout), MaxBatchItems: normalized.MaxBatchItems, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, + Egress: openAIEgressPolicyIdentity(normalized.EgressPolicy), + } + encoded, err := json.Marshal(identity, json.Deterministic(true)) + if err != nil { + return "", fmt.Errorf("openaiembed: encode policy identity: %w", err) + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +// New validates an immutable profile and isolates the supplied HTTP client +// from ambient cookies, timeouts, and redirect behavior. It performs no I/O. +func New(profile Profile, secrets SecretResolver, httpClient *http.Client) (*Client, error) { + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("openaiembed: invalid descriptor: %w", err) + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("openaiembed: descriptor policy fingerprint does not match profile") + } + if normalized.SecretBinding == "" { + if !nilValue(secrets) { + return nil, errors.New("openaiembed: secret resolver requires a named binding") + } + } else if nilValue(secrets) { + return nil, errors.New("openaiembed: named secret binding requires a resolver") + } + if httpClient == nil { + return nil, errors.New("openaiembed: HTTP client is required") + } + isolate := *httpClient + isolate.CheckRedirect = providerhttp.RefuseRedirects + isolate.Jar = nil + isolate.Timeout = 0 + normalized.Descriptor = cloneDescriptor(descriptor) + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, http: &isolate}, nil +} + +// Descriptor returns a defensive copy of the immutable provider contract. +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Embed validates and formats one text batch, performs one bounded request, +// and restores caller ordering from the response indices. +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if client == nil { + return document.EmbeddingResult{}, errors.New("openaiembed: client is required") + } + if ctx == nil { + return document.EmbeddingResult{}, errors.New("openaiembed: context is required") + } + if err := document.ValidateEmbeddingProviderRequest(client, inputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || + authorization.MaxInputBytes > client.profile.MaxInputBytes || + authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, errors.New("openaiembed: embedding authorization exceeds profile capacity") + } + if err := ctx.Err(); err != nil { + return document.EmbeddingResult{}, fmt.Errorf("openaiembed: embedding canceled: %w", err) + } + + rendered := make([]string, len(inputs)) + var renderedBytes int64 + for index, input := range inputs { + switch { + case input.Role == document.EmbeddingRoleDocument && input.Kind == document.EmbeddingInputRenditionChunk: + rendered[index] = client.profile.ModelInput.EncodeDocument(input.Text) + case input.Role == document.EmbeddingRoleQuery && input.Kind == document.EmbeddingInputQueryText: + rendered[index] = client.profile.ModelInput.EncodeQuery(input.Text) + default: + return document.EmbeddingResult{}, errors.New("openaiembed: unsupported non-text embedding input or role") + } + if int64(len(rendered[index])) > client.profile.MaxInputBytes-renderedBytes { + return document.EmbeddingResult{}, errors.New("openaiembed: embedding input exceeds profile byte capacity") + } + renderedBytes += int64(len(rendered[index])) + } + payload, err := json.Marshal(wireRequest{Input: rendered, Model: client.descriptor.Model, EncodingFormat: "float"}) + if err != nil { + return document.EmbeddingResult{}, errors.New("openaiembed: could not encode embedding request") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + return document.EmbeddingResult{}, fmt.Errorf("%w: embedding request byte limit exceeded", ErrCapacityResponse) + } + + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + request, err := http.NewRequestWithContext(requestCtx, http.MethodPost, client.profile.Origin+embeddingsPath, bytes.NewReader(payload)) + if err != nil { + return document.EmbeddingResult{}, errors.New("openaiembed: could not construct embedding request") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + if client.profile.SecretBinding != "" { + secret, resolveErr := client.secrets.ResolveSecret(requestCtx, client.profile.SecretBinding) + if resolveErr != nil { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("openaiembed: credential resolution canceled: %w", contextErr) + } + return document.EmbeddingResult{}, errors.New("openaiembed: could not resolve credential") + } + if !validSecret(secret) { + return document.EmbeddingResult{}, errors.New("openaiembed: resolved credential is invalid") + } + request.Header.Set("Authorization", "Bearer "+secret) + } + + response, err := client.http.Do(request) + if err != nil { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("openaiembed: embedding request canceled: %w", contextErr) + } + return document.EmbeddingResult{}, &ProviderError{Kind: ErrTransientResponse} + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode >= 300 && response.StatusCode < 400 { + return document.EmbeddingResult{}, fmt.Errorf("openaiembed: provider redirect refused: %w", + &ProviderError{Kind: ErrPermanentResponse, StatusCode: response.StatusCode}) + } + if response.StatusCode == http.StatusRequestEntityTooLarge { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrCapacityResponse, StatusCode: response.StatusCode} + } + if response.StatusCode == http.StatusTooManyRequests || response.StatusCode >= 500 { + delay, set := openAIRetryAfter(response.Header.Get("Retry-After"), time.Now().UTC()) + return document.EmbeddingResult{}, &ProviderError{Kind: ErrTransientResponse, StatusCode: response.StatusCode, + RetryDelay: delay, RetrySet: set} + } + if response.StatusCode != http.StatusOK { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse, StatusCode: response.StatusCode} + } + if err := validateResponseContentType(response.Header.Get("Content-Type")); err != nil { + return document.EmbeddingResult{}, err + } + if err := client.validateRevisionEcho(response.Header); err != nil { + return document.EmbeddingResult{}, err + } + body, err := readBounded(requestCtx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return document.EmbeddingResult{}, err + } + var decoded wireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response does not match the bounded embedding schema") + } + result, err := client.validateAndOrder(decoded, inputs) + if err != nil { + return document.EmbeddingResult{}, err + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, err + } + return result, nil +} + +func (client *Client) validateRevisionEcho(header http.Header) error { + if client.profile.ProviderRevisionHeader == "" { + return nil + } + values := header.Values(client.profile.ProviderRevisionHeader) + if len(values) != 1 || strings.TrimSpace(values[0]) != client.descriptor.ModelRevision { + return errors.New("openaiembed: provider revision echo does not match profile") + } + return nil +} + +func (client *Client) validateAndOrder(response wireResponse, inputs []document.EmbeddingInput) (document.EmbeddingResult, error) { + if response.Object != "list" || response.Model != client.descriptor.Model { + return document.EmbeddingResult{}, errors.New("openaiembed: provider model or response contract drifted") + } + if response.Usage != nil && (response.Usage.PromptTokens < 0 || response.Usage.TotalTokens < response.Usage.PromptTokens) { + return document.EmbeddingResult{}, errors.New("openaiembed: provider usage is invalid") + } + if len(response.Data) != len(inputs) { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response has a missing vector") + } + vectors := make([]document.EmbeddingVector, len(inputs)) + seen := make([]bool, len(inputs)) + for _, item := range response.Data { + if item.Object != "embedding" || item.Index == nil { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response item contract drifted") + } + index := *item.Index + if index < 0 || index >= len(inputs) { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response index is outside request bounds") + } + if seen[index] { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response has a duplicate vector index") + } + seen[index] = true + if err := client.validateVector(item.Embedding); err != nil { + return document.EmbeddingResult{}, err + } + vectors[index] = document.EmbeddingVector{Key: inputs[index].Key, Values: slices.Clone(item.Embedding)} + } + if slices.Contains(seen, false) { + return document.EmbeddingResult{}, errors.New("openaiembed: provider response has a missing vector index") + } + return document.EmbeddingResult{Vectors: vectors}, nil +} + +func (client *Client) validateVector(vector []float32) error { + if len(vector) != client.descriptor.Dimension { + return errors.New("openaiembed: provider vector dimension does not match profile") + } + var squaredNorm float64 + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("openaiembed: provider vector contains a non-finite value") + } + squaredNorm += float64(value) * float64(value) + } + if squaredNorm == 0 { + return errors.New("openaiembed: provider returned a zero vector") + } + if client.descriptor.Normalization == document.VectorNormalizationUnitLength && math.Abs(squaredNorm-1) > unitLengthTolerance { + return errors.New("openaiembed: provider vector normalization does not match profile") + } + return nil +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + if profile.DeploymentContract != nil { + contract := *profile.DeploymentContract + profile.DeploymentContract = &contract + } + profile.EgressPolicy = cloneEgressPolicy(profile.EgressPolicy) + origin, err := validateOrigin(profile.Origin) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + profile.Origin = origin + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultMaxBatchItems + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultMaxInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultMaxRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultMaxResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maxRequestTimeout || + profile.MaxBatchItems < 1 || profile.MaxBatchItems > maxBatchItems || + profile.MaxInputBytes < 1 || profile.MaxInputBytes > maxInputBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maxRequestBytes || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maxResponseBytes { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: execution bounds are invalid") + } + if profile.SecretBinding != "" && !validIdentityToken(profile.SecretBinding) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: secret binding is invalid") + } + if profile.EgressPolicy.Scheme != "" { + if profile.EgressPolicy.ConnectTimeout == 0 { + profile.EgressPolicy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if profile.EgressPolicy.KeepAlive == 0 { + profile.EgressPolicy.KeepAlive = providerhttp.DefaultKeepAlive + } + if profile.EgressPolicy.TLSHandshakeTimeout == 0 { + profile.EgressPolicy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if profile.EgressPolicy.ProxyMode == "" { + profile.EgressPolicy.ProxyMode = providerhttp.ProxyDisabled + } + if _, err := providerhttp.NewTransport(profile.EgressPolicy, nil); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, fmt.Errorf("openaiembed: invalid sealed egress policy: %w", err) + } + parsed, _ := url.Parse(profile.Origin) + port := parsed.Port() + if port == "" { + if parsed.Scheme == "https" { + port = "443" + } else { + port = "80" + } + } + if profile.EgressPolicy.Scheme != parsed.Scheme || !strings.EqualFold(profile.EgressPolicy.Host, parsed.Hostname()) || + strconv.Itoa(int(profile.EgressPolicy.Port)) != port || profile.EgressPolicy.ProxyMode != providerhttp.ProxyDisabled { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: origin and sealed egress authority differ") + } + slices.SortFunc(profile.EgressPolicy.AllowedCIDRs, func(a, b netip.Prefix) int { return strings.Compare(a.Masked().String(), b.Masked().String()) }) + slices.Sort(profile.EgressPolicy.TLS.SPKISHA256) + } + + descriptorIdentity := profile.Descriptor + descriptorIdentity.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptorIdentity.Fingerprint = "" + descriptorIdentity, err = document.NewEmbeddingDescriptor(descriptorIdentity) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, fmt.Errorf("openaiembed: invalid descriptor identity: %w", err) + } + descriptorIdentity.PolicyFingerprint = "" + descriptorIdentity.Fingerprint = "" + if err := validateDescriptorContract(descriptorIdentity, profile.ModelInput); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + if err := validateDeploymentContract(profile.DeploymentContract, descriptorIdentity, profile.ModelInput); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + if (profile.DeploymentEpoch == "") == (profile.ProviderRevisionHeader == "") { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: exactly one deployment epoch or provider revision header is required") + } + if profile.DeploymentEpoch != "" { + if !validIdentityToken(profile.DeploymentEpoch) || profile.DeploymentEpoch != descriptorIdentity.ModelRevision { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: deployment epoch must exactly match descriptor model revision") + } + } else if !validRevisionHeader(profile.ProviderRevisionHeader) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaiembed: provider revision header is not a canonical safe response header") + } + return profile, descriptorIdentity, nil +} + +func openAIEgressPolicyIdentity(policy providerhttp.EgressPolicy) *openAIEgressIdentity { + if policy.Scheme == "" { + return nil + } + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.Masked().String() + } + return &openAIEgressIdentity{Scheme: policy.Scheme, Host: strings.ToLower(policy.Host), Port: policy.Port, + AllowedCIDRs: cidrs, ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), + KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), + SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func cloneEgressPolicy(policy providerhttp.EgressPolicy) providerhttp.EgressPolicy { + policy.AllowedCIDRs = slices.Clone(policy.AllowedCIDRs) + policy.TLS.SPKISHA256 = slices.Clone(policy.TLS.SPKISHA256) + return policy +} + +func openAIRetryAfter(value string, now time.Time) (time.Duration, bool) { + value = strings.TrimSpace(value) + if seconds, err := strconv.ParseInt(value, 10, 64); err == nil && seconds >= 0 { + return min(time.Duration(seconds)*time.Second, time.Hour), true + } + when, err := http.ParseTime(value) + if err != nil { + return 0, false + } + return min(max(when.Sub(now), 0), time.Hour), true +} + +func validateDescriptorContract(descriptor document.EmbeddingDescriptor, modelInput document.ModelInputContract) error { + if descriptor.ID != ProviderID || descriptor.TrustBoundary != document.EmbeddingTrustOperatorNetwork || + descriptor.ScalarEncoding != ScalarEncodingFloat32 || descriptor.DocumentFormatter != DocumentFormatterV1 || + descriptor.QueryFormatter != QueryFormatterV1 || !descriptor.SupportsTextQuery || + !slices.Equal(descriptor.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || + !slices.Equal(descriptor.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeText}) { + return errors.New("openaiembed: descriptor does not match the text-only adapter contract") + } + if descriptor.ModelInput != modelInput || descriptor.CompatibilityID != modelInput.CompatibilityID { + return errors.New("openaiembed: descriptor and explicit model-input contract differ") + } + switch modelInput.Profile { + case document.ModelInputProfileNomic, document.ModelInputProfileE5, + document.ModelInputProfileBGEM3, document.ModelInputProfileGTE, + document.ModelInputProfileQwen3, document.ModelInputProfileQueryInstruction, + document.ModelInputProfileCustom: + default: + return errors.New("openaiembed: model-input contract must use a reviewed embedding family or complete custom profile") + } + return nil +} + +func validateOrigin(raw string) (string, error) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Opaque != "" || parsed.ForceQuery || parsed.Fragment != "" || + (parsed.Path != "" && parsed.Path != "/") || parsed.RawPath != "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return "", errors.New("openaiembed: origin must be one HTTP(S) origin without credentials, non-root path, query, or fragment") + } + hostname := parsed.Hostname() + if hostname == "" || !asciiHost(hostname) || strings.ToLower(hostname) != hostname { + return "", errors.New("openaiembed: origin host is not canonical ASCII") + } + port := parsed.Port() + if port != "" { + value, parseErr := strconv.ParseUint(port, 10, 16) + if parseErr != nil || value == 0 { + return "", errors.New("openaiembed: origin port is invalid") + } + } + authority := hostname + if strings.Contains(hostname, ":") { + authority = "[" + hostname + "]" + } + if port != "" { + authority = net.JoinHostPort(hostname, port) + } + return parsed.Scheme + "://" + authority, nil +} + +func asciiHost(host string) bool { + if address, err := netip.ParseAddr(host); err == nil { + return address.Zone() == "" && address.String() == host + } + if len(host) > 253 || strings.HasSuffix(host, ".") { + return false + } + for label := range strings.SplitSeq(host, ".") { + if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return false + } + for _, character := range label { + if (character < 'a' || character > 'z') && (character < '0' || character > '9') && character != '-' { + return false + } + } + } + return true +} + +func validIdentityToken(value string) bool { + return value != "" && len(value) <= maxIdentityBytes && value == strings.TrimSpace(value) && utf8.ValidString(value) && strings.IndexFunc(value, unicode.IsControl) < 0 +} + +func validRevisionHeader(name string) bool { + if name == "" || len(name) > maxIdentityBytes || http.CanonicalHeaderKey(name) != name { + return false + } + for _, character := range name { + if (character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || + (character >= '0' && character <= '9') || strings.ContainsRune("!#$%&'*+-.^_`|~", character) { + continue + } + return false + } + switch name { + case "Authorization", "Connection", "Content-Length", "Content-Type", "Cookie", "Date", "Location", "Server", "Set-Cookie", "Transfer-Encoding": + return false + default: + return true + } +} + +func validSecret(secret string) bool { + if secret == "" || len(secret) > maxSecretBytes { + return false + } + padding := false + content := 0 + for index := range len(secret) { + character := secret[index] + if character == '=' { + padding = true + continue + } + if padding || !validBearerCharacter(character) { + return false + } + content++ + } + return content > 0 +} + +func validBearerCharacter(character byte) bool { + return character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' || + character >= '0' && character <= '9' || strings.ContainsRune("-._~+/", rune(character)) +} + +func validateResponseContentType(value string) error { + mediaType, parameters, err := mime.ParseMediaType(value) + if err != nil || mediaType != "application/json" { + return errors.New("openaiembed: provider response content type is not application/json") + } + if len(parameters) == 0 { + return nil + } + charset, ok := parameters["charset"] + if len(parameters) != 1 || !ok || !strings.EqualFold(charset, "utf-8") { + return errors.New("openaiembed: provider response content type has unsupported parameters") + } + return nil +} + +func readBounded(ctx context.Context, reader io.Reader, maximum int64) ([]byte, error) { + body, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("openaiembed: response read canceled: %w", contextErr) + } + return nil, errors.New("openaiembed: could not read provider response") + } + if int64(len(body)) > maximum { + return nil, errors.New("openaiembed: provider response byte limit exceeded") + } + return body, nil +} + +func cloneDescriptor(descriptor document.EmbeddingDescriptor) document.EmbeddingDescriptor { + descriptor.InputKinds = slices.Clone(descriptor.InputKinds) + descriptor.SupportedRequestModes = slices.Clone(descriptor.SupportedRequestModes) + return descriptor +} + +func nilValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/openaiembed/client_test.go b/document/openaiembed/client_test.go new file mode 100644 index 00000000..93d4f294 --- /dev/null +++ b/document/openaiembed/client_test.go @@ -0,0 +1,640 @@ +package openaiembed + +import ( + "context" + _ "embed" + "encoding/json/v2" + "errors" + "io" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "net/netip" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +//go:embed testdata/success-indexed.json +var successIndexedResponse []byte + +//go:embed testdata/schema-drift.json +var schemaDriftResponse []byte + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("missing synthetic secret") + } + return value, nil +} + +type capturedRequest struct { + Input []string `json:"input"` + Model string `json:"model"` + EncodingFormat string `json:"encoding_format"` +} + +func TestEmbedAppliesExplicitModelInputProfilesAndRestoresResponseIndices(t *testing.T) { + tests := []struct { + name string + config document.ModelInputContractConfig + expected []string + }{ + {name: "nomic", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}, expected: []string{"search_document: passage", "search_query: question"}}, + {name: "e5", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileE5}, expected: []string{"passage: passage", "query: question"}}, + {name: "bge-m3", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileBGEM3}, expected: []string{"passage", "question"}}, + {name: "gte", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileGTE}, expected: []string{"passage", "question"}}, + {name: "qwen3 prefix free", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileQwen3}, expected: []string{"passage", "question"}}, + {name: "qwen3 instruction", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileQwen3, QueryInstruction: "Retrieve supporting passages"}, expected: []string{"passage", "Instruct: Retrieve supporting passages\nQuery:question"}}, + {name: "query-only instruction", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileQueryInstruction, QueryInstruction: "Find exact evidence"}, expected: []string{"passage", "Instruct: Find exact evidence\nQuery:question"}}, + {name: "complete custom", config: document.ModelInputContractConfig{Profile: document.ModelInputProfileCustom, CompatibilityID: "synthetic/custom-v1", Document: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "doc[{{content}}]"}, Query: document.ModelInputEncoder{Mode: document.ModelInputModeText, Template: "query[{{content}}]"}}, expected: []string{"doc[passage]", "query[question]"}}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + contract, err := document.NewModelInputContract(testCase.config) + require.NoError(t, err) + profile := testProfile(t, contract) + var seen capturedRequest + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + require.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "/v1/embeddings", request.URL.Path) + assert.Empty(t, request.URL.RawQuery) + assert.Equal(t, "application/json", request.Header.Get("Content-Type")) + assert.Equal(t, "application/json", request.Header.Get("Accept")) + require.NoError(t, json.UnmarshalRead(request.Body, &seen, json.RejectUnknownMembers(true))) + return jsonResponse(request, http.StatusOK, successIndexedResponse), nil + })) + + result, err := document.ExecuteEmbedding(t.Context(), client, testInputs(), testAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.Equal(t, testCase.expected, seen.Input) + assert.Equal(t, "synthetic-model", seen.Model) + assert.Equal(t, "float", seen.EncodingFormat) + require.Len(t, result.Vectors, 2) + assert.Equal(t, "document-1", result.Vectors[0].Key) + assert.Equal(t, []float32{1, 0, 0}, result.Vectors[0].Values) + assert.Nil(t, result.Vectors[0].Index) + assert.Equal(t, "query-1", result.Vectors[1].Key) + assert.Equal(t, []float32{0, 1, 0}, result.Vectors[1].Values) + assert.Nil(t, result.Vectors[1].Index) + }) + } +} + +func TestNewRequiresCanonicalProfileAndStableRevisionIdentity(t *testing.T) { + contract := modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}) + base := testProfile(t, contract) + + for name, mutate := range map[string]func(Profile) Profile{ + "credentials in origin": func(profile Profile) Profile { profile.Origin = "http://user@127.0.0.1:11434"; return profile }, + "query in origin": func(profile Profile) Profile { profile.Origin += "?model=private"; return profile }, + "fragment in origin": func(profile Profile) Profile { profile.Origin += "#fragment"; return profile }, + "non-root base path": func(profile Profile) Profile { profile.Origin += "/api"; return profile }, + "unsupported scheme": func(profile Profile) Profile { profile.Origin = "file:///tmp/provider"; return profile }, + "missing revision authority": func(profile Profile) Profile { profile.DeploymentEpoch = ""; return profile }, + "two revision authorities": func(profile Profile) Profile { profile.ProviderRevisionHeader = "X-Model-Revision"; return profile }, + "epoch differs from descriptor": func(profile Profile) Profile { profile.DeploymentEpoch = "epoch-other"; return profile }, + "non-canonical revision header": func(profile Profile) Profile { + profile.DeploymentEpoch = "" + profile.ProviderRevisionHeader = "x-model-revision" + return profile + }, + "unsafe revision header": func(profile Profile) Profile { + profile.DeploymentEpoch = "" + profile.ProviderRevisionHeader = "Authorization" + return profile + }, + "different explicit input contract": func(profile Profile) Profile { + profile.ModelInput = modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileE5}) + return profile + }, + "non-canonical descriptor": func(profile Profile) Profile { + profile.Descriptor.Fingerprint = strings.Repeat("0", 64) + return profile + }, + "wrong trust boundary": func(profile Profile) Profile { + profile.Descriptor = canonicalMutatedDescriptor(t, profile.Descriptor, func(descriptor *document.EmbeddingDescriptor) { + descriptor.TrustBoundary = document.EmbeddingTrustHostedProvider + }) + return profile + }, + "direct-file capability": func(profile Profile) Profile { + profile.Descriptor = canonicalMutatedDescriptor(t, profile.Descriptor, func(descriptor *document.EmbeddingDescriptor) { + descriptor.InputKinds = append(descriptor.InputKinds, document.EmbeddingInputOriginalFile) + }) + return profile + }, + "missing text query": func(profile Profile) Profile { + profile.Descriptor = canonicalMutatedDescriptor(t, profile.Descriptor, func(descriptor *document.EmbeddingDescriptor) { descriptor.SupportsTextQuery = false }) + return profile + }, + } { + t.Run(name, func(t *testing.T) { + _, err := New(mutate(base), nil, &http.Client{Transport: staticSuccessTransport()}) + require.Error(t, err) + }) + } + + for _, unsupported := range []document.ModelInputProfile{ + document.ModelInputProfileOpenAICompatible, document.ModelInputProfileMistral, + } { + t.Run("unsupported input profile "+string(unsupported), func(t *testing.T) { + profile := base + profile.ModelInput = modelInput(t, document.ModelInputContractConfig{Profile: unsupported}) + profile.Descriptor = canonicalMutatedDescriptor(t, profile.Descriptor, func(descriptor *document.EmbeddingDescriptor) { + descriptor.ModelInput = profile.ModelInput + descriptor.CompatibilityID = profile.ModelInput.CompatibilityID + }) + _, err := New(profile, nil, &http.Client{Transport: staticSuccessTransport()}) + require.ErrorContains(t, err, "model-input") + }) + } + + t.Run("empty input profile", func(t *testing.T) { + profile := base + profile.ModelInput = modelInput(t, document.ModelInputContractConfig{}) + _, err := New(profile, nil, &http.Client{Transport: staticSuccessTransport()}) + require.ErrorContains(t, err, "model-input") + }) + + t.Run("voyage request modes", func(t *testing.T) { + profile := base + profile.ModelInput = modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileVoyage}) + profile.Descriptor = canonicalMutatedDescriptor(t, profile.Descriptor, func(descriptor *document.EmbeddingDescriptor) { + descriptor.ModelInput = profile.ModelInput + descriptor.CompatibilityID = profile.ModelInput.CompatibilityID + descriptor.SupportedRequestModes = []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery} + }) + _, err := New(profile, nil, &http.Client{Transport: staticSuccessTransport()}) + require.ErrorContains(t, err, "text-only") + }) + + _, err := New(base, nil, nil) + require.ErrorContains(t, err, "HTTP client") +} + +func TestProfilePolicyFingerprintCoversEndpointIdentityAndBounds(t *testing.T) { + base := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + want, err := PolicyFingerprint(base) + require.NoError(t, err) + assert.Equal(t, want, base.Descriptor.PolicyFingerprint) + + mutations := map[string]func(*Profile){ + "origin": func(profile *Profile) { profile.Origin = "http://127.0.0.1:11435" }, + "deployment epoch": func(profile *Profile) { + profile.DeploymentEpoch = "epoch-other" + profile.Descriptor.ModelRevision = "epoch-other" + }, + "credential binding": func(profile *Profile) { profile.SecretBinding = "other-secret" }, + "batch capacity": func(profile *Profile) { profile.MaxBatchItems++ }, + "input capacity": func(profile *Profile) { profile.MaxInputBytes++ }, + "request bound": func(profile *Profile) { profile.MaxRequestBytes++ }, + "response bound": func(profile *Profile) { profile.MaxResponseBytes++ }, + "timeout": func(profile *Profile) { profile.RequestTimeout++ }, + "model": func(profile *Profile) { profile.Descriptor.Model = "different-model" }, + "dimension": func(profile *Profile) { profile.Descriptor.Dimension++ }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changed := base + mutate(&changed) + fingerprint, err := PolicyFingerprint(changed) + require.NoError(t, err) + assert.NotEqual(t, want, fingerprint) + changed.Descriptor = canonicalMutatedDescriptor(t, changed.Descriptor, func(*document.EmbeddingDescriptor) {}) + _, err = New(changed, nil, &http.Client{Transport: staticSuccessTransport()}) + require.ErrorContains(t, err, "policy fingerprint") + }) + } +} + +func TestProfilePolicyFingerprintCoversExactEgressAuthority(t *testing.T) { + base := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + base.EgressPolicy = providerhttp.EgressPolicy{ + Scheme: "http", Host: "127.0.0.1", Port: 11434, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, + } + want, err := PolicyFingerprint(base) + require.NoError(t, err) + changed := base + changed.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("127.0.0.1/32")} + got, err := PolicyFingerprint(changed) + require.NoError(t, err) + assert.NotEqual(t, want, got) + + mismatched := base + mismatched.EgressPolicy.Port++ + _, err = PolicyFingerprint(mismatched) + require.ErrorContains(t, err, "origin") +} + +func TestEmbedReturnsClassifiedProviderFailures(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + tests := []struct { + name string + status int + want error + retry string + wantDelay time.Duration + }{ + {name: "capacity", status: http.StatusRequestEntityTooLarge, want: ErrCapacityResponse}, + {name: "permanent", status: http.StatusBadRequest, want: ErrPermanentResponse}, + {name: "transient", status: http.StatusServiceUnavailable, want: ErrTransientResponse}, + {name: "retry after", status: http.StatusTooManyRequests, want: ErrTransientResponse, retry: "2", wantDelay: 2 * time.Second}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, testCase.status, []byte(`{"error":"synthetic"}`)) + if testCase.retry != "" { + response.Header.Set("Retry-After", testCase.retry) + } + return response, nil + })) + _, err := client.Embed(t.Context(), testInputs()[:1], testAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, testCase.want) + delay, set := RetryAfter(err) + assert.Equal(t, testCase.retry != "", set) + assert.Equal(t, testCase.wantDelay, delay) + }) + } +} + +func TestEmbedEnforcesProfileAuthorizationAndTextOnlyBoundsBeforeNetwork(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + profile.MaxBatchItems = 2 + profile.MaxInputBytes = 64 + profile.MaxRequestBytes = 256 + profile.MaxResponseBytes = 4096 + profile.Descriptor = descriptorFor(t, profile) + var requests atomic.Int64 + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + return jsonResponse(request, http.StatusOK, successIndexedResponse), nil + })) + + tests := []struct { + name string + inputs []document.EmbeddingInput + authorization document.EmbeddingAuthorization + }{ + {name: "batch authority above profile", inputs: testInputs(), authorization: mutateAuthorization(testAuthorization(profile.Descriptor), func(value *document.EmbeddingAuthorization) { value.MaxBatchItems = 3 })}, + {name: "input authority above profile", inputs: testInputs(), authorization: mutateAuthorization(testAuthorization(profile.Descriptor), func(value *document.EmbeddingAuthorization) { value.MaxInputBytes = 65 })}, + {name: "response authority above profile", inputs: testInputs(), authorization: mutateAuthorization(testAuthorization(profile.Descriptor), func(value *document.EmbeddingAuthorization) { value.MaxResponseBytes = 4097 })}, + {name: "direct file", inputs: []document.EmbeddingInput{{Key: "file", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputOriginalFile}}, authorization: testAuthorization(profile.Descriptor)}, + {name: "too many items", inputs: append(testInputs(), document.EmbeddingInput{Key: "extra", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "extra"}), authorization: testAuthorization(profile.Descriptor)}, + {name: "rendered text too large", inputs: []document.EmbeddingInput{{Key: "large", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: strings.Repeat("x", 64)}}, authorization: testAuthorization(profile.Descriptor)}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + _, err := client.Embed(t.Context(), testCase.inputs, testCase.authorization) + require.Error(t, err) + }) + } + assert.Zero(t, requests.Load()) +} + +func TestEmbedUsesOptionalNamedBearerSecretWithoutAmbientCookies(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileGTE})) + profile.SecretBinding = "local-embedding-key" + profile.Descriptor = descriptorFor(t, profile) + jar, err := cookiejar.New(nil) + require.NoError(t, err) + var seenAuthorization string + base := &http.Client{Jar: jar, Timeout: time.Hour, CheckRedirect: func(*http.Request, []*http.Request) error { return nil }} + base.Transport = roundTripFunc(func(request *http.Request) (*http.Response, error) { + seenAuthorization = request.Header.Get("Authorization") + return jsonResponse(request, http.StatusOK, successIndexedResponse), nil + }) + client, err := New(profile, testSecrets{"local-embedding-key": "synthetic-secret"}, base) + require.NoError(t, err) + assert.Nil(t, client.http.Jar) + assert.Zero(t, client.http.Timeout) + _, err = client.Embed(t.Context(), testInputs(), testAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.Equal(t, "Bearer synthetic-secret", seenAuthorization) + + _, err = New(profile, nil, base) + require.ErrorContains(t, err, "resolver") + withoutBinding := profile + withoutBinding.SecretBinding = "" + withoutBinding.Descriptor = descriptorFor(t, withoutBinding) + _, err = New(withoutBinding, testSecrets{}, base) + require.ErrorContains(t, err, "binding") + + for name, secret := range map[string]string{ + "line break": "bad\r\nvalue", + "space": "bad value", + "tab": "bad\tvalue", + "non-ASCII": "bad-välue", + "unsupported colon": "bad:value", + "padding in middle": "bad=value", + "padding only": "==", + } { + t.Run(name, func(t *testing.T) { + client = newTestClient(t, profile, testSecrets{"local-embedding-key": secret}, base.Transport) + _, err = client.Embed(t.Context(), testInputs(), testAuthorization(profile.Descriptor)) + require.ErrorContains(t, err, "credential") + assert.NotContains(t, err.Error(), "bad") + }) + } + + for name, secret := range map[string]string{ + "minimal": "a", + "full alphabet and padding": "aZ09-._~+/==", + } { + t.Run(name, func(t *testing.T) { + client = newTestClient(t, profile, testSecrets{"local-embedding-key": secret}, base.Transport) + _, err = client.Embed(t.Context(), testInputs(), testAuthorization(profile.Descriptor)) + require.NoError(t, err) + }) + } +} + +func TestEmbedAcceptsJSONUTF8ContentType(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileGTE})) + for _, contentType := range []string{"application/json; charset=utf-8", "application/json; charset=Utf-8"} { + t.Run(contentType, func(t *testing.T) { + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, http.StatusOK, singleVectorResponse("synthetic-model", []float32{1, 0, 0})) + response.Header.Set("Content-Type", contentType) + return response, nil + })) + _, err := client.Embed(t.Context(), testInputs()[:1], testAuthorization(profile.Descriptor)) + require.NoError(t, err) + }) + } + + for _, contentType := range []string{ + "application/json; charset=iso-8859-1", + "application/json; profile=embedding", + "application/json; charset=utf-8; profile=embedding", + } { + t.Run(contentType, func(t *testing.T) { + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, http.StatusOK, singleVectorResponse("synthetic-model", []float32{1, 0, 0})) + response.Header.Set("Content-Type", contentType) + return response, nil + })) + _, err := client.Embed(t.Context(), testInputs()[:1], testAuthorization(profile.Descriptor)) + require.ErrorContains(t, err, "content type") + }) + } +} + +func TestEmbedRequiresModelAndProviderRevisionEcho(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileBGEM3})) + profile.DeploymentEpoch = "" + profile.ProviderRevisionHeader = "X-Model-Revision" + profile.Descriptor = descriptorFor(t, profile) + + for name, testCase := range map[string]struct { + body []byte + revisions []string + }{ + "missing revision": {body: singleVectorResponse("synthetic-model", []float32{1, 0, 0})}, + "wrong revision": {body: singleVectorResponse("synthetic-model", []float32{1, 0, 0}), revisions: []string{"revision-other"}}, + "multiple revisions": {body: singleVectorResponse("synthetic-model", []float32{1, 0, 0}), revisions: []string{"epoch-2026-08-26", "epoch-2026-08-26"}}, + "wrong model": {body: singleVectorResponse("alias-drift", []float32{1, 0, 0}), revisions: []string{"epoch-2026-08-26"}}, + } { + t.Run(name, func(t *testing.T) { + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, http.StatusOK, testCase.body) + if testCase.revisions != nil { + response.Header["X-Model-Revision"] = testCase.revisions + } + return response, nil + })) + _, err := client.Embed(t.Context(), testInputs()[:1], testAuthorization(profile.Descriptor)) + require.Error(t, err) + }) + } + + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, http.StatusOK, singleVectorResponse("synthetic-model", []float32{1, 0, 0})) + response.Header.Set("X-Model-Revision", " epoch-2026-08-26 ") + return response, nil + })) + _, err := client.Embed(t.Context(), testInputs()[:1], testAuthorization(profile.Descriptor)) + require.NoError(t, err) +} + +func TestEmbedRejectsSchemaVectorAndTransportDriftWithoutLeakingBodies(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + one := testInputs()[:1] + tests := []struct { + name string + status int + contentType string + body []byte + }{ + {name: "unknown response member", status: http.StatusOK, contentType: "application/json", body: schemaDriftResponse}, + {name: "wrong media type", status: http.StatusOK, contentType: "text/plain", body: singleVectorResponse("synthetic-model", []float32{1, 0, 0})}, + {name: "provider error body", status: http.StatusTooManyRequests, contentType: "application/json", body: []byte(`{"error":"private request text and vector [9,8,7]"}`)}, + {name: "wrong list object", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"embedding","data":[{"object":"embedding","embedding":[1,0,0],"index":0}],"model":"synthetic-model"}`)}, + {name: "wrong item object", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[{"object":"vector","embedding":[1,0,0],"index":0}],"model":"synthetic-model"}`)}, + {name: "duplicate index", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[{"object":"embedding","embedding":[1,0,0],"index":0},{"object":"embedding","embedding":[0,1,0],"index":0}],"model":"synthetic-model"}`)}, + {name: "missing vector", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[],"model":"synthetic-model"}`)}, + {name: "index outside batch", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[{"object":"embedding","embedding":[1,0,0],"index":2}],"model":"synthetic-model"}`)}, + {name: "dimension drift", status: http.StatusOK, contentType: "application/json", body: singleVectorResponse("synthetic-model", []float32{1, 0})}, + {name: "non-finite scalar", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[{"object":"embedding","embedding":[1e1000,0,0],"index":0}],"model":"synthetic-model"}`)}, + {name: "zero vector", status: http.StatusOK, contentType: "application/json", body: singleVectorResponse("synthetic-model", []float32{0, 0, 0})}, + {name: "normalization drift", status: http.StatusOK, contentType: "application/json", body: singleVectorResponse("synthetic-model", []float32{2, 0, 0})}, + {name: "negative usage", status: http.StatusOK, contentType: "application/json", body: []byte(`{"object":"list","data":[{"object":"embedding","embedding":[1,0,0],"index":0}],"model":"synthetic-model","usage":{"prompt_tokens":-1,"total_tokens":-1}}`)}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := jsonResponse(request, testCase.status, testCase.body) + response.Header.Set("Content-Type", testCase.contentType) + return response, nil + })) + inputs := one + if testCase.name == "duplicate index" { + inputs = testInputs() + } + _, err := client.Embed(t.Context(), inputs, testAuthorization(profile.Descriptor)) + require.Error(t, err) + assert.NotContains(t, err.Error(), "private request text") + assert.NotContains(t, err.Error(), "[9,8,7]") + }) + } + + bounded := profile + bounded.MaxResponseBytes = 64 + bounded.Descriptor = descriptorFor(t, bounded) + client := newTestClient(t, bounded, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return jsonResponse(request, http.StatusOK, []byte(strings.Repeat("x", 65))), nil + })) + _, err := client.Embed(t.Context(), one, mutateAuthorization(testAuthorization(bounded.Descriptor), func(value *document.EmbeddingAuthorization) { value.MaxResponseBytes = 64 })) + require.ErrorContains(t, err, "response byte") +} + +func TestEmbedRefusesRedirectsAndHonorsCancellation(t *testing.T) { + contract := modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileGTE}) + var redirected atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.URL.Path == "/elsewhere" { + redirected.Add(1) + } + http.Redirect(response, request, "/elsewhere", http.StatusTemporaryRedirect) + })) + t.Cleanup(server.Close) + profile := testProfile(t, contract) + profile.Origin = server.URL + profile.Descriptor = descriptorFor(t, profile) + client, err := New(profile, nil, server.Client()) + require.NoError(t, err) + _, err = client.Embed(t.Context(), testInputs(), testAuthorization(profile.Descriptor)) + require.ErrorContains(t, err, "redirect") + assert.Zero(t, redirected.Load()) + + started := make(chan struct{}) + canceledProfile := testProfile(t, contract) + canceledClient := newTestClient(t, canceledProfile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + close(started) + <-request.Context().Done() + return nil, request.Context().Err() + })) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + _, embedErr := canceledClient.Embed(ctx, testInputs(), testAuthorization(canceledProfile.Descriptor)) + done <- embedErr + }() + <-started + cancel() + err = <-done + require.ErrorIs(t, err, context.Canceled) +} + +func TestEmbedBoundsRequestBodyAndPreservesDescriptorSnapshot(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileE5})) + profile.MaxRequestBytes = 128 + profile.Descriptor = descriptorFor(t, profile) + var requests atomic.Int64 + client := newTestClient(t, profile, nil, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + return jsonResponse(request, http.StatusOK, singleVectorResponse("synthetic-model", []float32{1, 0, 0})), nil + })) + input := []document.EmbeddingInput{{Key: "large", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: strings.Repeat("x", 100)}} + authorization := mutateAuthorization(testAuthorization(profile.Descriptor), func(value *document.EmbeddingAuthorization) { value.MaxInputBytes = 200 }) + _, err := client.Embed(t.Context(), input, authorization) + require.ErrorIs(t, err, ErrCapacityResponse) + require.ErrorContains(t, err, "request byte") + assert.Zero(t, requests.Load()) + + descriptor := client.Descriptor() + descriptor.InputKinds[0] = document.EmbeddingInputOriginalFile + descriptor.SupportedRequestModes[0] = document.ModelInputModeQuery + assert.Equal(t, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, client.Descriptor().InputKinds) + assert.Equal(t, []document.ModelInputMode{document.ModelInputModeText}, client.Descriptor().SupportedRequestModes) +} + +func testProfile(t *testing.T, contract document.ModelInputContract) Profile { + t.Helper() + profile := Profile{ + Origin: "http://127.0.0.1:11434", ModelInput: contract, DeploymentEpoch: "epoch-2026-08-26", + RequestTimeout: time.Second, MaxBatchItems: 8, MaxInputBytes: 4096, + MaxRequestBytes: 8192, MaxResponseBytes: 16384, + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustOperatorNetwork, Model: "synthetic-model", + ModelRevision: "epoch-2026-08-26", Dimension: 3, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }, + } + profile.Descriptor = descriptorFor(t, profile) + return profile +} + +func descriptorFor(t *testing.T, profile Profile) document.EmbeddingDescriptor { + t.Helper() + profile.Descriptor.PolicyFingerprint = "" + profile.Descriptor.Fingerprint = "" + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return descriptor +} + +func canonicalMutatedDescriptor(t *testing.T, descriptor document.EmbeddingDescriptor, mutate func(*document.EmbeddingDescriptor)) document.EmbeddingDescriptor { + t.Helper() + descriptor.Fingerprint = "" + mutate(&descriptor) + canonical, err := document.NewEmbeddingDescriptor(descriptor) + require.NoError(t, err) + return canonical +} + +func modelInput(t *testing.T, config document.ModelInputContractConfig) document.ModelInputContract { + t.Helper() + contract, err := document.NewModelInputContract(config) + require.NoError(t, err) + return contract +} + +func testInputs() []document.EmbeddingInput { + return []document.EmbeddingInput{ + {Key: "document-1", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage"}, + {Key: "query-1", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "question"}, + } +} + +func testAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 2, MaxInputBytes: 64, MaxResponseBytes: 4096} +} + +func mutateAuthorization(value document.EmbeddingAuthorization, mutate func(*document.EmbeddingAuthorization)) document.EmbeddingAuthorization { + mutate(&value) + return value +} + +func newTestClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, &http.Client{Transport: transport}) + require.NoError(t, err) + return client +} + +func staticSuccessTransport() http.RoundTripper { + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + return jsonResponse(request, http.StatusOK, successIndexedResponse), nil + }) +} + +func jsonResponse(request *http.Request, status int, body []byte) *http.Response { + return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(string(body))), Request: request} +} + +func singleVectorResponse(model string, vector []float32) []byte { + body, err := json.Marshal(map[string]any{"object": "list", "data": []map[string]any{{"object": "embedding", "embedding": vector, "index": 0}}, "model": model}) + if err != nil { + panic(err) + } + return body +} diff --git a/document/openaiembed/profiles.go b/document/openaiembed/profiles.go new file mode 100644 index 00000000..262496dd --- /dev/null +++ b/document/openaiembed/profiles.go @@ -0,0 +1,193 @@ +package openaiembed + +import ( + "errors" + "fmt" + "strings" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + bgeM3Model = "BAAI/bge-m3" + + // PoolingCLS selects the first-token vector used by BGE-M3. + PoolingCLS = "cls" + // PoolingLastToken selects the final non-padding token used by Qwen3. + PoolingLastToken = "last_token" + + // OutputDenseSingleVector excludes sparse and multi-vector outputs. + OutputDenseSingleVector = "dense_single_vector" + // DimensionTransformNone requires the deployment's native output width. + DimensionTransformNone = "none" +) + +// DeploymentContract pins the model-side behavior that an +// OpenAI-compatible transport cannot report or enforce itself. +type DeploymentContract struct { + ModelFamily string `json:"model_family"` + WeightsRevision string `json:"weights_revision"` + Tokenizer string `json:"tokenizer"` + TokenizerRevision string `json:"tokenizer_revision"` + Pooling string `json:"pooling"` + MaxSequenceTokens int `json:"max_sequence_tokens"` + OutputMode string `json:"output_mode"` + DimensionTransform string `json:"dimension_transform"` +} + +// ReviewedProfileConfig supplies deployment-specific identity and transport +// bounds for a reviewed self-hosted embedding profile. +type ReviewedProfileConfig struct { + Origin string + ServedModel string + DeploymentEpoch string + WeightsRevision string + TokenizerRevision string + SecretBinding string + RequestTimeout time.Duration + MaxBatchItems int + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + EgressPolicy providerhttp.EgressPolicy +} + +// Qwen3Model identifies one reviewed Qwen3 Embedding model size. +type Qwen3Model string + +const ( + // Qwen3Embedding06B is the 0.6B model with a native 1024-dimensional output. + Qwen3Embedding06B Qwen3Model = "Qwen/Qwen3-Embedding-0.6B" + // Qwen3Embedding4B is the 4B model with a native 2560-dimensional output. + Qwen3Embedding4B Qwen3Model = "Qwen/Qwen3-Embedding-4B" + // Qwen3Embedding8B is the 8B model with a native 4096-dimensional output. + Qwen3Embedding8B Qwen3Model = "Qwen/Qwen3-Embedding-8B" +) + +// Qwen3ProfileConfig adds the model size and required retrieval instruction. +type Qwen3ProfileConfig struct { + ReviewedProfileConfig + + Model Qwen3Model + QueryInstruction string +} + +// BGEM3Profile builds the reviewed dense-only BGE-M3 profile. +func BGEM3Profile(config ReviewedProfileConfig) (Profile, error) { + input, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileBGEM3}) + if err != nil { + return Profile{}, fmt.Errorf("openaiembed: build BGE-M3 input contract: %w", err) + } + return reviewedProfile(config, input, 1024, DeploymentContract{ + ModelFamily: bgeM3Model, WeightsRevision: config.WeightsRevision, + Tokenizer: bgeM3Model, TokenizerRevision: config.TokenizerRevision, + Pooling: PoolingCLS, MaxSequenceTokens: 8192, + OutputMode: OutputDenseSingleVector, DimensionTransform: DimensionTransformNone, + }) +} + +// Qwen3Profile builds one reviewed dense-only Qwen3 Embedding profile at its +// native output dimension. +func Qwen3Profile(config Qwen3ProfileConfig) (Profile, error) { + if strings.TrimSpace(config.QueryInstruction) == "" { + return Profile{}, errors.New("openaiembed: Qwen3 query instruction is required") + } + dimension := map[Qwen3Model]int{ + Qwen3Embedding06B: 1024, + Qwen3Embedding4B: 2560, + Qwen3Embedding8B: 4096, + }[config.Model] + if dimension == 0 { + return Profile{}, fmt.Errorf("openaiembed: unsupported Qwen3 embedding model %q", config.Model) + } + input, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileQwen3, QueryInstruction: config.QueryInstruction, + }) + if err != nil { + return Profile{}, fmt.Errorf("openaiembed: build Qwen3 input contract: %w", err) + } + model := string(config.Model) + return reviewedProfile(config.ReviewedProfileConfig, input, dimension, DeploymentContract{ + ModelFamily: model, WeightsRevision: config.WeightsRevision, + Tokenizer: model, TokenizerRevision: config.TokenizerRevision, + Pooling: PoolingLastToken, MaxSequenceTokens: 32768, + OutputMode: OutputDenseSingleVector, DimensionTransform: DimensionTransformNone, + }) +} + +func reviewedProfile(config ReviewedProfileConfig, input document.ModelInputContract, dimension int, deployment DeploymentContract) (Profile, error) { + profile := Profile{ + Origin: config.Origin, ModelInput: input, DeploymentContract: &deployment, + SecretBinding: config.SecretBinding, DeploymentEpoch: config.DeploymentEpoch, + RequestTimeout: config.RequestTimeout, MaxBatchItems: config.MaxBatchItems, + MaxInputBytes: config.MaxInputBytes, MaxRequestBytes: config.MaxRequestBytes, + MaxResponseBytes: config.MaxResponseBytes, EgressPolicy: cloneEgressPolicy(config.EgressPolicy), + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustOperatorNetwork, Model: config.ServedModel, + ModelRevision: config.DeploymentEpoch, Dimension: dimension, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: input.CompatibilityID, SupportsTextQuery: true, ModelInput: input, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }, + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return Profile{}, err + } + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil { + return Profile{}, fmt.Errorf("openaiembed: build reviewed descriptor: %w", err) + } + return profile, nil +} + +func validateDeploymentContract(contract *DeploymentContract, descriptor document.EmbeddingDescriptor, input document.ModelInputContract) error { + if contract == nil { + return nil + } + if !immutableGitRevision(contract.WeightsRevision) || !immutableGitRevision(contract.TokenizerRevision) { + return errors.New("openaiembed: reviewed deployment revisions must be lowercase 40-character Git commits") + } + if contract.OutputMode != OutputDenseSingleVector || contract.DimensionTransform != DimensionTransformNone { + return errors.New("openaiembed: reviewed deployment must return one native dense vector") + } + if descriptor.Metric != document.VectorMetricCosine || descriptor.Normalization != document.VectorNormalizationUnitLength { + return errors.New("openaiembed: reviewed deployment requires cosine metric and unit-length vectors") + } + switch input.Profile { + case document.ModelInputProfileBGEM3: + if contract.ModelFamily != bgeM3Model || contract.Tokenizer != bgeM3Model || + contract.Pooling != PoolingCLS || contract.MaxSequenceTokens != 8192 || descriptor.Dimension != 1024 { + return errors.New("openaiembed: deployment does not match the reviewed BGE-M3 dense contract") + } + case document.ModelInputProfileQwen3: + dimension := map[string]int{string(Qwen3Embedding06B): 1024, string(Qwen3Embedding4B): 2560, string(Qwen3Embedding8B): 4096}[contract.ModelFamily] + if dimension == 0 || contract.Tokenizer != contract.ModelFamily || contract.Pooling != PoolingLastToken || + contract.MaxSequenceTokens != 32768 || descriptor.Dimension != dimension || input.QueryInstruction == "" { + return errors.New("openaiembed: deployment does not match a reviewed Qwen3 dense contract") + } + default: + return errors.New("openaiembed: deployment contract requires a reviewed BGE-M3 or Qwen3 input profile") + } + return nil +} + +func immutableGitRevision(value string) bool { + if len(value) != 40 { + return false + } + for _, character := range value { + if character < '0' || character > '9' { + if character < 'a' || character > 'f' { + return false + } + } + } + return true +} diff --git a/document/openaiembed/profiles_test.go b/document/openaiembed/profiles_test.go new file mode 100644 index 00000000..b7b33915 --- /dev/null +++ b/document/openaiembed/profiles_test.go @@ -0,0 +1,170 @@ +package openaiembed + +import ( + "net/http" + "net/netip" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + testWeightsRevision = "0123456789abcdef0123456789abcdef01234567" + testTokenizerRevision = "89abcdef0123456789abcdef0123456789abcdef" +) + +func TestBGEM3ProfilePinsReviewedDenseContract(t *testing.T) { + profile, err := BGEM3Profile(reviewedProfileConfig()) + require.NoError(t, err) + assert.Equal(t, "BAAI/bge-m3", profile.DeploymentContract.ModelFamily) + assert.Equal(t, "BAAI/bge-m3", profile.DeploymentContract.Tokenizer) + assert.Equal(t, PoolingCLS, profile.DeploymentContract.Pooling) + assert.Equal(t, 8192, profile.DeploymentContract.MaxSequenceTokens) + assert.Equal(t, OutputDenseSingleVector, profile.DeploymentContract.OutputMode) + assert.Equal(t, DimensionTransformNone, profile.DeploymentContract.DimensionTransform) + assert.Equal(t, 1024, profile.Descriptor.Dimension) + assert.Equal(t, document.VectorNormalizationUnitLength, profile.Descriptor.Normalization) + assert.Equal(t, "passage", profile.ModelInput.EncodeDocument("passage")) + assert.Equal(t, "question", profile.ModelInput.EncodeQuery("question")) + _, err = New(profile, nil, &http.Client{}) + require.NoError(t, err) +} + +func TestQwen3ProfilePinsReviewedDenseContracts(t *testing.T) { + for _, testCase := range []struct { + model Qwen3Model + dimension int + }{ + {Qwen3Embedding06B, 1024}, + {Qwen3Embedding4B, 2560}, + {Qwen3Embedding8B, 4096}, + } { + t.Run(string(testCase.model), func(t *testing.T) { + profile, err := Qwen3Profile(Qwen3ProfileConfig{ + ReviewedProfileConfig: reviewedProfileConfig(), Model: testCase.model, + QueryInstruction: "Retrieve supporting passages", + }) + require.NoError(t, err) + assert.Equal(t, string(testCase.model), profile.DeploymentContract.ModelFamily) + assert.Equal(t, string(testCase.model), profile.DeploymentContract.Tokenizer) + assert.Equal(t, PoolingLastToken, profile.DeploymentContract.Pooling) + assert.Equal(t, 32768, profile.DeploymentContract.MaxSequenceTokens) + assert.Equal(t, testCase.dimension, profile.Descriptor.Dimension) + assert.Equal(t, "passage", profile.ModelInput.EncodeDocument("passage")) + assert.Equal(t, "Instruct: Retrieve supporting passages\nQuery:question", profile.ModelInput.EncodeQuery("question")) + _, err = New(profile, nil, &http.Client{}) + require.NoError(t, err) + }) + } +} + +func TestReviewedProfilesRejectAmbiguousOrMutableIdentity(t *testing.T) { + for name, build := range map[string]func() error{ + "mutable weights": func() error { + config := reviewedProfileConfig() + config.WeightsRevision = "main" + _, err := BGEM3Profile(config) + return err + }, + "mutable tokenizer": func() error { + config := reviewedProfileConfig() + config.TokenizerRevision = "main" + _, err := BGEM3Profile(config) + return err + }, + "missing qwen instruction": func() error { + _, err := Qwen3Profile(Qwen3ProfileConfig{ReviewedProfileConfig: reviewedProfileConfig(), Model: Qwen3Embedding06B}) + return err + }, + "unknown qwen model": func() error { + _, err := Qwen3Profile(Qwen3ProfileConfig{ReviewedProfileConfig: reviewedProfileConfig(), Model: "Qwen/Qwen3-Embedding-2B", QueryInstruction: "Retrieve passages"}) + return err + }, + } { + t.Run(name, func(t *testing.T) { require.Error(t, build()) }) + } +} + +func TestDeploymentContractIsFingerprintBoundAndDenseOnly(t *testing.T) { + base, err := BGEM3Profile(reviewedProfileConfig()) + require.NoError(t, err) + baseFingerprint := base.Descriptor.PolicyFingerprint + + changed := reviewedProfileConfig() + changed.TokenizerRevision = "fedcba9876543210fedcba9876543210fedcba98" + profile, err := BGEM3Profile(changed) + require.NoError(t, err) + assert.NotEqual(t, baseFingerprint, profile.Descriptor.PolicyFingerprint) + + for name, mutate := range map[string]func(*DeploymentContract){ + "sparse output": func(contract *DeploymentContract) { contract.OutputMode = "sparse" }, + "multi-vector output": func(contract *DeploymentContract) { contract.OutputMode = "multi_vector" }, + "wrong pooling": func(contract *DeploymentContract) { contract.Pooling = PoolingLastToken }, + "dimension transform": func(contract *DeploymentContract) { contract.DimensionTransform = "truncate" }, + } { + t.Run(name, func(t *testing.T) { + mutated := base + contract := *base.DeploymentContract + mutate(&contract) + mutated.DeploymentContract = &contract + _, err := PolicyFingerprint(mutated) + require.Error(t, err) + }) + } + + for name, mutate := range map[string]func(*document.EmbeddingDescriptor){ + "metric": func(descriptor *document.EmbeddingDescriptor) { descriptor.Metric = document.VectorMetricDotProduct }, + "normalization": func(descriptor *document.EmbeddingDescriptor) { + descriptor.Normalization = document.VectorNormalizationNone + }, + } { + t.Run(name, func(t *testing.T) { + mutated := base + mutated.Descriptor = canonicalMutatedDescriptor(t, base.Descriptor, mutate) + _, err := PolicyFingerprint(mutated) + require.Error(t, err) + }) + } +} + +func TestReviewedProfileDoesNotAliasCallerEgressPolicy(t *testing.T) { + config := reviewedProfileConfig() + config.Origin = "https://model.local" + config.EgressPolicy = providerhttp.EgressPolicy{ + Scheme: "https", Host: "model.local", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("10.1.0.0/16"), netip.MustParsePrefix("10.0.0.0/8")}, + TLS: providerhttp.TLSPolicy{SPKISHA256: []string{strings.Repeat("b", 64), strings.Repeat("a", 64)}}, + } + wantCIDRs := slices.Clone(config.EgressPolicy.AllowedCIDRs) + wantPins := slices.Clone(config.EgressPolicy.TLS.SPKISHA256) + + profile, err := BGEM3Profile(config) + require.NoError(t, err) + assert.Equal(t, wantCIDRs, config.EgressPolicy.AllowedCIDRs) + assert.Equal(t, wantPins, config.EgressPolicy.TLS.SPKISHA256) + + config.EgressPolicy.AllowedCIDRs[0] = netip.MustParsePrefix("192.0.2.0/24") + config.EgressPolicy.TLS.SPKISHA256[0] = strings.Repeat("c", 64) + assert.NotEqual(t, config.EgressPolicy.AllowedCIDRs, profile.EgressPolicy.AllowedCIDRs) + assert.NotEqual(t, config.EgressPolicy.TLS.SPKISHA256, profile.EgressPolicy.TLS.SPKISHA256) + _, err = New(profile, nil, &http.Client{}) + require.NoError(t, err) +} + +func TestLegacyGenericProfileFingerprintIsStable(t *testing.T) { + profile := testProfile(t, modelInput(t, document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic})) + assert.Equal(t, "8c5ca6d380549ccb86894384cd006d136481968ea3b84a5a762d3054c1d46848", profile.Descriptor.PolicyFingerprint) +} + +func reviewedProfileConfig() ReviewedProfileConfig { + return ReviewedProfileConfig{ + Origin: "http://127.0.0.1:11434", ServedModel: "operator-model", DeploymentEpoch: "deployment-2026-08-28", + WeightsRevision: testWeightsRevision, TokenizerRevision: testTokenizerRevision, + } +} diff --git a/document/openaiembed/testdata/schema-drift.json b/document/openaiembed/testdata/schema-drift.json new file mode 100644 index 00000000..8026fee0 --- /dev/null +++ b/document/openaiembed/testdata/schema-drift.json @@ -0,0 +1 @@ +{"object":"list","data":[{"object":"embedding","embedding":[1,0,0],"index":0}],"model":"synthetic-model","usage":{"prompt_tokens":3,"total_tokens":3},"provider_private_extension":"must fail closed"} diff --git a/document/openaiembed/testdata/success-indexed.json b/document/openaiembed/testdata/success-indexed.json new file mode 100644 index 00000000..3bec28cd --- /dev/null +++ b/document/openaiembed/testdata/success-indexed.json @@ -0,0 +1 @@ +{"object":"list","data":[{"object":"embedding","embedding":[0,1,0],"index":1},{"object":"embedding","embedding":[1,0,0],"index":0}],"model":"synthetic-model","usage":{"prompt_tokens":6,"total_tokens":6}} diff --git a/document/openaihosted/client.go b/document/openaihosted/client.go new file mode 100644 index 00000000..fac66cf1 --- /dev/null +++ b/document/openaihosted/client.go @@ -0,0 +1,210 @@ +package openaihosted + +import ( + "bytes" + "context" + "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "net/http" + "slices" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" +) + +const ( + maxSecretBytes = 64 << 10 + unitLengthTolerance = 1e-4 +) + +var _ document.EmbeddingProvider = (*Client)(nil) + +type wireRequest struct { + Input []string `json:"input"` + Model string `json:"model"` + Dimensions int `json:"dimensions"` + EncodingFormat string `json:"encoding_format"` +} + +type wireResponse struct { + Object string `json:"object"` + Data []wireEmbedding `json:"data"` + Model string `json:"model"` + Usage *wireUsage `json:"usage"` +} + +type wireEmbedding struct { + Object string `json:"object"` + Embedding []float32 `json:"embedding"` + Index *int `json:"index"` +} + +type wireUsage struct { + PromptTokens *int64 `json:"prompt_tokens"` + TotalTokens *int64 `json:"total_tokens"` +} + +// Embed sends one bounded text-only request to the fixed hosted endpoint. +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if client == nil || ctx == nil { + return document.EmbeddingResult{}, errors.New("openaihosted: client and context are required") + } + if err := document.ValidateEmbeddingProviderRequest(client, inputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || authorization.MaxInputBytes > client.profile.MaxInputBytes || + authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrCapacityResponse} + } + rendered := make([]string, len(inputs)) + var total int64 + for index, input := range inputs { + if input.Role == document.EmbeddingRoleDocument { + rendered[index] = client.descriptor.ModelInput.EncodeDocument(input.Text) + } else { + rendered[index] = client.descriptor.ModelInput.EncodeQuery(input.Text) + } + length := int64(len(rendered[index])) + if length > client.profile.MaxInputItemBytes || length > client.profile.MaxInputBytes-total { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrCapacityResponse} + } + total += length + } + payload, err := json.Marshal(wireRequest{Input: rendered, Model: Model, Dimensions: client.descriptor.Dimension, EncodingFormat: "float"}) + if err != nil { + return document.EmbeddingResult{}, errors.New("openaihosted: request encoding failed") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrCapacityResponse} + } + + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + secret, err := client.secrets.ResolveSecret(requestCtx, client.profile.SecretBinding) + if err != nil || !validSecret(secret) { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("openaihosted: credential resolution canceled: %w", contextErr) + } + return document.EmbeddingResult{}, errors.New("openaihosted: API-key resolution failed") + } + request, err := http.NewRequestWithContext(requestCtx, http.MethodPost, origin+embeddingsPath, bytes.NewReader(payload)) + if err != nil { + return document.EmbeddingResult{}, errors.New("openaihosted: request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+secret) + response, err := client.http.Do(request) + if err != nil { + if contextErr := requestCtx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("openaihosted: request canceled: %w", contextErr) + } + return document.EmbeddingResult{}, &ProviderError{Kind: ErrTransientResponse} + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return document.EmbeddingResult{}, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if err := requireJSONContentType(response.Header.Get("Content-Type")); err != nil { + return document.EmbeddingResult{}, err + } + body, err := readBounded(requestCtx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return document.EmbeddingResult{}, err + } + var decoded wireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + result, err := client.validateAndOrder(decoded, inputs) + if err != nil { + return document.EmbeddingResult{}, err + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + return result, nil +} + +func (client *Client) validateAndOrder(response wireResponse, inputs []document.EmbeddingInput) (document.EmbeddingResult, error) { + if response.Object != "list" || response.Model != Model || response.Usage == nil || + response.Usage.PromptTokens == nil || response.Usage.TotalTokens == nil || + *response.Usage.PromptTokens < 0 || *response.Usage.TotalTokens < *response.Usage.PromptTokens || len(response.Data) != len(inputs) { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + vectors := make([]document.EmbeddingVector, len(inputs)) + seen := make([]bool, len(inputs)) + for _, item := range response.Data { + if item.Object != "embedding" || item.Index == nil || *item.Index < 0 || *item.Index >= len(inputs) || seen[*item.Index] { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + if err := client.validateVector(item.Embedding); err != nil { + return document.EmbeddingResult{}, err + } + seen[*item.Index] = true + vectors[*item.Index] = document.EmbeddingVector{Key: inputs[*item.Index].Key, Values: slices.Clone(item.Embedding)} + } + if slices.Contains(seen, false) { + return document.EmbeddingResult{}, &ProviderError{Kind: ErrPermanentResponse} + } + return document.EmbeddingResult{Vectors: vectors}, nil +} + +func (client *Client) validateVector(vector []float32) error { + if len(vector) != client.descriptor.Dimension { + return &ProviderError{Kind: ErrPermanentResponse} + } + var norm float64 + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return &ProviderError{Kind: ErrPermanentResponse} + } + norm += float64(value) * float64(value) + } + if math.Abs(norm-1) > unitLengthTolerance { + return &ProviderError{Kind: ErrPermanentResponse} + } + return nil +} + +func validSecret(value string) bool { + return value != "" && len(value) <= maxSecretBytes && utf8.ValidString(value) && strings.IndexFunc(value, func(character rune) bool { + return unicode.IsControl(character) || unicode.IsSpace(character) + }) < 0 +} + +func requireJSONContentType(value string) error { + mediaType, parameters, err := mime.ParseMediaType(value) + if err != nil || mediaType != "application/json" { + return &ProviderError{Kind: ErrPermanentResponse} + } + if len(parameters) == 0 { + return nil + } + charset, ok := parameters["charset"] + if len(parameters) != 1 || !ok || !strings.EqualFold(charset, "utf-8") { + return &ProviderError{Kind: ErrPermanentResponse} + } + return nil +} + +func readBounded(ctx context.Context, reader io.Reader, maximum int64) ([]byte, error) { + body, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("openaihosted: response read canceled: %w", contextErr) + } + return nil, &ProviderError{Kind: ErrTransientResponse} + } + if int64(len(body)) > maximum { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + return body, nil +} diff --git a/document/openaihosted/client_test.go b/document/openaihosted/client_test.go new file mode 100644 index 00000000..dc83ac95 --- /dev/null +++ b/document/openaihosted/client_test.go @@ -0,0 +1,269 @@ +package openaihosted + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +func TestEmbedRejectsStrictHostedResponseDrift(t *testing.T) { + valid := `{"object":"list","data":[{"object":"embedding","embedding":[1,0,0],"index":0},{"object":"embedding","embedding":[0,1,0],"index":1}],"model":"text-embedding-3-large","usage":{"prompt_tokens":2,"total_tokens":2}}` + tests := map[string]string{ + "root object": strings.Replace(valid, `"object":"list"`, `"object":"collection"`, 1), + "item object": strings.Replace(valid, `"object":"embedding"`, `"object":"vector"`, 1), + "model echo": strings.Replace(valid, Model, "text-embedding-3-small", 1), + "usage missing": strings.Replace(valid, `,"usage":{"prompt_tokens":2,"total_tokens":2}`, "", 1), + "prompt usage missing": strings.Replace(valid, `"prompt_tokens":2,`, "", 1), + "total usage missing": strings.Replace(valid, `,"total_tokens":2`, "", 1), + "usage negative": strings.Replace(valid, `"prompt_tokens":2`, `"prompt_tokens":-1`, 1), + "usage inconsistent": strings.Replace(valid, `"total_tokens":2`, `"total_tokens":1`, 1), + "missing index": strings.Replace(valid, `,"index":0`, "", 1), + "duplicate index": strings.Replace(valid, `"index":1`, `"index":0`, 1), + "outside index": strings.Replace(valid, `"index":1`, `"index":2`, 1), + "missing vector": strings.Replace(valid, `,{"object":"embedding","embedding":[0,1,0],"index":1}`, "", 1), + "extra vector": strings.Replace(valid, `],"model"`, `,{"object":"embedding","embedding":[0,0,1],"index":2}],"model"`, 1), + "wrong dimension": strings.Replace(valid, `[1,0,0]`, `[1,0]`, 1), + "non finite": strings.Replace(valid, `[1,0,0]`, `[1e1000,0,0]`, 1), + "base64 vector": strings.Replace(valid, `[1,0,0]`, `"AQID"`, 1), + "unknown root": strings.Replace(valid, `"usage":`, `"future":true,"usage":`, 1), + "unknown item": strings.Replace(valid, `"index":0`, `"future":true,"index":0`, 1), + "unknown usage": strings.Replace(valid, `"prompt_tokens":2`, `"future":true,"prompt_tokens":2`, 1), + "duplicate member": strings.Replace(valid, `"object":"list"`, `"object":"list","object":"list"`, 1), + "trailing JSON": valid + `{}`, + } + for name, body := range tests { + t.Run(name, func(t *testing.T) { + client := hostedClientReturning(t, http.StatusOK, "application/json", body) + _, err := client.Embed(context.Background(), twoHostedInputs(), hostedAuthorization(client.descriptor, 2)) + require.Error(t, err) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.NotContains(t, err.Error(), "alpha") + assert.NotContains(t, err.Error(), "AQID") + }) + } + + client := hostedClientReturning(t, http.StatusOK, "text/plain", valid) + _, err := client.Embed(context.Background(), twoHostedInputs(), hostedAuthorization(client.descriptor, 2)) + require.ErrorIs(t, err, ErrPermanentResponse) +} + +func TestEmbedEnforcesItemTotalRequestResponseAndBatchBounds(t *testing.T) { + tests := []struct { + name string + mutateProfile func(*Profile) + inputs []document.EmbeddingInput + authorization func(document.EmbeddingDescriptor) document.EmbeddingAuthorization + response string + wantRequest bool + }{ + { + name: "per item", mutateProfile: func(profile *Profile) { profile.MaxInputItemBytes = 10 }, + inputs: []document.EmbeddingInput{{Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "xx"}}, + }, + { + name: "total input", mutateProfile: func(profile *Profile) { profile.MaxInputItemBytes = 10; profile.MaxInputBytes = 15 }, + inputs: []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "a"}, + {Key: "second", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "b"}, + }, + }, + { + name: "request bytes", mutateProfile: func(profile *Profile) { profile.MaxRequestBytes = 64 }, + inputs: []document.EmbeddingInput{{Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "a"}}, + }, + { + name: "batch", mutateProfile: func(profile *Profile) { profile.MaxBatchItems = 1 }, + inputs: twoHostedInputs(), + }, + { + name: "response bytes", mutateProfile: func(profile *Profile) { profile.MaxResponseBytes = 64 }, + inputs: []document.EmbeddingInput{{Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "a"}}, + authorization: func(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + authorization := hostedAuthorization(descriptor, 1) + authorization.MaxResponseBytes = 64 + return authorization + }, + response: strings.Repeat("x", 65), wantRequest: true, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + profile := hostedTestProfile(t) + test.mutateProfile(&profile) + profile.Descriptor = hostedDescriptorFor(t, profile) + var requests atomic.Int32 + client := newHostedTestClient(t, profile, hostedSecrets{"secret:openai": "sk-synthetic"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + return hostedJSONResponse(request, http.StatusOK, test.response), nil + })) + authorization := hostedAuthorization(client.descriptor, len(test.inputs)) + if test.authorization != nil { + authorization = test.authorization(client.descriptor) + } + _, err := client.Embed(context.Background(), test.inputs, authorization) + require.Error(t, err) + if test.name == "response bytes" { + require.ErrorIs(t, err, ErrCapacityResponse) + } + if test.wantRequest { + assert.Equal(t, int32(1), requests.Load()) + } else { + assert.Zero(t, requests.Load()) + } + }) + } +} + +func TestEmbedRequiresValidNamedAPIKeyWithoutLeakingResolverDetails(t *testing.T) { + profile := hostedTestProfile(t) + tests := map[string]SecretResolver{ + "empty": hostedSecrets{"secret:openai": ""}, + "whitespace": hostedSecrets{"secret:openai": "sk synthetic"}, + "control": hostedSecrets{"secret:openai": "sk-synthetic\nPRIVATE"}, + "invalid UTF-8": hostedSecrets{"secret:openai": string([]byte{'s', 'k', '-', 0xff})}, + "resolver": failingSecrets{err: errors.New("PRIVATE_RESOLVER_DETAIL")}, + } + for name, resolver := range tests { + t.Run(name, func(t *testing.T) { + var requests atomic.Int32 + client := newHostedTestClient(t, profile, resolver, roundTripFunc(func(request *http.Request) (*http.Response, error) { + requests.Add(1) + return hostedJSONResponse(request, http.StatusInternalServerError, "PRIVATE_BODY"), nil + })) + _, err := client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.Error(t, err) + assert.Zero(t, requests.Load()) + assert.NotContains(t, err.Error(), "PRIVATE") + assert.NotContains(t, err.Error(), "synthetic") + }) + } +} + +func TestEmbedClassifiesHTTPAndTransportFailuresWithoutLeakingBodies(t *testing.T) { + tests := []struct { + name string + status int + retryAfter string + want error + delay time.Duration + delaySet bool + }{ + {name: "request timeout", status: http.StatusRequestTimeout, want: ErrTransientResponse}, + {name: "rate limit", status: http.StatusTooManyRequests, retryAfter: "7200", want: ErrTransientResponse, delay: time.Hour, delaySet: true}, + {name: "server", status: http.StatusInternalServerError, want: ErrTransientResponse}, + {name: "capacity", status: http.StatusRequestEntityTooLarge, want: ErrCapacityResponse}, + {name: "permanent", status: http.StatusBadRequest, want: ErrPermanentResponse}, + {name: "redirect", status: http.StatusTemporaryRedirect, want: ErrPermanentResponse}, + {name: "nonstandard 600", status: 600, want: ErrPermanentResponse}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + profile := hostedTestProfile(t) + client := newHostedTestClient(t, profile, hostedSecrets{"secret:openai": "sk-PRIVATE"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := hostedJSONResponse(request, test.status, "PRIVATE_PROVIDER_BODY") + response.Header.Set("Retry-After", test.retryAfter) + if test.status == http.StatusTemporaryRedirect { + response.Header.Set("Location", "https://example.com/steal") + } + return response, nil + })) + _, err := client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.ErrorIs(t, err, test.want) + assert.NotContains(t, err.Error(), "PRIVATE") + delay, set := RetryAfter(err) + assert.Equal(t, test.delaySet, set) + assert.Equal(t, test.delay, delay) + }) + } + + transportErr := errors.New("PRIVATE_TRANSPORT_DETAIL") + client := newHostedTestClient(t, hostedTestProfile(t), hostedSecrets{"secret:openai": "sk-PRIVATE"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, transportErr + })) + _, err := client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.ErrorIs(t, err, ErrTransientResponse) + assert.NotContains(t, err.Error(), "PRIVATE") + + client = newHostedTestClient(t, hostedTestProfile(t), hostedSecrets{"secret:openai": "sk-PRIVATE"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &failingBody{err: errors.New("PRIVATE_RESPONSE_READ_DETAIL")}, Request: request, + }, nil + })) + _, err = client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.ErrorIs(t, err, ErrTransientResponse) + assert.NotContains(t, err.Error(), "PRIVATE") +} + +func TestEmbedPreservesCancellationAndRefusesRedirectReplay(t *testing.T) { + started := make(chan struct{}) + client := newHostedTestClient(t, hostedTestProfile(t), hostedSecrets{"secret:openai": "sk-synthetic"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + close(started) + <-request.Context().Done() + return nil, request.Context().Err() + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := client.Embed(ctx, oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + done <- err + }() + <-started + cancel() + require.ErrorIs(t, <-done, context.Canceled) + + var calls atomic.Int32 + client = newHostedTestClient(t, hostedTestProfile(t), hostedSecrets{"secret:openai": "sk-synthetic"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + calls.Add(1) + return &http.Response{ + StatusCode: http.StatusTemporaryRedirect, + Header: http.Header{"Location": []string{"https://example.com/steal"}}, + Body: io.NopCloser(strings.NewReader("PRIVATE_REDIRECT_BODY")), Request: request, + }, nil + })) + _, err := client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.Equal(t, int32(1), calls.Load()) + assert.NotContains(t, err.Error(), "PRIVATE") +} + +func hostedClientReturning(t *testing.T, status int, mediaType, body string) *Client { + t.Helper() + return newHostedTestClient(t, hostedTestProfile(t), hostedSecrets{"secret:openai": "sk-synthetic"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + response := hostedJSONResponse(request, status, body) + response.Header.Set("Content-Type", mediaType) + return response, nil + })) +} + +func oneHostedInput() []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}} +} + +func twoHostedInputs() []document.EmbeddingInput { + return []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}, + {Key: "second", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "beta"}, + } +} + +type failingSecrets struct{ err error } + +func (resolver failingSecrets) ResolveSecret(context.Context, string) (string, error) { + return "", resolver.err +} + +type failingBody struct{ err error } + +func (body *failingBody) Read([]byte) (int, error) { return 0, body.err } + +func (*failingBody) Close() error { return nil } diff --git a/document/openaihosted/errors.go b/document/openaihosted/errors.go new file mode 100644 index 00000000..9d465999 --- /dev/null +++ b/document/openaihosted/errors.go @@ -0,0 +1,72 @@ +package openaihosted + +import ( + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" +) + +var ( + // ErrTransientResponse identifies retryable transport, 408, 429, and 5xx failures. + ErrTransientResponse = errors.New("openaihosted: transient provider response") + // ErrCapacityResponse identifies request or response capacity failures. + ErrCapacityResponse = errors.New("openaihosted: provider capacity exceeded") + // ErrPermanentResponse identifies non-retryable provider or schema failures. + ErrPermanentResponse = errors.New("openaihosted: permanent provider response") +) + +// ProviderError contains only stable classification and bounded retry metadata. +type ProviderError struct { + Kind error + StatusCode int + RetryDelay time.Duration + RetrySet bool +} + +func (err *ProviderError) Error() string { + if err.StatusCode != 0 { + return fmt.Sprintf("openaihosted: HTTP %d: %v", err.StatusCode, err.Kind) + } + return err.Kind.Error() +} + +func (err *ProviderError) Unwrap() error { return err.Kind } + +// RetryAfter returns bounded provider retry guidance when one was valid. +func RetryAfter(err error) (time.Duration, bool) { + providerErr, ok := errors.AsType[*ProviderError](err) + if !ok || !providerErr.RetrySet { + return 0, false + } + return providerErr.RetryDelay, true +} + +func statusError(status int, retryAfter string, now time.Time) error { + switch { + case status == http.StatusRequestEntityTooLarge: + return &ProviderError{Kind: ErrCapacityResponse, StatusCode: status} + case status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500 && status <= 599: + delay, set := parseRetryAfter(retryAfter, now) + return &ProviderError{Kind: ErrTransientResponse, StatusCode: status, RetryDelay: delay, RetrySet: set} + default: + return &ProviderError{Kind: ErrPermanentResponse, StatusCode: status} + } +} + +func parseRetryAfter(value string, now time.Time) (time.Duration, bool) { + value = strings.TrimSpace(value) + if seconds, err := strconv.ParseInt(value, 10, 64); err == nil && seconds >= 0 { + if seconds >= int64(time.Hour/time.Second) { + return time.Hour, true + } + return time.Duration(seconds) * time.Second, true + } + when, err := http.ParseTime(value) + if err != nil { + return 0, false + } + return min(max(when.Sub(now), 0), time.Hour), true +} diff --git a/document/openaihosted/profile.go b/document/openaihosted/profile.go new file mode 100644 index 00000000..a39f9dd7 --- /dev/null +++ b/document/openaihosted/profile.go @@ -0,0 +1,319 @@ +// Package openaihosted implements the fixed hosted OpenAI embeddings contract. +package openaihosted + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "net/http" + "net/netip" + "reflect" + "slices" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + // ProviderID is the fixed E1 adapter identity for hosted OpenAI embeddings. + ProviderID = "openai.hosted.text-embedding-3-large-v1" + // Model is the only currently documented hosted model alias admitted here. + Model = "text-embedding-3-large" + // DocumentFormatterV1 identifies the document rendering path. + DocumentFormatterV1 = "openai-hosted/document/v1" + // QueryFormatterV1 identifies the query rendering path. + QueryFormatterV1 = "openai-hosted/query/v1" + // ScalarEncodingFloat32 is the only response representation admitted here. + ScalarEncodingFloat32 = "float32" + + host = "api.openai.com" + origin = "https://api.openai.com" + embeddingsPath = "/v1/embeddings" + adapterContract = "docbank-openai-hosted-embeddings/v1" + defaultTimeout = 30 * time.Second + defaultBatch = 128 + defaultItemBytes = int64(1 << 20) + defaultInputBytes = int64(16 << 20) + defaultRequest = int64(32 << 20) + defaultResponse = int64(64 << 20) + maximumTimeout = 5 * time.Minute + maximumBatch = 2_048 + maximumBytes = int64(1 << 30) + maximumTokenBytes = 128 +) + +// SecretResolver resolves only the configured named OpenAI API-key binding. +type SecretResolver interface { + ResolveSecret(ctx context.Context, binding string) (string, error) +} + +// Profile freezes the exact hosted contract and its execution bounds. +type Profile struct { + Descriptor document.EmbeddingDescriptor + CompatibilityEpoch string + SecretBinding string + RequestTimeout time.Duration + MaxBatchItems int + MaxInputItemBytes int64 + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + EgressPolicy providerhttp.EgressPolicy +} + +// Client calls only the fixed hosted OpenAI embeddings endpoint. +type Client struct { + profile Profile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + CompatibilityEpoch string `json:"compatibility_epoch"` + SecretBinding string `json:"secret_binding"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputItemBytes int64 `json:"max_input_item_bytes"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + Egress egressIdentity `json:"egress"` +} + +type egressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +// PolicyFingerprint returns the canonical hosted profile identity. +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptorIdentity, err := normalizeProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(policyIdentity{ + AdapterContract: adapterContract, Origin: origin, Route: embeddingsPath, + Descriptor: descriptorIdentity, CompatibilityEpoch: normalized.CompatibilityEpoch, + SecretBinding: normalized.SecretBinding, RequestTimeout: int64(normalized.RequestTimeout), + MaxBatchItems: normalized.MaxBatchItems, MaxInputItemBytes: normalized.MaxInputItemBytes, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, Egress: profileEgressIdentity(normalized.EgressPolicy), + }, json.Deterministic(true)) + if err != nil { + return "", errors.New("openaihosted: policy identity encoding failed") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +// New validates the fixed hosted profile and replaces every supplied transport +// authority with the sealed provider HTTP transport. +func New(profile Profile, secrets SecretResolver, resolver providerhttp.Resolver, supplied *http.Client) (*Client, error) { + if supplied == nil { + return nil, errors.New("openaihosted: HTTP client settings source is required") + } + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + return nil, errors.New("openaihosted: descriptor is not canonical") + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("openaihosted: descriptor policy fingerprint does not match profile") + } + if nilInterface(secrets) { + return nil, errors.New("openaihosted: named API-key resolver is required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, errors.New("openaihosted: sealed egress policy is invalid") + } + isolated := *supplied + isolated.Transport = transport + isolated.CheckRedirect = providerhttp.RefuseRedirects + isolated.Jar = nil + isolated.Timeout = 0 + normalized.Descriptor = cloneDescriptor(descriptor) + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, http: &isolated}, nil +} + +// Descriptor returns an immutable copy of the exact vector-space contract. +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + profile.EgressPolicy.TLS.SPKISHA256 = slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultTimeout + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultBatch + } + if profile.MaxInputItemBytes == 0 { + profile.MaxInputItemBytes = defaultItemBytes + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultRequest + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultResponse + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > maximumTimeout || + profile.MaxBatchItems < 1 || profile.MaxBatchItems > maximumBatch || + profile.MaxInputItemBytes < 1 || profile.MaxInputItemBytes > maximumBytes || + profile.MaxInputBytes < profile.MaxInputItemBytes || profile.MaxInputBytes > maximumBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maximumBytes || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maximumBytes { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaihosted: execution bounds are invalid") + } + if !validToken(profile.CompatibilityEpoch) || profile.Descriptor.ModelRevision != profile.CompatibilityEpoch { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaihosted: compatibility epoch must exactly match descriptor model revision") + } + if !validToken(profile.SecretBinding) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaihosted: named API-key binding is required") + } + if err := normalizeEgress(&profile.EgressPolicy); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + descriptorIdentity := cloneDescriptor(profile.Descriptor) + descriptorIdentity.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptorIdentity.Fingerprint = "" + var err error + descriptorIdentity, err = document.NewEmbeddingDescriptor(descriptorIdentity) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("openaihosted: descriptor identity is invalid") + } + descriptorIdentity.PolicyFingerprint = "" + descriptorIdentity.Fingerprint = "" + if err := validateDescriptorContract(descriptorIdentity, profile.CompatibilityEpoch); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + return profile, descriptorIdentity, nil +} + +func validateDescriptorContract(descriptor document.EmbeddingDescriptor, epoch string) error { + if descriptor.ID != ProviderID || descriptor.ContractVersion != document.EmbeddingProviderContractVersion || + descriptor.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptor.Model != Model || + descriptor.ModelRevision != epoch || descriptor.Dimension < 1 || descriptor.Metric != document.VectorMetricCosine || + descriptor.Normalization != document.VectorNormalizationUnitLength || descriptor.ScalarEncoding != ScalarEncodingFloat32 || + descriptor.DocumentFormatter != DocumentFormatterV1 || descriptor.QueryFormatter != QueryFormatterV1 || + !descriptor.SupportsTextQuery || !slices.Equal(descriptor.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || + !slices.Equal(descriptor.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeText}) || + descriptor.ModelInput.Document.Mode != document.ModelInputModeText || descriptor.ModelInput.Query.Mode != document.ModelInputModeText || + descriptor.CompatibilityID != descriptor.ModelInput.CompatibilityID { + return errors.New("openaihosted: descriptor does not match the fixed hosted text contract") + } + return nil +} + +func normalizeEgress(policy *providerhttp.EgressPolicy) error { + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = providerhttp.DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if policy.ProxyMode == "" { + policy.ProxyMode = providerhttp.ProxyDisabled + } + if policy.Scheme != "https" || policy.Host != host || policy.Port != 443 || + policy.ProxyMode != providerhttp.ProxyDisabled || policy.TLS.RootCAs != nil { + return errors.New("openaihosted: egress authority must be exactly api.openai.com:443 with system roots and no proxy") + } + for index := range policy.AllowedCIDRs { + policy.AllowedCIDRs[index] = policy.AllowedCIDRs[index].Masked() + } + slices.SortFunc(policy.AllowedCIDRs, func(left, right netip.Prefix) int { + return strings.Compare(left.String(), right.String()) + }) + for index := 1; index < len(policy.AllowedCIDRs); index++ { + if policy.AllowedCIDRs[index] == policy.AllowedCIDRs[index-1] { + return errors.New("openaihosted: egress policy has a duplicate CIDR") + } + } + for index := range policy.TLS.SPKISHA256 { + policy.TLS.SPKISHA256[index] = strings.ToLower(policy.TLS.SPKISHA256[index]) + } + slices.Sort(policy.TLS.SPKISHA256) + for index := 1; index < len(policy.TLS.SPKISHA256); index++ { + if policy.TLS.SPKISHA256[index] == policy.TLS.SPKISHA256[index-1] { + return errors.New("openaihosted: egress policy has a duplicate SPKI pin") + } + } + if _, err := providerhttp.NewTransport(*policy, nil); err != nil { + return errors.New("openaihosted: sealed egress policy is invalid") + } + return nil +} + +func profileEgressIdentity(policy providerhttp.EgressPolicy) egressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.String() + } + return egressIdentity{ + Scheme: policy.Scheme, Host: policy.Host, Port: policy.Port, AllowedCIDRs: cidrs, + ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), + KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), + SPKISHA256: slices.Clone(policy.TLS.SPKISHA256), + } +} + +func validToken(value string) bool { + return value != "" && len(value) <= maximumTokenBytes && utf8.ValidString(value) && value == strings.TrimSpace(value) && + strings.IndexFunc(value, func(character rune) bool { return unicode.IsControl(character) || unicode.IsSpace(character) }) < 0 +} + +func cloneDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/openaihosted/profile_test.go b/document/openaihosted/profile_test.go new file mode 100644 index 00000000..cd712778 --- /dev/null +++ b/document/openaihosted/profile_test.go @@ -0,0 +1,317 @@ +package openaihosted + +import ( + "bytes" + "context" + "crypto/x509" + "errors" + "io" + "net/http" + "net/netip" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +func TestNewRequiresTheFixedHostedDescriptorEpochAndCredential(t *testing.T) { + profile := hostedTestProfile(t) + resolver := hostedSecrets{"secret:openai": "sk-synthetic"} + + client, err := New(profile, resolver, staticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + assert.Equal(t, profile.Descriptor, client.Descriptor()) + + tests := map[string]func(*Profile){ + "provider id": func(value *Profile) { value.Descriptor.ID = "openai-compatible.embeddings-v1" }, + "contract version": func(value *Profile) { value.Descriptor.ContractVersion++ }, + "hosted trust": func(value *Profile) { value.Descriptor.TrustBoundary = document.EmbeddingTrustOperatorNetwork }, + "fixed model": func(value *Profile) { value.Descriptor.Model = "text-embedding-3-small" }, + "positive dimension": func(value *Profile) { value.Descriptor.Dimension = 0 }, + "metric": func(value *Profile) { value.Descriptor.Metric = document.VectorMetricL2 }, + "normalization": func(value *Profile) { value.Descriptor.Normalization = document.VectorNormalizationNone }, + "scalar": func(value *Profile) { value.Descriptor.ScalarEncoding = "float64" }, + "document formatter": func(value *Profile) { value.Descriptor.DocumentFormatter = "custom" }, + "query formatter": func(value *Profile) { value.Descriptor.QueryFormatter = "custom" }, + "input kinds": func(value *Profile) { + value.Descriptor.InputKinds = []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile} + }, + "request modes": func(value *Profile) { + value.Descriptor.SupportedRequestModes = []document.ModelInputMode{document.ModelInputModeDocument} + }, + "query support": func(value *Profile) { value.Descriptor.SupportsTextQuery = false }, + "compatibility": func(value *Profile) { value.Descriptor.CompatibilityID = "forged-space" }, + "revision epoch": func(value *Profile) { value.Descriptor.ModelRevision = "different-epoch" }, + "missing epoch": func(value *Profile) { value.CompatibilityEpoch = "" }, + "missing binding": func(value *Profile) { value.SecretBinding = "" }, + "whitespace binding": func(value *Profile) { value.SecretBinding = "secret: openai" }, + "control binding": func(value *Profile) { value.SecretBinding = "secret:openai\n" }, + } + for name, mutate := range tests { + t.Run(name, func(t *testing.T) { + changed := hostedTestProfile(t) + changed.Descriptor.Fingerprint = "" + mutate(&changed) + if changed.Descriptor.Dimension > 0 { + changed.Descriptor = canonicalDescriptorWithoutPolicyCheck(t, changed.Descriptor) + } + _, err := New(changed, resolver, staticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.Error(t, err) + }) + } + + _, err = New(profile, nil, staticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.Error(t, err, "a named API-key binding always requires a resolver") + _, err = New(profile, resolver, staticResolver{netip.MustParseAddr("192.0.2.10")}, nil) + require.Error(t, err, "the harmless client settings source is explicit") +} + +func TestHostedBatchLimitAccepts2048AndRejects2049BeforeRequestAuthority(t *testing.T) { + accepted := hostedTestProfile(t) + accepted.MaxBatchItems = 2048 + accepted.Descriptor = hostedDescriptorFor(t, accepted) + + _, err := New(accepted, hostedSecrets{"secret:openai": "sk-synthetic"}, staticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + + rejected := accepted + rejected.MaxBatchItems = 2049 + _, _, err = normalizeProfile(rejected) + require.Error(t, err) +} + +func TestPolicyFingerprintCoversFixedContractBoundsAndExactEgressWithoutMutatingCaller(t *testing.T) { + profile := hostedTestProfile(t) + originalCIDRs := slices.Clone(profile.EgressPolicy.AllowedCIDRs) + originalPins := slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + base, err := PolicyFingerprint(profile) + require.NoError(t, err) + + mutations := map[string]func(*Profile){ + "epoch": func(value *Profile) { + value.CompatibilityEpoch = "hosted-epoch-2" + value.Descriptor.ModelRevision = "hosted-epoch-2" + }, + "secret binding": func(value *Profile) { value.SecretBinding = "secret:other" }, + "dimension": func(value *Profile) { value.Descriptor.Dimension++ }, + "model input": func(value *Profile) { + contract, contractErr := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileNomic}) + require.NoError(t, contractErr) + value.Descriptor.ModelInput = contract + value.Descriptor.CompatibilityID = contract.CompatibilityID + }, + "batch": func(value *Profile) { value.MaxBatchItems++ }, + "per item": func(value *Profile) { value.MaxInputItemBytes++ }, + "total input": func(value *Profile) { value.MaxInputBytes++ }, + "request": func(value *Profile) { value.MaxRequestBytes++ }, + "response": func(value *Profile) { value.MaxResponseBytes++ }, + "timeout": func(value *Profile) { value.RequestTimeout += time.Second }, + "CIDR": func(value *Profile) { + value.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24")} + }, + "SPKI": func(value *Profile) { + value.EgressPolicy.TLS.SPKISHA256 = []string{"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"} + }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changed := hostedTestProfile(t) + mutate(&changed) + fingerprint, fingerprintErr := PolicyFingerprint(changed) + require.NoError(t, fingerprintErr) + assert.NotEqual(t, base, fingerprint) + }) + } + assert.Equal(t, originalCIDRs, profile.EgressPolicy.AllowedCIDRs) + assert.Equal(t, originalPins, profile.EgressPolicy.TLS.SPKISHA256) + + invalid := hostedTestProfile(t) + invalid.EgressPolicy.Host = "example.com" + _, err = PolicyFingerprint(invalid) + require.Error(t, err) + invalid = hostedTestProfile(t) + invalid.EgressPolicy.Port = 8443 + _, err = PolicyFingerprint(invalid) + require.Error(t, err) + invalid = hostedTestProfile(t) + invalid.EgressPolicy.TLS.RootCAs = x509.NewCertPool() + _, err = PolicyFingerprint(invalid) + require.Error(t, err, "custom roots expand hosted trust authority") +} + +func TestNewReplacesCallerTransportJarRedirectAndTimeoutAuthority(t *testing.T) { + profile := hostedTestProfile(t) + originalCIDRs := slices.Clone(profile.EgressPolicy.AllowedCIDRs) + originalPins := slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + ambient := &countingTransport{} + client, err := New(profile, hostedSecrets{"secret:openai": "sk-synthetic"}, + staticResolver{netip.MustParseAddr("203.0.113.7")}, + &http.Client{Transport: ambient, CheckRedirect: func(*http.Request, []*http.Request) error { return nil }, Timeout: time.Hour}) + require.NoError(t, err) + assert.Equal(t, originalCIDRs, profile.EgressPolicy.AllowedCIDRs) + assert.Equal(t, originalPins, profile.EgressPolicy.TLS.SPKISHA256) + assert.NotSame(t, ambient, client.http.Transport) + assert.Zero(t, client.http.Timeout) + assert.Nil(t, client.http.Jar) + require.NotNil(t, client.http.CheckRedirect) + require.ErrorIs(t, client.http.CheckRedirect(new(http.Request), nil), http.ErrUseLastResponse) + _, err = client.Embed(context.Background(), oneHostedInput(), hostedAuthorization(client.descriptor, 1)) + require.ErrorIs(t, err, ErrTransientResponse) + assert.Zero(t, ambient.calls, "the caller transport must never receive hosted authority") +} + +func TestEmbedSendsExactHostedRequestAndRestoresReturnedIndices(t *testing.T) { + profile := hostedTestProfile(t) + inputs := []document.EmbeddingInput{ + {Key: "document-a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}, + {Key: "query-a", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "beta"}, + } + originalInputs := slices.Clone(inputs) + var calls int + client := newHostedTestClient(t, profile, hostedSecrets{"secret:openai": "sk-synthetic"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + calls++ + assert.Equal(t, http.MethodPost, request.Method) + assert.Equal(t, "https://api.openai.com/v1/embeddings", request.URL.String()) + assert.Equal(t, "api.openai.com", request.URL.Host) + assert.Equal(t, "application/json", request.Header.Get("Accept")) + assert.Equal(t, "application/json", request.Header.Get("Content-Type")) + assert.Equal(t, "Bearer sk-synthetic", request.Header.Get("Authorization")) + assert.Empty(t, request.Header.Get("Cookie")) + payload, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.JSONEq(t, `{ + "input":["passage: alpha","query: beta"], + "model":"text-embedding-3-large", + "dimensions":3, + "encoding_format":"float" + }`, string(payload)) + return hostedJSONResponse(request, http.StatusOK, `{ + "object":"list", + "data":[ + {"object":"embedding","embedding":[0,1,0],"index":1}, + {"object":"embedding","embedding":[1,0,0],"index":0} + ], + "model":"text-embedding-3-large", + "usage":{"prompt_tokens":2,"total_tokens":2} + }`), nil + })) + + result, err := client.Embed(context.Background(), inputs, hostedAuthorization(profile.Descriptor, 2)) + require.NoError(t, err) + assert.Equal(t, 1, calls) + assert.Equal(t, originalInputs, inputs, "embedding must not mutate caller slices") + assert.Equal(t, document.EmbeddingResult{Vectors: []document.EmbeddingVector{ + {Key: "document-a", Values: []float32{1, 0, 0}}, + {Key: "query-a", Values: []float32{0, 1, 0}}, + }}, result) +} + +func hostedTestProfile(t *testing.T) Profile { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{Profile: document.ModelInputProfileE5}) + require.NoError(t, err) + profile := Profile{ + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustHostedProvider, Model: Model, + ModelRevision: "hosted-epoch-1", Dimension: 3, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeText}, + }, + CompatibilityEpoch: "hosted-epoch-1", SecretBinding: "secret:openai", + RequestTimeout: time.Second, MaxBatchItems: 8, MaxInputItemBytes: 2048, + MaxInputBytes: 4096, MaxRequestBytes: 8192, MaxResponseBytes: 16384, + EgressPolicy: providerhttp.EgressPolicy{ + Scheme: "https", Host: "api.openai.com", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24"), netip.MustParsePrefix("192.0.2.0/24")}, + ProxyMode: providerhttp.ProxyDisabled, ConnectTimeout: time.Second, + KeepAlive: time.Second, TLSHandshakeTimeout: time.Second, + TLS: providerhttp.TLSPolicy{SPKISHA256: []string{ + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + }}, + }, + } + profile.Descriptor = hostedDescriptorFor(t, profile) + return profile +} + +func hostedDescriptorFor(t *testing.T, profile Profile) document.EmbeddingDescriptor { + t.Helper() + profile.Descriptor.PolicyFingerprint = "" + profile.Descriptor.Fingerprint = "" + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return descriptor +} + +func canonicalDescriptorWithoutPolicyCheck(t *testing.T, value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + t.Helper() + canonical, err := document.NewEmbeddingDescriptor(value) + if err != nil { + return value + } + return canonical +} + +type hostedSecrets map[string]string + +func (secrets hostedSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + value, ok := secrets[name] + if !ok { + return "", errors.New("synthetic secret missing") + } + return value, nil +} + +type staticResolver []netip.Addr + +func (resolver staticResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return slices.Clone(resolver), nil +} + +type countingTransport struct{ calls int } + +func (transport *countingTransport) RoundTrip(*http.Request) (*http.Response, error) { + transport.calls++ + return nil, errors.New("ambient transport must not run") +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +func newHostedTestClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, staticResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + client.http.Transport = transport + return client +} + +func hostedAuthorization(descriptor document.EmbeddingDescriptor, batch int) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: batch, + MaxInputBytes: 4096, MaxResponseBytes: 4096, + } +} + +func hostedJSONResponse(request *http.Request, status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewBufferString(body)), Request: request, + } +} diff --git a/document/plaintext/provider.go b/document/plaintext/provider.go new file mode 100644 index 00000000..e4da2851 --- /dev/null +++ b/document/plaintext/provider.go @@ -0,0 +1,220 @@ +// Package plaintext implements the bounded in-process UTF-8 rendition provider. +package plaintext + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "slices" + "strconv" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" +) + +const ( + providerID = "plaintext.in-process-v1" + profileVersion = "docbank-plaintext-profile/v1" + timestampForm = "2006-01-02T15:04:05.000000000Z" + + // MaxDocumentBytes preserves the released bounded plain-text extraction limit. + MaxDocumentBytes = int64(16 << 20) +) + +// Profile fixes the maximum exact upload size accepted by one provider instance. +type Profile struct { + MaxDocumentBytes int64 +} + +// Provider renders verified UTF-8 bytes without network access. +type Provider struct { + descriptor document.RenditionDescriptor + maxDocumentBytes int64 +} + +// New constructs one immutable local provider profile. +func New(profile Profile) (*Provider, error) { + if profile.MaxDocumentBytes <= 0 || profile.MaxDocumentBytes > MaxDocumentBytes { + return nil, fmt.Errorf("plaintext: max document bytes must be between 1 and %d", MaxDocumentBytes) + } + policyDigest := sha256.Sum256([]byte(profileVersion + "\x00" + + strconv.FormatInt(profile.MaxDocumentBytes, 10))) + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, + ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: hex.EncodeToString(policyDigest[:]), + TrustBoundary: document.RenditionTrustLocalProcess, + SupportedFormats: supportedFormats(), + ReturnsStructured: true, + ArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + }) + if err != nil { + return nil, fmt.Errorf("plaintext: construct descriptor: %w", err) + } + return &Provider{descriptor: cloneDescriptor(descriptor), maxDocumentBytes: profile.MaxDocumentBytes}, nil +} + +// Descriptor returns the immutable provider identity fixed by the profile. +func (provider *Provider) Descriptor() document.RenditionDescriptor { + if provider == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(provider.descriptor) +} + +// Render reads and re-verifies one authorized upload, then emits one exact +// generic evidence unit. The provider never opens a path or performs I/O +// beyond the supplied read-once upload. +func (provider *Provider) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if provider == nil { + return document.RenditionResult{}, errors.New("plaintext: provider is required") + } + if _, err := document.ValidateRenditionProviderRequest(provider, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > provider.maxDocumentBytes { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "input exceeds the plain-text byte limit", nil) + } + if err := ctx.Err(); err != nil { + return document.RenditionResult{}, classifiedError(document.RenditionErrorCanceled, + "plain-text rendering canceled", err) + } + + startedAt := time.Now().UTC() + data, err := readExact(ctx, upload, metadata.ByteLength, provider.maxDocumentBytes) + if err != nil { + return document.RenditionResult{}, err + } + digest := sha256.Sum256(data) + if hex.EncodeToString(digest[:]) != metadata.SHA256 { + return document.RenditionResult{}, classifiedError(document.RenditionErrorPolicyRejected, + "authorized upload identity mismatch", nil) + } + if !utf8.Valid(data) || bytes.IndexByte(data, 0) >= 0 { + return document.RenditionResult{}, classifiedError(document.RenditionErrorUnsupportedInput, + "input is not UTF-8 plain text", nil) + } + completedAt := time.Now().UTC() + return document.RenditionResult{ + Evidence: document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceDegradedProvenance, + Family: authorization.MediaFamily, + UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_provenance", + Reason: "plain-text provider emits one generic unit", + }}, + Units: []document.SourceEvidenceUnitV1{{ + Order: 0, Text: string(data), + Locator: document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorGeneric, IndexOrigin: document.EvidenceIndexOriginNone, + }, + }}, + }, + Receipt: document.RenditionReceipt{ + ProviderID: provider.descriptor.ID, DescriptorFingerprint: provider.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "plaintext-" + authorization.RenditionRequestFingerprint[:24], + StartedAt: startedAt.Format(timestampForm), CompletedAt: completedAt.Format(timestampForm), + Warnings: []string{"degraded_provenance"}, + Usage: document.RenditionUsage{ + Requests: 1, InputBytes: int64(len(data)), OutputBytes: int64(len(data)), Units: 1, + }, + }, + }, nil +} + +func readExact( + ctx context.Context, reader io.Reader, expectedBytes, maxBytes int64, +) ([]byte, error) { + data := make([]byte, 0, expectedBytes) + buffer := make([]byte, 32<<10) + for { + if err := ctx.Err(); err != nil { + return nil, classifiedError(document.RenditionErrorCanceled, + "plain-text rendering canceled", err) + } + read, err := reader.Read(buffer) + if read > 0 { + if int64(len(data))+int64(read) > maxBytes { + return nil, classifiedError(document.RenditionErrorPolicyRejected, + "input exceeds the plain-text byte limit", nil) + } + data = append(data, buffer[:read]...) + } + switch { + case errors.Is(err, io.EOF): + if int64(len(data)) != expectedBytes { + return nil, classifiedError(document.RenditionErrorPolicyRejected, + "authorized upload identity mismatch", nil) + } + return data, nil + case err != nil: + if contextErr := ctx.Err(); contextErr != nil { + return nil, classifiedError(document.RenditionErrorCanceled, + "plain-text rendering canceled", contextErr) + } + return nil, classifiedError(document.RenditionErrorTransient, + "could not read the authorized upload", err) + case read == 0: + return nil, classifiedError(document.RenditionErrorTransient, + "authorized upload stopped making progress", io.ErrNoProgress) + } + } +} + +func classifiedError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("plaintext: classify provider error: %w", err) + } + return providerError +} + +func supportedFormats() []document.RenditionFormatCapability { + formats := []struct { + family string + mediaType string + }{ + {family: "mail", mediaType: "message/rfc822"}, + {family: "source", mediaType: "text/javascript"}, + {family: "source", mediaType: "text/x-go"}, + {family: "source", mediaType: "text/x-python"}, + {family: "spreadsheet", mediaType: "text/csv"}, + {family: "structured", mediaType: "application/json"}, + {family: "structured", mediaType: "application/x-ndjson"}, + {family: "structured", mediaType: "application/xml"}, + {family: "structured", mediaType: "application/yaml"}, + {family: "text", mediaType: "application/x-tex"}, + {family: "text", mediaType: "text/markdown"}, + {family: "text", mediaType: "text/plain"}, + {family: "text", mediaType: "text/x-rst"}, + } + result := make([]document.RenditionFormatCapability, 0, len(formats)) + for _, format := range formats { + result = append(result, document.RenditionFormatCapability{ + MediaFamily: format.family, MediaType: format.mediaType, + InputKind: document.RenditionInputOriginalFile, + }) + } + return result +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +var _ document.RenditionProvider = (*Provider)(nil) diff --git a/document/plaintext/provider_test.go b/document/plaintext/provider_test.go new file mode 100644 index 00000000..03767aa4 --- /dev/null +++ b/document/plaintext/provider_test.go @@ -0,0 +1,198 @@ +package plaintext + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +func TestProviderRendersExactUTF8AsOneGenericUnit(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + source := []byte("alpha\nβeta\n") + upload := newTestUpload(source) + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + + result, err := document.RenderRendition(t.Context(), provider, upload, authorization) + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 1) + assert.Equal(t, string(source), result.Evidence.Units[0].Text) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + assert.Equal(t, document.EvidenceLocatorGeneric, result.Evidence.Units[0].Locator.Kind) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, int64(len(source)), result.Receipt.Usage.InputBytes) + assert.Equal(t, authorization.SourceSHA256, result.Receipt.SourceSHA256) + + policy, err := document.NewEvidencePolicy(1024) + require.NoError(t, err) + normalized, err := document.NormalizeEvidenceV1(result.Evidence, policy) + require.NoError(t, err) + _, checksum, err := document.MarshalNormalizedEvidenceV1(normalized) + require.NoError(t, err) + assert.Equal(t, "unit_2bd7c59368d78ac033cd5ca2b3cf879c4023c4564f5b4f3b05548f3414e4271f", + normalized.Units[0].ID) + assert.Equal(t, "a7242bb9ba14c427bfafabdc43e8171ea58c839d7bfec2f26ec802c7406939c7", checksum) +} + +func TestProviderRejectsInvalidUTF8AndNUL(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + for _, testCase := range []struct { + name string + data []byte + }{ + {name: "invalid UTF-8", data: []byte{0xff, 0xfe}}, + {name: "NUL", data: []byte("alpha\x00beta")}, + } { + t.Run(testCase.name, func(t *testing.T) { + upload := newTestUpload(testCase.data) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorUnsupportedInput, providerErr.Code()) + }) + } +} + +func TestProviderRejectsEmptyInputBeforeReading(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + upload := newTestUpload(nil) + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + + _, err = provider.Render(t.Context(), upload, authorization) + require.ErrorContains(t, err, "byte length") + assert.Zero(t, upload.reads) +} + +func TestProviderEnforcesProfileSizeBeforeReading(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 4}) + require.NoError(t, err) + upload := newTestUpload([]byte("alpha")) + + _, err = provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorPolicyRejected, providerErr.Code()) + assert.Zero(t, upload.reads) +} + +func TestProviderHonorsCancellationBeforeReading(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + upload := newTestUpload([]byte("alpha")) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = provider.Render(ctx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.ErrorIs(t, err, context.Canceled) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorCanceled, providerErr.Code()) + assert.Zero(t, upload.reads) +} + +func TestProviderRejectsUploadByteSubstitution(t *testing.T) { + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + upload := newTestUpload([]byte("omega")) + metadata := upload.Metadata() + want := sha256.Sum256([]byte("alpha")) + metadata.SHA256 = hex.EncodeToString(want[:]) + upload.metadata = metadata + + _, err = provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, document.RenditionErrorPolicyRejected, providerErr.Code()) +} + +func TestNewRejectsInvalidBoundsAndDescriptorIsImmutable(t *testing.T) { + _, err := New(Profile{}) + require.ErrorContains(t, err, "max document bytes") + _, err = New(Profile{MaxDocumentBytes: MaxDocumentBytes + 1}) + require.ErrorContains(t, err, "max document bytes") + + provider, err := New(Profile{MaxDocumentBytes: 1024}) + require.NoError(t, err) + descriptor := provider.Descriptor() + assert.Equal(t, document.RenditionTrustLocalProcess, descriptor.TrustBoundary) + assert.True(t, descriptor.ReturnsStructured) + assert.False(t, descriptor.ReturnsMarkdown) + require.NotEmpty(t, descriptor.SupportedFormats) + descriptor.SupportedFormats[0].MediaType = "application/pdf" + assert.NotEqual(t, descriptor.SupportedFormats, provider.Descriptor().SupportedFormats) +} + +type testUpload struct { + reader *bytes.Reader + metadata document.AuthorizedUploadMetadata + reads int +} + +func newTestUpload(data []byte) *testUpload { + digest := sha256.Sum256(data) + return &testUpload{ + reader: bytes.NewReader(data), + metadata: document.AuthorizedUploadMetadata{ + Filename: "notes.txt", MediaFamily: "text", MediaType: "text/plain", + ByteLength: int64(len(data)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + }, + } +} + +func (upload *testUpload) Read(buffer []byte) (int, error) { + upload.reads++ + read, err := upload.reader.Read(buffer) + if err != nil { + return read, fmt.Errorf("read test upload: %w", err) + } + return read, nil +} + +func (*testUpload) Close() error { return nil } + +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +var _ document.AuthorizedUpload = (*testUpload)(nil) +var _ io.ReadCloser = (*testUpload)(nil) + +func testAuthorization( + descriptor document.RenditionDescriptor, metadata document.AuthorizedUploadMetadata, +) document.RenditionAuthorization { + started := time.Now().UTC().Add(-time.Minute) + return document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, + InputKind: metadata.InputKind, MaxTotalResultBytes: 1 << 20, + AuthorizedAt: started.Format("2006-01-02T15:04:05.000000000Z"), + ExpiresAt: started.Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + } +} diff --git a/document/profile.go b/document/profile.go index 4378be33..2c3550bd 100644 --- a/document/profile.go +++ b/document/profile.go @@ -121,6 +121,7 @@ type EmbeddingBindingV1 struct { MaxInputBytes int64 `json:"max_input_bytes"` MaxResponseBytes int64 `json:"max_response_bytes"` Metric string `json:"metric"` + ModelInput ModelInputContract `json:"model_input"` Model string `json:"model"` Name string `json:"name"` Normalization string `json:"normalization"` @@ -189,21 +190,23 @@ type renditionRequestIdentity struct { } type embeddingInputIdentity struct { - Chunk *EmbeddingChunkPolicyV1 `json:"chunk"` - EvidenceLexical string `json:"evidence_lexical"` - InputKind EmbeddingInputKind `json:"input_kind"` + Chunk *EmbeddingChunkPolicyV1 `json:"chunk"` + EvidenceLexical string `json:"evidence_lexical"` + InputKind EmbeddingInputKind `json:"input_kind"` + ModelInputFingerprint string `json:"model_input_fingerprint"` } type vectorSpaceIdentity struct { - CompatibilityID string `json:"compatibility_id"` - Descriptor ProviderDescriptorV1 `json:"descriptor"` - Dimensions int `json:"dimensions"` - DocumentFormatter string `json:"document_formatter"` - Metric string `json:"metric"` - Model string `json:"model"` - Normalization string `json:"normalization"` - QueryFormatter string `json:"query_formatter"` - ScalarEncoding string `json:"scalar_encoding"` + CompatibilityID string `json:"compatibility_id"` + Descriptor ProviderDescriptorV1 `json:"descriptor"` + Dimensions int `json:"dimensions"` + DocumentFormatter string `json:"document_formatter"` + Metric string `json:"metric"` + ModelInputFingerprint string `json:"model_input_fingerprint"` + Model string `json:"model"` + Normalization string `json:"normalization"` + QueryFormatter string `json:"query_formatter"` + ScalarEncoding string `json:"scalar_encoding"` } type providerDisclosureIdentity struct { @@ -258,7 +261,7 @@ func CanonicalProfile(profile ProcessingProfileV1) ([]byte, FingerprintSet, erro return nil, FingerprintSet{}, err } for _, binding := range canonical.Embeddings { - input := embeddingInputIdentity{InputKind: binding.InputKind, Chunk: binding.Chunk} + input := embeddingInputIdentity{InputKind: binding.InputKind, Chunk: binding.Chunk, ModelInputFingerprint: binding.ModelInput.Fingerprint} if binding.InputKind == EmbeddingInputRenditionChunk { input.EvidenceLexical = evidenceFingerprint } @@ -269,7 +272,8 @@ func CanonicalProfile(profile ProcessingProfileV1) ([]byte, FingerprintSet, erro result.VectorSpace[binding.Name], err = componentFingerprint("vector_space", vectorSpaceIdentity{ CompatibilityID: binding.CompatibilityID, Descriptor: binding.Descriptor, Dimensions: binding.Dimensions, DocumentFormatter: binding.DocumentFormatter, Metric: binding.Metric, Model: binding.Model, - Normalization: binding.Normalization, QueryFormatter: binding.QueryFormatter, ScalarEncoding: binding.ScalarEncoding, + ModelInputFingerprint: binding.ModelInput.Fingerprint, Normalization: binding.Normalization, + QueryFormatter: binding.QueryFormatter, ScalarEncoding: binding.ScalarEncoding, }) if err != nil { return nil, FingerprintSet{}, err @@ -316,6 +320,13 @@ func canonicalProcessingProfile(profile ProcessingProfileV1) (ProcessingProfileV chunk := *profile.Embeddings[index].Chunk canonical.Embeddings[index].Chunk = &chunk } + if canonical.Embeddings[index].ModelInput == (ModelInputContract{}) { + empty, err := NewModelInputContract(ModelInputContractConfig{}) + if err != nil { + return ProcessingProfileV1{}, err + } + canonical.Embeddings[index].ModelInput = empty + } } if err := normalizeProfileStrings(&canonical); err != nil { return ProcessingProfileV1{}, err @@ -525,6 +536,15 @@ func validateEmbeddingBinding(binding EmbeddingBindingV1, hasRendition bool) err if err := validateProviderDescriptor(binding.Descriptor); err != nil { return fmt.Errorf("descriptor: %w", err) } + if err := validateModelInputContract(binding.ModelInput); err != nil { + return fmt.Errorf("model input: %w", err) + } + if binding.ModelInput.Profile != "" && binding.ModelInput.CompatibilityID != binding.CompatibilityID { + return errors.New("model input compatibility ID does not match embedding compatibility ID") + } + if !validVectorNormalization(binding.Normalization) { + return errors.New("embedding normalization is invalid") + } for subject, value := range map[string]string{"authorization fingerprint": binding.AuthorizationFingerprint, "disclosure fingerprint": binding.DisclosureFingerprint} { if err := validateFingerprint(value, subject); err != nil { return err diff --git a/document/profile_test.go b/document/profile_test.go index b99bc011..ad32172d 100644 --- a/document/profile_test.go +++ b/document/profile_test.go @@ -26,16 +26,16 @@ func TestCanonicalProfileMatchesGoldenAndCanonicalizesInput(t *testing.T) { want = bytes.TrimSuffix(want, []byte("\n")) assert.Equal(t, want, encoded) assert.Equal(t, original, profile, "canonicalization must not mutate caller-owned policy") - assert.Equal(t, "389b38a03edb9c156e4804473d96accbef7d0e773f49387b84b1b7e6ade33b91", fingerprints.Profile) + assert.Equal(t, "38a06e593e2fe3588e88f76e842e3138184237c5a3d8052872fe865d117f0e1f", fingerprints.Profile) assert.Equal(t, "9d0a202be29b43e16684f74b540a41778443fe2f4757d4b1815d85994f5b2522", fingerprints.RenditionRequest) assert.Equal(t, "f5405d835d1bb377cade66ca6e18b488972fadc862bfd96faf746cd374184819", fingerprints.EvidenceLexical) assert.Equal(t, map[string]string{ - "direct": "59c57d1c5b5b3106ce6310660b0ec8aa038c58c20bdcaad9ef711c40b0ae9ff1", - "semantic": "dff3236afbf4b8810a1ff40a1c812f9ae7bef5dae28c08b6a80f78e811094872", + "direct": "6d0ee0f814c74992f6bcedfcab8b52f07c0024077526919ff5c2c336d9f353b2", + "semantic": "0a363d012a8caeb222de2a8eab8692fcf30a3d733a39f76199ebdd0bf364cf45", }, fingerprints.EmbeddingInput) assert.Equal(t, map[string]string{ - "direct": "dc2da50bb2d163a114834bbcee2ea439cd838da23f4e0ac8a504d9b007487357", - "semantic": "9bb1db44ac603c26f0f4028b009d50e2ca83cef5b61fa4fd2dad4108b31a541b", + "direct": "f731443d0c91d8b4a548ebf821090522fb3dbe772a4d3985ef9336f60c8f801e", + "semantic": "94b86507a61e14509a5c744f1db31e94d84cbfd61a62f2d8cf5ba893834f900d", }, fingerprints.VectorSpace) assert.Equal(t, "3debf7c7ff983edfebee3fe195f5db2a2d32dae3ff3f64eb834e382925e2a1ef", fingerprints.RetentionDisclosure) @@ -164,12 +164,34 @@ func TestCanonicalProfileFingerprintsTrackExactLayerFields(t *testing.T) { } } +// This test fails if changing the sealed document/query envelope leaves either +// the input-generation or vector-space authority unchanged. +func TestCanonicalProfileModelInputRotatesInputAndVectorSpace(t *testing.T) { + base := syntheticProcessingProfileV1() + _, before, err := document.CanonicalProfile(base) + require.NoError(t, err) + + modelInput, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "voyage-3-large/1024", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "document: {{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "query: {{content}}"}, + }) + require.NoError(t, err) + changed := cloneProcessingProfile(base) + changed.Embeddings[1].ModelInput = modelInput + _, after, err := document.CanonicalProfile(changed) + require.NoError(t, err) + assert.ElementsMatch(t, layers("profile", "input:semantic", "vector:semantic"), changedFingerprintLayers(before, after)) +} + func TestCanonicalProfileRejectsInvalidPolicy(t *testing.T) { tests := []struct { name string mutate func(*document.ProcessingProfileV1) want string }{ + {"l2 embedding normalization alias", func(p *document.ProcessingProfileV1) { p.Embeddings[0].Normalization = "l2" }, "normalization"}, + {"unknown embedding normalization", func(p *document.ProcessingProfileV1) { p.Embeddings[0].Normalization = "provider_default" }, "normalization"}, {"duplicate binding name", func(p *document.ProcessingProfileV1) { p.Embeddings[1].Name = "direct" }, "duplicated"}, {"unknown input kind", func(p *document.ProcessingProfileV1) { p.Embeddings[0].InputKind = "summary" }, "input kind"}, {"chunk embedding without rendition", func(p *document.ProcessingProfileV1) { diff --git a/document/provider.go b/document/provider.go new file mode 100644 index 00000000..8b822bc2 --- /dev/null +++ b/document/provider.go @@ -0,0 +1,1014 @@ +package document + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "mime" + "reflect" + "slices" + "strings" + "sync" + "time" +) + +const ( + // RenditionProviderContractVersion identifies the provider boundary defined here. + RenditionProviderContractVersion = 1 + + maxRenditionFormats = 64 + maxRenditionArtifactRoles = 16 + maxRenditionArtifacts = 64 + maxRenditionSourceBytes = int64(1 << 40) + maxRenditionMarkdownBytes = 64 << 20 + maxRenditionArtifactBytes = 256 << 20 + maxRenditionTotalResultBytes = 512 << 20 + maxRenditionWarnings = 64 + maxRenditionUsageValue = int64(1 << 50) + renditionTimestampForm = "2006-01-02T15:04:05.000000000Z" +) + +// RenditionProvider renders one authorized upload into provider-neutral evidence. +type RenditionProvider interface { + Descriptor() RenditionDescriptor + Render(ctx context.Context, upload AuthorizedUpload, authorization RenditionAuthorization) (RenditionResult, error) +} + +// RenditionResumeHandle is an opaque, provider-issued durable operation +// identity. Core persists it only after the provider checkpoints it; callers +// must never derive a handle from source or job identity. +type RenditionResumeHandle struct { + Value string +} + +// RenditionResumeCheckpoint durably records a provider-issued handle before +// the provider continues work whose outcome may otherwise become ambiguous. +type RenditionResumeCheckpoint func(RenditionResumeHandle) error + +// ResumableRenditionProvider is the narrow optional contract for providers +// that can continue a known durable operation without resubmitting source +// bytes. A nil handle starts new work; a non-nil handle resumes exactly that +// provider-issued operation. On resume, core validates the sealed upload +// metadata but passes a nil upload to the provider, making source resubmission +// unavailable through this contract. +type ResumableRenditionProvider interface { + RenditionProvider + RenderResumable( + ctx context.Context, upload AuthorizedUpload, authorization RenditionAuthorization, + resume *RenditionResumeHandle, checkpoint RenditionResumeCheckpoint, + ) (RenditionResult, error) +} + +// AuthorizedUpload is a read-once upload with immutable, authorization-bound metadata. +type AuthorizedUpload interface { + io.ReadCloser + Metadata() AuthorizedUploadMetadata +} + +// RenditionTrustBoundary identifies where provider-controlled processing occurs. +type RenditionTrustBoundary string + +const ( + // RenditionTrustLocalProcess keeps processing within a local child process. + RenditionTrustLocalProcess RenditionTrustBoundary = "local_process" + // RenditionTrustOperatorNetwork sends input to operator-controlled infrastructure. + RenditionTrustOperatorNetwork RenditionTrustBoundary = "operator_network" + // RenditionTrustHostedProvider sends input to a third-party provider. + RenditionTrustHostedProvider RenditionTrustBoundary = "hosted_provider" +) + +// RenditionInputKind identifies the exact representation accepted by a provider. +type RenditionInputKind string + +const ( + // RenditionInputOriginalFile sends the exact source file. + RenditionInputOriginalFile RenditionInputKind = "original_file" + // RenditionInputDerivedUpload sends a separately authorized derived representation. + RenditionInputDerivedUpload RenditionInputKind = "derived_upload" +) + +// RenditionFormatCapability is one exact supported media-family, media-type, and input-kind tuple. +type RenditionFormatCapability struct { + MediaFamily string `json:"media_family"` + MediaType string `json:"media_type"` + InputKind RenditionInputKind `json:"input_kind"` +} + +// RenditionDescriptor is the canonical immutable identity of a provider contract. +type RenditionDescriptor struct { + ID string `json:"id"` + ContractVersion int `json:"contract_version"` + PolicyFingerprint string `json:"policy_fingerprint"` + TrustBoundary RenditionTrustBoundary `json:"trust_boundary"` + SupportedFormats []RenditionFormatCapability `json:"supported_formats"` + ReturnsMarkdown bool `json:"returns_markdown"` + ReturnsStructured bool `json:"returns_structured"` + ArtifactRoles []EvidenceArtifactRole `json:"artifact_roles"` + Fingerprint string `json:"fingerprint"` +} + +type renditionDescriptorIdentity struct { + ID string `json:"id"` + ContractVersion int `json:"contract_version"` + PolicyFingerprint string `json:"policy_fingerprint"` + TrustBoundary RenditionTrustBoundary `json:"trust_boundary"` + SupportedFormats []RenditionFormatCapability `json:"supported_formats"` + ReturnsMarkdown bool `json:"returns_markdown"` + ReturnsStructured bool `json:"returns_structured"` + ArtifactRoles []EvidenceArtifactRole `json:"artifact_roles"` +} + +// NewRenditionDescriptor validates, canonicalizes, and fingerprints a descriptor. +func NewRenditionDescriptor(value RenditionDescriptor) (RenditionDescriptor, error) { + value.Fingerprint = "" + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + slices.SortFunc(value.SupportedFormats, compareRenditionFormats) + slices.Sort(value.ArtifactRoles) + if err := validateRenditionDescriptorFields(value); err != nil { + return RenditionDescriptor{}, err + } + encoded, err := canonicalJSON(descriptorIdentity(value)) + if err != nil { + return RenditionDescriptor{}, fmt.Errorf("encode rendition descriptor: %w", err) + } + value.Fingerprint = sha256Hex(encoded) + return value, nil +} + +// AuthorizedUploadMetadata describes the exact bytes presented to a provider. +type AuthorizedUploadMetadata struct { + Filename string `json:"filename"` + MediaFamily string `json:"media_family"` + MediaType string `json:"media_type"` + ByteLength int64 `json:"byte_length"` + SHA256 string `json:"sha256"` + CapabilityRecordChecksum string `json:"capability_record_checksum"` + ProviderMetadataChecksum string `json:"provider_metadata_checksum"` + InputKind RenditionInputKind `json:"input_kind"` +} + +// RenditionAuthorization binds one provider invocation to exact input and output limits. +type RenditionAuthorization struct { + ProviderID string `json:"provider_id"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + RenditionRequestFingerprint string `json:"rendition_request_fingerprint"` + SourceSHA256 string `json:"source_sha256"` + SourceBytes int64 `json:"source_bytes"` + CapabilityRecordChecksum string `json:"capability_record_checksum"` + ProviderMetadataChecksum string `json:"provider_metadata_checksum"` + MediaFamily string `json:"media_family"` + MediaType string `json:"media_type"` + InputKind RenditionInputKind `json:"input_kind"` + AllowedArtifactRoles []EvidenceArtifactRole `json:"allowed_artifact_roles"` + MaxProviderMarkdownBytes int `json:"max_provider_markdown_bytes"` + MaxArtifactBytes int `json:"max_artifact_bytes"` + MaxArtifacts int `json:"max_artifacts"` + MaxTotalResultBytes int `json:"max_total_result_bytes"` + AuthorizedAt string `json:"authorized_at"` + ExpiresAt string `json:"expires_at"` +} + +// RenditionArtifact is one bounded, checksum-addressed provider output. +type RenditionArtifact struct { + Role EvidenceArtifactRole `json:"role"` + MediaType string `json:"media_type"` + Payload []byte `json:"payload"` + SHA256 string `json:"sha256"` +} + +// RenditionUsage contains bounded numeric provider accounting only. +type RenditionUsage struct { + Requests int64 `json:"requests"` + Retries int64 `json:"retries"` + InputBytes int64 `json:"input_bytes"` + OutputBytes int64 `json:"output_bytes"` + Units int64 `json:"units"` +} + +// RenditionReceipt is a sanitized, bounded execution record without provider bodies or secrets. +type RenditionReceipt struct { + ProviderID string `json:"provider_id"` + DescriptorFingerprint string `json:"descriptor_fingerprint"` + PolicyFingerprint string `json:"policy_fingerprint"` + SourceSHA256 string `json:"source_sha256"` + OperationID string `json:"operation_id"` + StartedAt string `json:"started_at"` + CompletedAt string `json:"completed_at"` + Warnings []string `json:"warnings,omitempty"` + Usage RenditionUsage `json:"usage"` + RetryDelayMillis int64 `json:"retry_delay_millis,omitempty"` +} + +// RenditionResult contains bounded provider output and its sanitized receipt. +type RenditionResult struct { + Evidence SourceEvidenceV1 `json:"evidence"` + ProviderMarkdown []byte `json:"provider_markdown,omitempty"` + Artifacts []RenditionArtifact `json:"artifacts,omitempty"` + Receipt RenditionReceipt `json:"receipt"` +} + +// ValidateRenditionProviderRequest validates the provider snapshot and exact upload authorization. +func ValidateRenditionProviderRequest( + provider RenditionProvider, upload AuthorizedUpload, authorization RenditionAuthorization, +) (RenditionDescriptor, error) { + return ValidateRenditionProviderRequestAt(time.Now().UTC(), provider, upload, authorization) +} + +// ValidateRenditionProviderRequestAt validates a request against an explicit +// trusted clock. Callers with a transaction clock can avoid time-of-check +// drift while retaining the same expiry boundary. +func ValidateRenditionProviderRequestAt( + now time.Time, provider RenditionProvider, upload AuthorizedUpload, + authorization RenditionAuthorization, +) (RenditionDescriptor, error) { + if nilInterface(provider) { + return RenditionDescriptor{}, errors.New("rendition provider is required") + } + if nilInterface(upload) { + return RenditionDescriptor{}, errors.New("authorized upload is required") + } + authorization = cloneRenditionAuthorization(authorization) + descriptor := cloneRenditionDescriptor(provider.Descriptor()) + if err := validateRenditionDescriptor(descriptor); err != nil { + return RenditionDescriptor{}, err + } + if second := cloneRenditionDescriptor(provider.Descriptor()); !equalRenditionDescriptors(descriptor, second) { + return RenditionDescriptor{}, errors.New("rendition descriptor changed during validation") + } + metadata := upload.Metadata() + if err := validateAuthorizedUploadMetadata(metadata); err != nil { + return RenditionDescriptor{}, err + } + if second := upload.Metadata(); second != metadata { + return RenditionDescriptor{}, errors.New("authorized upload metadata changed during validation") + } + if err := validateRenditionAuthorization(descriptor, metadata, authorization); err != nil { + return RenditionDescriptor{}, err + } + if err := validateAuthorizationCurrentAt(authorization, now); err != nil { + return RenditionDescriptor{}, err + } + return cloneRenditionDescriptor(descriptor), nil +} + +// RenderRendition validates immutable boundary snapshots, gives the provider +// a separate authorization copy, and validates its result against the sealed +// copy. Provider mutation therefore cannot broaden its own output authority. +func RenderRendition( + ctx context.Context, provider RenditionProvider, upload AuthorizedUpload, + authorization RenditionAuthorization, +) (RenditionResult, error) { + sealed := cloneRenditionAuthorization(authorization) + descriptor, err := validateCurrentRenditionRequest(ctx, provider, upload, sealed) + if err != nil { + return RenditionResult{}, err + } + result, err := provider.Render(ctx, upload, cloneRenditionAuthorization(sealed)) + if err != nil { + if classified := ValidateRenditionProviderError(err); classified != nil { + return RenditionResult{}, classified + } + return RenditionResult{}, err + } + if err := ValidateRenditionResult(descriptor, sealed, result); err != nil { + return RenditionResult{}, err + } + return result, nil +} + +// RenderRenditionWithResume applies the ordinary sealed request/result +// contract while allowing an optional provider-neutral durable resume handle. +// Providers without the optional contract may start work but cannot consume a +// persisted handle. +func RenderRenditionWithResume( + ctx context.Context, provider RenditionProvider, upload AuthorizedUpload, + authorization RenditionAuthorization, resume *RenditionResumeHandle, + checkpoint RenditionResumeCheckpoint, +) (RenditionResult, error) { + sealed := cloneRenditionAuthorization(authorization) + descriptor, err := validateCurrentRenditionRequest(ctx, provider, upload, sealed) + if err != nil { + return RenditionResult{}, err + } + resumable, supportsResume := provider.(ResumableRenditionProvider) + if !supportsResume { + if resume != nil { + return RenditionResult{}, errors.New("rendition provider does not support durable resume") + } + result, renderErr := provider.Render(ctx, upload, cloneRenditionAuthorization(sealed)) + if renderErr != nil { + if classified := ValidateRenditionProviderError(renderErr); classified != nil { + return RenditionResult{}, classified + } + return RenditionResult{}, renderErr + } + if err := ValidateRenditionResult(descriptor, sealed, result); err != nil { + return RenditionResult{}, err + } + return result, nil + } + var resumeCopy *RenditionResumeHandle + if resume != nil { + if err := validateRenditionResumeHandle(*resume); err != nil { + return RenditionResult{}, err + } + resumeValue := *resume + resumeCopy = &resumeValue + } + if checkpoint == nil { + checkpoint = func(RenditionResumeHandle) error { return nil } + } + var checkpointMu sync.Mutex + var checkpointErr error + recordCheckpointError := func(err error) error { + if err == nil { + return nil + } + checkpointMu.Lock() + if checkpointErr == nil { + checkpointErr = err + } + checkpointMu.Unlock() + return err + } + checkedCheckpoint := func(handle RenditionResumeHandle) error { + if err := validateRenditionResumeHandle(handle); err != nil { + return recordCheckpointError(err) + } + return recordCheckpointError(checkpoint(handle)) + } + providerUpload := upload + if resumeCopy != nil { + providerUpload = nil + } + result, err := resumable.RenderResumable( + ctx, providerUpload, cloneRenditionAuthorization(sealed), resumeCopy, checkedCheckpoint, + ) + checkpointMu.Lock() + durableErr := checkpointErr + checkpointMu.Unlock() + if durableErr != nil { + return RenditionResult{}, durableErr + } + if err != nil { + if classified := ValidateRenditionProviderError(err); classified != nil { + return RenditionResult{}, classified + } + return RenditionResult{}, err + } + var validationErr error + if resumeCopy != nil { + validationErr = validateResumedRenditionResult(descriptor, sealed, result) + } else { + validationErr = ValidateRenditionResult(descriptor, sealed, result) + } + if validationErr != nil { + return RenditionResult{}, validationErr + } + return result, nil +} + +func validateCurrentRenditionRequest( + ctx context.Context, provider RenditionProvider, upload AuthorizedUpload, + authorization RenditionAuthorization, +) (RenditionDescriptor, error) { + if ctx != nil && ctx.Err() != nil { + return RenditionDescriptor{}, classifiedRenditionRequestError(RenditionErrorCanceled, "rendition request canceled", ctx.Err()) + } + descriptor, err := ValidateRenditionProviderRequest(provider, upload, authorization) + if err == nil { + return descriptor, nil + } + if ctx != nil && ctx.Err() != nil { + return RenditionDescriptor{}, classifiedRenditionRequestError(RenditionErrorCanceled, "rendition request canceled", ctx.Err()) + } + return RenditionDescriptor{}, classifiedRenditionRequestError(RenditionErrorPolicyRejected, "rendition request rejected", err) +} + +func classifiedRenditionRequestError(code RenditionErrorCode, message string, cause error) error { + providerError, err := NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return errors.Join(cause, err) + } + return providerError +} + +func validateRenditionResumeHandle(handle RenditionResumeHandle) error { + if handle.Value == "" || len(handle.Value) > 512 { + return errors.New("rendition resume handle must contain 1-512 characters") + } + for _, char := range handle.Value { + if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || + char >= '0' && char <= '9' || strings.ContainsRune("-._~", char) { + continue + } + return errors.New("rendition resume handle contains unsupported characters") + } + return nil +} + +// ValidateRenditionResult rejects provider output outside the authorized contract. +func ValidateRenditionResult( + descriptor RenditionDescriptor, authorization RenditionAuthorization, result RenditionResult, +) error { + return validateRenditionResult(descriptor, authorization, result, false) +} + +func validateResumedRenditionResult( + descriptor RenditionDescriptor, authorization RenditionAuthorization, result RenditionResult, +) error { + return validateRenditionResult(descriptor, authorization, result, true) +} + +func validateRenditionResult( + descriptor RenditionDescriptor, authorization RenditionAuthorization, result RenditionResult, + allowCompletionAfterExpiry bool, +) error { + if err := validateRenditionDescriptor(descriptor); err != nil { + return err + } + if err := validateAuthorizationWithoutUpload(descriptor, authorization); err != nil { + return err + } + if err := ValidateSourceEvidenceV1(result.Evidence); err != nil { + return fmt.Errorf("provider evidence: %w", err) + } + if result.Evidence.Family != authorization.MediaFamily { + return errors.New("provider evidence family does not match authorization") + } + if !descriptor.ReturnsMarkdown && len(result.ProviderMarkdown) != 0 { + return errors.New("provider Markdown is not declared by descriptor") + } + if len(result.ProviderMarkdown) > authorization.MaxProviderMarkdownBytes { + return errors.New("provider Markdown exceeds authorized byte limit") + } + if err := validateRenditionArtifacts(descriptor, authorization, result.Artifacts); err != nil { + return err + } + if err := validateEvidenceArtifactAuthorization( + authorization, result.Evidence.Artifacts, result.Artifacts); err != nil { + return err + } + if err := validateRenditionReceipt( + descriptor, authorization, result.Receipt, allowCompletionAfterExpiry); err != nil { + return err + } + encodedEvidence, err := canonicalJSON(result.Evidence) + if err != nil { + return fmt.Errorf("encode provider evidence: %w", err) + } + total := len(encodedEvidence) + len(result.ProviderMarkdown) + for _, artifact := range result.Artifacts { + total += len(artifact.Payload) + } + if total > authorization.MaxTotalResultBytes { + return errors.New("provider total result bytes exceed authorization") + } + return nil +} + +// RenditionErrorCode is a stable, bounded provider failure class. +type RenditionErrorCode string + +const ( + RenditionErrorUnsupportedInput RenditionErrorCode = "unsupported_input" + RenditionErrorPolicyRejected RenditionErrorCode = "policy_rejected" + RenditionErrorAuthentication RenditionErrorCode = "authentication" + RenditionErrorCapacity RenditionErrorCode = "capacity" + RenditionErrorRateLimited RenditionErrorCode = "rate_limited" + RenditionErrorTransient RenditionErrorCode = "transient" + RenditionErrorMalformedEvidence RenditionErrorCode = "malformed_evidence" + RenditionErrorUnknownJob RenditionErrorCode = "unknown_job" + RenditionErrorCanceled RenditionErrorCode = "canceled" + RenditionErrorAmbiguousSubmission RenditionErrorCode = "ambiguous_submission" +) + +// RenditionProviderError preserves a private cause behind a sanitized error class and message. +type RenditionProviderError struct { + code RenditionErrorCode + message string + retryAfter time.Duration + cause error +} + +// NewRenditionProviderError constructs a classified provider failure. +func NewRenditionProviderError( + code RenditionErrorCode, message string, retryAfter time.Duration, cause error, +) (*RenditionProviderError, error) { + providerError := &RenditionProviderError{code: code, message: message, retryAfter: retryAfter, cause: cause} + if err := validateClassifiedProviderError(providerError); err != nil { + return nil, err + } + return providerError, nil +} + +// Error returns only the sanitized failure class and message. +func (providerError *RenditionProviderError) Error() string { + if providerError == nil { + return "rendition provider error" + } + return fmt.Sprintf("rendition provider %s: %s", providerError.code, providerError.message) +} + +// Unwrap exposes the private cause for programmatic matching without rendering it. +func (providerError *RenditionProviderError) Unwrap() error { + if providerError == nil { + return nil + } + return providerError.cause +} + +// Code returns the stable provider failure class. +func (providerError *RenditionProviderError) Code() RenditionErrorCode { + if providerError == nil { + return "" + } + return providerError.code +} + +// RetryAfter returns the bounded provider-supplied delay. +func (providerError *RenditionProviderError) RetryAfter() time.Duration { + if providerError == nil { + return 0 + } + return providerError.retryAfter +} + +// ValidateRenditionProviderError rejects raw, unclassified provider failures. +func ValidateRenditionProviderError(err error) error { + providerError, ok := err.(*RenditionProviderError) //nolint:errorlint // wrappers may expose unsanitized text + if !ok { + return errors.New("unclassified rendition provider error") + } + return validateClassifiedProviderError(providerError) +} + +// IsRenditionProviderErrorRetryable reports whether an explicitly classified failure is retryable. +func IsRenditionProviderErrorRetryable(err error) bool { + providerError, ok := err.(*RenditionProviderError) //nolint:errorlint // only the top-level safe error is retryable + if !ok || validateClassifiedProviderError(providerError) != nil { + return false + } + switch providerError.code { + case RenditionErrorCapacity, RenditionErrorRateLimited, RenditionErrorTransient: + return true + default: + return false + } +} + +func validateRenditionDescriptor(descriptor RenditionDescriptor) error { + if err := validateRenditionDescriptorFields(descriptor); err != nil { + return err + } + if err := validateFingerprint(descriptor.Fingerprint, "descriptor fingerprint"); err != nil { + return err + } + canonical, err := NewRenditionDescriptor(descriptor) + if err != nil { + return err + } + if canonical.Fingerprint != descriptor.Fingerprint || !equalRenditionDescriptors(canonical, descriptor) { + return errors.New("descriptor fingerprint or canonical ordering is invalid") + } + return nil +} + +func cloneRenditionAuthorization(value RenditionAuthorization) RenditionAuthorization { + value.AllowedArtifactRoles = slices.Clone(value.AllowedArtifactRoles) + return value +} + +func validateRenditionDescriptorFields(descriptor RenditionDescriptor) error { + if err := validateStableToken(descriptor.ID, "descriptor ID", 128); err != nil { + return err + } + if descriptor.ContractVersion != RenditionProviderContractVersion { + return fmt.Errorf("descriptor contract version must be %d", RenditionProviderContractVersion) + } + if err := validateFingerprint(descriptor.PolicyFingerprint, "descriptor policy fingerprint"); err != nil { + return err + } + switch descriptor.TrustBoundary { + case RenditionTrustLocalProcess, RenditionTrustOperatorNetwork, RenditionTrustHostedProvider: + default: + return errors.New("descriptor trust boundary is invalid") + } + if len(descriptor.SupportedFormats) == 0 || len(descriptor.SupportedFormats) > maxRenditionFormats { + return fmt.Errorf("descriptor supported formats must contain 1-%d entries", maxRenditionFormats) + } + seenFormats := make(map[RenditionFormatCapability]struct{}, len(descriptor.SupportedFormats)) + for _, format := range descriptor.SupportedFormats { + if err := validateRenditionFormat(format); err != nil { + return err + } + if _, exists := seenFormats[format]; exists { + return errors.New("descriptor contains duplicate supported format") + } + seenFormats[format] = struct{}{} + } + if len(descriptor.ArtifactRoles) > maxRenditionArtifactRoles { + return errors.New("descriptor has too many artifact roles") + } + seenRoles := make(map[EvidenceArtifactRole]struct{}, len(descriptor.ArtifactRoles)) + for _, role := range descriptor.ArtifactRoles { + if !validProfileArtifactRole(role) { + return fmt.Errorf("descriptor artifact role %q is invalid", role) + } + if _, exists := seenRoles[role]; exists { + return fmt.Errorf("descriptor artifact role %q is duplicated", role) + } + seenRoles[role] = struct{}{} + } + if !descriptor.ReturnsMarkdown && !descriptor.ReturnsStructured && len(descriptor.ArtifactRoles) == 0 { + return errors.New("descriptor must declare at least one result kind") + } + return nil +} + +func validateRenditionFormat(format RenditionFormatCapability) error { + if err := validateStableToken(format.MediaFamily, "media family", 63); err != nil { + return err + } + if err := validateCanonicalMediaType(format.MediaType); err != nil { + return err + } + if !validRenditionInputKind(format.InputKind) { + return errors.New("rendition input kind is invalid") + } + return nil +} + +func validateAuthorizedUploadMetadata(metadata AuthorizedUploadMetadata) error { + if metadata.Filename == "" || len(metadata.Filename) > 255 || strings.ContainsAny(metadata.Filename, "/\\\x00") { + return errors.New("upload filename must be a safe basename of at most 255 bytes") + } + if err := validateRenditionFormat(RenditionFormatCapability{ + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, InputKind: metadata.InputKind, + }); err != nil { + return fmt.Errorf("upload metadata: %w", err) + } + if metadata.ByteLength <= 0 || metadata.ByteLength > maxRenditionSourceBytes { + return errors.New("upload byte length is outside the supported bound") + } + for subject, value := range map[string]string{ + "upload SHA-256": metadata.SHA256, "capability record checksum": metadata.CapabilityRecordChecksum, + "provider metadata checksum": metadata.ProviderMetadataChecksum, + } { + if err := validateFingerprint(value, subject); err != nil { + return err + } + } + return nil +} + +func validateRenditionAuthorization( + descriptor RenditionDescriptor, metadata AuthorizedUploadMetadata, authorization RenditionAuthorization, +) error { + if err := validateAuthorizationWithoutUpload(descriptor, authorization); err != nil { + return err + } + if authorization.SourceSHA256 != metadata.SHA256 || authorization.SourceBytes != metadata.ByteLength { + return errors.New("authorization does not match exact upload bytes") + } + if authorization.CapabilityRecordChecksum != metadata.CapabilityRecordChecksum { + return errors.New("authorization capability record checksum does not match upload") + } + if authorization.ProviderMetadataChecksum != metadata.ProviderMetadataChecksum { + return errors.New("authorization provider metadata checksum does not match upload") + } + if authorization.MediaFamily != metadata.MediaFamily || authorization.MediaType != metadata.MediaType || + authorization.InputKind != metadata.InputKind { + return errors.New("authorization format does not match upload metadata") + } + return nil +} + +func validateAuthorizationWithoutUpload(descriptor RenditionDescriptor, authorization RenditionAuthorization) error { + if authorization.ProviderID != descriptor.ID { + return errors.New("authorization provider ID does not match descriptor") + } + if authorization.DescriptorFingerprint != descriptor.Fingerprint { + return errors.New("authorization descriptor fingerprint does not match descriptor") + } + if authorization.PolicyFingerprint != descriptor.PolicyFingerprint { + return errors.New("authorization policy fingerprint does not match descriptor") + } + for subject, value := range map[string]string{ + "rendition request fingerprint": authorization.RenditionRequestFingerprint, + "source SHA-256": authorization.SourceSHA256, "capability record checksum": authorization.CapabilityRecordChecksum, + "provider metadata checksum": authorization.ProviderMetadataChecksum, + } { + if err := validateFingerprint(value, subject); err != nil { + return err + } + } + format := RenditionFormatCapability{ + MediaFamily: authorization.MediaFamily, MediaType: authorization.MediaType, InputKind: authorization.InputKind, + } + if !slices.Contains(descriptor.SupportedFormats, format) { + return errors.New("authorization requests an unsupported format") + } + if authorization.SourceBytes <= 0 || authorization.SourceBytes > maxRenditionSourceBytes { + return errors.New("authorization source bytes are outside the supported bound") + } + if err := validateAuthorizedRoles(descriptor, authorization.AllowedArtifactRoles); err != nil { + return err + } + if authorization.MaxProviderMarkdownBytes < 0 || authorization.MaxProviderMarkdownBytes > maxRenditionMarkdownBytes || + (!descriptor.ReturnsMarkdown && authorization.MaxProviderMarkdownBytes != 0) { + return errors.New("authorization provider Markdown bytes are invalid") + } + if authorization.MaxArtifactBytes < 0 || authorization.MaxArtifactBytes > maxRenditionArtifactBytes { + return errors.New("authorization artifact bytes are invalid") + } + if authorization.MaxArtifacts < 0 || authorization.MaxArtifacts > maxRenditionArtifacts || + authorization.MaxArtifacts > len(authorization.AllowedArtifactRoles) { + return errors.New("authorization artifact count is invalid") + } + if authorization.MaxTotalResultBytes <= 0 || authorization.MaxTotalResultBytes > maxRenditionTotalResultBytes { + return errors.New("authorization total result bytes are invalid") + } + authorizedAt, err := parseRenditionTimestamp(authorization.AuthorizedAt) + if err != nil { + return errors.New("authorization time must be canonical RFC3339Nano") + } + expiresAt, err := parseRenditionTimestamp(authorization.ExpiresAt) + if err != nil || !expiresAt.After(authorizedAt) { + return errors.New("authorization expiry must be canonical and after authorization time") + } + return nil +} + +func validateAuthorizationCurrentAt(authorization RenditionAuthorization, now time.Time) error { + authorizedAt, err := parseRenditionTimestamp(authorization.AuthorizedAt) + if err != nil { + return errors.New("authorization time must be canonical") + } + expiresAt, err := parseRenditionTimestamp(authorization.ExpiresAt) + if err != nil { + return errors.New("authorization expiry must be canonical") + } + if now.Before(authorizedAt) || !now.Before(expiresAt) { + return errors.New("authorization is not current") + } + return nil +} + +func validateAuthorizedRoles(descriptor RenditionDescriptor, roles []EvidenceArtifactRole) error { + if len(roles) > maxRenditionArtifactRoles { + return errors.New("authorization has too many artifact roles") + } + seen := make(map[EvidenceArtifactRole]struct{}, len(roles)) + for _, role := range roles { + if !validProfileArtifactRole(role) || !slices.Contains(descriptor.ArtifactRoles, role) { + return fmt.Errorf("authorization artifact role %q is not declared", role) + } + if _, exists := seen[role]; exists { + return fmt.Errorf("authorization artifact role %q is duplicated", role) + } + seen[role] = struct{}{} + } + return nil +} + +func validateRenditionArtifacts( + descriptor RenditionDescriptor, authorization RenditionAuthorization, artifacts []RenditionArtifact, +) error { + seen := make(map[EvidenceArtifactRole]struct{}, len(artifacts)) + for _, artifact := range artifacts { + if !validProfileArtifactRole(artifact.Role) || !slices.Contains(descriptor.ArtifactRoles, artifact.Role) || + !slices.Contains(authorization.AllowedArtifactRoles, artifact.Role) { + return fmt.Errorf("provider artifact role %q is not authorized", artifact.Role) + } + if _, exists := seen[artifact.Role]; exists { + return fmt.Errorf("provider artifact role %q is duplicated", artifact.Role) + } + seen[artifact.Role] = struct{}{} + if err := validateCanonicalMediaType(artifact.MediaType); err != nil { + return fmt.Errorf("provider artifact: %w", err) + } + if len(artifact.Payload) > authorization.MaxArtifactBytes { + return errors.New("provider artifact exceeds authorized byte limit") + } + digest := sha256.Sum256(artifact.Payload) + if artifact.SHA256 != hex.EncodeToString(digest[:]) { + return errors.New("provider artifact checksum does not match payload") + } + } + if len(artifacts) > authorization.MaxArtifacts { + return errors.New("provider artifact count exceeds authorization") + } + return nil +} + +func validateEvidenceArtifactAuthorization( + authorization RenditionAuthorization, evidence []SourceEvidenceArtifactV1, + retained []RenditionArtifact, +) error { + type identity struct { + role EvidenceArtifactRole + sha256 string + } + retainedSet := make(map[identity]struct{}, len(retained)) + for _, artifact := range retained { + retainedSet[identity{role: artifact.Role, sha256: artifact.SHA256}] = struct{}{} + } + referenced := make(map[identity]struct{}, len(evidence)) + for _, artifact := range evidence { + if !slices.Contains(authorization.AllowedArtifactRoles, artifact.Role) { + return fmt.Errorf("source evidence artifact role %q is not authorized", artifact.Role) + } + key := identity{role: artifact.Role, sha256: artifact.SHA256} + if _, ok := retainedSet[key]; !ok { + return errors.New("source evidence artifact does not match a retained provider artifact") + } + referenced[key] = struct{}{} + } + for key := range retainedSet { + if _, ok := referenced[key]; !ok { + return errors.New("retained provider artifact is absent from source evidence") + } + } + return nil +} + +func validateRenditionReceipt( + descriptor RenditionDescriptor, authorization RenditionAuthorization, receipt RenditionReceipt, + allowCompletionAfterExpiry bool, +) error { + if receipt.ProviderID != descriptor.ID || receipt.DescriptorFingerprint != descriptor.Fingerprint || + receipt.PolicyFingerprint != authorization.PolicyFingerprint || receipt.SourceSHA256 != authorization.SourceSHA256 { + return errors.New("provider receipt does not match authorization") + } + if err := validateStableToken(receipt.OperationID, "operation ID", 128); err != nil { + return err + } + startedAt, err := parseRenditionTimestamp(receipt.StartedAt) + if err != nil { + return errors.New("receipt start time must be canonical RFC3339Nano") + } + completedAt, err := parseRenditionTimestamp(receipt.CompletedAt) + if err != nil || completedAt.Before(startedAt) { + return errors.New("receipt completion time must be canonical and not precede start") + } + authorizedAt, _ := parseRenditionTimestamp(authorization.AuthorizedAt) + expiresAt, _ := parseRenditionTimestamp(authorization.ExpiresAt) + if startedAt.Before(authorizedAt) || !startedAt.Before(expiresAt) || + !allowCompletionAfterExpiry && completedAt.After(expiresAt) { + return errors.New("receipt execution is outside the authorization interval") + } + if len(receipt.Warnings) > maxRenditionWarnings { + return errors.New("receipt has too many warning codes") + } + seenWarnings := make(map[string]struct{}, len(receipt.Warnings)) + for _, warning := range receipt.Warnings { + if err := validateStableToken(warning, "receipt warning", 63); err != nil { + return err + } + if _, exists := seenWarnings[warning]; exists { + return errors.New("receipt warning codes must be unique") + } + seenWarnings[warning] = struct{}{} + } + if err := validateRenditionUsage(receipt.Usage); err != nil { + return err + } + if receipt.Usage.Retries > receipt.Usage.Requests { + return errors.New("receipt retries cannot exceed requests") + } + if receipt.RetryDelayMillis < 0 || receipt.RetryDelayMillis > int64((24*time.Hour)/time.Millisecond) { + return errors.New("receipt retry delay is outside the supported bound") + } + return nil +} + +func validateRenditionUsage(usage RenditionUsage) error { + for subject, value := range map[string]int64{ + "requests": usage.Requests, "retries": usage.Retries, "input bytes": usage.InputBytes, + "output bytes": usage.OutputBytes, "units": usage.Units, + } { + if value < 0 || value > maxRenditionUsageValue { + return fmt.Errorf("receipt usage %s is outside the supported bound", subject) + } + } + return nil +} + +func validateClassifiedProviderError(providerError *RenditionProviderError) error { + if providerError == nil { + return errors.New("classified rendition provider error is nil") + } + switch providerError.code { + case RenditionErrorUnsupportedInput, RenditionErrorPolicyRejected, RenditionErrorAuthentication, + RenditionErrorCapacity, RenditionErrorRateLimited, RenditionErrorTransient, + RenditionErrorMalformedEvidence, RenditionErrorUnknownJob, RenditionErrorCanceled, + RenditionErrorAmbiguousSubmission: + default: + return errors.New("rendition provider error code is invalid") + } + if err := validateSafeMessage(providerError.message); err != nil { + return err + } + if providerError.retryAfter < 0 || providerError.retryAfter > 24*time.Hour { + return errors.New("rendition provider retry delay is outside the supported bound") + } + return nil +} + +func validateSafeMessage(value string) error { + if value == "" || len(value) > 160 || strings.ContainsAny(value, "\r\n{}[]<>=\"") { + return errors.New("rendition provider error message must be a bounded safe summary") + } + lower := strings.ToLower(value) + for _, unsafe := range []string{"authorization:", "bearer ", "api_key", "apikey", "secret", "token="} { + if strings.Contains(lower, unsafe) { + return errors.New("rendition provider error message contains credential-shaped content") + } + } + return nil +} + +func validateStableToken(value, subject string, maxLength int) error { + if value == "" || len(value) > maxLength { + return fmt.Errorf("%s must contain 1-%d characters", subject, maxLength) + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= '0' && char <= '9' || char == '_' || char == '-' || char == '.' { + continue + } + return fmt.Errorf("%s contains unsupported characters", subject) + } + return nil +} + +func validateCanonicalMediaType(value string) error { + parsed, parameters, err := mime.ParseMediaType(value) + if err != nil || parsed != value || len(parameters) != 0 || !strings.Contains(value, "/") { + return fmt.Errorf("media type %q must be canonical and contain no parameters", value) + } + return nil +} + +func validRenditionInputKind(kind RenditionInputKind) bool { + return kind == RenditionInputOriginalFile || kind == RenditionInputDerivedUpload +} + +func parseRenditionTimestamp(value string) (time.Time, error) { + parsed, err := time.Parse(renditionTimestampForm, value) + if err != nil || parsed.Format(renditionTimestampForm) != value { + return time.Time{}, errors.New("timestamp is not canonical UTC RFC3339Nano") + } + return parsed, nil +} + +func compareRenditionFormats(left, right RenditionFormatCapability) int { + if comparison := strings.Compare(left.MediaFamily, right.MediaFamily); comparison != 0 { + return comparison + } + if comparison := strings.Compare(left.MediaType, right.MediaType); comparison != 0 { + return comparison + } + return strings.Compare(string(left.InputKind), string(right.InputKind)) +} + +func descriptorIdentity(value RenditionDescriptor) renditionDescriptorIdentity { + return renditionDescriptorIdentity{ + ID: value.ID, ContractVersion: value.ContractVersion, PolicyFingerprint: value.PolicyFingerprint, + TrustBoundary: value.TrustBoundary, SupportedFormats: value.SupportedFormats, + ReturnsMarkdown: value.ReturnsMarkdown, ReturnsStructured: value.ReturnsStructured, + ArtifactRoles: value.ArtifactRoles, + } +} + +func cloneRenditionDescriptor(value RenditionDescriptor) RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +func equalRenditionDescriptors(left, right RenditionDescriptor) bool { + return left.ID == right.ID && left.ContractVersion == right.ContractVersion && + left.PolicyFingerprint == right.PolicyFingerprint && left.TrustBoundary == right.TrustBoundary && + left.ReturnsMarkdown == right.ReturnsMarkdown && left.ReturnsStructured == right.ReturnsStructured && + left.Fingerprint == right.Fingerprint && slices.Equal(left.SupportedFormats, right.SupportedFormats) && + slices.Equal(left.ArtifactRoles, right.ArtifactRoles) +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/provider_test.go b/document/provider_test.go new file mode 100644 index 00000000..89270c34 --- /dev/null +++ b/document/provider_test.go @@ -0,0 +1,687 @@ +package document + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "math" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type syntheticRenditionProvider struct { + descriptor RenditionDescriptor + result RenditionResult + err error +} + +type countingRenditionProvider struct { + descriptor RenditionDescriptor + renderCalls int +} + +func (provider *countingRenditionProvider) Descriptor() RenditionDescriptor { + return provider.descriptor +} + +func (provider *countingRenditionProvider) Render( + context.Context, AuthorizedUpload, RenditionAuthorization, +) (RenditionResult, error) { + provider.renderCalls++ + return RenditionResult{}, errors.New("unexpected render call") +} + +type syntheticResumableRenditionProvider struct { + syntheticRenditionProvider + + resume *RenditionResumeHandle + checkpoint RenditionResumeCheckpoint + calls int + ignoreCheckpointError bool + uploadWasNil bool +} + +func (provider *syntheticResumableRenditionProvider) RenderResumable( + _ context.Context, upload AuthorizedUpload, _ RenditionAuthorization, + resume *RenditionResumeHandle, checkpoint RenditionResumeCheckpoint, +) (RenditionResult, error) { + provider.calls++ + provider.uploadWasNil = upload == nil + if resume != nil { + resumeValue := *resume + provider.resume = &resumeValue + } + provider.checkpoint = checkpoint + if resume == nil { + if err := checkpoint(RenditionResumeHandle{Value: "remote-job-1"}); err != nil && + !provider.ignoreCheckpointError { + return RenditionResult{}, err + } + } + return provider.result, provider.err +} + +var _ ResumableRenditionProvider = (*syntheticResumableRenditionProvider)(nil) + +func (provider syntheticRenditionProvider) Descriptor() RenditionDescriptor { + return provider.descriptor +} + +func (provider syntheticRenditionProvider) Render( + context.Context, AuthorizedUpload, RenditionAuthorization, +) (RenditionResult, error) { + return provider.result, provider.err +} + +var _ RenditionProvider = syntheticRenditionProvider{} + +type syntheticAuthorizedUpload struct { + io.ReadCloser + + metadata AuthorizedUploadMetadata +} + +func (upload *syntheticAuthorizedUpload) Metadata() AuthorizedUploadMetadata { + return upload.metadata +} + +var _ AuthorizedUpload = (*syntheticAuthorizedUpload)(nil) + +type changingRenditionProvider struct { + descriptors []RenditionDescriptor + calls int +} + +type mutatingRenditionProvider struct { + descriptor RenditionDescriptor + render func(RenditionAuthorization) RenditionResult + mutateDescriptor bool + calls int +} + +func (provider *mutatingRenditionProvider) Descriptor() RenditionDescriptor { + provider.calls++ + if provider.mutateDescriptor && provider.calls == 2 { + provider.descriptor.ArtifactRoles[0] = EvidenceArtifactImage + } + return provider.descriptor +} + +func (provider *mutatingRenditionProvider) Render( + _ context.Context, _ AuthorizedUpload, authorization RenditionAuthorization, +) (RenditionResult, error) { + return provider.render(authorization), nil +} + +func (provider *changingRenditionProvider) Descriptor() RenditionDescriptor { + descriptor := provider.descriptors[provider.calls] + provider.calls++ + return descriptor +} + +func (*changingRenditionProvider) Render( + context.Context, AuthorizedUpload, RenditionAuthorization, +) (RenditionResult, error) { + return RenditionResult{}, nil +} + +type changingAuthorizedUpload struct { + io.ReadCloser + + metadata []AuthorizedUploadMetadata + calls int +} + +func (upload *changingAuthorizedUpload) Metadata() AuthorizedUploadMetadata { + metadata := upload.metadata[upload.calls] + upload.calls++ + return metadata +} + +func TestRenditionProviderContractAcceptsExactBoundedResult(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + result := validRenditionResult(descriptor, authorization) + provider := syntheticRenditionProvider{descriptor: descriptor, result: result} + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + + validated, err := ValidateRenditionProviderRequest(provider, upload, authorization) + require.NoError(t, err) + assert.Equal(t, descriptor, validated) + produced, err := provider.Render(t.Context(), upload, authorization) + require.NoError(t, err) + require.NoError(t, ValidateRenditionResult(validated, authorization, produced)) +} + +func TestRenditionProviderContractRejectsInvalidDescriptorAndAuthorization(t *testing.T) { + metadata := validAuthorizedUploadMetadata() + for _, testCase := range []struct { + name string + mutate func(*RenditionDescriptor, *RenditionAuthorization) + want string + }{ + { + name: "empty descriptor identity", + mutate: func(descriptor *RenditionDescriptor, _ *RenditionAuthorization) { + descriptor.ID = "" + }, + want: "descriptor ID", + }, + { + name: "unsupported format", + mutate: func(_ *RenditionDescriptor, authorization *RenditionAuthorization) { + authorization.MediaFamily = "image" + authorization.MediaType = "image/png" + }, + want: "unsupported format", + }, + { + name: "mismatched descriptor authorization", + mutate: func(_ *RenditionDescriptor, authorization *RenditionAuthorization) { + authorization.DescriptorFingerprint = strings.Repeat("9", 64) + }, + want: "descriptor fingerprint", + }, + { + name: "oversized result declaration", + mutate: func(_ *RenditionDescriptor, authorization *RenditionAuthorization) { + authorization.MaxTotalResultBytes = math.MaxInt64 + }, + want: "total result bytes", + }, + } { + t.Run(testCase.name, func(t *testing.T) { + descriptor := validRenditionDescriptor(t) + authorization := validRenditionAuthorization(descriptor, metadata) + testCase.mutate(&descriptor, &authorization) + provider := syntheticRenditionProvider{descriptor: descriptor} + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader(nil)), metadata: metadata, + } + _, err := ValidateRenditionProviderRequest(provider, upload, authorization) + require.ErrorContains(t, err, testCase.want) + }) + } +} + +func TestRenditionProviderContractRejectsDuplicateAndOversizedOutputs(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + for _, testCase := range []struct { + name string + mutate func(*RenditionResult) + want string + }{ + { + name: "duplicate artifact role", + mutate: func(result *RenditionResult) { + result.Artifacts = append(result.Artifacts, result.Artifacts[0]) + }, + want: "artifact role", + }, + { + name: "oversized provider Markdown", + mutate: func(result *RenditionResult) { + result.ProviderMarkdown = bytes.Repeat([]byte("x"), authorization.MaxProviderMarkdownBytes+1) + }, + want: "provider Markdown", + }, + { + name: "artifact checksum mismatch", + mutate: func(result *RenditionResult) { + result.Artifacts[0].SHA256 = strings.Repeat("0", 64) + }, + want: "artifact checksum", + }, + } { + t.Run(testCase.name, func(t *testing.T) { + result := validRenditionResult(descriptor, authorization) + testCase.mutate(&result) + require.ErrorContains(t, + ValidateRenditionResult(descriptor, authorization, result), testCase.want) + }) + } +} + +func TestRenditionProviderContractRejectsUnauthorizedOrUnmatchedEvidenceArtifacts(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + + result := validRenditionResult(descriptor, authorization) + result.Evidence.Artifacts[0].Role = EvidenceArtifactImage + require.ErrorContains(t, ValidateRenditionResult(descriptor, authorization, result), + "not authorized") + + result = validRenditionResult(descriptor, authorization) + result.Evidence.Artifacts[0].SHA256 = strings.Repeat("9", 64) + require.ErrorContains(t, ValidateRenditionResult(descriptor, authorization, result), + "does not match") + + result = validRenditionResult(descriptor, authorization) + result.Evidence.Artifacts = nil + require.ErrorContains(t, ValidateRenditionResult(descriptor, authorization, result), + "absent from source evidence") +} + +func TestRenditionProviderContractRejectsUnsafeReceiptFields(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + for _, testCase := range []struct { + name string + mutate func(*RenditionReceipt) + want string + }{ + { + name: "provider body in warning", + mutate: func(receipt *RenditionReceipt) { + receipt.Warnings = []string{"raw body: {\"document\":\"secret\"}"} + }, + want: "warning", + }, + { + name: "credential-shaped operation identity", + mutate: func(receipt *RenditionReceipt) { + receipt.OperationID = "Authorization: Bearer secret" + }, + want: "operation ID", + }, + } { + t.Run(testCase.name, func(t *testing.T) { + result := validRenditionResult(descriptor, authorization) + testCase.mutate(&result.Receipt) + require.ErrorContains(t, + ValidateRenditionResult(descriptor, authorization, result), testCase.want) + }) + } +} + +func TestRenditionProviderContractRejectsUnclassifiedErrors(t *testing.T) { + require.ErrorContains(t, + ValidateRenditionProviderError(errors.New("raw provider body with secret")), + "unclassified", + ) + cause := errors.New("raw provider body with secret") + providerError, err := NewRenditionProviderError( + RenditionErrorRateLimited, "provider rate limit", 30*time.Second, cause, + ) + require.NoError(t, err) + require.NoError(t, ValidateRenditionProviderError(providerError)) + assert.True(t, IsRenditionProviderErrorRetryable(providerError)) + wrapped := fmt.Errorf("unsafe wrapper includes provider body: %w", providerError) + require.ErrorContains(t, ValidateRenditionProviderError(wrapped), "unclassified") + assert.False(t, IsRenditionProviderErrorRetryable(wrapped)) + assert.NotContains(t, providerError.Error(), "secret") + assert.ErrorIs(t, providerError, cause) +} + +func TestRenditionProviderContractRejectsExpiredAuthorizationAndOutOfWindowReceipt(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader([]byte("synthetic exact source"))), metadata: metadata, + } + provider := syntheticRenditionProvider{descriptor: descriptor} + authorization := validRenditionAuthorization(descriptor, metadata) + now, err := parseRenditionTimestamp(authorization.ExpiresAt) + require.NoError(t, err) + _, err = ValidateRenditionProviderRequestAt(now, provider, upload, authorization) + require.ErrorContains(t, err, "not current") + authorizedAt, err := parseRenditionTimestamp(authorization.AuthorizedAt) + require.NoError(t, err) + _, err = ValidateRenditionProviderRequestAt( + authorizedAt.Add(-time.Nanosecond), + provider, upload, authorization) + require.ErrorContains(t, err, "not current") + + result := validRenditionResult(descriptor, authorization) + result.Receipt.StartedAt = authorizedAt.Add(-time.Nanosecond).Format(renditionTimestampForm) + require.ErrorContains(t, ValidateRenditionResult(descriptor, authorization, result), + "outside the authorization interval") + result = validRenditionResult(descriptor, authorization) + result.Receipt.CompletedAt = now.Add(time.Nanosecond).Format(renditionTimestampForm) + require.ErrorContains(t, ValidateRenditionResult(descriptor, authorization, result), + "outside the authorization interval") +} + +func TestRenderRenditionClassifiesExpiredRequestBeforeProviderCall(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader([]byte("synthetic exact source"))), metadata: metadata, + } + provider := &countingRenditionProvider{descriptor: descriptor} + authorization := validRenditionAuthorization(descriptor, metadata) + authorization.ExpiresAt = time.Now().UTC().Add(-time.Second).Format(renditionTimestampForm) + + _, err := RenderRendition(t.Context(), provider, upload, authorization) + + providerError, ok := errors.AsType[*RenditionProviderError](err) + require.True(t, ok, "%T: %v", err, err) + assert.Equal(t, RenditionErrorPolicyRejected, providerError.Code()) + assert.Zero(t, provider.renderCalls) +} + +func TestRenderRenditionWithResumeClassifiesCanceledRequestBeforeProviderCall(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader([]byte("synthetic exact source"))), metadata: metadata, + } + provider := &countingRenditionProvider{descriptor: descriptor} + authorization := validRenditionAuthorization(descriptor, metadata) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err := RenderRenditionWithResume(ctx, provider, upload, authorization, nil, nil) + + providerError, ok := errors.AsType[*RenditionProviderError](err) + require.True(t, ok, "%T: %v", err, err) + assert.Equal(t, RenditionErrorCanceled, providerError.Code()) + assert.Zero(t, provider.renderCalls) +} + +func TestRenditionDescriptorCanonicalizesAndOwnsCollections(t *testing.T) { + formats := []RenditionFormatCapability{ + {MediaFamily: "text", MediaType: "text/plain", InputKind: RenditionInputOriginalFile}, + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: RenditionInputOriginalFile}, + } + roles := []EvidenceArtifactRole{EvidenceArtifactTranscript, EvidenceArtifactStructured} + descriptor, err := NewRenditionDescriptor(RenditionDescriptor{ + ID: "synthetic-rendition", ContractVersion: RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), TrustBoundary: RenditionTrustHostedProvider, + SupportedFormats: formats, ReturnsStructured: true, ArtifactRoles: roles, + }) + require.NoError(t, err) + formats[0].MediaFamily = "mutated" + roles[0] = EvidenceArtifactImage + assert.Equal(t, "pdf", descriptor.SupportedFormats[0].MediaFamily) + assert.Equal(t, EvidenceArtifactTranscript, descriptor.ArtifactRoles[0]) + + reordered, err := NewRenditionDescriptor(RenditionDescriptor{ + ID: "synthetic-rendition", ContractVersion: RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), TrustBoundary: RenditionTrustHostedProvider, + SupportedFormats: []RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: RenditionInputOriginalFile}, + }, + ReturnsStructured: true, + ArtifactRoles: []EvidenceArtifactRole{EvidenceArtifactStructured, EvidenceArtifactTranscript}, + }) + require.NoError(t, err) + assert.Equal(t, descriptor, reordered) +} + +func TestRenditionProviderContractRejectsMutableBoundarySnapshots(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + + changedDescriptor := descriptor + changedDescriptor.ID = "changed-provider" + provider := &changingRenditionProvider{descriptors: []RenditionDescriptor{descriptor, changedDescriptor}} + upload := &syntheticAuthorizedUpload{ReadCloser: io.NopCloser(bytes.NewReader(nil)), metadata: metadata} + _, err := ValidateRenditionProviderRequest(provider, upload, authorization) + require.ErrorContains(t, err, "descriptor changed") + + changedMetadata := metadata + changedMetadata.Filename = "changed.pdf" + provider = &changingRenditionProvider{descriptors: []RenditionDescriptor{descriptor, descriptor}} + changingUpload := &changingAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader(nil)), + metadata: []AuthorizedUploadMetadata{metadata, changedMetadata}, + } + _, err = ValidateRenditionProviderRequest(provider, changingUpload, authorization) + require.ErrorContains(t, err, "upload metadata changed") +} + +func TestRenditionProviderContractFreezesAliasedDescriptorSlices(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + provider := &mutatingRenditionProvider{descriptor: descriptor, mutateDescriptor: true} + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(bytes.NewReader(nil)), metadata: metadata, + } + + _, err := ValidateRenditionProviderRequest(provider, upload, authorization) + require.ErrorContains(t, err, "descriptor changed") +} + +func TestRenderRenditionKeepsSealedAuthorizationSeparateFromProvider(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + provider := &mutatingRenditionProvider{descriptor: descriptor} + provider.render = func(received RenditionAuthorization) RenditionResult { + received.AllowedArtifactRoles[0] = EvidenceArtifactImage + result := validRenditionResult(descriptor, received) + result.Artifacts[0].Role = EvidenceArtifactImage + result.Evidence.Artifacts[0].Role = EvidenceArtifactImage + return result + } + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + + _, err := RenderRendition(t.Context(), provider, upload, authorization) + require.ErrorContains(t, err, "not authorized") + assert.Equal(t, []EvidenceArtifactRole{EvidenceArtifactStructured}, authorization.AllowedArtifactRoles, + "provider mutation must not reach the caller or sealed validation snapshot") +} + +func TestRenderRenditionWithResumePersistsProviderIssuedHandleAndResumesIt(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + provider := &syntheticResumableRenditionProvider{ + descriptor: descriptor, result: validRenditionResult(descriptor, authorization)} + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + var persisted []RenditionResumeHandle + + result, err := RenderRenditionWithResume(t.Context(), provider, upload, authorization, nil, + func(handle RenditionResumeHandle) error { + persisted = append(persisted, handle) + return nil + }) + require.NoError(t, err) + assert.Equal(t, "operation-synthetic-1", result.Receipt.OperationID) + assert.Equal(t, []RenditionResumeHandle{{Value: "remote-job-1"}}, persisted) + assert.Equal(t, 1, provider.calls) + + provider.result = validRenditionResult(descriptor, authorization) + expiresAt, err := parseRenditionTimestamp(authorization.ExpiresAt) + require.NoError(t, err) + provider.result.Receipt.CompletedAt = expiresAt.Add(time.Second).Format(renditionTimestampForm) + resume := persisted[0] + upload = &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + _, err = RenderRenditionWithResume(t.Context(), provider, upload, authorization, &resume, + func(RenditionResumeHandle) error { return errors.New("resume must not invent a replacement handle") }) + require.NoError(t, err) + require.NotNil(t, provider.resume) + assert.Equal(t, resume, *provider.resume) + assert.Equal(t, 2, provider.calls) + assert.True(t, provider.uploadWasNil, "resume must not expose source bytes for resubmission") +} + +func TestRenderRenditionWithResumeRejectsUnsafeOrUnsupportedHandles(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + upload := func() *syntheticAuthorizedUpload { + return &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + } + + provider := syntheticRenditionProvider{ + descriptor: descriptor, result: validRenditionResult(descriptor, authorization), + } + _, err := RenderRenditionWithResume(t.Context(), provider, upload(), authorization, + &RenditionResumeHandle{Value: "remote-job-1"}, nil) + require.ErrorContains(t, err, "does not support durable resume") + + resumable := &syntheticResumableRenditionProvider{syntheticRenditionProvider: provider} + _, err = RenderRenditionWithResume(t.Context(), resumable, upload(), authorization, + &RenditionResumeHandle{Value: "unsafe\nprovider-body"}, nil) + require.ErrorContains(t, err, "resume handle") + assert.Zero(t, resumable.calls, "an unsafe persisted handle must be rejected before provider work") +} + +func TestRenderRenditionWithResumeSurfacesIgnoredCheckpointFailure(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + provider := &syntheticResumableRenditionProvider{ + descriptor: descriptor, result: validRenditionResult(descriptor, authorization), + ignoreCheckpointError: true, + } + upload := &syntheticAuthorizedUpload{ + ReadCloser: io.NopCloser(strings.NewReader("synthetic exact source")), metadata: metadata, + } + checkpointErr := errors.New("durable checkpoint failed") + + _, err := RenderRenditionWithResume(t.Context(), provider, upload, authorization, nil, + func(RenditionResumeHandle) error { return checkpointErr }) + require.ErrorIs(t, err, checkpointErr) +} + +func TestRenditionProviderContractRejectsTypedNilBoundaryValues(t *testing.T) { + descriptor := validRenditionDescriptor(t) + metadata := validAuthorizedUploadMetadata() + authorization := validRenditionAuthorization(descriptor, metadata) + + var provider *changingRenditionProvider + upload := &syntheticAuthorizedUpload{ReadCloser: io.NopCloser(bytes.NewReader(nil)), metadata: metadata} + assert.NotPanics(t, func() { + _, err := ValidateRenditionProviderRequest(provider, upload, authorization) + require.ErrorContains(t, err, "provider is required") + }) + + validProvider := syntheticRenditionProvider{descriptor: descriptor} + var nilUpload *syntheticAuthorizedUpload + assert.NotPanics(t, func() { + _, err := ValidateRenditionProviderRequest(validProvider, nilUpload, authorization) + require.ErrorContains(t, err, "upload is required") + }) +} + +func validRenditionDescriptor(t *testing.T) RenditionDescriptor { + t.Helper() + descriptor, err := NewRenditionDescriptor(RenditionDescriptor{ + ID: "synthetic-rendition", + ContractVersion: RenditionProviderContractVersion, + PolicyFingerprint: strings.Repeat("1", 64), + TrustBoundary: RenditionTrustHostedProvider, + SupportedFormats: []RenditionFormatCapability{{ + MediaFamily: "pdf", MediaType: "application/pdf", InputKind: RenditionInputOriginalFile, + }}, + ReturnsMarkdown: true, + ReturnsStructured: true, + ArtifactRoles: []EvidenceArtifactRole{EvidenceArtifactStructured}, + }) + require.NoError(t, err) + return descriptor +} + +func validAuthorizedUploadMetadata() AuthorizedUploadMetadata { + source := []byte("synthetic exact source") + digest := sha256.Sum256(source) + return AuthorizedUploadMetadata{ + Filename: "document.pdf", + MediaFamily: "pdf", + MediaType: "application/pdf", + ByteLength: int64(len(source)), + SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: RenditionInputOriginalFile, + } +} + +func validRenditionAuthorization( + descriptor RenditionDescriptor, metadata AuthorizedUploadMetadata, + +) RenditionAuthorization { + authorizedAt := time.Now().UTC().Add(-time.Minute) + expiresAt := authorizedAt.Add(10 * time.Minute) + return RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, InputKind: metadata.InputKind, + AllowedArtifactRoles: []EvidenceArtifactRole{EvidenceArtifactStructured}, + MaxProviderMarkdownBytes: 1_024, MaxArtifactBytes: 1_024, + MaxArtifacts: 1, MaxTotalResultBytes: 4_096, + AuthorizedAt: authorizedAt.Format(renditionTimestampForm), + ExpiresAt: expiresAt.Format(renditionTimestampForm), + } +} + +func validRenditionResult( + descriptor RenditionDescriptor, authorization RenditionAuthorization, +) RenditionResult { + authorizedAt, _ := parseRenditionTimestamp(authorization.AuthorizedAt) + payload := []byte(`{"synthetic":"structured"}`) + digest := sha256.Sum256(payload) + return RenditionResult{ + Evidence: SourceEvidenceV1{ + ContractVersion: SourceEvidenceContractV1, + Completeness: EvidenceDegradedProvenance, + Family: "pdf", + Artifacts: []SourceEvidenceArtifactV1{{ + ProviderID: "provider-artifact-1", Pointer: "provider/structured.json", + Role: EvidenceArtifactStructured, SHA256: hex.EncodeToString(digest[:]), + }}, + UnitKind: EvidenceUnitGeneric, + Omissions: []SourceEvidenceOmissionV1{{ + Kind: EvidenceOmissionField, Field: "natural_provenance", + Reason: "synthetic provider returned generic evidence", + }}, + Units: []SourceEvidenceUnitV1{{ + Order: 0, Text: "synthetic evidence", + Locator: SourceEvidenceLocatorV1{ + Kind: EvidenceLocatorGeneric, IndexOrigin: EvidenceIndexOriginNone, + }, + }}, + }, + ProviderMarkdown: []byte("synthetic evidence\n"), + Artifacts: []RenditionArtifact{{ + Role: EvidenceArtifactStructured, MediaType: "application/json", + Payload: payload, SHA256: hex.EncodeToString(digest[:]), + }}, + Receipt: RenditionReceipt{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, + SourceSHA256: authorization.SourceSHA256, + OperationID: "operation-synthetic-1", + StartedAt: authorizedAt.Add(time.Second).Format(renditionTimestampForm), + CompletedAt: authorizedAt.Add(2 * time.Second).Format(renditionTimestampForm), + Warnings: []string{"degraded_provenance"}, + Usage: RenditionUsage{ + Requests: 1, InputBytes: authorization.SourceBytes, + OutputBytes: int64(len(payload) + len("synthetic evidence\n")), Units: 1, + }, + }, + } +} diff --git a/document/providerhttp/policy.go b/document/providerhttp/policy.go new file mode 100644 index 00000000..3761be64 --- /dev/null +++ b/document/providerhttp/policy.go @@ -0,0 +1,154 @@ +// Package providerhttp constructs HTTP transports whose network destination is +// bound to an explicit provider egress policy. +package providerhttp + +import ( + "crypto/x509" + "encoding/hex" + "errors" + "fmt" + "net/netip" + "strings" + "time" +) + +const ( + // DefaultConnectTimeout bounds one TCP connection attempt. + DefaultConnectTimeout = 30 * time.Second + // DefaultKeepAlive preserves the standard library's TCP keepalive cadence. + DefaultKeepAlive = 30 * time.Second + // DefaultTLSHandshakeTimeout bounds a provider TLS handshake. + DefaultTLSHandshakeTimeout = 10 * time.Second + maxTransportTimeout = 5 * time.Minute +) + +// ProxyMode describes whether a provider transport may use a proxy. +type ProxyMode string + +const ( + // ProxyDisabled makes all connections directly to a policy-approved IP. + // Environment proxy variables are never consulted. + ProxyDisabled ProxyMode = "disabled" +) + +// TLSPolicy contains optional certificate authority and SPKI pin restrictions. +// Normal hostname and chain verification always runs before a pin is checked. +type TLSPolicy struct { + RootCAs *x509.CertPool + SPKISHA256 []string +} + +// EgressPolicy authorizes one exact scheme, host, and port plus every network +// prefix to which that hostname is allowed to resolve. +type EgressPolicy struct { + Scheme string + Host string + Port uint16 + AllowedCIDRs []netip.Prefix + ProxyMode ProxyMode + ConnectTimeout time.Duration + KeepAlive time.Duration + TLSHandshakeTimeout time.Duration + TLS TLSPolicy +} + +type validatedPolicy struct { + scheme string + host string + port uint16 + allowedCIDRs []netip.Prefix + rootCAs *x509.CertPool + spkiPins [][32]byte + connect time.Duration + keepAlive time.Duration + tlsHandshake time.Duration +} + +func validatePolicy(policy EgressPolicy) (validatedPolicy, error) { + if policy.Scheme != "http" && policy.Scheme != "https" { + return validatedPolicy{}, errors.New("provider egress scheme must be http or https") + } + host, err := normalizeHost(policy.Host) + if err != nil { + return validatedPolicy{}, err + } + if policy.Port == 0 { + return validatedPolicy{}, errors.New("provider egress port is required") + } + if policy.ProxyMode != "" && policy.ProxyMode != ProxyDisabled { + return validatedPolicy{}, errors.New("provider egress supports only direct connections") + } + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = DefaultTLSHandshakeTimeout + } + if policy.ConnectTimeout < 0 || policy.ConnectTimeout > maxTransportTimeout || + policy.KeepAlive < 0 || policy.KeepAlive > maxTransportTimeout || + policy.TLSHandshakeTimeout < 0 || policy.TLSHandshakeTimeout > maxTransportTimeout { + return validatedPolicy{}, errors.New("provider egress transport timeout is outside package bounds") + } + if len(policy.AllowedCIDRs) == 0 { + return validatedPolicy{}, errors.New("provider egress requires at least one allowed CIDR") + } + allowed := make([]netip.Prefix, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + if !prefix.IsValid() || prefix.Addr().Zone() != "" { + return validatedPolicy{}, fmt.Errorf("provider egress CIDR %d is invalid", index) + } + allowed[index] = prefix.Masked() + } + if policy.Scheme != "https" && (policy.TLS.RootCAs != nil || len(policy.TLS.SPKISHA256) != 0) { + return validatedPolicy{}, errors.New("provider egress TLS policy requires an https scheme") + } + pins := make([][32]byte, len(policy.TLS.SPKISHA256)) + for index, encoded := range policy.TLS.SPKISHA256 { + decoded, err := hex.DecodeString(encoded) + if err != nil || len(decoded) != len(pins[index]) { + return validatedPolicy{}, fmt.Errorf("provider egress SPKI pin %d must be a SHA-256 hex digest", index) + } + copy(pins[index][:], decoded) + } + var roots *x509.CertPool + if policy.TLS.RootCAs != nil { + roots = policy.TLS.RootCAs.Clone() + } + return validatedPolicy{ + scheme: policy.Scheme, host: host, port: policy.Port, + allowedCIDRs: allowed, rootCAs: roots, spkiPins: pins, + connect: policy.ConnectTimeout, keepAlive: policy.KeepAlive, + tlsHandshake: policy.TLSHandshakeTimeout, + }, nil +} + +func normalizeHost(host string) (string, error) { + if host == "" || host != strings.TrimSpace(host) { + return "", errors.New("provider egress host is required without surrounding whitespace") + } + if address, err := netip.ParseAddr(host); err == nil { + if address.Zone() != "" { + return "", errors.New("provider egress IP host cannot contain a zone") + } + return address.Unmap().String(), nil + } + if len(host) > 253 || strings.ContainsAny(host, ":/\\?#@[]") { + return "", errors.New("provider egress host must be a bare DNS name or IP address") + } + host = strings.ToLower(host) + for label := range strings.SplitSeq(host, ".") { + if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return "", errors.New("provider egress host is not a valid DNS name") + } + for _, character := range label { + if (character < 'a' || character > 'z') && + (character < '0' || character > '9') && character != '-' { + return "", errors.New("provider egress host is not a valid ASCII DNS name") + } + } + } + return host, nil +} diff --git a/document/providerhttp/transport.go b/document/providerhttp/transport.go new file mode 100644 index 00000000..3012975e --- /dev/null +++ b/document/providerhttp/transport.go @@ -0,0 +1,252 @@ +package providerhttp + +import ( + "context" + "crypto/sha256" + "crypto/subtle" + "crypto/tls" + "errors" + "fmt" + "net" + "net/http" + "net/netip" + "net/url" + "strconv" +) + +var ( + // ErrDestinationDenied means the requested scheme, host, or port was not + // the exact destination authorized by the egress policy. + ErrDestinationDenied = errors.New("provider egress destination denied") + // ErrAddressDenied means DNS returned at least one address outside every + // network prefix authorized by the egress policy. + ErrAddressDenied = errors.New("provider egress address denied") + // ErrCertificatePin means normal TLS verification succeeded but the leaf + // certificate did not match an authorized SHA-256 SPKI digest. + ErrCertificatePin = errors.New("provider egress certificate pin mismatch") +) + +// Resolver performs the one DNS lookup allowed for each connection attempt. +type Resolver interface { + LookupNetIP(ctx context.Context, network string, host string) ([]netip.Addr, error) +} + +// RefuseRedirects is an http.Client CheckRedirect function that prevents a +// request body or credential from being replayed, including to the same host. +func RefuseRedirects(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse +} + +// NewTransport returns a transport that resolves once per connection attempt, +// validates every answer, and connects to one exact allowed IP without an +// ambient proxy or a second hostname lookup. +func NewTransport(policy EgressPolicy, resolver Resolver) (*http.Transport, error) { + validated, err := validatePolicy(policy) + if err != nil { + return nil, err + } + if resolver == nil { + resolver = net.DefaultResolver + } + dialer := &policyDialer{ + policy: validated, resolver: resolver, + dialer: net.Dialer{Timeout: validated.connect, KeepAlive: validated.keepAlive}, + } + tlsConfig := &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: validated.rootCAs, + ServerName: validated.host, + } + if len(validated.spkiPins) != 0 { + pins := append([][32]byte(nil), validated.spkiPins...) + tlsConfig.VerifyConnection = func(state tls.ConnectionState) error { + if len(state.PeerCertificates) == 0 { + return ErrCertificatePin + } + digest := sha256.Sum256(state.PeerCertificates[0].RawSubjectPublicKeyInfo) + for _, pin := range pins { + if subtle.ConstantTimeCompare(digest[:], pin[:]) == 1 { + return nil + } + } + return ErrCertificatePin + } + } + defaultTransport, ok := http.DefaultTransport.(*http.Transport) + if !ok { + return nil, errors.New("default HTTP transport has an unsupported implementation") + } + transport := defaultTransport.Clone() + transport.Proxy = func(request *http.Request) (*url.URL, error) { + if err := validated.validateRequest(request); err != nil { + return nil, err + } + return nil, nil //nolint:nilnil // nil is the http.Transport contract for a direct connection. + } + transport.TLSClientConfig = tlsConfig + transport.TLSHandshakeTimeout = validated.tlsHandshake + transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + if validated.scheme != "http" { + return nil, fmt.Errorf("%w: plaintext connection", ErrDestinationDenied) + } + return dialer.dial(ctx, network, address) + } + transport.DialTLSContext = func(ctx context.Context, network, address string) (net.Conn, error) { + if validated.scheme != "https" { + return nil, fmt.Errorf("%w: TLS connection", ErrDestinationDenied) + } + connection, err := dialer.dial(ctx, network, address) + if err != nil { + return nil, err + } + tlsConnection := tls.Client(connection, tlsConfig.Clone()) + handshakeContext, cancel := context.WithTimeout(ctx, validated.tlsHandshake) + defer cancel() + if err := tlsConnection.HandshakeContext(handshakeContext); err != nil { + handshakeErr := fmt.Errorf("handshake with provider: %w", err) + if closeErr := connection.Close(); closeErr != nil { + return nil, errors.Join(handshakeErr, fmt.Errorf("close failed provider connection: %w", closeErr)) + } + return nil, handshakeErr + } + return tlsConnection, nil + } + return transport, nil +} + +func (policy validatedPolicy) validateRequest(request *http.Request) error { + if request == nil || request.URL == nil || request.URL.Scheme != policy.scheme || + request.URL.User != nil { + return fmt.Errorf("%w: request scheme or authority", ErrDestinationDenied) + } + host, port, err := parseAuthority(request.URL.Host, policy.scheme) + if err != nil || host != policy.host || port != policy.port { + return fmt.Errorf("%w: request host or port", ErrDestinationDenied) + } + if request.Host != "" { + host, port, err = parseAuthority(request.Host, policy.scheme) + if err != nil || host != policy.host || port != policy.port { + return fmt.Errorf("%w: Host header", ErrDestinationDenied) + } + } + return nil +} + +func parseAuthority(authority, scheme string) (string, uint16, error) { + parsed, err := url.Parse(scheme + "://" + authority) + if err != nil || parsed.User != nil || parsed.Host == "" || parsed.Path != "" || + parsed.RawQuery != "" || parsed.Fragment != "" { + return "", 0, errors.New("invalid provider egress authority") + } + host, err := normalizeHost(parsed.Hostname()) + if err != nil { + return "", 0, err + } + port := parsed.Port() + if port == "" { + switch scheme { + case "http": + port = "80" + case "https": + port = "443" + default: + return "", 0, errors.New("invalid provider egress scheme") + } + } + parsedPort, err := strconv.ParseUint(port, 10, 16) + if err != nil || parsedPort == 0 { + return "", 0, errors.New("invalid provider egress port") + } + return host, uint16(parsedPort), nil +} + +type policyDialer struct { + policy validatedPolicy + resolver Resolver + dialer net.Dialer +} + +func (dialer *policyDialer) dial( + ctx context.Context, + network string, + address string, +) (net.Conn, error) { + if network != "tcp" && network != "tcp4" && network != "tcp6" { + return nil, fmt.Errorf("%w: unsupported network", ErrDestinationDenied) + } + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, fmt.Errorf("%w: malformed address", ErrDestinationDenied) + } + host, err = normalizeHost(host) + if err != nil || host != dialer.policy.host || port != strconv.Itoa(int(dialer.policy.port)) { + return nil, fmt.Errorf("%w: host or port", ErrDestinationDenied) + } + addresses, err := dialer.resolve(ctx) + if err != nil { + return nil, err + } + var failures []error + for _, selected := range addresses { + exactNetwork := "tcp6" + if selected.Is4() { + exactNetwork = "tcp4" + } + connection, dialErr := dialer.dialer.DialContext( + ctx, exactNetwork, net.JoinHostPort(selected.String(), port), + ) + if dialErr == nil { + return connection, nil + } + failures = append(failures, dialErr) + if err := ctx.Err(); err != nil { + return nil, err + } + } + return nil, fmt.Errorf("dial provider IPs: %w", errors.Join(failures...)) +} + +func (dialer *policyDialer) resolve(ctx context.Context) ([]netip.Addr, error) { + if address, err := netip.ParseAddr(dialer.policy.host); err == nil { + addresses := []netip.Addr{address.Unmap()} + if err := dialer.validateAddresses(addresses); err != nil { + return nil, err + } + return addresses, nil + } + addresses, err := dialer.resolver.LookupNetIP(ctx, "ip", dialer.policy.host) + if err != nil { + return nil, fmt.Errorf("resolve provider egress host: %w", err) + } + if err := dialer.validateAddresses(addresses); err != nil { + return nil, err + } + result := make([]netip.Addr, len(addresses)) + for index, address := range addresses { + result[index] = address.Unmap() + } + return result, nil +} + +func (dialer *policyDialer) validateAddresses(addresses []netip.Addr) error { + if len(addresses) == 0 { + return errors.New("provider egress DNS returned no addresses") + } + for _, address := range addresses { + if !address.IsValid() || address.Zone() != "" { + return fmt.Errorf("%w: invalid DNS answer", ErrAddressDenied) + } + address = address.Unmap() + allowed := false + for _, prefix := range dialer.policy.allowedCIDRs { + if prefix.Contains(address) { + allowed = true + break + } + } + if !allowed { + return ErrAddressDenied + } + } + return nil +} diff --git a/document/providerhttp/transport_test.go b/document/providerhttp/transport_test.go new file mode 100644 index 00000000..e6741f37 --- /dev/null +++ b/document/providerhttp/transport_test.go @@ -0,0 +1,450 @@ +package providerhttp_test + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/hex" + "encoding/pem" + "io" + "math/big" + "net" + "net/http" + "net/http/httptest" + "net/netip" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document/providerhttp" +) + +func TestEgressDialsResolvedIPWithoutASecondLookup(t *testing.T) { + var hostHeader string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hostHeader = r.Host + _, _ = io.WriteString(w, "direct") + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + resolver := &recordingResolver{answers: [][]netip.Addr{{netip.MustParseAddr(host)}}} + transport := newHTTPTransport(t, port, resolver) + client := &http.Client{Transport: transport, CheckRedirect: providerhttp.RefuseRedirects} + + response, err := client.Get("http://provider.invalid:" + strconv.Itoa(int(port)) + "/source") + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + + assert.Equal(t, "direct", string(body)) + assert.Equal(t, "provider.invalid:"+strconv.Itoa(int(port)), hostHeader) + assert.Equal(t, []string{"provider.invalid"}, resolver.hosts()) +} + +func TestEgressIgnoresAmbientProxy(t *testing.T) { + var proxyRequests atomic.Int32 + proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + proxyRequests.Add(1) + http.Error(w, "proxy used", http.StatusBadGateway) + })) + defer proxy.Close() + t.Setenv("HTTP_PROXY", proxy.URL) + t.Setenv("HTTPS_PROXY", proxy.URL) + t.Setenv("NO_PROXY", "") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "origin") + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + transport := newHTTPTransport(t, port, &recordingResolver{ + answers: [][]netip.Addr{{netip.MustParseAddr(host)}}, + }) + + response, err := (&http.Client{Transport: transport}).Get( + "http://provider.invalid:" + strconv.Itoa(int(port)), + ) + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + + assert.Equal(t, "origin", string(body)) + assert.Zero(t, proxyRequests.Load()) +} + +func TestEgressRejectsMixedAllowedAndDeniedDNSAnswers(t *testing.T) { + resolver := &recordingResolver{answers: [][]netip.Addr{{ + netip.MustParseAddr("127.0.0.1"), + netip.MustParseAddr("203.0.113.10"), + }}} + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: "http", Host: "provider.invalid", Port: 8080, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, + }, resolver) + require.NoError(t, err) + + response, err := (&http.Client{Transport: transport}).Get("http://provider.invalid:8080/") + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, providerhttp.ErrAddressDenied) + assert.Equal(t, []string{"provider.invalid"}, resolver.hosts()) +} + +func TestEgressRevalidatesDNSOnEachConnection(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "first") + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + resolver := &recordingResolver{answers: [][]netip.Addr{ + {netip.MustParseAddr(host)}, + {netip.MustParseAddr("203.0.113.10")}, + }} + transport := newHTTPTransport(t, port, resolver) + transport.DisableKeepAlives = true + client := &http.Client{Transport: transport} + url := "http://provider.invalid:" + strconv.Itoa(int(port)) + "/" + + first, err := client.Get(url) + require.NoError(t, err) + require.NoError(t, first.Body.Close()) + second, err := client.Get(url) + if second != nil { + require.NoError(t, second.Body.Close()) + } + + require.ErrorIs(t, err, providerhttp.ErrAddressDenied) + assert.Equal(t, []string{"provider.invalid", "provider.invalid"}, resolver.hosts()) +} + +func TestEgressFallsBackAcrossApprovedDNSAnswers(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "fallback") + })) + defer server.Close() + _, port := endpoint(t, server.Listener.Addr()) + resolver := &recordingResolver{answers: [][]netip.Addr{{ + netip.MustParseAddr("127.0.0.2"), + netip.MustParseAddr("127.0.0.1"), + }}} + transport := newHTTPTransport(t, port, resolver) + + response, err := (&http.Client{Transport: transport}).Get( + "http://provider.invalid:" + strconv.Itoa(int(port)), + ) + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + assert.Equal(t, "fallback", string(body)) + assert.Equal(t, []string{"provider.invalid"}, resolver.hosts()) +} + +func TestEgressRejectsUnapprovedDestinationBeforeResolution(t *testing.T) { + resolver := &recordingResolver{answers: [][]netip.Addr{{netip.MustParseAddr("127.0.0.1")}}} + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: "http", Host: "provider.invalid", Port: 8080, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + }, resolver) + require.NoError(t, err) + client := &http.Client{Transport: transport} + + for _, rawURL := range []string{ + "http://other.invalid:8080/", + "http://provider.invalid:8081/", + "https://provider.invalid:8080/", + } { + response, err := client.Get(rawURL) + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, providerhttp.ErrDestinationDenied, rawURL) + } + request, err := http.NewRequest(http.MethodGet, "http://provider.invalid:8080/", nil) + require.NoError(t, err) + request.Host = "other.invalid:8080" + response, err := client.Do(request) + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, providerhttp.ErrDestinationDenied) + assert.Empty(t, resolver.hosts()) +} + +func TestEgressUsesLiteralIPWithoutDNS(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "literal") + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + resolver := &recordingResolver{} + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: "http", Host: host, Port: port, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + }, resolver) + require.NoError(t, err) + + response, err := (&http.Client{Transport: transport}).Get(server.URL) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + assert.Empty(t, resolver.hosts()) +} + +func TestEgressPreservesTLSIdentityAndEnforcesSPKIPin(t *testing.T) { + certificate, roots, pin := providerCertificate(t, "provider.invalid") + var sni string + var hostHeader string + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hostHeader = r.Host + _, _ = io.WriteString(w, "secure") + })) + server.TLS = &tls.Config{ + Certificates: []tls.Certificate{certificate}, + GetConfigForClient: func(hello *tls.ClientHelloInfo) (*tls.Config, error) { + sni = hello.ServerName + return nil, nil //nolint:nilnil // nil keeps this server's configured certificate. + }, + } + server.StartTLS() + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + resolver := &recordingResolver{answers: [][]netip.Addr{{netip.MustParseAddr(host)}}} + policy := providerhttp.EgressPolicy{ + Scheme: "https", Host: "provider.invalid", Port: port, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + TLS: providerhttp.TLSPolicy{RootCAs: roots, SPKISHA256: []string{pin}}, + } + transport, err := providerhttp.NewTransport(policy, resolver) + require.NoError(t, err) + client := &http.Client{Transport: transport} + url := "https://provider.invalid:" + strconv.Itoa(int(port)) + "/result" + + response, err := client.Get(url) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + assert.Equal(t, "provider.invalid", sni) + assert.Equal(t, "provider.invalid:"+strconv.Itoa(int(port)), hostHeader) + + policy.TLS.SPKISHA256 = []string{hex.EncodeToString(make([]byte, sha256.Size))} + transport, err = providerhttp.NewTransport(policy, resolver) + require.NoError(t, err) + response, err = (&http.Client{Transport: transport}).Get(url) + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, providerhttp.ErrCertificatePin) +} + +func TestEgressBoundsTLSHandshake(t *testing.T) { + listener, err := net.Listen("tcp4", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, listener.Close()) }) + accepted := make(chan net.Conn, 1) + go func() { + connection, acceptErr := listener.Accept() + if acceptErr == nil { + accepted <- connection + } + }() + _, port := endpoint(t, listener.Addr()) + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: "https", Host: "127.0.0.1", Port: port, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + TLSHandshakeTimeout: 25 * time.Millisecond, + }, &recordingResolver{}) + require.NoError(t, err) + started := time.Now() + + response, err := (&http.Client{Transport: transport}).Get( + "https://127.0.0.1:" + strconv.Itoa(int(port)), + ) + if response != nil { + require.NoError(t, response.Body.Close()) + } + require.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(started), time.Second) + select { + case connection := <-accepted: + require.NoError(t, connection.Close()) + case <-time.After(time.Second): + require.Fail(t, "provider connection was not accepted") + } +} + +func TestEgressPolicyIsImmutableAfterTransportCreation(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "immutable") + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + allowed := []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")} + policy := providerhttp.EgressPolicy{ + Scheme: "http", Host: "provider.invalid", Port: port, AllowedCIDRs: allowed, + } + transport, err := providerhttp.NewTransport(policy, &recordingResolver{ + answers: [][]netip.Addr{{netip.MustParseAddr(host)}}, + }) + require.NoError(t, err) + allowed[0] = netip.MustParsePrefix("203.0.113.0/24") + + response, err := (&http.Client{Transport: transport}).Get( + "http://provider.invalid:" + strconv.Itoa(int(port)), + ) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) +} + +func TestEgressRefusesRedirectReplay(t *testing.T) { + var redirected atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/redirected" { + redirected.Add(1) + return + } + http.Redirect(w, r, "/redirected", http.StatusTemporaryRedirect) + })) + defer server.Close() + host, port := endpoint(t, server.Listener.Addr()) + transport := newHTTPTransport(t, port, &recordingResolver{ + answers: [][]netip.Addr{{netip.MustParseAddr(host)}}, + }) + client := &http.Client{Transport: transport, CheckRedirect: providerhttp.RefuseRedirects} + + response, err := client.Get("http://provider.invalid:" + strconv.Itoa(int(port))) + require.NoError(t, err) + require.NoError(t, response.Body.Close()) + assert.Equal(t, http.StatusTemporaryRedirect, response.StatusCode) + assert.Zero(t, redirected.Load()) +} + +func TestEgressValidatesPolicy(t *testing.T) { + valid := providerhttp.EgressPolicy{ + Scheme: "https", Host: "provider.invalid", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("203.0.113.0/24")}, + ProxyMode: providerhttp.ProxyDisabled, + } + for _, test := range []struct { + name string + mutate func(*providerhttp.EgressPolicy) + }{ + {name: "scheme", mutate: func(policy *providerhttp.EgressPolicy) { policy.Scheme = "ftp" }}, + {name: "host", mutate: func(policy *providerhttp.EgressPolicy) { policy.Host = "" }}, + {name: "host with port", mutate: func(policy *providerhttp.EgressPolicy) { policy.Host = "provider.invalid:443" }}, + {name: "port", mutate: func(policy *providerhttp.EgressPolicy) { policy.Port = 0 }}, + {name: "CIDRs", mutate: func(policy *providerhttp.EgressPolicy) { policy.AllowedCIDRs = nil }}, + {name: "proxy", mutate: func(policy *providerhttp.EgressPolicy) { policy.ProxyMode = "environment" }}, + {name: "pin", mutate: func(policy *providerhttp.EgressPolicy) { policy.TLS.SPKISHA256 = []string{"bad"} }}, + {name: "connect timeout", mutate: func(policy *providerhttp.EgressPolicy) { policy.ConnectTimeout = -time.Second }}, + {name: "keepalive", mutate: func(policy *providerhttp.EgressPolicy) { policy.KeepAlive = -time.Second }}, + {name: "TLS handshake timeout", mutate: func(policy *providerhttp.EgressPolicy) { policy.TLSHandshakeTimeout = -time.Second }}, + } { + t.Run(test.name, func(t *testing.T) { + policy := valid + test.mutate(&policy) + _, err := providerhttp.NewTransport(policy, &recordingResolver{}) + require.Error(t, err) + }) + } +} + +func newHTTPTransport( + t *testing.T, + port uint16, + resolver providerhttp.Resolver, +) *http.Transport { + t.Helper() + transport, err := providerhttp.NewTransport(providerhttp.EgressPolicy{ + Scheme: "http", Host: "provider.invalid", Port: port, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, + }, resolver) + require.NoError(t, err) + return transport +} + +func endpoint(t *testing.T, address net.Addr) (string, uint16) { + t.Helper() + host, rawPort, err := net.SplitHostPort(address.String()) + require.NoError(t, err) + port, err := strconv.ParseUint(rawPort, 10, 16) + require.NoError(t, err) + return host, uint16(port) +} + +type recordingResolver struct { + mu sync.Mutex + answers [][]netip.Addr + calls []string +} + +func (resolver *recordingResolver) LookupNetIP( + _ context.Context, + network string, + host string, +) ([]netip.Addr, error) { + resolver.mu.Lock() + defer resolver.mu.Unlock() + resolver.calls = append(resolver.calls, host) + if network != "ip" { + return nil, assert.AnError + } + if len(resolver.answers) == 0 { + return nil, assert.AnError + } + answer := resolver.answers[0] + if len(resolver.answers) > 1 { + resolver.answers = resolver.answers[1:] + } + return append([]netip.Addr(nil), answer...), nil +} + +func (resolver *recordingResolver) hosts() []string { + resolver.mu.Lock() + defer resolver.mu.Unlock() + return append([]string(nil), resolver.calls...) +} + +func providerCertificate(t *testing.T, hostname string) (tls.Certificate, *x509.CertPool, string) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120)) + require.NoError(t, err) + template := &x509.Certificate{ + SerialNumber: serial, + Subject: pkix.Name{CommonName: hostname}, + DNSNames: []string{hostname}, + NotBefore: time.Now().Add(-time.Minute), + NotAfter: time.Now().Add(time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + } + der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key) + require.NoError(t, err) + keyDER, err := x509.MarshalPKCS8PrivateKey(key) + require.NoError(t, err) + certificate, err := tls.X509KeyPair( + pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}), + pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}), + ) + require.NoError(t, err) + parsed, err := x509.ParseCertificate(der) + require.NoError(t, err) + roots := x509.NewCertPool() + roots.AddCert(parsed) + digest := sha256.Sum256(parsed.RawSubjectPublicKeyInfo) + return certificate, roots, hex.EncodeToString(digest[:]) +} diff --git a/document/pymupdf/doc.go b/document/pymupdf/doc.go new file mode 100644 index 00000000..86c4a64d --- /dev/null +++ b/document/pymupdf/doc.go @@ -0,0 +1,2 @@ +// Package pymupdf implements the bounded local-process PyMuPDF PDF-text provider. +package pymupdf diff --git a/document/pymupdf/provider.go b/document/pymupdf/provider.go new file mode 100644 index 00000000..96694628 --- /dev/null +++ b/document/pymupdf/provider.go @@ -0,0 +1,494 @@ +package pymupdf + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "slices" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/formatdetect" +) + +const ( + providerID = "pymupdf.local-v1" + protocolVersion = "docbank-pymupdf/v1" + profileVersion = "docbank-pymupdf-profile/v1" + timestampForm = "2006-01-02T15:04:05.000000000Z" + childDrainWindow = 250 * time.Millisecond + + // MaxDocumentBytes is the largest PDF accepted by a local provider profile. + MaxDocumentBytes = formatdetect.MaxDocumentBytes + // MaxResponseBytes is the largest structured child response accepted by a profile. + MaxResponseBytes = int64(64 << 20) + // MaxPages is the largest locally verified page sequence accepted by a profile. + MaxPages = 100_000 + // MaxTimeout is the largest local child deadline accepted by a profile. + MaxTimeout = 30 * time.Minute +) + +var ( + errInputIdentity = errors.New("authorized input identity changed") + errOutputTooLarge = errors.New("child output exceeds limit") +) + +// Profile fixes one executable, immutable runtime identity, and all local bounds. +type Profile struct { + Executable string + RuntimeIdentity string + MaxDocumentBytes int64 + MaxResponseBytes int64 + MaxPages int + Timeout time.Duration +} + +// Provider renders authorized PDF bytes through one directly configured executable. +type Provider struct { + descriptor document.RenditionDescriptor + executable string + runtimeIdentity string + maxDocumentBytes int64 + maxResponseBytes int64 + maxPages int + timeout time.Duration +} + +// New constructs one immutable local-process provider profile. +func New(profile Profile) (*Provider, error) { + if !filepath.IsAbs(profile.Executable) || filepath.Clean(profile.Executable) != profile.Executable { + return nil, errors.New("pymupdf: executable must be an absolute clean path") + } + if pythonInterpreter(filepath.Base(profile.Executable)) { + return nil, errors.New("pymupdf: executable must not be a Python interpreter") + } + info, err := os.Lstat(profile.Executable) + if err != nil { + return nil, errors.New("pymupdf: configured executable is unavailable") + } + if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 { + return nil, errors.New("pymupdf: executable must be a regular non-symlink file") + } + if err := validateRuntimeIdentity(profile.RuntimeIdentity); err != nil { + return nil, err + } + if profile.MaxDocumentBytes <= 0 || profile.MaxDocumentBytes > MaxDocumentBytes { + return nil, fmt.Errorf("pymupdf: max document bytes must be between 1 and %d", MaxDocumentBytes) + } + if profile.MaxResponseBytes <= 0 || profile.MaxResponseBytes > MaxResponseBytes { + return nil, fmt.Errorf("pymupdf: max response bytes must be between 1 and %d", MaxResponseBytes) + } + if profile.MaxPages <= 0 || profile.MaxPages > MaxPages { + return nil, fmt.Errorf("pymupdf: max pages must be between 1 and %d", MaxPages) + } + if profile.Timeout <= 0 || profile.Timeout > MaxTimeout { + return nil, fmt.Errorf("pymupdf: timeout must be between 1ns and %s", MaxTimeout) + } + identity := strings.Join([]string{ + profileVersion, protocolVersion, profile.Executable, profile.RuntimeIdentity, + strconv.FormatInt(profile.MaxDocumentBytes, 10), strconv.FormatInt(profile.MaxResponseBytes, 10), + strconv.Itoa(profile.MaxPages), strconv.FormatInt(int64(profile.Timeout), 10), + }, "\x00") + policyDigest := sha256.Sum256([]byte(identity)) + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: hex.EncodeToString(policyDigest[:]), + TrustBoundary: document.RenditionTrustLocalProcess, + SupportedFormats: []document.RenditionFormatCapability{{ + MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile, + }}, + ReturnsStructured: true, + }) + if err != nil { + return nil, fmt.Errorf("pymupdf: construct descriptor: %w", err) + } + return &Provider{ + descriptor: cloneDescriptor(descriptor), executable: profile.Executable, + runtimeIdentity: profile.RuntimeIdentity, maxDocumentBytes: profile.MaxDocumentBytes, + maxResponseBytes: profile.MaxResponseBytes, maxPages: profile.MaxPages, timeout: profile.Timeout, + }, nil +} + +// Descriptor returns the immutable provider identity fixed by the profile. +func (provider *Provider) Descriptor() document.RenditionDescriptor { + if provider == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(provider.descriptor) +} + +// Render re-verifies one PDF and sends its exact authorized bytes through stdin only. +func (provider *Provider) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if provider == nil { + return document.RenditionResult{}, errors.New("pymupdf: provider is required") + } + if _, err := document.ValidateRenditionProviderRequest(provider, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > provider.maxDocumentBytes { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF input exceeds the configured byte limit", nil) + } + expiresAt, err := time.Parse(timestampForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF authorization is invalid", err) + } + startedAt := time.Now().UTC() + timeoutAt := startedAt.Add(provider.timeout) + deadline := timeoutAt + expiryDeadline := false + if expiresAt.Before(deadline) { + deadline = expiresAt + expiryDeadline = true + } + operationCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + if err := operationCtx.Err(); err != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, err) + } + + source, err := readExact(operationCtx, upload, metadata.ByteLength, provider.maxDocumentBytes) + if err != nil { + if operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + if errors.Is(err, errInputIdentity) { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF input identity does not match authorization", err) + } + return document.RenditionResult{}, providerError(document.RenditionErrorTransient, + "PyMuPDF input could not be read", err) + } + defer clear(source) + digest := sha256.Sum256(source) + if int64(len(source)) != metadata.ByteLength || hex.EncodeToString(digest[:]) != metadata.SHA256 { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF input identity does not match authorization", nil) + } + candidate, err := formatdetect.DetectFormat(bytes.NewReader(source), int64(len(source)), metadata.MediaType) + if err != nil || candidate.ID != "pdf" || candidate.Family != "pdf" { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, + "PyMuPDF input is not a locally verified PDF", err) + } + localPages, err := formatdetect.CountPDFPages(source) + if err != nil || localPages <= 0 { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, + "PyMuPDF page count could not be verified", err) + } + if localPages > int64(provider.maxPages) { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF PDF exceeds the configured page limit", nil) + } + + stdout := &boundedBuffer{limit: provider.maxResponseBytes} + command := exec.CommandContext( //nolint:gosec // the operator pins one direct executable; source bytes never select it + operationCtx, provider.executable, "--protocol", protocolVersion, + ) + stdout.overflow = func() { + if command.Process != nil { + _ = command.Process.Kill() + } + } + command.Dir = filepath.Dir(provider.executable) + command.Env = cleanEnvironment() + command.Stdin = bytes.NewReader(source) + command.Stdout = stdout + command.Stderr = io.Discard + command.WaitDelay = childDrainWindow + runErr := command.Run() + if operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + if stdout.exceeded { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "PyMuPDF output exceeds the configured byte limit", nil) + } + if runErr != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorTransient, + "PyMuPDF executable failed", runErr) + } + + wire, err := parseResponse(stdout.Bytes(), provider.runtimeIdentity, metadata, int(localPages), provider.maxPages) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "PyMuPDF output is malformed", err) + } + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: "pdf", UnitKind: document.EvidenceUnitPage, + Units: make([]document.SourceEvidenceUnitV1, len(wire.Pages)), + } + for index, page := range wire.Pages { + evidence.Units[index] = document.SourceEvidenceUnitV1{ + Order: index, ProviderID: fmt.Sprintf("pymupdf-page-%d", page.Number), Text: page.Text, + Locator: document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginOne, + Start: int64(page.Number), End: int64(page.Number), + }, + } + } + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "PyMuPDF output is malformed", err) + } + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + completedAt := time.Now().UTC() + result := document.RenditionResult{ + Evidence: evidence, + Receipt: document.RenditionReceipt{ + ProviderID: provider.descriptor.ID, DescriptorFingerprint: provider.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "pymupdf-" + authorization.RenditionRequestFingerprint[:24], + StartedAt: startedAt.Format(timestampForm), CompletedAt: completedAt.Format(timestampForm), + Usage: document.RenditionUsage{ + Requests: 1, InputBytes: metadata.ByteLength, OutputBytes: int64(stdout.Len()), Units: localPages, + }, + }, + } + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + return result, nil +} + +type response struct { + ContractVersion string `json:"contract_version"` + RuntimeIdentity string `json:"runtime_identity"` + SourceSHA256 string `json:"source_sha256"` + SourceBytes int64 `json:"source_bytes"` + Complete bool `json:"complete"` + PageCount int `json:"page_count"` + Pages []page `json:"pages"` +} + +type page struct { + Number int `json:"number"` + Text string `json:"text"` + EmptyReason string `json:"empty_reason,omitempty"` +} + +func parseResponse( + raw []byte, runtimeIdentity string, metadata document.AuthorizedUploadMetadata, + localPages, maxPages int, +) (response, error) { + if len(raw) == 0 { + return response{}, errors.New("empty response") + } + var wire response + if err := json.Unmarshal(raw, &wire, json.RejectUnknownMembers(true)); err != nil { + return response{}, err + } + if wire.ContractVersion != protocolVersion { + return response{}, errors.New("protocol version changed") + } + if wire.RuntimeIdentity != runtimeIdentity { + return response{}, errors.New("runtime identity changed") + } + if wire.SourceSHA256 != metadata.SHA256 || wire.SourceBytes != metadata.ByteLength { + return response{}, errors.New("source identity changed") + } + if !wire.Complete { + return response{}, errors.New("response is partial") + } + if wire.PageCount <= 0 || wire.PageCount > maxPages || wire.PageCount != localPages || len(wire.Pages) != wire.PageCount { + return response{}, errors.New("page count changed") + } + for index, page := range wire.Pages { + if page.Number != index+1 { + return response{}, errors.New("page sequence is not complete and contiguous") + } + if !utf8.ValidString(page.Text) || strings.ContainsRune(page.Text, '\x00') { + return response{}, errors.New("page text is invalid") + } + empty := strings.TrimSpace(page.Text) == "" + if empty == (strings.TrimSpace(page.EmptyReason) == "") { + return response{}, errors.New("page emptiness is not explained exactly once") + } + if page.EmptyReason != "" && !validExplanation(page.EmptyReason) { + return response{}, errors.New("empty page explanation is invalid") + } + } + return wire, nil +} + +func readExact(ctx context.Context, reader io.Reader, expected, maximum int64) ([]byte, error) { + data := make([]byte, 0, expected) + buffer := make([]byte, 32<<10) + for { + if err := ctx.Err(); err != nil { + return nil, err + } + read, err := reader.Read(buffer) + if read > 0 { + if int64(len(data))+int64(read) > maximum { + return nil, errInputIdentity + } + data = append(data, buffer[:read]...) + } + switch { + case errors.Is(err, io.EOF): + if int64(len(data)) != expected { + return nil, errInputIdentity + } + return data, nil + case err != nil: + return nil, err + case read == 0: + return nil, io.ErrNoProgress + } + } +} + +type boundedBuffer struct { + data bytes.Buffer + limit int64 + exceeded bool + overflow func() +} + +func (buffer *boundedBuffer) Write(data []byte) (int, error) { + remaining := buffer.limit - int64(buffer.data.Len()) + if remaining <= 0 { + buffer.failOverflow() + return 0, errOutputTooLarge + } + if int64(len(data)) > remaining { + written, _ := buffer.data.Write(data[:remaining]) + buffer.failOverflow() + return written, errOutputTooLarge + } + written, _ := buffer.data.Write(data) + return written, nil +} + +func (buffer *boundedBuffer) Bytes() []byte { return buffer.data.Bytes() } + +func (buffer *boundedBuffer) Len() int { return buffer.data.Len() } + +func (buffer *boundedBuffer) failOverflow() { + if buffer.exceeded { + return + } + buffer.exceeded = true + if buffer.overflow != nil { + buffer.overflow() + } +} + +func (provider *Provider) contextError(parent context.Context, expiry bool, cause error) error { + if parent.Err() != nil { + return providerError(document.RenditionErrorCanceled, "PyMuPDF rendering canceled", parent.Err()) + } + if expiry { + return providerError(document.RenditionErrorPolicyRejected, + "PyMuPDF authorization expired during rendering", cause) + } + return providerError(document.RenditionErrorTransient, "PyMuPDF rendering timed out", cause) +} + +func (provider *Provider) postProcessError( + parent, operation context.Context, deadline time.Time, expiry bool, +) error { + if parentErr := parent.Err(); parentErr != nil { + return provider.contextError(parent, expiry, parentErr) + } + if operationErr := operation.Err(); operationErr != nil { + return provider.contextError(parent, expiry, operationErr) + } + if !time.Now().UTC().Before(deadline) { + return provider.contextError(parent, expiry, context.DeadlineExceeded) + } + return nil +} + +func providerError(code document.RenditionErrorCode, message string, cause error) error { + classified, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("pymupdf: classify provider error: %w", err) + } + return classified +} + +func cleanEnvironment() []string { + return []string{ + "LANG=C.UTF-8", "LC_ALL=C.UTF-8", "TZ=UTC", "PYTHONHASHSEED=0", + "PYTHONNOUSERSITE=1", "PYTHONDONTWRITEBYTECODE=1", + } +} + +func validateRuntimeIdentity(value string) error { + if value == "" || len(value) > 512 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return errors.New("pymupdf: runtime identity must be non-empty bounded UTF-8") + } + for _, char := range value { + if unicode.IsControl(char) { + return errors.New("pymupdf: runtime identity contains a control character") + } + } + return nil +} + +func validExplanation(value string) bool { + if len(value) > 256 || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return false + } + for _, char := range value { + if unicode.IsControl(char) { + return false + } + } + return true +} + +func pythonInterpreter(base string) bool { + name := strings.TrimSuffix(strings.ToLower(base), ".exe") + if name == "py" { + return true + } + for _, prefix := range []string{"python", "pypy"} { + if !strings.HasPrefix(name, prefix) { + continue + } + suffix := strings.TrimPrefix(name, prefix) + if suffix == "" { + return true + } + for _, char := range suffix { + if (char < '0' || char > '9') && char != '.' { + return false + } + } + return true + } + return false +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +var _ document.RenditionProvider = (*Provider)(nil) diff --git a/document/pymupdf/provider_test.go b/document/pymupdf/provider_test.go new file mode 100644 index 00000000..67ef600c --- /dev/null +++ b/document/pymupdf/provider_test.go @@ -0,0 +1,387 @@ +package pymupdf + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +const testRuntimeIdentity = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + +var testHelperBinary string + +func TestMain(m *testing.M) { + directory, err := os.MkdirTemp("", "docbank-pymupdf-test-") + if err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + extension := "" + if runtime.GOOS == "windows" { + extension = ".exe" + } + testHelperBinary = filepath.Join(directory, "renderer-base"+extension) + command := exec.Command("go", "build", "-o", testHelperBinary, "./testdata/helper") + if output, buildErr := command.CombinedOutput(); buildErr != nil { + _, _ = fmt.Fprintf(os.Stderr, "build PyMuPDF test helper: %v\n%s", buildErr, output) + _ = os.RemoveAll(directory) + os.Exit(1) + } + code := m.Run() + _ = os.RemoveAll(directory) + os.Exit(code) +} + +func TestProviderRendersExactPDFPagesThroughFixedProtocol(t *testing.T) { + t.Setenv("DOCBANK_PYMUPDF_AMBIENT_SECRET", "must-not-reach-child") + executable := helperExecutable(t, "success") + provider, err := New(Profile{ + Executable: executable, RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1 << 20, MaxResponseBytes: 1 << 20, + MaxPages: 10, Timeout: time.Second, + }) + require.NoError(t, err) + source := testPDF(2) + upload := newTestUpload(source) + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + + result, err := document.RenderRendition(t.Context(), provider, upload, authorization) + require.NoError(t, err) + assert.Equal(t, document.SourceEvidenceContractV1, result.Evidence.ContractVersion) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, "pdf", result.Evidence.Family) + assert.Equal(t, document.EvidenceUnitPage, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, "first page", result.Evidence.Units[0].Text) + assert.Equal(t, "second page", result.Evidence.Units[1].Text) + assert.Equal(t, document.SourceEvidenceLocatorV1{ + Kind: document.EvidenceLocatorPage, IndexOrigin: document.EvidenceIndexOriginOne, + Start: 1, End: 1, + }, result.Evidence.Units[0].Locator) + assert.Equal(t, int64(len(source)), result.Receipt.Usage.InputBytes) + assert.Equal(t, int64(2), result.Receipt.Usage.Units) + assert.Equal(t, upload.metadata.SHA256, result.Receipt.SourceSHA256) +} + +func TestNewPinsExecutableAndRuntimeIdentityInDescriptor(t *testing.T) { + first, err := New(Profile{ + Executable: helperExecutable(t, "success"), RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1024, MaxResponseBytes: 1024, MaxPages: 2, Timeout: time.Second, + }) + require.NoError(t, err) + second, err := New(Profile{ + Executable: helperExecutable(t, "success-copy"), RuntimeIdentity: testRuntimeIdentity + ".revision", + MaxDocumentBytes: 1024, MaxResponseBytes: 1024, MaxPages: 2, Timeout: time.Second, + }) + require.NoError(t, err) + + descriptor := first.Descriptor() + assert.Equal(t, document.RenditionTrustLocalProcess, descriptor.TrustBoundary) + assert.Equal(t, []document.RenditionFormatCapability{{ + MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile, + }}, descriptor.SupportedFormats) + assert.True(t, descriptor.ReturnsStructured) + assert.False(t, descriptor.ReturnsMarkdown) + assert.NotEqual(t, descriptor.Fingerprint, second.Descriptor().Fingerprint) + descriptor.SupportedFormats[0].MediaType = "text/plain" + assert.Equal(t, "application/pdf", first.Descriptor().SupportedFormats[0].MediaType) +} + +func TestProviderRejectsMalformedPartialAndDriftedOutput(t *testing.T) { + for _, mode := range []string{ + "malformed", "unknown-field", "version-drift", "runtime-drift", "source-hash-drift", + "source-size-drift", "partial", "page-count-drift", "gap", "duplicate", "empty-unexplained", + } { + t.Run(mode, func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, mode), time.Second, 1<<20) + upload := newTestUpload(testPDF(2)) + + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestProviderAcceptsExplicitlyExplainedEmptyPage(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "empty-explained"), time.Second, 1<<20) + upload := newTestUpload(testPDF(2)) + + result, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err) + require.Len(t, result.Evidence.Units, 2) + assert.Empty(t, result.Evidence.Units[0].Text) +} + +func TestProviderBoundsOutputAndSanitizesProcessFailure(t *testing.T) { + t.Run("oversized stdout", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "oversized"), time.Second, 1024) + upload := newTestUpload(testPDF(2)) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + t.Run("private stderr", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "failure"), time.Second, 1<<20) + upload := newTestUpload(testPDF(2)) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorTransient) + assert.NotContains(t, err.Error(), "private-stderr-token") + assert.NotContains(t, err.Error(), provider.executable) + }) +} + +func TestProviderTerminatesChildThatNeverStopsOversizedOutput(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "unbounded-output"), 2*time.Second, 1024) + upload := newTestUpload(testPDF(2)) + started := time.Now() + + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Less(t, time.Since(started), time.Second) + assert.NotContains(t, err.Error(), "broken pipe") +} + +func TestProviderEnforcesTimeoutAndCancellation(t *testing.T) { + t.Run("timeout", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "wait"), 20*time.Millisecond, 1<<20) + upload := newTestUpload(testPDF(2)) + started := time.Now() + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorTransient) + assert.Less(t, time.Since(started), time.Second) + }) + t.Run("cancellation", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "wait"), time.Second, 1<<20) + upload := newTestUpload(testPDF(2)) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, err := provider.Render(ctx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorCanceled) + assert.ErrorIs(t, err, context.Canceled) + }) + t.Run("authorization expiry", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "wait"), time.Second, 1<<20) + upload := newTestUpload(testPDF(2)) + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + authorization.ExpiresAt = time.Now().UTC().Add(20 * time.Millisecond).Format(timestampForm) + _, err := provider.Render(t.Context(), upload, authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.ErrorIs(t, err, context.DeadlineExceeded) + }) +} + +func TestProviderRejectsCancellationObservedAfterLargeChildResponse(t *testing.T) { + const pages = 20_000 + executable := helperExecutable(t, "many-pages") + provider, err := New(Profile{ + Executable: executable, RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 8 << 20, MaxResponseBytes: 8 << 20, + MaxPages: pages, Timeout: 30 * time.Second, + }) + require.NoError(t, err) + upload := newTestUpload(testPDF(pages)) + ctx := newCancelWhenCheckedContext(t.Context()) + defer ctx.cancel() + + _, err = provider.Render(ctx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorCanceled) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestProviderRejectsUnverifiedOrSubstitutedInputBeforeExecution(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "failure"), time.Second, 1<<20) + t.Run("not PDF", func(t *testing.T) { + upload := newTestUpload([]byte("not a PDF")) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorUnsupportedInput) + }) + t.Run("substituted bytes", func(t *testing.T) { + upload := newTestUpload(testPDF(2)) + upload.Reader = bytes.NewReader(testPDF(1)) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) +} + +func TestNewRejectsInterpreterAndInvalidBounds(t *testing.T) { + python := filepath.Join(t.TempDir(), "python3") + require.NoError(t, os.WriteFile(python, []byte("synthetic"), 0o700)) + _, err := New(Profile{ + Executable: python, RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1, MaxResponseBytes: 1, MaxPages: 1, Timeout: time.Second, + }) + require.ErrorContains(t, err, "must not be a Python interpreter") + + valid := Profile{ + Executable: helperExecutable(t, "success"), RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1024, MaxResponseBytes: 1024, MaxPages: 2, Timeout: time.Second, + } + for _, mutate := range []func(*Profile){ + func(profile *Profile) { profile.Executable = "renderer" }, + func(profile *Profile) { profile.RuntimeIdentity = "" }, + func(profile *Profile) { profile.MaxDocumentBytes = 0 }, + func(profile *Profile) { profile.MaxResponseBytes = 0 }, + func(profile *Profile) { profile.MaxPages = 0 }, + func(profile *Profile) { profile.MaxPages = 100_001 }, + func(profile *Profile) { profile.Timeout = 0 }, + } { + profile := valid + mutate(&profile) + _, err := New(profile) + require.Error(t, err) + } +} + +func newTestProvider(t *testing.T, executable string, timeout time.Duration, maxResponse int64) *Provider { + t.Helper() + provider, err := New(Profile{ + Executable: executable, RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1 << 20, MaxResponseBytes: maxResponse, + MaxPages: 10, Timeout: timeout, + }) + require.NoError(t, err) + return provider +} + +func assertProviderCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok) + assert.Equal(t, want, providerErr.Code()) +} + +type cancelWhenCheckedContext struct { + context.Context + + done chan struct{} + once sync.Once +} + +func newCancelWhenCheckedContext(parent context.Context) *cancelWhenCheckedContext { + return &cancelWhenCheckedContext{Context: parent, done: make(chan struct{})} +} + +func (ctx *cancelWhenCheckedContext) Done() <-chan struct{} { return ctx.done } + +func (ctx *cancelWhenCheckedContext) Err() error { + ctx.cancel() + return context.Canceled +} + +func (ctx *cancelWhenCheckedContext) cancel() { + ctx.once.Do(func() { close(ctx.done) }) +} + +func helperExecutable(t *testing.T, mode string) string { + t.Helper() + extension := filepath.Ext(testHelperBinary) + target := filepath.Join(filepath.Dir(testHelperBinary), "renderer-"+mode+extension) + data, err := os.ReadFile(testHelperBinary) + require.NoError(t, err) + require.NoError(t, os.WriteFile(target, data, 0o700)) + return target +} + +type testUpload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata +} + +func newTestUpload(data []byte) *testUpload { + digest := sha256.Sum256(data) + return &testUpload{ + Reader: bytes.NewReader(data), + metadata: document.AuthorizedUploadMetadata{ + Filename: "document.pdf", MediaFamily: "pdf", MediaType: "application/pdf", + ByteLength: int64(len(data)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), + ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + }, + } +} + +func (*testUpload) Close() error { return nil } + +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func testAuthorization( + descriptor document.RenditionDescriptor, metadata document.AuthorizedUploadMetadata, +) document.RenditionAuthorization { + started := time.Now().UTC().Add(-time.Minute) + return document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, + InputKind: metadata.InputKind, MaxTotalResultBytes: 1 << 20, + AuthorizedAt: started.Format("2006-01-02T15:04:05.000000000Z"), + ExpiresAt: started.Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + } +} + +func testPDF(pageCount int) []byte { + objects := []string{ + "<< /Type /Catalog /Pages 2 0 R >>", + fmt.Sprintf("<< /Type /Pages /Kids [%s] /Count %d >>", pdfPageReferences(pageCount), pageCount), + } + for page := range pageCount { + objects = append(objects, fmt.Sprintf("<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Rotate %d >>", page*90)) + } + var output bytes.Buffer + _, _ = output.WriteString("%PDF-1.4\n%synthetic\n") + offsets := make([]int, len(objects)) + for index, object := range objects { + offsets[index] = output.Len() + _, _ = fmt.Fprintf(&output, "%d 0 obj\n%s\nendobj\n", index+1, object) + } + xref := output.Len() + _, _ = fmt.Fprintf(&output, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1) + for _, offset := range offsets { + _, _ = fmt.Fprintf(&output, "%010d 00000 n \n", offset) + } + _, _ = fmt.Fprintf(&output, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref) + return output.Bytes() +} + +func pdfPageReferences(pageCount int) string { + references := make([]string, pageCount) + for page := range pageCount { + references[page] = fmt.Sprintf("%d 0 R", page+3) + } + return strings.Join(references, " ") +} + +var _ document.AuthorizedUpload = (*testUpload)(nil) +var _ io.ReadCloser = (*testUpload)(nil) diff --git a/document/pymupdf/testdata/helper/main.go b/document/pymupdf/testdata/helper/main.go new file mode 100644 index 00000000..5c46bb4f --- /dev/null +++ b/document/pymupdf/testdata/helper/main.go @@ -0,0 +1,136 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" +) + +const ( + protocol = "docbank-pymupdf/v1" + runtimeIdentity = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" +) + +func main() { + if len(os.Args) != 3 || os.Args[1] != "--protocol" || os.Args[2] != protocol { + fail("unexpected arguments") + } + // The operating system may synthesize bookkeeping variables that Cmd.Env + // cannot control. Verify the security boundary directly: ambient caller + // state is absent and every provider-controlled value is exact. + if os.Getenv("DOCBANK_PYMUPDF_AMBIENT_SECRET") != "" { + fail("ambient environment reached child") + } + for name, expected := range map[string]string{ + "LANG": "C.UTF-8", "LC_ALL": "C.UTF-8", "TZ": "UTC", "PYTHONHASHSEED": "0", + "PYTHONNOUSERSITE": "1", "PYTHONDONTWRITEBYTECODE": "1", + } { + if os.Getenv(name) != expected { + fail("controlled environment changed") + } + } + source, err := io.ReadAll(os.Stdin) + if err != nil { + fail("read stdin") + } + digest := sha256.Sum256(source) + mode := strings.TrimSuffix(filepath.Base(os.Args[0]), filepath.Ext(os.Args[0])) + mode = strings.TrimPrefix(mode, "renderer-") + if mode == "failure" { + _, _ = fmt.Fprintln(os.Stderr, "private-stderr-token") + os.Exit(7) + } + if mode == "wait" { + time.Sleep(10 * time.Second) + } + if mode == "unbounded-output" { + chunk := []byte(strings.Repeat("x", 32<<10)) + for { + if _, writeErr := os.Stdout.Write(chunk); writeErr != nil { + time.Sleep(10 * time.Second) + } + } + } + if mode == "malformed" { + _, _ = io.WriteString(os.Stdout, "{") + return + } + if mode == "oversized" { + _, _ = io.WriteString(os.Stdout, strings.Repeat("x", 2048)) + return + } + type outputPage struct { + Number int `json:"number"` + Text string `json:"text"` + EmptyReason string `json:"empty_reason,omitempty"` + } + response := struct { + ContractVersion string `json:"contract_version"` + RuntimeIdentity string `json:"runtime_identity"` + SourceSHA256 string `json:"source_sha256"` + SourceBytes int64 `json:"source_bytes"` + Complete bool `json:"complete"` + PageCount int `json:"page_count"` + Pages []outputPage `json:"pages"` + }{ + ContractVersion: protocol, RuntimeIdentity: runtimeIdentity, + SourceSHA256: hex.EncodeToString(digest[:]), SourceBytes: int64(len(source)), + Complete: true, PageCount: 2, + } + response.Pages = append(response.Pages, + outputPage{Number: 1, Text: "first page"}, + outputPage{Number: 2, Text: "second page"}, + ) + if mode == "many-pages" { + const pages = 20_000 + response.PageCount = pages + response.Pages = make([]outputPage, pages) + for index := range pages { + response.Pages[index] = outputPage{Number: index + 1, Text: "synthetic page text"} + } + } + switch mode { + case "version-drift": + response.ContractVersion = "docbank-pymupdf/v2" + case "runtime-drift": + response.RuntimeIdentity = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + case "source-hash-drift": + response.SourceSHA256 = strings.Repeat("b", 64) + case "source-size-drift": + response.SourceBytes++ + case "partial": + response.Complete = false + case "page-count-drift": + response.PageCount++ + case "gap": + response.Pages[1].Number = 3 + case "duplicate": + response.Pages[1].Number = 1 + case "empty-unexplained": + response.Pages[0].Text = "" + case "empty-explained": + response.Pages[0].Text = "" + response.Pages[0].EmptyReason = "blank page" + } + encoded, err := json.Marshal(response) + if err != nil { + fail("encode stdout") + } + if mode == "unknown-field" { + encoded = append(encoded[:len(encoded)-1], []byte(`,"unexpected":true}`)...) + } + if _, err := os.Stdout.Write(encoded); err != nil { + fail("write stdout") + } +} + +func fail(message string) { + _, _ = fmt.Fprintln(os.Stderr, message) + os.Exit(7) +} diff --git a/document/reducto/client.go b/document/reducto/client.go new file mode 100644 index 00000000..646e2353 --- /dev/null +++ b/document/reducto/client.go @@ -0,0 +1,1342 @@ +package reducto + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json/jsontext" + json "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "mime/multipart" + "net/http" + "net/textproto" + "reflect" + "slices" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" +) + +const ( + apiHost = "platform.reducto.ai" + apiOrigin = "https://" + apiHost + uploadPath = "/upload" + parsePath = "/parse_async" + providerID = "reducto.parse-v1" + timeForm = "2006-01-02T15:04:05.000000000Z" + + // The official Go SDK does not expose a selectable model. These constants + // bind the adapter to its immutable reviewed parse operation and schema. + sdkSchemaCommit = "a000fe10d8db34a4bde067f3c55a04f208c68e1e" + modelIdentity = "reducto-parse" + profileID = "go-sdk-" + sdkSchemaCommit + "-page-v2" + + defaultMaxUploadBytes = int64(100 << 20) + defaultRequestOverhead = int64(1 << 20) + defaultMaxControlBytes = int64(64 << 10) + defaultMaxResultBytes = int64(64 << 20) + defaultMaxArtifactBytes = int64(64 << 20) + defaultMaxPolls = 900 + defaultPollInterval = 2 * time.Second + defaultRequestTimeout = 30 * time.Second + defaultMaxWallTime = 30 * time.Minute + + maxConfiguredBytes = int64(512 << 20) + maxConfiguredPolls = 100_000 + maxConfiguredDuration = 24 * time.Hour + maxCredentialBytes = 4096 + maxUploadTokenBytes = 256 + maxJobTokenBytes = 120 + resumeHandlePrefix = "r2." + resumeHandleVersion = "reducto-resume/v2" + maxResumeUsageValue = int64(1 << 50) + maxNaturalUnits = 100_000 + markdownUnitSeparator = "\n\n---\n\n" +) + +var ( + _ document.RenditionProvider = (*Client)(nil) + _ document.ResumableRenditionProvider = (*Client)(nil) + + errResponseTooLarge = errors.New("reducto response exceeds byte limit") +) + +// SecretResolver resolves the one credential named by the frozen profile. +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +// Profile freezes all request, polling, result, artifact, and wall-clock +// bounds. The provider origin, routes, model identity, and parse options are +// deliberately not configurable. +type Profile struct { + SecretBinding string + MaxUploadBytes int64 + MaxRequestBytes int64 + MaxControlBytes int64 + MaxPolls int + PollInterval time.Duration + RequestTimeout time.Duration + MaxResultBytes int64 + MaxArtifactBytes int64 + MaxArtifacts int + MaxWallTime time.Duration + RetainStructured bool +} + +// Client is a fixed-origin hosted Reducto rendition provider. +type Client struct { + descriptor document.RenditionDescriptor + profile Profile + secrets SecretResolver + transport http.RoundTripper +} + +// NewProvider constructs a provider around an injected hardened transport. +// It does not consult proxy environment variables, resolve credentials, or +// perform network access during construction. +func NewProvider(profile Profile, secrets SecretResolver, transport http.RoundTripper) (*Client, error) { + profile = defaultProfile(profile) + if err := validateProfile(profile); err != nil { + return nil, err + } + if nilValue(secrets) { + return nil, errors.New("reducto named credential resolver is required") + } + if nilValue(transport) { + return nil, errors.New("reducto hardened transport is required") + } + identity := fmt.Sprintf( + "reducto-profile/v1\x00%s\x00%s\x00%s\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%t", + modelIdentity, profileID, profile.SecretBinding, profile.MaxUploadBytes, + profile.MaxRequestBytes, profile.MaxControlBytes, profile.MaxPolls, + profile.PollInterval, profile.RequestTimeout, profile.MaxResultBytes, + profile.MaxArtifactBytes, profile.MaxWallTime, profile.RetainStructured, + ) + digest := sha256.Sum256([]byte(identity)) + roles := []document.EvidenceArtifactRole(nil) + if profile.RetainStructured { + roles = []document.EvidenceArtifactRole{document.EvidenceArtifactStructured} + } + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, + ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: hex.EncodeToString(digest[:]), + TrustBoundary: document.RenditionTrustHostedProvider, + SupportedFormats: []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + }, + ReturnsMarkdown: true, ReturnsStructured: true, ArtifactRoles: roles, + }) + if err != nil { + return nil, fmt.Errorf("reducto descriptor: %w", err) + } + return &Client{ + descriptor: cloneDescriptor(descriptor), profile: profile, secrets: secrets, + transport: fixedOriginTransport{base: transport}, + }, nil +} + +// Descriptor returns a defensive copy of the immutable provider identity. +func (client *Client) Descriptor() document.RenditionDescriptor { + if client == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +// Render starts and completes one hosted asynchronous parse operation. +func (client *Client) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + return client.RenderResumable(ctx, upload, authorization, nil, nil) +} + +// RenderResumable submits exact authorized bytes or resumes one known opaque +// job handle. A newly accepted job is checkpointed before its first poll. +func (client *Client) RenderResumable( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, resume *document.RenditionResumeHandle, + checkpoint document.RenditionResumeCheckpoint, +) (document.RenditionResult, error) { + if client == nil { + return document.RenditionResult{}, errors.New("reducto client is required") + } + startedAt := time.Now().UTC() + expiresAt, resumeFacts, err := client.validateInvocation(startedAt, upload, authorization, resume) + if err != nil { + return document.RenditionResult{}, err + } + operationExpiry := expiresAt + if resume != nil { + operationExpiry = time.Time{} + } + operationCtx, cancel := boundedOperationContext(ctx, operationExpiry, client.profile.MaxWallTime) + defer cancel() + state := operationState{startedAt: startedAt, inputBytes: authorization.SourceBytes} + var jobID string + if resume == nil { + fileID, submitErr := client.upload(operationCtx, ctx, expiresAt, upload, &state) + if submitErr != nil { + return document.RenditionResult{}, submitErr + } + jobID, submitErr = client.submit(operationCtx, ctx, expiresAt, fileID, &state) + if submitErr != nil { + return document.RenditionResult{}, submitErr + } + checkpointedAt := time.Now().UTC() + if !checkpointedAt.Before(expiresAt) { + return document.RenditionResult{}, renditionError(document.RenditionErrorAmbiguousSubmission, + "Reducto job was accepted after the durable authorization interval", nil) + } + state.submittedAt = checkpointedAt + if checkpointErr := client.checkpoint(checkpoint, jobID, &state); checkpointErr != nil { + return document.RenditionResult{}, checkpointErr + } + } else { + jobID = resumeFacts.JobID + state.startedAt, _ = time.Parse(timeForm, resumeFacts.StartedAt) + state.submittedAt, _ = time.Parse(timeForm, resumeFacts.SubmittedAt) + state.requests = resumeFacts.Requests + state.retries = resumeFacts.Retries + state.inputBytes = resumeFacts.InputBytes + state.outputBytes = resumeFacts.OutputBytes + state.pollDelay = time.Duration(resumeFacts.RetryDelayMillis) * time.Millisecond + state.durableCheckpoint = true + } + result, err := client.poll( + operationCtx, ctx, operationExpiry, jobID, authorization, checkpoint, &state) + if err != nil { + return document.RenditionResult{}, err + } + if err := operationCtx.Err(); err != nil { + return document.RenditionResult{}, client.resultFailure( + ctx, operationCtx, operationExpiry, state.durableCheckpoint, + "Reducto operation ended after its lifecycle boundary", err) + } + completedAt := time.Now().UTC() + if resume == nil && !completedAt.Before(expiresAt) { + if state.durableCheckpoint { + return document.RenditionResult{}, renditionError(document.RenditionErrorTransient, + "Reducto durable result reached its authorization boundary", nil) + } + return document.RenditionResult{}, expiredError(nil) + } + result.Receipt = document.RenditionReceipt{ + ProviderID: client.descriptor.ID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: authorization.SourceSHA256, + OperationID: "reducto-" + jobID, + StartedAt: state.startedAt.Format(timeForm), CompletedAt: completedAt.Format(timeForm), + Warnings: state.warnings, + Usage: document.RenditionUsage{ + Requests: state.requests, Retries: state.retries, InputBytes: state.inputBytes, + OutputBytes: state.outputBytes, Units: int64(len(result.Evidence.Units)), + }, + RetryDelayMillis: state.pollDelay.Milliseconds(), + } + return result, nil +} + +func (client *Client) validateInvocation( + now time.Time, upload document.AuthorizedUpload, authorization document.RenditionAuthorization, + resume *document.RenditionResumeHandle, +) (time.Time, resumePayload, error) { + expiresAt, err := time.Parse(timeForm, authorization.ExpiresAt) + if err != nil { + return time.Time{}, resumePayload{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto authorization expiry is invalid", err) + } + if resume == nil { + if nilValue(upload) { + return time.Time{}, resumePayload{}, errors.New("reducto authorized upload is required for submission") + } + if _, err := document.ValidateRenditionProviderRequestAt(now, client, upload, authorization); err != nil { + return time.Time{}, resumePayload{}, err + } + return expiresAt, resumePayload{}, nil + } + if !nilValue(upload) { + return time.Time{}, resumePayload{}, errors.New("reducto resume must not receive source bytes") + } + facts, decodeErr := decodeResumeHandle(resume.Value, authorization) + if decodeErr != nil { + return time.Time{}, resumePayload{}, renditionError(document.RenditionErrorUnknownJob, + "Reducto resume handle is invalid", decodeErr) + } + submittedAt, _ := time.Parse(timeForm, facts.SubmittedAt) + sealed := resumeAuthorizationUpload{metadata: document.AuthorizedUploadMetadata{ + Filename: "resume.input", MediaFamily: authorization.MediaFamily, + MediaType: authorization.MediaType, ByteLength: authorization.SourceBytes, + SHA256: authorization.SourceSHA256, + CapabilityRecordChecksum: authorization.CapabilityRecordChecksum, + ProviderMetadataChecksum: authorization.ProviderMetadataChecksum, + InputKind: authorization.InputKind, + }} + if _, err := document.ValidateRenditionProviderRequestAt(submittedAt, client, sealed, authorization); err != nil { + return time.Time{}, resumePayload{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto resume authority is invalid", err) + } + if authorization.ProviderID != client.descriptor.ID || + authorization.DescriptorFingerprint != client.descriptor.Fingerprint || + authorization.PolicyFingerprint != client.descriptor.PolicyFingerprint { + return time.Time{}, resumePayload{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto resume authority changed", nil) + } + return expiresAt, facts, nil +} + +func (client *Client) upload( + operationCtx, callerCtx context.Context, expiresAt time.Time, + upload document.AuthorizedUpload, state *operationState, +) (string, error) { + metadata := upload.Metadata() + if metadata.ByteLength > client.profile.MaxUploadBytes { + return "", renditionError(document.RenditionErrorPolicyRejected, + "Reducto upload exceeds profile limit", nil) + } + source, err := readExactUpload( + operationCtx, callerCtx, expiresAt, upload, metadata, client.profile.MaxUploadBytes) + if err != nil { + return "", err + } + defer clear(source) + body := new(bytes.Buffer) + writer := multipart.NewWriter(body) + header := make(textproto.MIMEHeader) + header.Set("Content-Disposition", mime.FormatMediaType("form-data", map[string]string{ + "name": "file", "filename": metadata.Filename, + })) + header.Set("Content-Type", metadata.MediaType) + part, err := writer.CreatePart(header) + if err == nil { + _, err = part.Write(source) + } + if closeErr := writer.Close(); err == nil { + err = closeErr + } + if err != nil { + return "", renditionError(document.RenditionErrorTransient, + "Reducto upload request could not be built", err) + } + if int64(body.Len()) > client.profile.MaxRequestBytes { + return "", renditionError(document.RenditionErrorPolicyRejected, + "Reducto upload request exceeds profile limit", nil) + } + requestBody := bytes.Clone(body.Bytes()) + defer clear(requestBody) + raw, status, err := client.do(operationCtx, http.MethodPost, uploadPath, + writer.FormDataContentType(), requestBody, client.profile.MaxControlBytes, state) + if err != nil { + return "", client.submissionFailure(callerCtx, operationCtx, expiresAt, + "Reducto upload outcome is ambiguous", err) + } + if status < 200 || status >= 300 { + return "", client.httpError(status, "upload") + } + var result uploadResponse + if err := strictJSON(raw, &result); err != nil || validateUploadToken(result.FileID) != nil { + return "", renditionError(document.RenditionErrorPolicyRejected, + "Reducto upload schema changed", err) + } + return result.FileID, nil +} + +func (client *Client) submit( + operationCtx, callerCtx context.Context, expiresAt time.Time, fileID string, + state *operationState, +) (string, error) { + request := parseRequest{DocumentURL: fileID, Priority: false} + request.AdvancedOptions.AddPageMarkers = true + request.AdvancedOptions.ReturnOCRData = false + request.Options.Chunking.ChunkMode = "page" + request.Options.ForceURLResult = false + body, err := json.Marshal(request) + if err != nil { + return "", renditionError(document.RenditionErrorTransient, + "Reducto parse request could not be built", err) + } + if int64(len(body)) > client.profile.MaxRequestBytes { + return "", renditionError(document.RenditionErrorPolicyRejected, + "Reducto parse request exceeds profile limit", nil) + } + defer clear(body) + raw, status, err := client.do(operationCtx, http.MethodPost, parsePath, + "application/json", body, client.profile.MaxControlBytes, state) + if err != nil { + return "", client.submissionFailure(callerCtx, operationCtx, expiresAt, + "Reducto parse submission outcome is ambiguous", err) + } + if status < 200 || status >= 300 { + return "", client.httpError(status, "submission") + } + var result parseResponse + if err := strictJSON(raw, &result); err != nil || validateJobToken(result.JobID) != nil { + return "", renditionError(document.RenditionErrorAmbiguousSubmission, + "Reducto submission schema changed", err) + } + return result.JobID, nil +} + +func (client *Client) poll( + operationCtx, callerCtx context.Context, expiresAt time.Time, jobID string, + authorization document.RenditionAuthorization, checkpoint document.RenditionResumeCheckpoint, + state *operationState, +) (document.RenditionResult, error) { + limit := min64(client.profile.MaxResultBytes, int64(authorization.MaxTotalResultBytes)) + for attempt := range client.profile.MaxPolls { + raw, status, err := client.do(operationCtx, http.MethodGet, jobPath(jobID), "", nil, limit, state) + if checkpointErr := client.checkpoint(checkpoint, jobID, state); checkpointErr != nil { + return document.RenditionResult{}, checkpointErr + } + if err != nil { + return document.RenditionResult{}, client.resultFailure( + callerCtx, operationCtx, expiresAt, state.durableCheckpoint, + "Reducto job request failed", err) + } + if status < 200 || status >= 300 { + return document.RenditionResult{}, client.httpError(status, "job") + } + var job jobResponse + if err := requireMembers(raw, "status"); err != nil { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto job schema changed", err) + } + if err := strictJSON(raw, &job); err != nil || !validProgress(job.Progress) { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto job schema changed", err) + } + switch job.Status { + case "Pending", "Idle": + if !isNullJSON(job.Result) { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto pending job schema changed", nil) + } + if attempt+1 == client.profile.MaxPolls { + break + } + if err := waitContext(operationCtx, client.profile.PollInterval); err != nil { + return document.RenditionResult{}, client.pollFailure( + callerCtx, operationCtx, expiresAt, state.durableCheckpoint, err) + } + state.pollDelay += client.profile.PollInterval + if checkpointErr := client.checkpoint(checkpoint, jobID, state); checkpointErr != nil { + return document.RenditionResult{}, checkpointErr + } + case "Completed": + return client.completed(job.Result, jobID, authorization) + case "Failed": + return document.RenditionResult{}, renditionError(document.RenditionErrorUnsupportedInput, + "Reducto could not parse the input", nil) + default: + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto job status schema changed", nil) + } + } + return document.RenditionResult{}, renditionError(document.RenditionErrorTransient, + "Reducto polling limit was reached", nil) +} + +func (client *Client) completed( + raw []byte, jobID string, authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if isNullJSON(raw) { + return document.RenditionResult{}, malformedError("Reducto completed without a result", nil) + } + if err := requireMembers(raw, "result", "usage", "duration", "job_id"); err != nil { + return document.RenditionResult{}, malformedError("Reducto completed result schema changed", err) + } + var completed completedResult + if err := strictJSON(raw, &completed); err != nil { + return document.RenditionResult{}, malformedError("Reducto completed result schema changed", err) + } + if completed.JobID != jobID { + return document.RenditionResult{}, renditionError(document.RenditionErrorPolicyRejected, + "Reducto completed job identity changed", nil) + } + if math.IsNaN(completed.Duration) || math.IsInf(completed.Duration, 0) || completed.Duration < 0 { + return document.RenditionResult{}, malformedError("Reducto result duration is invalid", nil) + } + if completed.Usage.NumPages == nil || *completed.Usage.NumPages < 0 { + return document.RenditionResult{}, malformedError("Reducto page usage is invalid", nil) + } + if err := requireMembers(completed.Result, "type"); err != nil { + return document.RenditionResult{}, malformedError("Reducto parse result schema changed", err) + } + var parsed parseResult + if err := strictJSON(completed.Result, &parsed); err != nil { + return document.RenditionResult{}, malformedError("Reducto parse result schema changed", err) + } + if parsed.Type == "url" { + if err := requireMembers(completed.Result, "result_id", "url"); err != nil || + parsed.ResultID == "" || parsed.URL == "" { + return document.RenditionResult{}, malformedError("Reducto URL result schema changed", err) + } + return document.RenditionResult{}, malformedError("Reducto returned a provider URL instead of inline output", nil) + } + if parsed.Type != "full" || len(parsed.Chunks) == 0 { + return document.RenditionResult{}, malformedError("Reducto parse result type changed", nil) + } + if parsed.Custom != nil || parsed.OCR != nil { + return document.RenditionResult{}, malformedError("Reducto returned unrequested nested output", nil) + } + var rawChunks []jsontext.Value + if err := strictJSON(parsed.Chunks, &rawChunks); err != nil || len(rawChunks) == 0 { + return document.RenditionResult{}, malformedError("Reducto chunk schema changed", err) + } + chunks := make([]resultChunk, len(rawChunks)) + for index, rawChunk := range rawChunks { + if err := requireMembers(rawChunk, "blocks", "content", "embed", "enriched"); err != nil { + return document.RenditionResult{}, malformedError("Reducto chunk schema changed", err) + } + if err := strictJSON(rawChunk, &chunks[index]); err != nil { + return document.RenditionResult{}, malformedError("Reducto chunk schema changed", err) + } + } + evidence, markdown, err := naturalEvidence(chunks, authorization.MediaFamily, *completed.Usage.NumPages) + if err != nil { + return document.RenditionResult{}, malformedError("Reducto output is partial or malformed", err) + } + if len(markdown) > authorization.MaxProviderMarkdownBytes { + return document.RenditionResult{}, malformedError("Reducto Markdown exceeds authorization", nil) + } + result := document.RenditionResult{Evidence: evidence, ProviderMarkdown: markdown} + if client.profile.RetainStructured && allowsStructured(authorization) { + maximum := min64(client.profile.MaxArtifactBytes, int64(authorization.MaxArtifactBytes)) + if int64(len(raw)) > maximum { + return document.RenditionResult{}, malformedError("Reducto structured result exceeds authorization", nil) + } + payload := bytes.Clone(raw) + digest := sha256.Sum256(payload) + checksum := hex.EncodeToString(digest[:]) + result.Artifacts = []document.RenditionArtifact{{ + Role: document.EvidenceArtifactStructured, MediaType: "application/json", + Payload: payload, SHA256: checksum, + }} + result.Evidence.Artifacts = []document.SourceEvidenceArtifactV1{{ + Pointer: "result", ProviderID: "reducto-result", Role: document.EvidenceArtifactStructured, + SHA256: checksum, + }} + } + encodedEvidence, err := json.Marshal(result.Evidence) + if err != nil { + return document.RenditionResult{}, malformedError("Reducto evidence could not be bounded", err) + } + total := len(encodedEvidence) + len(result.ProviderMarkdown) + for _, artifact := range result.Artifacts { + total += len(artifact.Payload) + } + if total > authorization.MaxTotalResultBytes { + return document.RenditionResult{}, malformedError("Reducto total result exceeds authorization", nil) + } + if err := document.ValidateSourceEvidenceV1(result.Evidence); err != nil { + return document.RenditionResult{}, malformedError("Reducto evidence is unrepresentable", err) + } + return result, nil +} + +func naturalEvidence( + chunks []resultChunk, family string, reportedUnits int64, +) (document.SourceEvidenceV1, []byte, error) { + if family == "spreadsheet" { + return document.SourceEvidenceV1{}, nil, errors.New("provider did not report a stable sheet name") + } + unitKind, locatorKind, ok := familyUnit(family) + if !ok || reportedUnits <= 0 { + return document.SourceEvidenceV1{}, nil, errors.New("natural unit family or count is invalid") + } + if reportedUnits > maxNaturalUnits { + return document.SourceEvidenceV1{}, nil, errors.New("natural unit count exceeds provider bound") + } + units := make([][]string, int(reportedUnits)) + lastPage := int64(0) + for _, chunk := range chunks { + if !utf8.ValidString(chunk.Content) || len(chunk.Blocks) == 0 { + return document.SourceEvidenceV1{}, nil, errors.New("chunk is incomplete") + } + chunkPage := int64(0) + for _, block := range chunk.Blocks { + if err := validateBlockLocation(block); err != nil { + return document.SourceEvidenceV1{}, nil, err + } + if chunk.Content == "" && block.Content != "" { + return document.SourceEvidenceV1{}, nil, errors.New("contentless chunk contains unassigned block text") + } + if chunk.Content != "" && (block.Content == "" || !utf8.ValidString(block.Content)) { + return document.SourceEvidenceV1{}, nil, errors.New("block content is invalid") + } + page := *block.BBox.Page + if chunkPage == 0 { + chunkPage = page + } else if page != chunkPage { + return document.SourceEvidenceV1{}, nil, errors.New("page chunk spans multiple natural units") + } + } + if chunkPage < lastPage || chunkPage > reportedUnits { + return document.SourceEvidenceV1{}, nil, errors.New("block unit sequence is incomplete") + } + lastPage = chunkPage + if chunk.Content != "" { + units[chunkPage-1] = append(units[chunkPage-1], chunk.Content) + } + } + evidence := document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: family, UnitKind: unitKind, + Units: make([]document.SourceEvidenceUnitV1, 0, len(units)), + } + markdown := make([]string, 0, len(units)) + for index, parts := range units { + if len(parts) == 0 { + locator := document.SourceEvidenceLocatorV1{ + Kind: locatorKind, IndexOrigin: document.EvidenceIndexOriginOne, + Start: int64(index + 1), End: int64(index + 1), + } + evidence.Omissions = append(evidence.Omissions, document.SourceEvidenceOmissionV1{ + Kind: document.EvidenceOmissionUnit, Locator: &locator, + Reason: "Reducto returned no content for the reported natural unit", + }) + continue + } + text := strings.Join(parts, "\n\n") + locator := document.SourceEvidenceLocatorV1{ + Kind: locatorKind, IndexOrigin: document.EvidenceIndexOriginOne, + Start: int64(index + 1), End: int64(index + 1), + } + evidence.Units = append(evidence.Units, document.SourceEvidenceUnitV1{ + Order: len(evidence.Units), Text: text, Locator: locator, + }) + markdown = append(markdown, text) + } + if len(evidence.Units) == 0 { + return document.SourceEvidenceV1{}, nil, errors.New("all natural units are missing") + } + if len(evidence.Omissions) != 0 { + evidence.Completeness = document.EvidencePartial + } + return evidence, []byte(strings.Join(markdown, markdownUnitSeparator)), nil +} + +func validateBlockLocation(block resultBlock) error { + if !knownBlockType(block.Type) { + return errors.New("block content or type is invalid") + } + box := block.BBox + if box.Height == nil || box.Left == nil || box.Page == nil || box.Top == nil || box.Width == nil || + !finite(*box.Height) || !finite(*box.Left) || !finite(*box.Top) || !finite(*box.Width) || + *box.Height < 0 || *box.Width < 0 || *box.Page <= 0 || + (box.OriginalPage != nil && *box.OriginalPage <= 0) { + return errors.New("block bounding box is invalid") + } + return nil +} + +func knownBlockType(value string) bool { + switch value { + case "Header", "Footer", "Title", "Section Header", "Page Number", "List Item", + "Figure", "Table", "Key Value", "Text", "Comment", "Discard": + return true + default: + return false + } +} + +func (client *Client) do( + ctx context.Context, method, path, contentType string, body []byte, limit int64, + state *operationState, +) ([]byte, int, error) { + if err := ctx.Err(); err != nil { + return nil, 0, &preEgressContextError{cause: err} + } + credential, err := client.secrets.ResolveSecret(ctx, client.profile.SecretBinding) + if err != nil && ctx.Err() != nil { + return nil, 0, &preEgressContextError{cause: ctx.Err()} + } + if err != nil || !validCredential(credential) { + return nil, 0, renditionError(document.RenditionErrorAuthentication, + "Reducto credential is unavailable", err) + } + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + var reader io.Reader + if body != nil { + reader = bytes.NewReader(body) + } + request, err := http.NewRequestWithContext(requestCtx, method, apiOrigin+path, reader) + if err != nil { + return nil, 0, &preEgressContextError{cause: err} + } + request.Header.Set("Authorization", "Bearer "+credential) + request.Header.Set("Accept", "application/json") + if contentType != "" { + request.Header.Set("Content-Type", contentType) + } + state.requests++ + response, err := client.transport.RoundTrip(request) + if err != nil { + return nil, 0, err + } + if response == nil || response.Body == nil { + return nil, 0, errors.New("reducto returned an empty HTTP response") + } + raw, readErr := readBounded(response.Body, limit) + closeErr := response.Body.Close() + if readErr != nil { + return nil, response.StatusCode, readErr + } + if closeErr != nil { + return nil, response.StatusCode, fmt.Errorf("close Reducto response: %w", closeErr) + } + state.outputBytes += int64(len(raw)) + return raw, response.StatusCode, nil +} + +func (client *Client) httpError(status int, stage string) error { + switch status { + case http.StatusUnauthorized, http.StatusForbidden: + return renditionError(document.RenditionErrorAuthentication, + "Reducto authentication was rejected", nil) + case http.StatusTooManyRequests: + return renditionError(document.RenditionErrorRateLimited, + "Reducto rate limit was exhausted", nil) + case http.StatusServiceUnavailable, http.StatusInsufficientStorage: + return renditionError(document.RenditionErrorCapacity, + "Reducto capacity is unavailable", nil) + case http.StatusNotFound, http.StatusGone: + if stage == "job" { + return renditionError(document.RenditionErrorUnknownJob, + "Reducto job is unknown or expired", nil) + } + case http.StatusBadRequest, http.StatusUnprocessableEntity, http.StatusUnsupportedMediaType, + http.StatusRequestEntityTooLarge: + return renditionError(document.RenditionErrorUnsupportedInput, + "Reducto rejected the input", nil) + case http.StatusRequestTimeout, http.StatusInternalServerError, http.StatusBadGateway, + http.StatusGatewayTimeout: + return renditionError(document.RenditionErrorTransient, + "Reducto service request failed", nil) + } + if status >= 300 && status < 400 { + return renditionError(document.RenditionErrorPolicyRejected, + "Reducto redirects are refused", nil) + } + if status >= 500 { + return renditionError(document.RenditionErrorTransient, + "Reducto service request failed", nil) + } + return renditionError(document.RenditionErrorPolicyRejected, + "Reducto rejected the request", nil) +} + +func (client *Client) submissionFailure( + callerCtx, operationCtx context.Context, expiresAt time.Time, message string, cause error, +) error { + if local, ok := errors.AsType[*preEgressContextError](cause); ok { + return classifyContextFailure(callerCtx, operationCtx, expiresAt, local.cause) + } + if providerError, ok := errors.AsType[*document.RenditionProviderError](cause); ok { + return providerError + } + return renditionError(document.RenditionErrorAmbiguousSubmission, message, cause) +} + +func (client *Client) resultFailure( + callerCtx, operationCtx context.Context, expiresAt time.Time, durableCheckpoint bool, + message string, cause error, +) error { + if errors.Is(callerCtx.Err(), context.Canceled) { + return renditionError(document.RenditionErrorCanceled, "Reducto rendering was canceled", callerCtx.Err()) + } + if callerCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, "Reducto caller deadline was reached", cause) + } + if !expiresAt.IsZero() && !time.Now().UTC().Before(expiresAt) { + if durableCheckpoint { + return renditionError(document.RenditionErrorTransient, + "Reducto durable operation reached its authorization boundary", cause) + } + return expiredError(cause) + } + if providerError, ok := errors.AsType[*document.RenditionProviderError](cause); ok { + return providerError + } + if errors.Is(cause, errResponseTooLarge) { + return malformedError("Reducto job response exceeds profile limit", cause) + } + if operationCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, "Reducto operation timed out", cause) + } + return renditionError(document.RenditionErrorTransient, message, cause) +} + +func (client *Client) pollFailure( + callerCtx, operationCtx context.Context, expiresAt time.Time, durableCheckpoint bool, cause error, +) error { + return client.resultFailure( + callerCtx, operationCtx, expiresAt, durableCheckpoint, "Reducto polling failed", cause) +} + +type operationState struct { + startedAt time.Time + submittedAt time.Time + durableCheckpoint bool + requests int64 + retries int64 + inputBytes int64 + outputBytes int64 + pollDelay time.Duration + warnings []string +} + +type preEgressContextError struct{ cause error } + +func (failure *preEgressContextError) Error() string { return "reducto request stopped before egress" } +func (failure *preEgressContextError) Unwrap() error { return failure.cause } + +type resumePayload struct { + Version string `json:"v"` + JobID string `json:"j"` + StartedAt string `json:"s"` + SubmittedAt string `json:"a"` + Requests int64 `json:"q"` + Retries int64 `json:"r"` + InputBytes int64 `json:"i"` + OutputBytes int64 `json:"o"` + RetryDelayMillis int64 `json:"d"` +} + +func (client *Client) checkpoint( + checkpoint document.RenditionResumeCheckpoint, jobID string, state *operationState, +) error { + if checkpoint == nil { + return nil + } + resumeValue, err := encodeResumeHandle(resumePayload{ + Version: resumeHandleVersion, JobID: jobID, + StartedAt: state.startedAt.Format(timeForm), SubmittedAt: state.submittedAt.Format(timeForm), + Requests: state.requests, Retries: state.retries, InputBytes: state.inputBytes, + OutputBytes: state.outputBytes, RetryDelayMillis: state.pollDelay.Milliseconds(), + }) + if err != nil { + return renditionError(document.RenditionErrorAmbiguousSubmission, + "Reducto job resume state could not be encoded", err) + } + if err := checkpoint(document.RenditionResumeHandle{Value: resumeValue}); err != nil { + return renditionError(document.RenditionErrorAmbiguousSubmission, + "Reducto job could not be durably checkpointed", err) + } + state.durableCheckpoint = true + return nil +} + +type uploadResponse struct { + FileID string `json:"file_id"` + PresignedURL *string `json:"presigned_url"` +} + +type parseRequest struct { + DocumentURL string `json:"document_url"` + AdvancedOptions parseAdvancedOptions `json:"advanced_options"` + Options parseProcessingOptions `json:"options"` + Priority bool `json:"priority"` +} + +type parseAdvancedOptions struct { + AddPageMarkers bool `json:"add_page_markers"` + ReturnOCRData bool `json:"return_ocr_data"` +} + +type parseProcessingOptions struct { + Chunking parseChunking `json:"chunking"` + ForceURLResult bool `json:"force_url_result"` +} + +type parseChunking struct { + ChunkMode string `json:"chunk_mode"` +} + +type parseResponse struct { + JobID string `json:"job_id"` +} + +type jobResponse struct { + Status string `json:"status"` + Progress *float64 `json:"progress"` + Reason *string `json:"reason"` + Result jsontext.Value `json:"result"` +} + +type completedResult struct { + Result jsontext.Value `json:"result"` + Usage parseUsage `json:"usage"` + Duration float64 `json:"duration"` + JobID string `json:"job_id"` + PDFURL *string `json:"pdf_url"` +} + +type parseUsage struct { + NumPages *int64 `json:"num_pages"` +} + +type parseResult struct { + Type string `json:"type"` + Chunks jsontext.Value `json:"chunks"` + Custom any `json:"custom"` + OCR *parseOCR `json:"ocr"` + ResultID string `json:"result_id"` + URL string `json:"url"` +} + +type resultChunk struct { + Blocks []resultBlock `json:"blocks"` + Content string `json:"content"` + Embed string `json:"embed"` + Enriched *string `json:"enriched"` + EnrichmentSuccess bool `json:"enrichment_success"` +} + +type resultBlock struct { + BBox resultBoundingBox `json:"bbox"` + Content string `json:"content"` + Type string `json:"type"` + ImageURL *string `json:"image_url"` +} + +type resultBoundingBox struct { + Height *float64 `json:"height"` + Left *float64 `json:"left"` + Page *int64 `json:"page"` + Top *float64 `json:"top"` + Width *float64 `json:"width"` + OriginalPage *int64 `json:"original_page"` +} + +type parseOCR struct { + Lines []parseOCRLine `json:"lines"` + Words []parseOCRWord `json:"words"` +} + +type parseOCRLine struct { + BBox resultBoundingBox `json:"bbox"` + Text string `json:"text"` +} + +type parseOCRWord struct { + BBox resultBoundingBox `json:"bbox"` + Text string `json:"text"` +} + +type fixedOriginTransport struct{ base http.RoundTripper } + +func (transport fixedOriginTransport) RoundTrip(request *http.Request) (*http.Response, error) { + if request == nil || request.URL == nil || request.URL.Scheme != "https" || + request.URL.Host != apiHost || request.URL.User != nil || request.URL.RawQuery != "" || + request.URL.Fragment != "" || request.Host != "" && request.Host != apiHost || + !fixedRequestPath(request.Method, request.URL.Path) { + return nil, errors.New("reducto request destination is not fixed") + } + return transport.base.RoundTrip(request) +} + +func fixedRequestPath(method, path string) bool { + if method == http.MethodPost && (path == uploadPath || path == parsePath) { + return true + } + if method != http.MethodGet || !strings.HasPrefix(path, "/job/") { + return false + } + return validateJobToken(strings.TrimPrefix(path, "/job/")) == nil +} + +func jobPath(jobID string) string { return "/job/" + jobID } + +func defaultProfile(profile Profile) Profile { + if profile.MaxUploadBytes == 0 { + profile.MaxUploadBytes = defaultMaxUploadBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = profile.MaxUploadBytes + defaultRequestOverhead + } + if profile.MaxControlBytes == 0 { + profile.MaxControlBytes = defaultMaxControlBytes + } + if profile.MaxPolls == 0 { + profile.MaxPolls = defaultMaxPolls + } + if profile.PollInterval == 0 { + profile.PollInterval = defaultPollInterval + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultRequestTimeout + } + if profile.MaxResultBytes == 0 { + profile.MaxResultBytes = defaultMaxResultBytes + } + if profile.MaxArtifactBytes == 0 { + profile.MaxArtifactBytes = defaultMaxArtifactBytes + } + if profile.MaxWallTime == 0 { + profile.MaxWallTime = defaultMaxWallTime + } + if profile.RetainStructured && profile.MaxArtifacts == 0 { + profile.MaxArtifacts = 1 + } + return profile +} + +func validateProfile(profile Profile) error { + if err := validateCanonicalToken(profile.SecretBinding, 128); err != nil { + return fmt.Errorf("reducto credential binding: %w", err) + } + if profile.MaxUploadBytes <= 0 || profile.MaxUploadBytes > defaultMaxUploadBytes || + profile.MaxRequestBytes <= profile.MaxUploadBytes || + profile.MaxRequestBytes > maxConfiguredBytes+defaultRequestOverhead || + profile.MaxControlBytes <= 0 || profile.MaxControlBytes > 1<<20 || + profile.MaxResultBytes <= 0 || profile.MaxResultBytes > maxConfiguredBytes || + profile.MaxArtifactBytes <= 0 || profile.MaxArtifactBytes > maxConfiguredBytes || + profile.MaxPolls <= 0 || profile.MaxPolls > maxConfiguredPolls || + profile.PollInterval <= 0 || profile.PollInterval > time.Hour || + profile.RequestTimeout <= 0 || profile.RequestTimeout > maxConfiguredDuration || + profile.MaxWallTime <= 0 || profile.MaxWallTime > maxConfiguredDuration { + return errors.New("reducto profile bounds are invalid") + } + if profile.RetainStructured && profile.MaxArtifacts != 1 || + !profile.RetainStructured && profile.MaxArtifacts != 0 { + return errors.New("reducto structured artifact count is invalid") + } + return nil +} + +func validateUploadToken(value string) error { + if len(value) <= len("reducto://") || len(value) > maxUploadTokenBytes || + !strings.HasPrefix(value, "reducto://") { + return errors.New("upload token is invalid") + } + return validateOpaqueToken(strings.TrimPrefix(value, "reducto://"), maxUploadTokenBytes-len("reducto://")) +} + +func validateJobToken(value string) error { + if value == "" || len(value) > maxJobTokenBytes || value == "." || value == ".." { + return errors.New("job token is invalid") + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= '0' && char <= '9' || strings.ContainsRune("-._", char) { + continue + } + return errors.New("job token contains unsupported characters") + } + return nil +} + +func encodeResumeHandle(payload resumePayload) (string, error) { + raw, err := json.Marshal(payload) + if err != nil { + return "", err + } + value := resumeHandlePrefix + base64.RawURLEncoding.EncodeToString(raw) + if len(value) > 512 { + return "", errors.New("resume handle exceeds core bound") + } + if err := validateOpaqueToken(value, 512); err != nil { + return "", err + } + return value, nil +} + +func decodeResumeHandle(value string, authorization document.RenditionAuthorization) (resumePayload, error) { + if err := validateOpaqueToken(value, 512); err != nil || !strings.HasPrefix(value, resumeHandlePrefix) { + return resumePayload{}, errors.New("resume handle envelope is invalid") + } + raw, err := base64.RawURLEncoding.DecodeString(strings.TrimPrefix(value, resumeHandlePrefix)) + if err != nil { + return resumePayload{}, errors.New("resume handle payload is invalid") + } + var payload resumePayload + if err := strictJSON(raw, &payload); err != nil { + return resumePayload{}, errors.New("resume handle schema is invalid") + } + canonical, err := json.Marshal(payload) + if err != nil || !bytes.Equal(raw, canonical) || payload.Version != resumeHandleVersion || + validateJobToken(payload.JobID) != nil { + return resumePayload{}, errors.New("resume handle identity is invalid") + } + startedAt, err := parseCanonicalTimestamp(payload.StartedAt) + if err != nil { + return resumePayload{}, err + } + submittedAt, err := parseCanonicalTimestamp(payload.SubmittedAt) + if err != nil || submittedAt.Before(startedAt) { + return resumePayload{}, errors.New("resume handle submission is invalid") + } + authorizedAt, err := parseCanonicalTimestamp(authorization.AuthorizedAt) + if err != nil { + return resumePayload{}, errors.New("resume authorization start is invalid") + } + expiresAt, err := parseCanonicalTimestamp(authorization.ExpiresAt) + if err != nil || startedAt.Before(authorizedAt) || !startedAt.Before(expiresAt) || + !submittedAt.Before(expiresAt) { + return resumePayload{}, errors.New("resume handle is outside the sealed authorization interval") + } + if payload.Requests < 2 || payload.Requests > maxResumeUsageValue || + payload.Retries < 0 || payload.Retries > payload.Requests || + payload.InputBytes != authorization.SourceBytes || + payload.OutputBytes < 0 || payload.OutputBytes > maxResumeUsageValue || + payload.RetryDelayMillis < 0 || payload.RetryDelayMillis > int64((24*time.Hour)/time.Millisecond) { + return resumePayload{}, errors.New("resume handle accounting is invalid") + } + return payload, nil +} + +func parseCanonicalTimestamp(value string) (time.Time, error) { + parsed, err := time.Parse(timeForm, value) + if err != nil || parsed.Format(timeForm) != value { + return time.Time{}, errors.New("timestamp is not canonical UTC RFC3339Nano") + } + return parsed, nil +} + +func validateCanonicalToken(value string, maximum int) error { + if value == "" || len(value) > maximum || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return errors.New("value is not a canonical token") + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || + char >= '0' && char <= '9' || strings.ContainsRune("_.-", char) { + continue + } + return errors.New("value contains unsupported token characters") + } + return nil +} + +func validateOpaqueToken(value string, maximum int) error { + if value == "" || len(value) > maximum || value == "." || value == ".." || + value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return errors.New("opaque token is invalid") + } + for _, char := range value { + if char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || + char >= '0' && char <= '9' || strings.ContainsRune("-._~", char) { + continue + } + return errors.New("opaque token contains unsupported characters") + } + return nil +} + +func validCredential(value string) bool { + if value == "" || len(value) > maxCredentialBytes || value != strings.TrimSpace(value) || !utf8.ValidString(value) { + return false + } + for _, char := range value { + if char < 0x21 || char == 0x7f { + return false + } + } + return true +} + +func readExactUpload( + ctx, callerCtx context.Context, expiresAt time.Time, + upload document.AuthorizedUpload, metadata document.AuthorizedUploadMetadata, limit int64, +) ([]byte, error) { + readDone := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = upload.Close() + case <-readDone: + } + }() + data, err := io.ReadAll(io.LimitReader(contextReader{ctx: ctx, reader: upload}, limit+1)) + close(readDone) + if ctx.Err() != nil { + clear(data) + return nil, classifyContextFailure(callerCtx, ctx, expiresAt, ctx.Err()) + } + if err != nil { + return nil, renditionError(document.RenditionErrorTransient, + "Reducto upload could not be read", err) + } + digest := sha256.Sum256(data) + if int64(len(data)) != metadata.ByteLength || int64(len(data)) > limit || + hex.EncodeToString(digest[:]) != metadata.SHA256 { + clear(data) + return nil, renditionError(document.RenditionErrorPolicyRejected, + "Reducto upload identity changed", nil) + } + return data, nil +} + +func classifyContextFailure( + callerCtx, operationCtx context.Context, expiresAt time.Time, cause error, +) error { + if errors.Is(callerCtx.Err(), context.Canceled) { + return renditionError(document.RenditionErrorCanceled, "Reducto rendering was canceled", callerCtx.Err()) + } + if callerCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, "Reducto caller deadline was reached", cause) + } + if !expiresAt.IsZero() && !time.Now().UTC().Before(expiresAt) { + return expiredError(cause) + } + if operationCtx.Err() != nil { + return renditionError(document.RenditionErrorTransient, "Reducto operation timed out", cause) + } + return renditionError(document.RenditionErrorTransient, "Reducto operation context failed", cause) +} + +type contextReader struct { + ctx context.Context + reader io.Reader +} + +type resumeAuthorizationUpload struct { + metadata document.AuthorizedUploadMetadata +} + +func (upload resumeAuthorizationUpload) Metadata() document.AuthorizedUploadMetadata { + return upload.metadata +} + +func (resumeAuthorizationUpload) Read([]byte) (int, error) { return 0, io.EOF } +func (resumeAuthorizationUpload) Close() error { return nil } + +func (reader contextReader) Read(buffer []byte) (int, error) { + if err := reader.ctx.Err(); err != nil { + return 0, err + } + return reader.reader.Read(buffer) +} + +func readBounded(reader io.Reader, limit int64) ([]byte, error) { + if limit <= 0 { + return nil, errResponseTooLarge + } + data, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if err != nil { + return nil, err + } + if int64(len(data)) > limit { + return nil, errResponseTooLarge + } + return data, nil +} + +func strictJSON(raw []byte, target any) error { + return json.Unmarshal(raw, target, json.RejectUnknownMembers(true)) +} + +func requireMembers(raw []byte, names ...string) error { + var members map[string]jsontext.Value + if err := json.Unmarshal(raw, &members); err != nil { + return err + } + for _, name := range names { + if _, ok := members[name]; !ok { + return fmt.Errorf("required JSON member %q is missing", name) + } + } + return nil +} + +func validProgress(value *float64) bool { + return value == nil || finite(*value) && *value >= 0 && *value <= 1 +} + +func finite(value float64) bool { return !math.IsNaN(value) && !math.IsInf(value, 0) } + +func isNullJSON(raw []byte) bool { + trimmed := bytes.TrimSpace(raw) + return len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) +} + +func familyUnit(family string) (document.EvidenceUnitKind, document.EvidenceLocatorKind, bool) { + switch family { + case "pdf": + return document.EvidenceUnitPage, document.EvidenceLocatorPage, true + case "presentation": + return document.EvidenceUnitSlide, document.EvidenceLocatorSlide, true + case "spreadsheet": + return document.EvidenceUnitSheet, document.EvidenceLocatorSheet, true + default: + return "", "", false + } +} + +func allowsStructured(authorization document.RenditionAuthorization) bool { + return authorization.MaxArtifacts >= 1 && authorization.MaxArtifactBytes > 0 && + slices.Contains(authorization.AllowedArtifactRoles, document.EvidenceArtifactStructured) +} + +func waitContext(ctx context.Context, delay time.Duration) error { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func boundedOperationContext( + ctx context.Context, expiresAt time.Time, wall time.Duration, +) (context.Context, context.CancelFunc) { + deadline := time.Now().Add(wall) + if !expiresAt.IsZero() && expiresAt.Before(deadline) { + deadline = expiresAt + } + return context.WithDeadline(ctx, deadline) +} + +func renditionError(code document.RenditionErrorCode, message string, cause error) error { + providerError, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("reducto error classification failed: %w", err) + } + return providerError +} + +func malformedError(message string, cause error) error { + return renditionError(document.RenditionErrorMalformedEvidence, message, cause) +} + +func expiredError(cause error) error { + return renditionError(document.RenditionErrorPolicyRejected, "Reducto authorization expired", cause) +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +func nilValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} + +func min64(left, right int64) int64 { + if left < right { + return left + } + return right +} diff --git a/document/reducto/client_test.go b/document/reducto/client_test.go new file mode 100644 index 00000000..3d0d0d53 --- /dev/null +++ b/document/reducto/client_test.go @@ -0,0 +1,1045 @@ +package reducto + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +const ( + testFileID = "reducto://123e4567-e89b-12d3-a456-426614174000.pdf" + testJobID = "123e4567-e89b-12d3-a456-426614174001" +) + +var _ document.ResumableRenditionProvider = (*Client)(nil) + +func TestClientUploadsExactBytesAndMapsNaturalEvidence(t *testing.T) { + tests := []struct { + family, mediaType, filename string + unitKind document.EvidenceUnitKind + locatorKind document.EvidenceLocatorKind + }{ + {family: "pdf", mediaType: "application/pdf", filename: "synthetic.pdf", unitKind: document.EvidenceUnitPage, locatorKind: document.EvidenceLocatorPage}, + {family: "presentation", mediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", filename: "synthetic.pptx", unitKind: document.EvidenceUnitSlide, locatorKind: document.EvidenceLocatorSlide}, + } + for _, testCase := range tests { + t.Run(testCase.family, func(t *testing.T) { + source := []byte("synthetic exact " + testCase.family + " bytes") + fixture := newFixture(t, testCase.family, testCase.mediaType, testCase.filename, source) + fixture.profile.RetainStructured = true + fixture.profile.MaxArtifacts = 1 + var paths []string + polls := 0 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + paths = append(paths, request.URL.Path) + assert.Equal(t, "https", request.URL.Scheme) + assert.Equal(t, apiHost, request.URL.Host) + assert.Empty(t, request.URL.RawQuery) + assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) + switch request.URL.Path { + case uploadPath: + assertUpload(t, request, fixture.metadata, source) + return response(request, http.StatusOK, `{"file_id":"`+testFileID+`","presigned_url":"https://attacker.invalid/upload"}`), nil + case parsePath: + assert.Equal(t, http.MethodPost, request.Method) + var body map[string]any + require.NoError(t, json.NewDecoder(request.Body).Decode(&body)) + assert.Equal(t, map[string]any{ + "document_url": testFileID, + "advanced_options": map[string]any{"add_page_markers": true, "return_ocr_data": false}, + "options": map[string]any{ + "chunking": map[string]any{"chunk_mode": "page"}, + "force_url_result": false, + }, + "priority": false, + }, body) + return response(request, http.StatusOK, `{"job_id":"`+testJobID+`"}`), nil + case jobPath(testJobID): + assert.Equal(t, http.MethodGet, request.Method) + polls++ + if polls == 1 { + return response(request, http.StatusOK, `{"status":"Pending","progress":0.25,"reason":null,"result":null}`), nil + } + return response(request, http.StatusOK, completedBody()), nil + default: + t.Fatalf("unexpected route %s", request.URL.String()) + return nil, errors.New("unexpected route") + } + }) + client := fixture.client(t, transport) + var checkpoint document.RenditionResumeHandle + checkpointCount := 0 + result, err := document.RenderRenditionWithResume( + t.Context(), client, fixture.upload(), fixture.authorization, nil, + func(handle document.RenditionResumeHandle) error { + checkpoint = handle + checkpointCount++ + return nil + }, + ) + require.NoError(t, err) + assert.NotEqual(t, testJobID, checkpoint.Value) + assert.LessOrEqual(t, len(checkpoint.Value), 512) + issued := decodeTestResumeHandle(t, checkpoint.Value) + assert.Equal(t, testJobID, issued.JobID) + assert.Equal(t, int64(4), issued.Requests) + assert.Equal(t, fixture.metadata.ByteLength, issued.InputBytes) + assert.Positive(t, issued.OutputBytes) + assert.GreaterOrEqual(t, checkpointCount, 4) + started, err := time.Parse(timeForm, issued.StartedAt) + require.NoError(t, err) + assert.False(t, time.Now().UTC().Before(started)) + assert.Equal(t, []string{uploadPath, parsePath, jobPath(testJobID), jobPath(testJobID)}, paths) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, testCase.unitKind, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 2) + assert.Equal(t, testCase.locatorKind, result.Evidence.Units[0].Locator.Kind) + assert.Equal(t, document.EvidenceIndexOriginOne, result.Evidence.Units[0].Locator.IndexOrigin) + assert.Equal(t, int64(1), result.Evidence.Units[0].Locator.Start) + assert.Empty(t, result.Evidence.Units[0].ProviderID) + assert.Equal(t, "# First unit\n\nFirst body", result.Evidence.Units[0].Text) + assert.Equal(t, "# First unit\n\nFirst body\n\n---\n\nSecond unit", string(result.ProviderMarkdown)) + require.Len(t, result.Artifacts, 1) + assert.Equal(t, document.EvidenceArtifactStructured, result.Artifacts[0].Role) + assert.Contains(t, string(result.Artifacts[0].Payload), `"pdf_url":"https://attacker.invalid/converted.pdf"`) + assert.Equal(t, int64(4), result.Receipt.Usage.Requests) + assert.Equal(t, int64(2), result.Receipt.Usage.Units) + }) + } +} + +func TestNaturalEvidenceRefusesUnreportedSpreadsheetNames(t *testing.T) { + chunks := []resultChunk{{ + Content: "sheet content", + Blocks: []resultBlock{{ + Content: "sheet content", Type: "Text", + BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(1)), + Top: new(0.0), Width: new(1.0), OriginalPage: new(int64(1)), + }, + }}, + }} + _, _, err := naturalEvidence(chunks, "spreadsheet", 1) + require.ErrorContains(t, err, "stable sheet name") + + fixture := newFixture(t, "spreadsheet", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "synthetic.xlsx", []byte("spreadsheet source")) + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("unsupported spreadsheet reached egress") + return nil, errors.New("spreadsheet reached egress") + })) + _, err = document.RenderRenditionWithResume(t.Context(), client, + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorPolicyRejected) +} + +func TestClientResumesOnlyTheFixedJobRoute(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("resume source")) + started := time.Now().UTC().Add(-30 * time.Second) + checkpointed := started.Add(time.Second) + resumeValue := testResumeValue(t, testJobID, started, checkpointed, fixture.metadata.ByteLength) + var paths []string + var transport http.RoundTripper = roundTripFunc(func(request *http.Request) (*http.Response, error) { + paths = append(paths, request.URL.Path) + return response(request, http.StatusOK, completedBody()), nil + }) + result, err := fixture.client(t, transport).RenderResumable( + t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: resumeValue}, nil, + ) + require.NoError(t, err) + assert.Equal(t, []string{jobPath(testJobID)}, paths) + assert.Equal(t, "reducto-"+testJobID, result.Receipt.OperationID) + + for _, handle := range []string{"", testJobID, "../job", "https://attacker.invalid", strings.Repeat("a", 513)} { + _, err = fixture.client(t, transport).RenderResumable( + t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: handle}, nil) + assertCode(t, err, document.RenditionErrorUnknownJob) + } + + invalidAuthorization := fixture.authorization + invalidAuthorization.MaxTotalResultBytes = 0 + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("invalid resume authorization reached egress") + return nil, errors.New("invalid resume reached egress") + })) + _, err = client.RenderResumable(t.Context(), nil, invalidAuthorization, + &document.RenditionResumeHandle{Value: resumeValue}, nil) + require.Error(t, err) +} + +func TestClientResumesHistoricalAuthorizationWithRecordedReceiptInterval(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("historical source")) + authorizedAt := time.Now().UTC().Add(-2 * time.Hour) + expiresAt := authorizedAt.Add(time.Minute) + startedAt := authorizedAt.Add(10 * time.Second) + checkpointedAt := startedAt.Add(5 * time.Second) + fixture.authorization.AuthorizedAt = authorizedAt.Format(timeForm) + fixture.authorization.ExpiresAt = expiresAt.Format(timeForm) + client := fixture.client(t, completedTransport(t)) + snapshot := executionSnapshot(t, startedAt, client, fixture) + result, err := document.ResumeRendition(t.Context(), client, snapshot, document.RenditionResumeHandle{ + Value: testResumeValue(t, testJobID, startedAt, checkpointedAt, fixture.metadata.ByteLength), + }, nil) + require.NoError(t, err) + assert.Equal(t, startedAt.Format(timeForm), result.Receipt.StartedAt) + completedAt, err := time.Parse(timeForm, result.Receipt.CompletedAt) + require.NoError(t, err) + assert.True(t, completedAt.After(expiresAt)) + assert.Equal(t, int64(3), result.Receipt.Usage.Requests) +} + +func TestClientAcceptsOnlyPinnedSDKAsyncLifecycleStates(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("lifecycle source")) + polls := 0 + var transport http.RoundTripper = roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath: + return response(request, http.StatusOK, `{"file_id":"`+testFileID+`","presigned_url":null}`), nil + case parsePath: + return response(request, http.StatusOK, `{"job_id":"`+testJobID+`"}`), nil + case jobPath(testJobID): + polls++ + if polls == 1 { + return response(request, http.StatusOK, `{"status":"Pending","progress":0.5,"reason":null,"result":null}`), nil + } + if polls == 2 { + return response(request, http.StatusOK, `{"status":"Idle","progress":0.9,"reason":null,"result":null}`), nil + } + return response(request, http.StatusOK, completedBody()), nil + default: + t.Fatalf("unexpected route %s", request.URL.String()) + return nil, errors.New("unexpected route") + } + }) + result, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + require.NoError(t, err) + assert.Equal(t, int64(2), result.Receipt.Usage.Units) + + for _, status := range []string{"InProgress", "Completing"} { + fixture = newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("status drift")) + transport = routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): `{"status":"` + status + `","progress":0.5,"reason":null,"result":null}`, + }) + _, err = document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorPolicyRejected) + } +} + +func TestClientNeverFollowsProviderAuthoredURLs(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("url source")) + var paths []string + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + paths = append(paths, request.URL.Path) + switch request.URL.Path { + case uploadPath: + return response(request, http.StatusOK, `{"file_id":"`+testFileID+`","presigned_url":"https://attacker.invalid/upload"}`), nil + case parsePath: + return response(request, http.StatusOK, `{"job_id":"`+testJobID+`"}`), nil + case jobPath(testJobID): + return response(request, http.StatusOK, `{ + "status":"Completed","progress":1,"reason":null, + "result":{"result":{"type":"url","result_id":"result-1","url":"https://attacker.invalid/result"},"usage":{"num_pages":2},"duration":1,"job_id":"`+testJobID+`","pdf_url":"https://attacker.invalid/pdf"} + }`), nil + default: + t.Fatalf("provider-authored URL was followed: %s", request.URL.String()) + return nil, errors.New("provider URL followed") + } + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Equal(t, []string{uploadPath, parsePath, jobPath(testJobID)}, paths) +} + +func TestClientRetainsRawResultOnlyWhenProfileAndAuthorizationPermit(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("artifact source")) + transport := completedTransport(t) + + result, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + require.NoError(t, err) + assert.Empty(t, result.Artifacts) + assert.Empty(t, result.Evidence.Artifacts) + + fixture.profile.RetainStructured = true + fixture.profile.MaxArtifacts = 1 + client := fixture.client(t, completedTransport(t)) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + fixture.authorization.AllowedArtifactRoles = nil + fixture.authorization.MaxArtifacts = 0 + fixture.authorization.MaxArtifactBytes = 0 + result, err = document.RenderRenditionWithResume(t.Context(), client, + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + require.NoError(t, err) + assert.Empty(t, result.Artifacts) +} + +func TestClientRejectsPartialOutputAndSchemaDrift(t *testing.T) { + tests := []struct { + name, body string + want document.RenditionErrorCode + }{ + {name: "failed", body: `{"status":"Failed","progress":0.5,"reason":"private source detail","result":null}`, want: document.RenditionErrorUnsupportedInput}, + {name: "unknown status", body: `{"status":"Paused","progress":0.5,"reason":null,"result":null}`, want: document.RenditionErrorPolicyRejected}, + {name: "unknown model field", body: `{"status":"Completed","progress":1,"reason":null,"model":"next-model","result":null}`, want: document.RenditionErrorPolicyRejected}, + {name: "newer API usage drift", body: strings.Replace(completedBody(), + `"num_pages":2`, `"num_pages":2,"credits":2.5`, 1), want: document.RenditionErrorMalformedEvidence}, + {name: "empty chunk", body: completedBodyWithChunks(testJobID, `[]`, 0), want: document.RenditionErrorMalformedEvidence}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("partial source")) + transport := routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): testCase.body, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "private source detail") + assert.NotContains(t, err.Error(), "next-model") + }) + } +} + +func TestNaturalEvidenceRepresentsProviderIdentifiedMissingPages(t *testing.T) { + chunks := []resultChunk{ + { + Content: "first", Embed: "first", Blocks: []resultBlock{{ + Content: "first", Type: "Text", BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(1)), Top: new(0.0), Width: new(1.0), + }, + }}, + }, + { + Content: "", Embed: "", Blocks: []resultBlock{{ + Content: "", Type: "Text", BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(2)), Top: new(0.0), Width: new(1.0), + }, + }}, + }, + { + Content: "third", Embed: "third", Blocks: []resultBlock{{ + Content: "third", Type: "Text", BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(3)), Top: new(0.0), Width: new(1.0), + }, + }}, + }, + } + + evidence, markdown, err := naturalEvidence(chunks, "pdf", 4) + require.NoError(t, err) + assert.Equal(t, document.EvidencePartial, evidence.Completeness) + require.Len(t, evidence.Units, 2) + assert.Equal(t, int64(1), evidence.Units[0].Locator.Start) + assert.Equal(t, int64(3), evidence.Units[1].Locator.Start) + require.Len(t, evidence.Omissions, 2) + assert.Equal(t, document.EvidenceOmissionUnit, evidence.Omissions[0].Kind) + assert.Equal(t, int64(2), evidence.Omissions[0].Locator.Start) + assert.Equal(t, int64(4), evidence.Omissions[1].Locator.Start) + assert.Equal(t, "first\n\n---\n\nthird", string(markdown)) +} + +func TestNaturalEvidenceRejectsUnlocatedContentlessChunk(t *testing.T) { + _, _, err := naturalEvidence([]resultChunk{{Content: "", Blocks: nil}}, "pdf", 1) + require.ErrorContains(t, err, "incomplete") +} + +func TestClientClassifiesHTTPAndSubmissionFailures(t *testing.T) { + tests := []struct { + name string + status int + want document.RenditionErrorCode + }{ + {name: "authentication", status: http.StatusUnauthorized, want: document.RenditionErrorAuthentication}, + {name: "rate", status: http.StatusTooManyRequests, want: document.RenditionErrorRateLimited}, + {name: "capacity", status: http.StatusServiceUnavailable, want: document.RenditionErrorCapacity}, + {name: "terminal input", status: http.StatusUnsupportedMediaType, want: document.RenditionErrorUnsupportedInput}, + {name: "redirect", status: http.StatusTemporaryRedirect, want: document.RenditionErrorPolicyRejected}, + {name: "transient", status: http.StatusBadGateway, want: document.RenditionErrorTransient}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("http source")) + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + result := response(request, testCase.status, `{"detail":"private provider data"}`) + result.Header.Set("Location", "https://attacker.invalid/redirect") + return result, nil + }) + _, err := fixture.client(t, transport).Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, testCase.want) + assert.NotContains(t, err.Error(), "private provider data") + assert.NotContains(t, err.Error(), "attacker.invalid") + }) + } + + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("ambiguous source")) + var transport http.RoundTripper = roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("private transport detail") + }) + _, err := fixture.client(t, transport).Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) + assert.NotContains(t, err.Error(), "private transport detail") + + for _, body := range []string{`{}`, `{"job_id":"UPPER"}`, `{"job_id":"job~suffix"}`} { + transport = routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: body, + }) + _, err = fixture.client(t, transport).Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) + } + + for _, status := range []int{http.StatusNotFound, http.StatusGone} { + transport = roundTripFunc(func(request *http.Request) (*http.Response, error) { + return response(request, status, `{"detail":"private expired job"}`), nil + }) + _, err = fixture.client(t, transport).RenderResumable(t.Context(), nil, fixture.authorization, + &document.RenditionResumeHandle{Value: testResumeValue(t, testJobID, + time.Now().UTC().Add(-time.Second), time.Now().UTC(), fixture.metadata.ByteLength)}, nil) + assertCode(t, err, document.RenditionErrorUnknownJob) + } + + client, newErr := NewProvider(testProfile(), failingSecrets{}, completedTransport(t)) + require.NoError(t, newErr) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorAuthentication) + assert.NotContains(t, err.Error(), "private secret backend") +} + +func TestClientEnforcesInputControlPollResultAndArtifactBounds(t *testing.T) { + t.Run("input identity", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("authorized")) + bad := &testUpload{Reader: strings.NewReader("different"), metadata: fixture.metadata} + transport := roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("identity mismatch reached egress") + return nil, errors.New("reached egress") + }) + _, err := fixture.client(t, transport).Render(t.Context(), bad, fixture.authorization) + assertCode(t, err, document.RenditionErrorPolicyRejected) + }) + + t.Run("control response", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("control")) + fixture.profile.MaxControlBytes = 32 + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + return response(request, http.StatusOK, strings.Repeat("x", 33)), nil + }) + _, err := fixture.client(t, transport).Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) + }) + + t.Run("poll limit", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("poll")) + fixture.profile.MaxPolls = 2 + transport := routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): `{"status":"Idle","progress":0,"reason":null,"result":null}`, + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorTransient) + }) + + t.Run("result response", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("result")) + fixture.profile.MaxResultBytes = 128 + transport := routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): strings.Repeat("x", 129), + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) + + t.Run("structured artifact", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("artifact")) + fixture.profile.RetainStructured = true + fixture.profile.MaxArtifacts = 1 + fixture.profile.MaxArtifactBytes = 64 + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, completedTransport(t)), + fixture.upload(), fixture.authorization, nil, + func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) +} + +func TestNaturalEvidenceRejectsUnboundedReportedUnitAllocation(t *testing.T) { + chunks := []resultChunk{{ + Content: "one", + Blocks: []resultBlock{{ + Content: "one", Type: "Text", + BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(1)), + Top: new(0.0), Width: new(1.0), OriginalPage: new(int64(1)), + }, + }}, + }} + _, _, err := naturalEvidence(chunks, "pdf", 100_001) + require.EqualError(t, err, "natural unit count exceeds provider bound") +} + +func TestNaturalEvidencePreservesCompleteChunkContent(t *testing.T) { + chunks := []resultChunk{{ + Content: "Complete provider Markdown", + Blocks: []resultBlock{{ + Content: "geometry text", Type: "Text", + BBox: resultBoundingBox{ + Height: new(1.0), Left: new(0.0), Page: new(int64(1)), + Top: new(0.0), Width: new(1.0), OriginalPage: new(int64(1)), + }, + }}, + }} + evidence, markdown, err := naturalEvidence(chunks, "pdf", 1) + require.NoError(t, err) + assert.Equal(t, "Complete provider Markdown", evidence.Units[0].Text) + assert.Equal(t, "Complete provider Markdown", string(markdown)) +} + +func TestNaturalEvidenceRejectsBlockTypeOutsidePinnedSDK(t *testing.T) { + chunks := []resultChunk{{ + Content: "Signed by Synthetic Person", + Blocks: []resultBlock{{ + Content: "Signed by Synthetic Person", Type: "Signature", + BBox: resultBoundingBox{ + Height: new(0.1), Left: new(0.1), Page: new(int64(1)), + Top: new(0.1), Width: new(0.2), OriginalPage: new(int64(1)), + }, + }}, + }} + _, _, err := naturalEvidence(chunks, "pdf", 1) + require.ErrorContains(t, err, "type") +} + +func TestStrictJSONRejectsDuplicateMembers(t *testing.T) { + var job jobResponse + err := strictJSON([]byte(`{"status":"Pending","status":"Completed","progress":0,"reason":null,"result":null}`), &job) + require.ErrorContains(t, err, "duplicate") +} + +func TestCompletedRequiresPinnedSDKFields(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("required fields")) + client := fixture.client(t, completedTransport(t)) + valid := completedResultJSON() + + for _, testCase := range []struct { + name, body string + }{ + {name: "duration", body: strings.Replace(valid, `,"duration":1.25`, "", 1)}, + {name: "job id", body: strings.Replace(valid, `,"job_id":"`+testJobID+`"`, "", 1)}, + {name: "usage", body: strings.Replace(valid, `"usage":{"num_pages":2},`, "", 1)}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := client.completed(json.RawMessage(testCase.body), testJobID, fixture.authorization) + assertCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestClientClassifiesCancellationAfterEgressAsAmbiguous(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("cancel")) + ctx, cancel := context.WithCancel(t.Context()) + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + cancel() + <-request.Context().Done() + return nil, request.Context().Err() + }) + _, err := fixture.client(t, transport).Render(ctx, fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) +} + +func TestSubmissionFailureRemainsAmbiguousWhenWallTimeoutOccursAfterEgress(t *testing.T) { + operationCtx, cancel := context.WithDeadline(t.Context(), time.Now().Add(-time.Second)) + defer cancel() + client := &Client{} + err := client.submissionFailure(t.Context(), operationCtx, time.Time{}, + "Reducto submission outcome is ambiguous", context.DeadlineExceeded) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) +} + +func TestSubmissionFailureRemainsAmbiguousWhenAuthorizationExpiresAfterEgress(t *testing.T) { + client := &Client{} + err := client.submissionFailure(t.Context(), t.Context(), time.Now().UTC().Add(-time.Second), + "Reducto submission outcome is ambiguous", errors.New("synthetic transport failure")) + assertCode(t, err, document.RenditionErrorAmbiguousSubmission) +} + +func TestClientDistinguishesPreEgressContextFailures(t *testing.T) { + t.Run("wall timeout while reading source", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("wall source")) + fixture.profile.MaxWallTime = 10 * time.Millisecond + transport := roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("wall timeout reached egress") + return nil, errors.New("unexpected egress") + }) + upload := newBlockingTestUpload(fixture.metadata) + _, err := fixture.client(t, transport).Render(t.Context(), upload, fixture.authorization) + assertCode(t, err, document.RenditionErrorTransient) + assert.True(t, upload.wasClosed()) + }) + + t.Run("wall timeout resolving secret", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("secret wall source")) + fixture.profile.MaxWallTime = time.Millisecond + client, err := NewProvider(fixture.profile, blockingSecrets{}, completedTransport(t)) + require.NoError(t, err) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + fixture.authorization.AllowedArtifactRoles = nil + fixture.authorization.MaxArtifacts = 0 + fixture.authorization.MaxArtifactBytes = 0 + _, err = client.Render(t.Context(), fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorTransient) + }) + + t.Run("caller cancellation resolving secret", func(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("secret cancel source")) + ctx, cancel := context.WithCancel(t.Context()) + client, err := NewProvider(fixture.profile, cancelingSecrets{cancel: cancel}, completedTransport(t)) + require.NoError(t, err) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + fixture.authorization.AllowedArtifactRoles = nil + fixture.authorization.MaxArtifacts = 0 + fixture.authorization.MaxArtifactBytes = 0 + _, err = client.Render(ctx, fixture.upload(), fixture.authorization) + assertCode(t, err, document.RenditionErrorCanceled) + }) +} + +func TestClientTreatsExpiryAfterDurableCheckpointAsResumableTimeout(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("durable expiry source")) + fixture.profile.PollInterval = 50 * time.Millisecond + fixture.authorization.ExpiresAt = time.Now().UTC().Add(20 * time.Millisecond).Format(timeForm) + transport := routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): `{"status":"Pending","progress":0.5,"reason":null,"result":null}`, + }) + checkpointed := false + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { + checkpointed = true + return nil + }) + assert.True(t, checkpointed) + assertCode(t, err, document.RenditionErrorTransient) +} + +func TestClientTreatsExpiryAfterCompletedResultAsResumableTimeout(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("completed expiry")) + fixture.authorization.ExpiresAt = time.Now().UTC().Add(100 * time.Millisecond).Format(timeForm) + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath: + return response(request, http.StatusOK, `{"file_id":"`+testFileID+`","presigned_url":null}`), nil + case parsePath: + return response(request, http.StatusOK, `{"job_id":"`+testJobID+`"}`), nil + case jobPath(testJobID): + time.Sleep(150 * time.Millisecond) + return response(request, http.StatusOK, completedBody()), nil + default: + return nil, errors.New("unexpected route") + } + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorTransient) +} + +func TestClientClosesBlockedUploadOnWallTimeout(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("blocked upload")) + fixture.profile.MaxWallTime = 10 * time.Millisecond + upload := newBlockingTestUpload(fixture.metadata) + client := fixture.client(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("blocked upload reached egress") + return nil, errors.New("blocked upload reached egress") + })) + + _, err := client.Render(t.Context(), upload, fixture.authorization) + assertCode(t, err, document.RenditionErrorTransient) + assert.True(t, upload.wasClosed()) +} + +func TestClientAppliesWallDeadlineAfterCompletedResultProcessing(t *testing.T) { + fixture := newFixture(t, "pdf", "application/pdf", "synthetic.pdf", []byte("final wall")) + fixture.profile.MaxWallTime = 10 * time.Millisecond + transport := roundTripFunc(func(request *http.Request) (*http.Response, error) { + switch request.URL.Path { + case uploadPath: + return response(request, http.StatusOK, `{"file_id":"`+testFileID+`","presigned_url":null}`), nil + case parsePath: + return response(request, http.StatusOK, `{"job_id":"`+testJobID+`"}`), nil + case jobPath(testJobID): + time.Sleep(20 * time.Millisecond) + return response(request, http.StatusOK, completedBody()), nil + default: + return nil, errors.New("unexpected route") + } + }) + _, err := document.RenderRenditionWithResume(t.Context(), fixture.client(t, transport), + fixture.upload(), fixture.authorization, nil, func(document.RenditionResumeHandle) error { return nil }) + assertCode(t, err, document.RenditionErrorTransient) +} + +func TestProviderRequiresFrozenNamedProfileAndHardenedTransport(t *testing.T) { + profile := testProfile() + _, err := NewProvider(profile, nil, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("unused transport") + })) + require.ErrorContains(t, err, "credential") + _, err = NewProvider(profile, testSecrets{value: "secret"}, nil) + require.ErrorContains(t, err, "transport") + profile.SecretBinding = "bad/binding" + _, err = NewProvider(profile, testSecrets{value: "secret"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("unused transport") + })) + require.ErrorContains(t, err, "binding") + + profile = testProfile() + client, err := NewProvider(profile, testSecrets{value: "secret"}, roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("unused transport") + })) + require.NoError(t, err) + descriptor := client.Descriptor() + descriptor.SupportedFormats[0].MediaFamily = "changed" + assert.NotEqual(t, descriptor, client.Descriptor()) + assert.Equal(t, providerID, client.Descriptor().ID) +} + +type fixture struct { + profile Profile + metadata document.AuthorizedUploadMetadata + authorization document.RenditionAuthorization + source []byte +} + +func newFixture(t *testing.T, family, mediaType, filename string, source []byte) *fixture { + t.Helper() + profile := testProfile() + client, err := NewProvider(profile, testSecrets{value: "synthetic-secret"}, + roundTripFunc(func(*http.Request) (*http.Response, error) { return nil, errors.New("unused") })) + require.NoError(t, err) + descriptor := client.Descriptor() + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: filename, MediaFamily: family, MediaType: mediaType, + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("1", 64), + ProviderMetadataChecksum: strings.Repeat("2", 64), + InputKind: document.RenditionInputOriginalFile, + } + now := time.Now().UTC() + return &fixture{ + profile: profile, metadata: metadata, source: bytes.Clone(source), + authorization: document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, + RenditionRequestFingerprint: strings.Repeat("3", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: family, MediaType: mediaType, InputKind: metadata.InputKind, + AllowedArtifactRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + MaxProviderMarkdownBytes: 16 << 10, MaxArtifactBytes: 32 << 10, + MaxArtifacts: 1, MaxTotalResultBytes: 128 << 10, + AuthorizedAt: now.Add(-time.Minute).Format(timeForm), + ExpiresAt: now.Add(time.Minute).Format(timeForm), + }, + } +} + +func testProfile() Profile { + return Profile{ + SecretBinding: "reducto-production", MaxUploadBytes: 1 << 20, + MaxRequestBytes: 2 << 20, MaxControlBytes: 4 << 10, + MaxPolls: 3, PollInterval: time.Millisecond, RequestTimeout: time.Second, + MaxResultBytes: 64 << 10, MaxArtifactBytes: 32 << 10, + MaxWallTime: time.Second, + } +} + +func (fixture *fixture) client(t *testing.T, transport http.RoundTripper) *Client { + t.Helper() + client, err := NewProvider(fixture.profile, testSecrets{value: "synthetic-secret"}, transport) + require.NoError(t, err) + descriptor := client.Descriptor() + fixture.authorization.ProviderID = descriptor.ID + fixture.authorization.DescriptorFingerprint = descriptor.Fingerprint + fixture.authorization.PolicyFingerprint = descriptor.PolicyFingerprint + if fixture.profile.RetainStructured { + fixture.authorization.AllowedArtifactRoles = []document.EvidenceArtifactRole{document.EvidenceArtifactStructured} + fixture.authorization.MaxArtifacts = 1 + } else { + fixture.authorization.AllowedArtifactRoles = nil + fixture.authorization.MaxArtifacts = 0 + fixture.authorization.MaxArtifactBytes = 0 + } + return client +} + +func (fixture *fixture) upload() document.AuthorizedUpload { + return &testUpload{Reader: bytes.NewReader(fixture.source), metadata: fixture.metadata} +} + +func assertUpload(t *testing.T, request *http.Request, metadata document.AuthorizedUploadMetadata, source []byte) { + t.Helper() + assert.Equal(t, http.MethodPost, request.Method) + mediaType, parameters, err := mime.ParseMediaType(request.Header.Get("Content-Type")) + require.NoError(t, err) + assert.Equal(t, "multipart/form-data", mediaType) + reader := multipart.NewReader(request.Body, parameters["boundary"]) + part, err := reader.NextPart() + require.NoError(t, err) + assert.Equal(t, "file", part.FormName()) + assert.Equal(t, metadata.Filename, part.FileName()) + assert.Equal(t, metadata.MediaType, part.Header.Get("Content-Type")) + payload, err := io.ReadAll(part) + require.NoError(t, err) + assert.Equal(t, source, payload) + _, err = reader.NextPart() + assert.ErrorIs(t, err, io.EOF) +} + +func completedBody() string { + return completedBodyWithChunks(testJobID, `[ + {"blocks":[ + {"bbox":{"height":1,"left":0,"page":1,"top":0,"width":1,"original_page":1},"content":"# First unit","type":"Title","image_url":null}, + {"bbox":{"height":1,"left":0,"page":1,"top":1,"width":1,"original_page":1},"content":"First body","type":"Text","image_url":"https://attacker.invalid/image.png"} + ],"content":"# First unit\n\nFirst body","embed":"First unit First body","enriched":null,"enrichment_success":false}, + {"blocks":[ + {"bbox":{"height":1,"left":0,"page":2,"top":0,"width":1,"original_page":2},"content":"Second unit","type":"Text","image_url":null} + ],"content":"Second unit","embed":"Second unit","enriched":null,"enrichment_success":false} + ]`, 2) +} + +func completedBodyWithChunks(jobID, chunks string, pages int) string { + return `{"status":"Completed","progress":1,"reason":null,"result":{` + + `"result":{"type":"full","chunks":` + chunks + `,"custom":null,"ocr":null},` + + `"usage":{"num_pages":` + jsonInt(pages) + `},"duration":1.25,` + + `"job_id":"` + jobID + `","pdf_url":"https://attacker.invalid/converted.pdf"}}` +} + +func completedResultJSON() string { + body := completedBody() + var job struct { + Result json.RawMessage `json:"result"` + } + if err := json.Unmarshal([]byte(body), &job); err != nil { + panic(err) + } + return string(job.Result) +} + +func jsonInt(value int) string { + return strconv.Itoa(value) +} + +func completedTransport(t *testing.T) http.RoundTripper { + t.Helper() + return routeTransport(t, map[string]string{ + uploadPath: `{"file_id":"` + testFileID + `","presigned_url":null}`, + parsePath: `{"job_id":"` + testJobID + `"}`, + jobPath(testJobID): completedBody(), + }) +} + +func routeTransport(t *testing.T, routes map[string]string) http.RoundTripper { + t.Helper() + return roundTripFunc(func(request *http.Request) (*http.Response, error) { + body, ok := routes[request.URL.Path] + if !ok { + t.Fatalf("unexpected route %s", request.URL.String()) + } + return response(request, http.StatusOK, body), nil + }) +} + +type testResumePayload struct { + Version string `json:"v"` + JobID string `json:"j"` + StartedAt string `json:"s"` + SubmittedAt string `json:"a"` + Requests int64 `json:"q"` + Retries int64 `json:"r"` + InputBytes int64 `json:"i"` + OutputBytes int64 `json:"o"` + RetryDelayMillis int64 `json:"d"` +} + +func testResumeValue(t *testing.T, jobID string, startedAt, submittedAt time.Time, inputBytes int64) string { + t.Helper() + raw, err := json.Marshal(testResumePayload{ + Version: "reducto-resume/v2", JobID: jobID, + StartedAt: startedAt.UTC().Format(timeForm), SubmittedAt: submittedAt.UTC().Format(timeForm), + Requests: 2, InputBytes: inputBytes, OutputBytes: 1, + }) + require.NoError(t, err) + return "r2." + base64.RawURLEncoding.EncodeToString(raw) +} + +func decodeTestResumeHandle(t *testing.T, value string) testResumePayload { + t.Helper() + require.True(t, strings.HasPrefix(value, "r2.")) + raw, err := base64.RawURLEncoding.DecodeString(strings.TrimPrefix(value, "r2.")) + require.NoError(t, err) + var payload testResumePayload + require.NoError(t, json.Unmarshal(raw, &payload)) + return payload +} + +func executionSnapshot( + t *testing.T, at time.Time, client *Client, fixture *fixture, +) document.RenditionExecutionSnapshotV1 { + t.Helper() + evidence, err := document.NewEvidencePolicy(100_000) + require.NoError(t, err) + normalization, err := document.NewNormalizePolicy(100_000) + require.NoError(t, err) + rendition, err := document.NewRenditionPolicy(normalization, 1_000) + require.NoError(t, err) + snapshot, err := document.SealRenditionExecutionAt( + at, client, fixture.upload(), fixture.authorization, evidence, rendition) + require.NoError(t, err) + return snapshot +} + +type testUpload struct { + io.Reader + + metadata document.AuthorizedUploadMetadata +} + +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } +func (*testUpload) Close() error { return nil } + +type blockingTestUpload struct { + metadata document.AuthorizedUploadMetadata + closed chan struct{} +} + +func newBlockingTestUpload(metadata document.AuthorizedUploadMetadata) *blockingTestUpload { + return &blockingTestUpload{metadata: metadata, closed: make(chan struct{})} +} + +func (upload *blockingTestUpload) Metadata() document.AuthorizedUploadMetadata { + return upload.metadata +} + +func (upload *blockingTestUpload) Read([]byte) (int, error) { + <-upload.closed + return 0, io.EOF +} + +func (upload *blockingTestUpload) Close() error { + select { + case <-upload.closed: + default: + close(upload.closed) + } + return nil +} + +func (upload *blockingTestUpload) wasClosed() bool { + select { + case <-upload.closed: + return true + default: + return false + } +} + +type testSecrets struct { + value string +} + +type failingSecrets struct{} + +type blockingSecrets struct{} + +func (blockingSecrets) ResolveSecret(ctx context.Context, _ string) (string, error) { + <-ctx.Done() + return "", ctx.Err() +} + +type cancelingSecrets struct{ cancel context.CancelFunc } + +func (secrets cancelingSecrets) ResolveSecret(ctx context.Context, _ string) (string, error) { + secrets.cancel() + <-ctx.Done() + return "", ctx.Err() +} + +func (failingSecrets) ResolveSecret(context.Context, string) (string, error) { + return "", errors.New("private secret backend detail") +} + +func (secrets testSecrets) ResolveSecret(context.Context, string) (string, error) { + return secrets.value, nil +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +func response(request *http.Request, status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, Status: http.StatusText(status), Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), Request: request, + } +} + +func assertCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerError, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok, "expected classified provider error, got %T: %v", err, err) + assert.Equal(t, want, providerError.Code()) +} diff --git a/document/reducto/doc.go b/document/reducto/doc.go new file mode 100644 index 00000000..cd84c30b --- /dev/null +++ b/document/reducto/doc.go @@ -0,0 +1,3 @@ +// Package reducto adapts Reducto's fixed hosted asynchronous parse API to +// Docbank's storage-neutral, resumable rendition provider contract. +package reducto diff --git a/document/rendition_frontmatter.go b/document/rendition_frontmatter.go new file mode 100644 index 00000000..cf065657 --- /dev/null +++ b/document/rendition_frontmatter.go @@ -0,0 +1,338 @@ +package document + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "slices" + "strconv" + "strings" + "unicode/utf8" + + "gopkg.in/yaml.v3" +) + +const ( + RenditionMarkdownContractV1 = "docbank-sanitized-markdown/v1" + RenditionNavigationOffsetBody = "body" + maxRenditionNavigationEntries = 1024 + maxRenditionFrontMatterBytes = 256 << 10 +) + +type RenditionMarkdownSourceV1 struct { + SHA256 string `yaml:"sha256"` + Format string `yaml:"format"` + MediaType string `yaml:"media_type"` +} + +type RenditionMarkdownBuildV1 struct { + BuildID string `yaml:"build_id"` + RenditionRequestFingerprint string `yaml:"rendition_request_fingerprint"` + EvidenceLexicalFingerprint string `yaml:"evidence_lexical_fingerprint"` + NormalizedEvidenceContract string `yaml:"normalized_evidence_contract"` + BodySHA256 string `yaml:"body_sha256"` + Completeness EvidenceCompleteness `yaml:"completeness"` + Truncated bool `yaml:"truncated"` +} + +type RenditionMarkdownDocumentV1 struct { + Title string `yaml:"title,omitempty"` + Language string `yaml:"language,omitempty"` + UnitKind EvidenceUnitKind `yaml:"unit_kind"` + UnitCount int `yaml:"unit_count"` +} + +type RenditionNavigationEntryV1 struct { + Key string `yaml:"key"` + Kind EvidenceLocatorKind `yaml:"kind"` + Title string `yaml:"title,omitempty"` + Line int `yaml:"line"` + Byte int `yaml:"byte"` +} + +type RenditionMarkdownNavigationV1 struct { + OffsetBase string `yaml:"offset_base"` + Complete bool `yaml:"complete"` + Entries []RenditionNavigationEntryV1 `yaml:"entries"` +} + +type RenditionFrontMatterV1 struct { + Contract string `yaml:"contract"` + Source RenditionMarkdownSourceV1 `yaml:"source"` + Rendition RenditionMarkdownBuildV1 `yaml:"rendition"` + Document RenditionMarkdownDocumentV1 `yaml:"document"` + Navigation RenditionMarkdownNavigationV1 `yaml:"navigation"` +} + +type RenditionEnvelopeV1 struct { + BuildID string + SourceSHA256 string + SourceFormat string + SourceMediaType string + RenditionRequestFingerprint string + EvidenceLexicalFingerprint string + NormalizedEvidenceContract string + UnitKind EvidenceUnitKind + Title string + Language string +} + +// EnvelopeRenditionV1 adds deterministic build-scoped YAML frontmatter to the +// retained Markdown only. Units and lexical segments remain body-derived and +// therefore never ingest frontmatter metadata. +func EnvelopeRenditionV1(rendition RenditionV1, envelope RenditionEnvelopeV1) (RenditionV1, RenditionFrontMatterV1, error) { + if len(rendition.Markdown) == 0 || rendition.MarkdownChecksum != checksumBytes(rendition.Markdown) { + return RenditionV1{}, RenditionFrontMatterV1{}, errors.New("rendition Markdown body authority is invalid") + } + for name, value := range map[string]string{ + "build ID": envelope.BuildID, "source SHA-256": envelope.SourceSHA256, + "rendition request fingerprint": envelope.RenditionRequestFingerprint, + "evidence lexical fingerprint": envelope.EvidenceLexicalFingerprint, + } { + if len(value) != 64 { + return RenditionV1{}, RenditionFrontMatterV1{}, fmt.Errorf("rendition frontmatter %s is invalid", name) + } + } + if envelope.SourceFormat == "" || envelope.SourceMediaType == "" || + envelope.NormalizedEvidenceContract != NormalizedEvidenceContractV1 || envelope.UnitKind == "" { + return RenditionV1{}, RenditionFrontMatterV1{}, errors.New("rendition frontmatter identity is incomplete") + } + navigation, err := renditionNavigation(rendition) + if err != nil { + return RenditionV1{}, RenditionFrontMatterV1{}, err + } + frontmatter := RenditionFrontMatterV1{Contract: RenditionMarkdownContractV1, + Source: RenditionMarkdownSourceV1{SHA256: envelope.SourceSHA256, + Format: envelope.SourceFormat, MediaType: envelope.SourceMediaType}, + Rendition: RenditionMarkdownBuildV1{BuildID: envelope.BuildID, + RenditionRequestFingerprint: envelope.RenditionRequestFingerprint, + EvidenceLexicalFingerprint: envelope.EvidenceLexicalFingerprint, + NormalizedEvidenceContract: envelope.NormalizedEvidenceContract, + BodySHA256: rendition.MarkdownChecksum, Completeness: rendition.Completeness, + Truncated: slices.ContainsFunc(rendition.Warnings, func(w RenditionWarningV1) bool { return w.Code == "truncated" })}, + Document: RenditionMarkdownDocumentV1{Title: envelope.Title, Language: envelope.Language, + UnitKind: envelope.UnitKind, UnitCount: len(rendition.Units)}, Navigation: navigation} + header, err := marshalRenditionFrontMatterV1(frontmatter) + if err != nil { + return RenditionV1{}, RenditionFrontMatterV1{}, err + } + if len(header) > maxRenditionFrontMatterBytes { + return RenditionV1{}, RenditionFrontMatterV1{}, errors.New("rendition frontmatter exceeds its byte bound") + } + body := slices.Clone(rendition.Markdown) + rendition.Markdown = make([]byte, 0, len(header)+len(body)) + rendition.Markdown = append(rendition.Markdown, header...) + rendition.Markdown = append(rendition.Markdown, body...) + rendition.MarkdownChecksum = checksumBytes(rendition.Markdown) + rendition.Checksum = renditionChecksum(rendition) + return rendition, frontmatter, nil +} + +// ParseRenditionFrontMatterV1 validates the exact deterministic envelope and +// returns its body as a view into data. Navigation offsets are body-relative. +func ParseRenditionFrontMatterV1(data []byte) (RenditionFrontMatterV1, []byte, error) { + const opening = "---\n" + closing := []byte("\n---\n") + if !bytes.HasPrefix(data, []byte(opening)) { + return RenditionFrontMatterV1{}, nil, errors.New("rendition frontmatter opening delimiter is missing") + } + relativeEnd := bytes.Index(data[len(opening):], closing) + if relativeEnd < 0 { + return RenditionFrontMatterV1{}, nil, errors.New("rendition frontmatter closing delimiter is missing") + } + headerLength := len(opening) + relativeEnd + len(closing) + if headerLength > maxRenditionFrontMatterBytes { + return RenditionFrontMatterV1{}, nil, errors.New("rendition frontmatter exceeds its byte bound") + } + yamlEnd := len(opening) + relativeEnd + 1 + var envelope struct { + Docbank RenditionFrontMatterV1 `yaml:"docbank"` + } + decoder := yaml.NewDecoder(bytes.NewReader(data[len(opening):yamlEnd])) + decoder.KnownFields(true) + if err := decoder.Decode(&envelope); err != nil { + return RenditionFrontMatterV1{}, nil, fmt.Errorf("decoding rendition frontmatter: %w", err) + } + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + return RenditionFrontMatterV1{}, nil, errors.New("rendition frontmatter contains multiple YAML documents") + } + frontmatter, body := envelope.Docbank, data[headerLength:] + if len(body) == 0 || !utf8.Valid(body) { + return RenditionFrontMatterV1{}, nil, errors.New("rendition Markdown body is empty or not UTF-8") + } + if err := validateRenditionFrontMatterV1(frontmatter, body); err != nil { + return RenditionFrontMatterV1{}, nil, err + } + canonical, err := marshalRenditionFrontMatterV1(frontmatter) + if err != nil { + return RenditionFrontMatterV1{}, nil, err + } + if !bytes.Equal(canonical, data[:headerLength]) { + return RenditionFrontMatterV1{}, nil, errors.New("rendition frontmatter is not canonical") + } + return frontmatter, body, nil +} + +func validateRenditionFrontMatterV1(value RenditionFrontMatterV1, body []byte) error { + validDigest := func(value string) bool { + if len(value) != sha256.Size*2 || value != strings.ToLower(value) { + return false + } + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == sha256.Size + } + for name, digest := range map[string]string{ + "source SHA-256": value.Source.SHA256, "build ID": value.Rendition.BuildID, + "rendition request fingerprint": value.Rendition.RenditionRequestFingerprint, + "evidence lexical fingerprint": value.Rendition.EvidenceLexicalFingerprint, + "body SHA-256": value.Rendition.BodySHA256, + } { + if !validDigest(digest) { + return fmt.Errorf("rendition frontmatter %s is invalid", name) + } + } + if value.Contract != RenditionMarkdownContractV1 || value.Source.Format == "" || + value.Source.MediaType == "" || value.Rendition.NormalizedEvidenceContract != NormalizedEvidenceContractV1 || + !validEvidenceCompleteness(value.Rendition.Completeness) || !validEvidenceUnitKind(value.Document.UnitKind) || + value.Document.UnitCount < 1 || value.Navigation.OffsetBase != RenditionNavigationOffsetBody || + len(value.Navigation.Entries) > maxRenditionNavigationEntries || + len(value.Navigation.Entries) > value.Document.UnitCount { + return errors.New("rendition frontmatter contract is invalid") + } + if got := checksumBytes(body); got != value.Rendition.BodySHA256 { + return fmt.Errorf("rendition body SHA-256 %s differs from frontmatter %s", got, value.Rendition.BodySHA256) + } + seen := make(map[string]struct{}, len(value.Navigation.Entries)) + priorByte := -1 + for _, entry := range value.Navigation.Entries { + if entry.Key == "" || !renditionFrontMatterLocatorKind(entry.Kind) || entry.Byte < 0 || + entry.Byte >= len(body) || entry.Byte < priorByte || !utf8.RuneStart(body[entry.Byte]) || + entry.Line != 1+bytes.Count(body[:entry.Byte], []byte{'\n'}) { + return errors.New("rendition frontmatter navigation is invalid") + } + if _, exists := seen[entry.Key]; exists { + return errors.New("rendition frontmatter navigation contains a duplicate key") + } + seen[entry.Key] = struct{}{} + priorByte = entry.Byte + } + return nil +} + +func renditionFrontMatterLocatorKind(value EvidenceLocatorKind) bool { + switch value { + case EvidenceLocatorGeneric, EvidenceLocatorLine, EvidenceLocatorMessage, + EvidenceLocatorPage, EvidenceLocatorRecord, EvidenceLocatorSection, + EvidenceLocatorSheet, EvidenceLocatorSlide, EvidenceLocatorSpine: + return true + default: + return false + } +} + +func renditionNavigation(rendition RenditionV1) (RenditionMarkdownNavigationV1, error) { + navigation := RenditionMarkdownNavigationV1{OffsetBase: RenditionNavigationOffsetBody, + Complete: len(rendition.Units) <= maxRenditionNavigationEntries} + maximum := min(len(rendition.Units), maxRenditionNavigationEntries) + navigation.Entries = make([]RenditionNavigationEntryV1, 0, maximum) + cursor := 0 + for _, unit := range rendition.Units[:maximum] { + if unit.Text == "" { + continue + } + relative := bytes.Index(rendition.Markdown[cursor:], []byte(unit.Text)) + if relative < 0 { + return RenditionMarkdownNavigationV1{}, errors.New("rendition unit is absent from its Markdown body") + } + offset := cursor + relative + title := unit.Locator.Name + if len(unit.HeadingPath) != 0 { + title = unit.HeadingPath[len(unit.HeadingPath)-1] + } + navigation.Entries = append(navigation.Entries, RenditionNavigationEntryV1{ + Key: unit.EvidenceUnitID, Kind: unit.Locator.Kind, Title: title, + Line: 1 + bytes.Count(rendition.Markdown[:offset], []byte{'\n'}), Byte: offset}) + cursor = offset + len(unit.Text) + } + return navigation, nil +} + +func marshalRenditionFrontMatterV1(value RenditionFrontMatterV1) ([]byte, error) { + if value.Contract != RenditionMarkdownContractV1 || value.Navigation.OffsetBase != RenditionNavigationOffsetBody || + value.Document.UnitCount < 1 || len(value.Navigation.Entries) > maxRenditionNavigationEntries { + return nil, errors.New("rendition frontmatter is invalid") + } + var builder strings.Builder + builder.WriteString("---\ndocbank:\n") + writeYAMLString(&builder, 2, "contract", value.Contract) + builder.WriteString(" source:\n") + writeYAMLString(&builder, 4, "sha256", value.Source.SHA256) + writeYAMLString(&builder, 4, "format", value.Source.Format) + writeYAMLString(&builder, 4, "media_type", value.Source.MediaType) + builder.WriteString(" rendition:\n") + writeYAMLString(&builder, 4, "build_id", value.Rendition.BuildID) + writeYAMLString(&builder, 4, "rendition_request_fingerprint", value.Rendition.RenditionRequestFingerprint) + writeYAMLString(&builder, 4, "evidence_lexical_fingerprint", value.Rendition.EvidenceLexicalFingerprint) + writeYAMLString(&builder, 4, "normalized_evidence_contract", value.Rendition.NormalizedEvidenceContract) + writeYAMLString(&builder, 4, "body_sha256", value.Rendition.BodySHA256) + writeYAMLString(&builder, 4, "completeness", string(value.Rendition.Completeness)) + writeYAMLBool(&builder, 4, "truncated", value.Rendition.Truncated) + builder.WriteString(" document:\n") + if value.Document.Title != "" { + writeYAMLString(&builder, 4, "title", value.Document.Title) + } + if value.Document.Language != "" { + writeYAMLString(&builder, 4, "language", value.Document.Language) + } + writeYAMLString(&builder, 4, "unit_kind", string(value.Document.UnitKind)) + writeYAMLInt(&builder, 4, "unit_count", value.Document.UnitCount) + builder.WriteString(" navigation:\n") + writeYAMLString(&builder, 4, "offset_base", value.Navigation.OffsetBase) + writeYAMLBool(&builder, 4, "complete", value.Navigation.Complete) + builder.WriteString(" entries:\n") + for _, entry := range value.Navigation.Entries { + builder.WriteString(" - key: ") + builder.WriteString(strconv.Quote(entry.Key)) + builder.WriteByte('\n') + writeYAMLString(&builder, 8, "kind", string(entry.Kind)) + if entry.Title != "" { + writeYAMLString(&builder, 8, "title", entry.Title) + } + writeYAMLInt(&builder, 8, "line", entry.Line) + writeYAMLInt(&builder, 8, "byte", entry.Byte) + } + builder.WriteString("---\n") + encoded := builder.String() + if !utf8.ValidString(encoded) { + return nil, errors.New("rendition frontmatter is not UTF-8") + } + return []byte(encoded), nil +} + +func writeYAMLString(builder *strings.Builder, indent int, key, value string) { + builder.WriteString(strings.Repeat(" ", indent)) + builder.WriteString(key) + builder.WriteString(": ") + builder.WriteString(strconv.Quote(value)) + builder.WriteByte('\n') +} + +func writeYAMLBool(builder *strings.Builder, indent int, key string, value bool) { + builder.WriteString(strings.Repeat(" ", indent)) + builder.WriteString(key) + builder.WriteString(": ") + builder.WriteString(strconv.FormatBool(value)) + builder.WriteByte('\n') +} + +func writeYAMLInt(builder *strings.Builder, indent int, key string, value int) { + builder.WriteString(strings.Repeat(" ", indent)) + builder.WriteString(key) + builder.WriteString(": ") + builder.WriteString(strconv.Itoa(value)) + builder.WriteByte('\n') +} diff --git a/document/rendition_frontmatter_test.go b/document/rendition_frontmatter_test.go new file mode 100644 index 00000000..2b8f7af3 --- /dev/null +++ b/document/rendition_frontmatter_test.go @@ -0,0 +1,85 @@ +package document + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestEnvelopeRenditionV1KeepsNavigationBodyRelative(t *testing.T) { + body := []byte("# First\n\nAlpha\n\n---\n\n# Second\n\nBeta\n") + rendition := RenditionV1{ContractVersion: RenditionContractV1, + Completeness: EvidenceComplete, EvidenceChecksum: frontmatterHash("evidence"), + Markdown: body, MarkdownChecksum: checksumBytes(body), + Units: []NormalizedUnitV1{ + {EvidenceUnitID: "page:000000", Order: 0, Text: "# First\n\nAlpha", + Locator: EvidenceLocatorV1{Kind: EvidenceLocatorPage, IndexOrigin: EvidenceIndexOriginZero}}, + {EvidenceUnitID: "page:000001", Order: 1, Text: "# Second\n\nBeta", + Locator: EvidenceLocatorV1{Kind: EvidenceLocatorPage, IndexOrigin: EvidenceIndexOriginZero}}, + }} + rendition.Checksum = renditionChecksum(rendition) + + got, frontmatter, err := EnvelopeRenditionV1(rendition, RenditionEnvelopeV1{ + BuildID: frontmatterHash("build"), SourceSHA256: frontmatterHash("source"), + SourceFormat: "pdf", SourceMediaType: "application/pdf", + RenditionRequestFingerprint: frontmatterHash("request"), + EvidenceLexicalFingerprint: frontmatterHash("lexical"), + NormalizedEvidenceContract: NormalizedEvidenceContractV1, UnitKind: EvidenceUnitPage, + }) + require.NoError(t, err) + require.True(t, bytes.HasPrefix(got.Markdown, + []byte("---\ndocbank:\n contract: \"docbank-sanitized-markdown/v1\"\n"))) + parts := bytes.SplitN(got.Markdown, []byte("---\n"), 3) + require.Len(t, parts, 3) + require.Equal(t, body, parts[2]) + require.Equal(t, checksumBytes(body), frontmatter.Rendition.BodySHA256) + for _, entry := range frontmatter.Navigation.Entries { + require.Less(t, entry.Byte, len(body)) + require.Equal(t, entry.Line, 1+bytes.Count(body[:entry.Byte], []byte{'\n'})) + marker := strings.TrimPrefix(entry.Key, "page:") + if marker == "000000" { + require.True(t, bytes.HasPrefix(body[entry.Byte:], []byte("# First"))) + } else { + require.True(t, bytes.HasPrefix(body[entry.Byte:], []byte("# Second"))) + } + } + + parsed, parsedBody, err := ParseRenditionFrontMatterV1(got.Markdown) + require.NoError(t, err) + require.Equal(t, frontmatter, parsed) + require.Equal(t, body, parsedBody) +} + +func TestParseRenditionFrontMatterV1RejectsCorruptBodyAndNavigation(t *testing.T) { + body := []byte("# First\n\nAlpha\n") + rendition := RenditionV1{ContractVersion: RenditionContractV1, + Completeness: EvidenceComplete, EvidenceChecksum: frontmatterHash("evidence"), + Markdown: body, MarkdownChecksum: checksumBytes(body), + Units: []NormalizedUnitV1{{EvidenceUnitID: "page:000000", Order: 0, Text: string(body), + Locator: EvidenceLocatorV1{Kind: EvidenceLocatorPage, IndexOrigin: EvidenceIndexOriginZero}}}} + rendered, _, err := EnvelopeRenditionV1(rendition, RenditionEnvelopeV1{ + BuildID: frontmatterHash("build"), SourceSHA256: frontmatterHash("source"), + SourceFormat: "pdf", SourceMediaType: "application/pdf", + RenditionRequestFingerprint: frontmatterHash("request"), + EvidenceLexicalFingerprint: frontmatterHash("lexical"), + NormalizedEvidenceContract: NormalizedEvidenceContractV1, UnitKind: EvidenceUnitPage}) + require.NoError(t, err) + + corrupt := append([]byte(nil), rendered.Markdown...) + corrupt[len(corrupt)-2] ^= 1 + _, _, err = ParseRenditionFrontMatterV1(corrupt) + require.ErrorContains(t, err, "body SHA-256") + + badNavigation := bytes.Replace(rendered.Markdown, []byte("byte: 0"), []byte("byte: 999999"), 1) + _, _, err = ParseRenditionFrontMatterV1(badNavigation) + require.Error(t, err) +} + +func frontmatterHash(value string) string { + digest := sha256.Sum256([]byte(value)) + return hex.EncodeToString(digest[:]) +} diff --git a/document/rendition_normalize_test.go b/document/rendition_normalize_test.go index 893d2d52..e020285e 100644 --- a/document/rendition_normalize_test.go +++ b/document/rendition_normalize_test.go @@ -325,7 +325,7 @@ func TestBuildRenditionV1PreservesParserGeneratedTaskCheckboxes(t *testing.T) { rendered, err := BuildRenditionV1(evidence, policy) require.NoError(t, err) - assert.Equal(t, "- \\[x\\] shipped\n- \\[ \\] pending\n", string(rendered.Markdown)) + assert.Equal(t, "- [x] shipped\n- [ ] pending\n", string(rendered.Markdown)) } func TestBuildRenditionV1DropsProviderSuppliedCheckboxControls(t *testing.T) { diff --git a/document/source_metadata.go b/document/source_metadata.go new file mode 100644 index 00000000..006424ab --- /dev/null +++ b/document/source_metadata.go @@ -0,0 +1,93 @@ +package document + +const ( + // SourceMetadataContractV1 identifies canonical embedded source claims. + SourceMetadataContractV1 = "source-metadata/v1" + + MaxSourceMetadataFields = 512 + MaxSourceMetadataWarnings = 256 + MaxSourceMetadataListValues = 256 + MaxSourceMetadataValueBytes = 64 << 10 + MaxSourceMetadataEncodedBytes = 8 << 20 +) + +// SourceMetadataValueKind identifies the one typed payload carried by a field. +type SourceMetadataValueKind string + +const ( + SourceMetadataString SourceMetadataValueKind = "string" + SourceMetadataStringList SourceMetadataValueKind = "string_list" + SourceMetadataInteger SourceMetadataValueKind = "integer" + SourceMetadataNumber SourceMetadataValueKind = "number" + SourceMetadataBoolean SourceMetadataValueKind = "boolean" + SourceMetadataTimestamp SourceMetadataValueKind = "timestamp" +) + +// SourceMetadataTimestampPrecision preserves what the source actually stated. +type SourceMetadataTimestampPrecision string + +const ( + SourceMetadataPrecisionDate SourceMetadataTimestampPrecision = "date" + SourceMetadataPrecisionHour SourceMetadataTimestampPrecision = "hour" + SourceMetadataPrecisionMinute SourceMetadataTimestampPrecision = "minute" + SourceMetadataPrecisionSecond SourceMetadataTimestampPrecision = "second" + SourceMetadataPrecisionFraction SourceMetadataTimestampPrecision = "fraction" +) + +// SourceMetadataTimezoneKind distinguishes an explicit zone from an omitted one. +type SourceMetadataTimezoneKind string + +const ( + SourceMetadataTimezoneOmitted SourceMetadataTimezoneKind = "omitted" + SourceMetadataTimezoneUTC SourceMetadataTimezoneKind = "utc" + SourceMetadataTimezoneOffset SourceMetadataTimezoneKind = "offset" +) + +// SourceMetadataTimestampV1 retains the raw value and a normalization only +// when the source syntax is unambiguous. Offset is empty for UTC or omission. +type SourceMetadataTimestampV1 struct { + Normalized string `json:"normalized"` + Offset string `json:"offset,omitempty"` + Precision SourceMetadataTimestampPrecision `json:"precision"` + Raw string `json:"raw"` + Timezone SourceMetadataTimezoneKind `json:"timezone"` +} + +// SourceMetadataValueV1 is a closed typed union. Exactly one payload matching +// Kind is populated. +type SourceMetadataValueV1 struct { + Boolean *bool `json:"boolean,omitempty"` + Integer *int64 `json:"integer,omitempty"` + Kind SourceMetadataValueKind `json:"kind"` + Number *float64 `json:"number,omitempty"` + String string `json:"string,omitempty"` + Strings []string `json:"strings,omitempty"` + Timestamp *SourceMetadataTimestampV1 `json:"timestamp,omitempty"` +} + +// SourceMetadataFieldV1 is one canonical claim with its exact source label. +// Sensitive fields remain local unless a future explicit disclosure policy +// selects them. +type SourceMetadataFieldV1 struct { + Key string `json:"key"` + Namespace string `json:"namespace"` + Sensitive bool `json:"sensitive"` + SourceField string `json:"source_field"` + Value SourceMetadataValueV1 `json:"value"` +} + +// SourceMetadataWarningV1 preserves extraction uncertainty without guessing. +type SourceMetadataWarningV1 struct { + Code string `json:"code"` + Detail string `json:"detail"` + Namespace string `json:"namespace"` + SourceField string `json:"source_field"` +} + +// SourceMetadataV1 is content-scoped embedded evidence. It intentionally has +// no filename, path, ingest, filesystem, or user-authored metadata fields. +type SourceMetadataV1 struct { + ContractVersion string `json:"contract_version"` + Fields []SourceMetadataFieldV1 `json:"fields"` + Warnings []SourceMetadataWarningV1 `json:"warnings"` +} diff --git a/document/source_metadata_codec.go b/document/source_metadata_codec.go new file mode 100644 index 00000000..d46f2a6e --- /dev/null +++ b/document/source_metadata_codec.go @@ -0,0 +1,356 @@ +package document + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "math" + "regexp" + "slices" + "strings" + "time" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +var sourceMetadataKeyPattern = regexp.MustCompile(`^[a-z][a-z0-9_]*(\.[a-z][a-z0-9_]*)*$`) +var sourceMetadataOffsetPattern = regexp.MustCompile(`^[+-](0[0-9]|1[0-4]):[0-5][0-9]$`) + +var sourceMetadataCommonKeys = map[string]bool{ + "attachment_count": true, "calendar.end": true, "calendar.start": true, + "created": true, "creators": true, "description": true, + "email.bcc": true, "email.cc": true, "email.from": true, + "email.received": true, "email.sent": true, "email.subject": true, + "email.to": true, "keywords": true, "language": true, "modified": true, + "page_count": true, "subject": true, "title": true, +} + +var sourceMetadataNamespaces = []string{ + "calendar.", "email.", "image.exif.", "image.iptc.", "image.xmp.", + "media.container.", "media.id3.", "office.core.", "office.custom.", + "pdf.info.", "xmp.", +} + +var sourceMetadataSourceNamespaces = map[string]bool{ + "calendar": true, "email": true, "image.exif": true, "image.iptc": true, + "image.xmp": true, "media.container": true, "media.id3": true, + "office.core": true, "office.custom": true, "pdf.info": true, "xmp": true, +} + +var forbiddenSourceMetadataSegments = map[string]bool{ + "bates": true, "bates_start": true, "bates_end": true, "collection_processing_timezone": true, + "custodian": true, "duplicate_path": true, "extension": true, "family_date": true, + "filename": true, "filesystem_mtime": true, "ingest_time": true, + "produced_document_link": true, "produced_text_link": true, "production_volume": true, + "redaction_state": true, "source_path": true, +} + +// SourceMetadataCanonicalKeyAllowed reports whether a key belongs to the +// content-scoped contract rather than attachment/provenance or legal workflow. +func SourceMetadataCanonicalKeyAllowed(key string) bool { + if !sourceMetadataKeyPattern.MatchString(key) { + return false + } + for segment := range strings.SplitSeq(key, ".") { + if forbiddenSourceMetadataSegments[segment] { + return false + } + } + if sourceMetadataCommonKeys[key] { + return true + } + for _, prefix := range sourceMetadataNamespaces { + if strings.HasPrefix(key, prefix) { + return true + } + } + return false +} + +// MarshalSourceMetadataV1 validates, canonicalizes, and hashes one record. +func MarshalSourceMetadataV1(value SourceMetadataV1) ([]byte, string, error) { + canonical, err := canonicalSourceMetadataV1(value) + if err != nil { + return nil, "", err + } + encoded, err := json.Marshal(canonical, json.Deterministic(true)) + if err != nil { + return nil, "", fmt.Errorf("encoding source metadata: %w", err) + } + if len(encoded) > MaxSourceMetadataEncodedBytes { + return nil, "", errors.New("source metadata record is too large") + } + digest := sha256.Sum256(encoded) + return encoded, hex.EncodeToString(digest[:]), nil +} + +// DecodeSourceMetadataV1 accepts the exact v1 typed schema and returns its +// canonical value and checksum. Unknown members and major contracts fail. +func DecodeSourceMetadataV1(encoded []byte) (SourceMetadataV1, string, error) { + if len(encoded) > MaxSourceMetadataEncodedBytes { + return SourceMetadataV1{}, "", errors.New("source metadata record is too large") + } + var value SourceMetadataV1 + if err := json.Unmarshal(encoded, &value, json.RejectUnknownMembers(true)); err != nil { + return SourceMetadataV1{}, "", fmt.Errorf("decoding source metadata: %w", err) + } + canonical, checksum, err := MarshalSourceMetadataV1(value) + if err != nil { + return SourceMetadataV1{}, "", err + } + if !slices.Equal(encoded, canonical) { + return SourceMetadataV1{}, "", errors.New("source metadata bytes are not canonical") + } + return value, checksum, nil +} + +func canonicalSourceMetadataV1(value SourceMetadataV1) (SourceMetadataV1, error) { + if value.ContractVersion != SourceMetadataContractV1 { + return SourceMetadataV1{}, fmt.Errorf( + "source metadata contract version must be %q", SourceMetadataContractV1) + } + if len(value.Fields) > MaxSourceMetadataFields { + return SourceMetadataV1{}, errors.New("source metadata has too many fields") + } + if len(value.Warnings) > MaxSourceMetadataWarnings { + return SourceMetadataV1{}, errors.New("source metadata has too many warnings") + } + value.Fields = append([]SourceMetadataFieldV1(nil), value.Fields...) + seen := make(map[string]bool, len(value.Fields)) + for index := range value.Fields { + field := &value.Fields[index] + field.Key = norm.NFC.String(strings.TrimSpace(field.Key)) + field.Namespace = norm.NFC.String(strings.TrimSpace(field.Namespace)) + field.SourceField = norm.NFC.String(strings.TrimSpace(field.SourceField)) + if !SourceMetadataCanonicalKeyAllowed(field.Key) { + return SourceMetadataV1{}, fmt.Errorf("source metadata field %d has forbidden canonical key %q", index, field.Key) + } + if seen[field.Key] { + return SourceMetadataV1{}, fmt.Errorf("source metadata has duplicate canonical key %q", field.Key) + } + seen[field.Key] = true + if err := validateSourceMetadataLabel(field.Namespace, "namespace"); err != nil { + return SourceMetadataV1{}, fmt.Errorf("source metadata field %q: %w", field.Key, err) + } + if !sourceMetadataSourceNamespaces[field.Namespace] { + return SourceMetadataV1{}, fmt.Errorf("source metadata field %q has unknown namespace %q", field.Key, field.Namespace) + } + if err := validateSourceMetadataLabel(field.SourceField, "source field"); err != nil { + return SourceMetadataV1{}, fmt.Errorf("source metadata field %q: %w", field.Key, err) + } + if err := canonicalSourceMetadataValue(&field.Value); err != nil { + return SourceMetadataV1{}, fmt.Errorf("source metadata field %q: %w", field.Key, err) + } + } + slices.SortFunc(value.Fields, func(a, b SourceMetadataFieldV1) int { + return strings.Compare(a.Key, b.Key) + }) + value.Warnings = append([]SourceMetadataWarningV1(nil), value.Warnings...) + for index := range value.Warnings { + warning := &value.Warnings[index] + warning.Code = norm.NFC.String(strings.TrimSpace(warning.Code)) + warning.Namespace = norm.NFC.String(strings.TrimSpace(warning.Namespace)) + warning.SourceField = norm.NFC.String(strings.TrimSpace(warning.SourceField)) + warning.Detail = norm.NFC.String(warning.Detail) + for label, name := range map[string]string{ + warning.Code: "warning code", warning.Namespace: "warning namespace", + warning.SourceField: "warning source field", + } { + if err := validateSourceMetadataLabel(label, name); err != nil { + return SourceMetadataV1{}, fmt.Errorf("source metadata warning %d: %w", index, err) + } + } + if err := validateSourceMetadataString(warning.Detail); err != nil { + return SourceMetadataV1{}, fmt.Errorf("source metadata warning %d detail: %w", index, err) + } + } + slices.SortFunc(value.Warnings, func(a, b SourceMetadataWarningV1) int { + return strings.Compare(a.Namespace+"\x00"+a.SourceField+"\x00"+a.Code+"\x00"+a.Detail, + b.Namespace+"\x00"+b.SourceField+"\x00"+b.Code+"\x00"+b.Detail) + }) + if value.Fields == nil { + value.Fields = []SourceMetadataFieldV1{} + } + if value.Warnings == nil { + value.Warnings = []SourceMetadataWarningV1{} + } + return value, nil +} + +func validateSourceMetadataLabel(value, name string) error { + if value == "" || len(value) > 256 || !utf8.ValidString(value) { + return fmt.Errorf("%s must be bounded UTF-8", name) + } + return nil +} + +func validateSourceMetadataString(value string) error { + if !utf8.ValidString(value) { + return errors.New("value must be UTF-8") + } + if len(value) > MaxSourceMetadataValueBytes { + return errors.New("value is too large") + } + return nil +} + +func canonicalSourceMetadataValue(value *SourceMetadataValueV1) error { + payloads := 0 + if value.String != "" { + payloads++ + } + if value.Strings != nil { + payloads++ + } + if value.Integer != nil { + payloads++ + } + if value.Number != nil { + payloads++ + } + if value.Boolean != nil { + payloads++ + } + if value.Timestamp != nil { + payloads++ + } + switch value.Kind { + case SourceMetadataString: + if payloads != 1 || value.Strings != nil || value.Integer != nil || value.Number != nil || + value.Boolean != nil || value.Timestamp != nil { + return errors.New("string value has conflicting payloads") + } + value.String = norm.NFC.String(value.String) + return validateSourceMetadataString(value.String) + case SourceMetadataStringList: + if payloads != 1 || value.Strings == nil { + return errors.New("string-list value has conflicting payloads") + } + if len(value.Strings) > MaxSourceMetadataListValues { + return errors.New("string-list value has too many entries") + } + value.Strings = append([]string(nil), value.Strings...) + for index := range value.Strings { + value.Strings[index] = norm.NFC.String(value.Strings[index]) + if err := validateSourceMetadataString(value.Strings[index]); err != nil { + return fmt.Errorf("string-list entry %d: %w", index, err) + } + } + return nil + case SourceMetadataInteger: + if payloads != 1 || value.Integer == nil { + return errors.New("integer value has conflicting payloads") + } + return nil + case SourceMetadataNumber: + if payloads != 1 || value.Number == nil { + return errors.New("number value has conflicting payloads") + } + if math.IsNaN(*value.Number) || math.IsInf(*value.Number, 0) { + return errors.New("number value must be finite") + } + return nil + case SourceMetadataBoolean: + if payloads != 1 || value.Boolean == nil { + return errors.New("boolean value has conflicting payloads") + } + return nil + case SourceMetadataTimestamp: + if payloads != 1 || value.Timestamp == nil { + return errors.New("timestamp value has conflicting payloads") + } + return canonicalSourceMetadataTimestamp(value.Timestamp) + default: + return fmt.Errorf("unknown source metadata value kind %q", value.Kind) + } +} + +func canonicalSourceMetadataTimestamp(value *SourceMetadataTimestampV1) error { + value.Raw = norm.NFC.String(value.Raw) + value.Normalized = norm.NFC.String(value.Normalized) + if err := validateSourceMetadataString(value.Raw); err != nil || value.Raw == "" { + if err == nil { + err = errors.New("value is empty") + } + return fmt.Errorf("raw timestamp: %w", err) + } + if err := validateSourceMetadataString(value.Normalized); err != nil || value.Normalized == "" { + if err == nil { + err = errors.New("value is empty") + } + return fmt.Errorf("normalized timestamp: %w", err) + } + switch value.Precision { + case SourceMetadataPrecisionDate, SourceMetadataPrecisionHour, + SourceMetadataPrecisionMinute, SourceMetadataPrecisionSecond, + SourceMetadataPrecisionFraction: + default: + return errors.New("timestamp precision is invalid") + } + switch value.Timezone { + case SourceMetadataTimezoneOmitted: + if value.Offset != "" || strings.HasSuffix(value.Normalized, "Z") || + hasRFC3339Offset(value.Normalized) { + return errors.New("timestamp timezone omission conflicts with normalized value") + } + return validateLocalTimestampPrecision(value.Normalized, value.Precision) + case SourceMetadataTimezoneUTC: + if value.Offset != "" || !strings.HasSuffix(value.Normalized, "Z") { + return errors.New("UTC timestamp timezone is inconsistent") + } + case SourceMetadataTimezoneOffset: + if !sourceMetadataOffsetPattern.MatchString(value.Offset) || + !strings.HasSuffix(value.Normalized, value.Offset) { + return errors.New("timestamp timezone offset is inconsistent") + } + default: + return errors.New("timestamp timezone is invalid") + } + if value.Precision == SourceMetadataPrecisionDate || value.Precision == SourceMetadataPrecisionHour { + return errors.New("timestamp timezone requires minute or finer precision") + } + parsed, err := time.Parse(time.RFC3339Nano, value.Normalized) + if err != nil { + return fmt.Errorf("timestamp timezone value is not RFC3339: %w", err) + } + _ = parsed + return validateZonedTimestampPrecision(value.Normalized, value.Precision) +} + +func hasRFC3339Offset(value string) bool { + return len(value) >= 6 && (value[len(value)-6] == '+' || value[len(value)-6] == '-') && + value[len(value)-3] == ':' +} + +func validateLocalTimestampPrecision(value string, precision SourceMetadataTimestampPrecision) error { + layout := map[SourceMetadataTimestampPrecision]string{ + SourceMetadataPrecisionDate: "2006-01-02", SourceMetadataPrecisionHour: "2006-01-02T15", + SourceMetadataPrecisionMinute: "2006-01-02T15:04", SourceMetadataPrecisionSecond: "2006-01-02T15:04:05", + SourceMetadataPrecisionFraction: "2006-01-02T15:04:05.999999999", + }[precision] + if _, err := time.Parse(layout, value); err != nil { + return fmt.Errorf("timestamp without timezone does not match precision: %w", err) + } + return nil +} + +func validateZonedTimestampPrecision(value string, precision SourceMetadataTimestampPrecision) error { + base := strings.TrimSuffix(value, "Z") + if hasRFC3339Offset(value) { + base = value[:len(value)-6] + } + hasFraction := strings.Contains(base, ".") + if (precision == SourceMetadataPrecisionFraction) != hasFraction { + return errors.New("timestamp fraction does not match precision") + } + if precision == SourceMetadataPrecisionMinute && len(base) != len("2006-01-02T15:04") { + return errors.New("timestamp seconds do not match minute precision") + } + if precision == SourceMetadataPrecisionSecond && len(base) != len("2006-01-02T15:04:05") { + return errors.New("timestamp does not match second precision") + } + return nil +} diff --git a/document/source_metadata_test.go b/document/source_metadata_test.go new file mode 100644 index 00000000..113820bf --- /dev/null +++ b/document/source_metadata_test.go @@ -0,0 +1,125 @@ +package document + +import ( + "bytes" + "math" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSourceMetadataV1CanonicalGolden(t *testing.T) { + pageCount := int64(7) + record := SourceMetadataV1{ + ContractVersion: SourceMetadataContractV1, + Fields: []SourceMetadataFieldV1{ + {Key: "pdf.info.page_count", Namespace: "pdf.info", SourceField: "PageCount", + Value: SourceMetadataValueV1{Kind: SourceMetadataInteger, Integer: &pageCount}}, + {Key: "creators", Namespace: "pdf.info", SourceField: "Author", + Value: SourceMetadataValueV1{Kind: SourceMetadataStringList, Strings: []string{"Ada Lovelace", "Grace Hopper"}}}, + {Key: "created", Namespace: "xmp", SourceField: "CreateDate", + Value: SourceMetadataValueV1{Kind: SourceMetadataTimestamp, Timestamp: &SourceMetadataTimestampV1{ + Raw: "2024-01-02T03:04:05-07:00", Normalized: "2024-01-02T03:04:05-07:00", + Offset: "-07:00", Precision: SourceMetadataPrecisionSecond, Timezone: SourceMetadataTimezoneOffset, + }}}, + {Key: "title", Namespace: "pdf.info", SourceField: "Title", + Value: SourceMetadataValueV1{Kind: SourceMetadataString, String: "Synthetic report"}}, + }, + Warnings: []SourceMetadataWarningV1{{ + Code: "unparseable_timestamp", Namespace: "pdf.info", SourceField: "ModDate", + Detail: "value retained without coercion", + }}, + } + + canonical, checksum, err := MarshalSourceMetadataV1(record) + require.NoError(t, err) + assert.Len(t, checksum, 64) + want, err := os.ReadFile(filepath.Join("testdata", "source-metadata-v1.golden.json")) + require.NoError(t, err) + want = bytes.TrimSuffix(want, []byte("\n")) + assert.Equal(t, want, canonical) + + decoded, decodedChecksum, err := DecodeSourceMetadataV1(canonical) + require.NoError(t, err) + assert.Equal(t, checksum, decodedChecksum) + assert.Equal(t, []string{"created", "creators", "pdf.info.page_count", "title"}, + []string{decoded.Fields[0].Key, decoded.Fields[1].Key, decoded.Fields[2].Key, decoded.Fields[3].Key}) +} + +func TestSourceMetadataV1RejectsAmbiguityAndUnboundedValues(t *testing.T) { + base := SourceMetadataV1{ContractVersion: SourceMetadataContractV1, + Fields: []SourceMetadataFieldV1{{Key: "title", Namespace: "pdf.info", SourceField: "Title", + Value: SourceMetadataValueV1{Kind: SourceMetadataString, String: "safe"}}}} + for _, testCase := range []struct { + name string + mutate func(*SourceMetadataV1) + want string + }{ + {name: "duplicate canonical key", want: "duplicate canonical key", mutate: func(v *SourceMetadataV1) { + v.Fields = append(v.Fields, v.Fields[0]) + }}, + {name: "invalid UTF-8", want: "UTF-8", mutate: func(v *SourceMetadataV1) { + v.Fields[0].Value.String = string([]byte{0xff}) + }}, + {name: "oversized value", want: "too large", mutate: func(v *SourceMetadataV1) { + v.Fields[0].Value.String = string(make([]byte, MaxSourceMetadataValueBytes+1)) + }}, + {name: "unbounded array", want: "too many", mutate: func(v *SourceMetadataV1) { + v.Fields[0].Value = SourceMetadataValueV1{Kind: SourceMetadataStringList, + Strings: make([]string, MaxSourceMetadataListValues+1)} + }}, + {name: "non-finite number", want: "finite", mutate: func(v *SourceMetadataV1) { + number := math.Inf(1) + v.Fields[0].Value = SourceMetadataValueV1{Kind: SourceMetadataNumber, Number: &number} + }}, + {name: "ambiguous timestamp", want: "timezone", mutate: func(v *SourceMetadataV1) { + v.Fields[0].Value = SourceMetadataValueV1{Kind: SourceMetadataTimestamp, + Timestamp: &SourceMetadataTimestampV1{Raw: "01/02/03", Normalized: "2003-01-02", + Precision: SourceMetadataPrecisionDate, Timezone: SourceMetadataTimezoneOffset}} + }}, + {name: "unknown major", want: "contract version", mutate: func(v *SourceMetadataV1) { + v.ContractVersion = "source-metadata/v2" + }}, + {name: "unknown namespace", want: "unknown namespace", mutate: func(v *SourceMetadataV1) { + v.Fields[0].Namespace = "provider.guess" + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + value := base + value.Fields = append([]SourceMetadataFieldV1(nil), base.Fields...) + testCase.mutate(&value) + _, _, err := MarshalSourceMetadataV1(value) + require.ErrorContains(t, err, testCase.want) + }) + } +} + +func TestSourceMetadataV1KeepsAttachmentFactsOutOfContentRecord(t *testing.T) { + for _, key := range []string{ + "filename", "extension", "source_path", "ingest_time", "filesystem_mtime", + "bates_start", "bates_end", "production_volume", "custodian", "redaction_state", + "family_date", "collection_processing_timezone", "duplicate_path", "produced_document_link", + } { + assert.False(t, SourceMetadataCanonicalKeyAllowed(key), key) + } + assert.True(t, SourceMetadataCanonicalKeyAllowed("title")) + assert.True(t, SourceMetadataCanonicalKeyAllowed("image.exif.gps_latitude")) +} + +func TestSourceMetadataV1PreservesTimestampPrecisionAndTimezoneEvidence(t *testing.T) { + for _, stamp := range []SourceMetadataTimestampV1{ + {Raw: "2024-01-02", Normalized: "2024-01-02", Precision: SourceMetadataPrecisionDate, Timezone: SourceMetadataTimezoneOmitted}, + {Raw: "2024-01-02T03:04:05Z", Normalized: "2024-01-02T03:04:05Z", Precision: SourceMetadataPrecisionSecond, Timezone: SourceMetadataTimezoneUTC}, + {Raw: "2024-01-02T03:04:05.123-07:00", Normalized: "2024-01-02T03:04:05.123-07:00", Offset: "-07:00", Precision: SourceMetadataPrecisionFraction, Timezone: SourceMetadataTimezoneOffset}, + } { + record := SourceMetadataV1{ContractVersion: SourceMetadataContractV1, Fields: []SourceMetadataFieldV1{{Key: "created", Namespace: "xmp", SourceField: "CreateDate", Value: SourceMetadataValueV1{Kind: SourceMetadataTimestamp, Timestamp: &stamp}}}} + encoded, _, err := MarshalSourceMetadataV1(record) + require.NoError(t, err) + decoded, _, err := DecodeSourceMetadataV1(encoded) + require.NoError(t, err) + assert.Equal(t, stamp, *decoded.Fields[0].Value.Timestamp) + } +} diff --git a/document/testdata/chunks-bge-m3-v1.golden.json b/document/testdata/chunks-bge-m3-v1.golden.json new file mode 100644 index 00000000..2ed359ed --- /dev/null +++ b/document/testdata/chunks-bge-m3-v1.golden.json @@ -0,0 +1,66 @@ +{ + "version": 2, + "checksum": "e61984509c8f9506140a3b3e0bd4628a29d697f497f98f04af314f88b33df4d9", + "policy_fingerprint": "6b8a4f17ccce3fd4d68d9a4fcff5cf82dd9593ef25e3780e6b6b4b1df8addfa4", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "c6721247ba66fcc388b97a93ee88bce53a58f2d8e13ad13500adfde69f32e40d", + "content_token_budget": 6, + "overlap_tokens": 1, + "truncation_policy": "reject_indivisible", + "context_fingerprint": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "total_content_tokens": 12, + "total_rendered_tokens": 12, + "total_content_bytes": 23, + "total_rendered_bytes": 23, + "inputs": [ + { + "key": "chunk-000000-6d09c2dd9f6a", + "content": "Alpha beta.", + "rendered": "Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "6d09c2dd9f6acb7404974f19cfc7e9036ac59abb930e505351a55e94b8b1701c", + "heading_paths": [ + [ + "Intro" + ] + ], + "source_spans": [ + { + "unit_index": 0, + "char_start": 0, + "char_end": 11 + } + ], + "truncated": false + }, + { + "key": "chunk-000001-750e49755ff9", + "content": "Gamma delta.", + "rendered": "Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "750e49755ff93394b3ad9ff170f45d3e2ebc528e427347d6f08939db37835113", + "heading_paths": [ + [ + "Details" + ] + ], + "source_spans": [ + { + "unit_index": 1, + "char_start": 0, + "char_end": 12 + } + ], + "truncated": false + } + ] +} diff --git a/document/testdata/chunks-e5-v1.golden.json b/document/testdata/chunks-e5-v1.golden.json new file mode 100644 index 00000000..9cd07eef --- /dev/null +++ b/document/testdata/chunks-e5-v1.golden.json @@ -0,0 +1,66 @@ +{ + "version": 2, + "checksum": "ae69cde35cd8a675856f3d2c473692a63c71ed2def56402efc64ff87c6e9da6e", + "policy_fingerprint": "c399a944a459d1f78ef06d9a08aa043abc255c3233082ba8a0531267a0d7b39d", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "6ae96a49ec1dc2d2e9e2df1a63af01d5bd47980022e5128d16bc8ffe8ed95e7e", + "content_token_budget": 6, + "overlap_tokens": 1, + "truncation_policy": "reject_indivisible", + "context_fingerprint": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "total_content_tokens": 12, + "total_rendered_tokens": 21, + "total_content_bytes": 23, + "total_rendered_bytes": 41, + "inputs": [ + { + "key": "chunk-000000-266a1f658b28", + "content": "Alpha beta.", + "rendered": "passage: Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 10, + "checksum": "266a1f658b286585d3f73b81025b643c25e87c7554960cf2cd4e0246772af1e7", + "heading_paths": [ + [ + "Intro" + ] + ], + "source_spans": [ + { + "unit_index": 0, + "char_start": 0, + "char_end": 11 + } + ], + "truncated": false + }, + { + "key": "chunk-000001-f72ba775cd56", + "content": "Gamma delta.", + "rendered": "passage: Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 11, + "checksum": "f72ba775cd56fda7cad7a6c1e9ec43b907e3a2a5a3aa329c1b19d7e34bb1038f", + "heading_paths": [ + [ + "Details" + ] + ], + "source_spans": [ + { + "unit_index": 1, + "char_start": 0, + "char_end": 12 + } + ], + "truncated": false + } + ] +} diff --git a/document/testdata/chunks-gte-v1.golden.json b/document/testdata/chunks-gte-v1.golden.json new file mode 100644 index 00000000..76dcd675 --- /dev/null +++ b/document/testdata/chunks-gte-v1.golden.json @@ -0,0 +1,66 @@ +{ + "version": 2, + "checksum": "6669ae0498397dfd2cdb28df1529a0f4d423c4e63a051f18372f3f96286b01e7", + "policy_fingerprint": "90a21ad0d1787d9f88bdf7d7a1e09d32ac50e7954b214ae4249811fd6b675522", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "535de15db9409ee11c011dc473f9183513d3efd9423e6d45c75c64b3ef31ad02", + "content_token_budget": 6, + "overlap_tokens": 1, + "truncation_policy": "reject_indivisible", + "context_fingerprint": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "total_content_tokens": 12, + "total_rendered_tokens": 12, + "total_content_bytes": 23, + "total_rendered_bytes": 23, + "inputs": [ + { + "key": "chunk-000000-6d09c2dd9f6a", + "content": "Alpha beta.", + "rendered": "Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "6d09c2dd9f6acb7404974f19cfc7e9036ac59abb930e505351a55e94b8b1701c", + "heading_paths": [ + [ + "Intro" + ] + ], + "source_spans": [ + { + "unit_index": 0, + "char_start": 0, + "char_end": 11 + } + ], + "truncated": false + }, + { + "key": "chunk-000001-750e49755ff9", + "content": "Gamma delta.", + "rendered": "Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "750e49755ff93394b3ad9ff170f45d3e2ebc528e427347d6f08939db37835113", + "heading_paths": [ + [ + "Details" + ] + ], + "source_spans": [ + { + "unit_index": 1, + "char_start": 0, + "char_end": 12 + } + ], + "truncated": false + } + ] +} diff --git a/document/testdata/chunks-nomic-v1.golden.json b/document/testdata/chunks-nomic-v1.golden.json new file mode 100644 index 00000000..236b55ab --- /dev/null +++ b/document/testdata/chunks-nomic-v1.golden.json @@ -0,0 +1,66 @@ +{ + "version": 2, + "checksum": "53ff67cdb1e076cc2af4db3a51adcf4d2fb7b657c0752f812f4ff7218f5ba04c", + "policy_fingerprint": "5b75326f6b5bd3cd4515bb151ec978e6f3218f6b3aa62dfe38691d272e0be0c0", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "c8286ca7895bcc8dbe3f761b94807abc30ec2b800d062561be3d3a5662fcab6a", + "content_token_budget": 6, + "overlap_tokens": 1, + "truncation_policy": "reject_indivisible", + "context_fingerprint": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "total_content_tokens": 12, + "total_rendered_tokens": 29, + "total_content_bytes": 23, + "total_rendered_bytes": 57, + "inputs": [ + { + "key": "chunk-000000-30be6a24a77a", + "content": "Alpha beta.", + "rendered": "search_document: Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 14, + "checksum": "30be6a24a77ac2dae8990bac594a34fcdd03ba6a7a643ecdf68fe7900453fa99", + "heading_paths": [ + [ + "Intro" + ] + ], + "source_spans": [ + { + "unit_index": 0, + "char_start": 0, + "char_end": 11 + } + ], + "truncated": false + }, + { + "key": "chunk-000001-0e192444c494", + "content": "Gamma delta.", + "rendered": "search_document: Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 15, + "checksum": "0e192444c494cb0c4ad53d9fa290948c3035c8aff4573d68b4035a24991df282", + "heading_paths": [ + [ + "Details" + ] + ], + "source_spans": [ + { + "unit_index": 1, + "char_start": 0, + "char_end": 12 + } + ], + "truncated": false + } + ] +} diff --git a/document/testdata/chunks-qwen3-v1.golden.json b/document/testdata/chunks-qwen3-v1.golden.json new file mode 100644 index 00000000..19975208 --- /dev/null +++ b/document/testdata/chunks-qwen3-v1.golden.json @@ -0,0 +1,66 @@ +{ + "version": 2, + "checksum": "d423df2193723b79d18b43f1775d1093a64e02a189e3f891b6a112074fb0b10d", + "policy_fingerprint": "b605e26a5585df6f6a809e234abbf5c900217025418584ccf7a6ae3940c8e457", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "f280138c305dc66fcea0eca352898efba5f94513cef80e4475186d596d6f8cce", + "content_token_budget": 6, + "overlap_tokens": 1, + "truncation_policy": "reject_indivisible", + "context_fingerprint": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "total_content_tokens": 12, + "total_rendered_tokens": 12, + "total_content_bytes": 23, + "total_rendered_bytes": 23, + "inputs": [ + { + "key": "chunk-000000-6d09c2dd9f6a", + "content": "Alpha beta.", + "rendered": "Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "6d09c2dd9f6acb7404974f19cfc7e9036ac59abb930e505351a55e94b8b1701c", + "heading_paths": [ + [ + "Intro" + ] + ], + "source_spans": [ + { + "unit_index": 0, + "char_start": 0, + "char_end": 11 + } + ], + "truncated": false + }, + { + "key": "chunk-000001-750e49755ff9", + "content": "Gamma delta.", + "rendered": "Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 6, + "checksum": "750e49755ff93394b3ad9ff170f45d3e2ebc528e427347d6f08939db37835113", + "heading_paths": [ + [ + "Details" + ] + ], + "source_spans": [ + { + "unit_index": 1, + "char_start": 0, + "char_end": 12 + } + ], + "truncated": false + } + ] +} diff --git a/document/testdata/embedding-input-generation-v1.golden.json b/document/testdata/embedding-input-generation-v1.golden.json new file mode 100644 index 00000000..b7503b3f --- /dev/null +++ b/document/testdata/embedding-input-generation-v1.golden.json @@ -0,0 +1,42 @@ +{ + "version": 1, + "checksum": "007666b8ded9ec8d15734c8a3da9e4bd140bbddef27866835480a929fb4bfcf6", + "policy_fingerprint": "3339d8944fdbf65a726a83f65226a86ed84255edfc8cfd811ca6db5de1cb9481", + "evidence_checksum": "be5eaa990540eb4b5c5512ef07b63333b29897716b01283017f46f1ac6ab8b04", + "tokenizer_identity": { + "name": "synthetic-pairs", + "revision": "v1", + "prefix_token_counts_monotonic": true + }, + "lexical_evidence_fingerprint": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "formatter": "evidence-text/v1", + "model_input_fingerprint": "c8286ca7895bcc8dbe3f761b94807abc30ec2b800d062561be3d3a5662fcab6a", + "total_content_tokens": 12, + "total_rendered_tokens": 29, + "total_content_bytes": 23, + "total_rendered_bytes": 57, + "inputs": [ + { + "key": "chunk-000000-30be6a24a77a", + "content": "Alpha beta.", + "rendered": "search_document: Alpha beta.", + "content_tokens": 6, + "rendered_tokens": 14, + "checksum": "30be6a24a77ac2dae8990bac594a34fcdd03ba6a7a643ecdf68fe7900453fa99", + "heading_paths": [["Intro"]], + "source_spans": [{"unit_index": 0, "char_start": 0, "char_end": 11}], + "truncated": false + }, + { + "key": "chunk-000001-0e192444c494", + "content": "Gamma delta.", + "rendered": "search_document: Gamma delta.", + "content_tokens": 6, + "rendered_tokens": 15, + "checksum": "0e192444c494cb0c4ad53d9fa290948c3035c8aff4573d68b4035a24991df282", + "heading_paths": [["Details"]], + "source_spans": [{"unit_index": 1, "char_start": 0, "char_end": 12}], + "truncated": false + } + ] +} diff --git a/document/testdata/profile-v1.golden.json b/document/testdata/profile-v1.golden.json index 3cfaf8d8..d04e0546 100644 --- a/document/testdata/profile-v1.golden.json +++ b/document/testdata/profile-v1.golden.json @@ -1 +1 @@ -{"contract_version":"processing-profile/v1","embeddings":[{"activation":"optional","authorization_fingerprint":"7777777777777777777777777777777777777777777777777777777777777777","chunk":null,"compatibility_id":"voyage-multimodal-3/1024","credential_binding":"credential:embedding-primary","descriptor":{"fingerprint":"8888888888888888888888888888888888888888888888888888888888888888","id":"voyage-multimodal-v1"},"dimensions":1024,"disclosure_fingerprint":"9999999999999999999999999999999999999999999999999999999999999999","document_formatter":"direct-file/v1","input_kind":"original_file","max_batch_items":8,"max_input_bytes":8388608,"max_response_bytes":1048576,"metric":"cosine","model":"voyage-Café\nmodel","name":"direct","normalization":"unit_length","query_formatter":"query/v1","scalar_encoding":"float32","trust_boundary":"processor-primary"},{"activation":"required","authorization_fingerprint":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","chunk":{"context_fingerprint":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","formatter":"rendition-chunk/v1","max_tokens":800,"overlap_tokens":80,"tokenizer":"voyage-3","truncation_policy":"reject"},"compatibility_id":"voyage-3-large/1024","credential_binding":"credential:embedding-primary","descriptor":{"fingerprint":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","id":"voyage-text-v1"},"dimensions":1024,"disclosure_fingerprint":"dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd","document_formatter":"document/v1","input_kind":"rendition_chunk","max_batch_items":32,"max_input_bytes":1048576,"max_response_bytes":1048576,"metric":"cosine","model":"voyage-3-large","name":"semantic","normalization":"unit_length","query_formatter":"query/v1","scalar_encoding":"float32","trust_boundary":"processor-primary"}],"evidence_lexical":{"completeness_fingerprint":"6666666666666666666666666666666666666666666666666666666666666666","lexical_segmenter_fingerprint":"5555555555555555555555555555555555555555555555555555555555555555","max_segment_runes":2000,"max_unit_runes":100000,"normalized_evidence_contract":"normalized-evidence/v1","normalizer_fingerprint":"3333333333333333333333333333333333333333333333333333333333333333","rendition_contract":"rendition/v1","sanitizer_fingerprint":"4444444444444444444444444444444444444444444444444444444444444444","source_evidence_contract":"source-evidence/v1"},"rendition":{"adapter_contract":"mistral-ocr-adapter/v1","authorization_fingerprint":"2222222222222222222222222222222222222222222222222222222222222222","credential_binding":"credential:ocr-primary","deployment_fingerprint":"dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd","descriptor":{"fingerprint":"1111111111111111111111111111111111111111111111111111111111111111","id":"mistral-Café\nruntime"},"disclose_filename":false,"disclosure_fingerprint":"eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee","max_document_bytes":10485760,"max_response_bytes":4194304,"max_units":256,"name":"primary","requested_artifacts":["provider_markdown","structured_evidence"],"trust_boundary":"processor-primary","upload_options_fingerprint":"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"},"retention_disclosure":{"attachment_policy_fingerprint":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","consent_fingerprint":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","retain_provider_markdown":false,"retain_sanitized_markdown":true,"retain_typed_artifacts":true,"trust_boundary":"vault-primary"},"retrieval":{"lexical_limit":40,"vector_limit":60}} +{"contract_version":"processing-profile/v1","embeddings":[{"activation":"optional","authorization_fingerprint":"7777777777777777777777777777777777777777777777777777777777777777","chunk":null,"compatibility_id":"voyage-multimodal-3/1024","credential_binding":"credential:embedding-primary","descriptor":{"fingerprint":"8888888888888888888888888888888888888888888888888888888888888888","id":"voyage-multimodal-v1"},"dimensions":1024,"disclosure_fingerprint":"9999999999999999999999999999999999999999999999999999999999999999","document_formatter":"direct-file/v1","input_kind":"original_file","max_batch_items":8,"max_input_bytes":8388608,"max_response_bytes":1048576,"metric":"cosine","model":"voyage-Café\nmodel","model_input":{"compatibility_id":"","document":{"mode":"","template":""},"fingerprint":"2079965fdab2f527ef5fdef9f15668005cadd309efa577c8ee98e7668d6ddb5e","profile":"","query":{"mode":"","template":""},"version":1},"name":"direct","normalization":"unit_length","query_formatter":"query/v1","scalar_encoding":"float32","trust_boundary":"processor-primary"},{"activation":"required","authorization_fingerprint":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","chunk":{"context_fingerprint":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","formatter":"rendition-chunk/v1","max_tokens":800,"overlap_tokens":80,"tokenizer":"voyage-3","truncation_policy":"reject"},"compatibility_id":"voyage-3-large/1024","credential_binding":"credential:embedding-primary","descriptor":{"fingerprint":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","id":"voyage-text-v1"},"dimensions":1024,"disclosure_fingerprint":"dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd","document_formatter":"document/v1","input_kind":"rendition_chunk","max_batch_items":32,"max_input_bytes":1048576,"max_response_bytes":1048576,"metric":"cosine","model":"voyage-3-large","model_input":{"compatibility_id":"","document":{"mode":"","template":""},"fingerprint":"2079965fdab2f527ef5fdef9f15668005cadd309efa577c8ee98e7668d6ddb5e","profile":"","query":{"mode":"","template":""},"version":1},"name":"semantic","normalization":"unit_length","query_formatter":"query/v1","scalar_encoding":"float32","trust_boundary":"processor-primary"}],"evidence_lexical":{"completeness_fingerprint":"6666666666666666666666666666666666666666666666666666666666666666","lexical_segmenter_fingerprint":"5555555555555555555555555555555555555555555555555555555555555555","max_segment_runes":2000,"max_unit_runes":100000,"normalized_evidence_contract":"normalized-evidence/v1","normalizer_fingerprint":"3333333333333333333333333333333333333333333333333333333333333333","rendition_contract":"rendition/v1","sanitizer_fingerprint":"4444444444444444444444444444444444444444444444444444444444444444","source_evidence_contract":"source-evidence/v1"},"rendition":{"adapter_contract":"mistral-ocr-adapter/v1","authorization_fingerprint":"2222222222222222222222222222222222222222222222222222222222222222","credential_binding":"credential:ocr-primary","deployment_fingerprint":"dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd","descriptor":{"fingerprint":"1111111111111111111111111111111111111111111111111111111111111111","id":"mistral-Café\nruntime"},"disclose_filename":false,"disclosure_fingerprint":"eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee","max_document_bytes":10485760,"max_response_bytes":4194304,"max_units":256,"name":"primary","requested_artifacts":["provider_markdown","structured_evidence"],"trust_boundary":"processor-primary","upload_options_fingerprint":"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"},"retention_disclosure":{"attachment_policy_fingerprint":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","consent_fingerprint":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","retain_provider_markdown":false,"retain_sanitized_markdown":true,"retain_typed_artifacts":true,"trust_boundary":"vault-primary"},"retrieval":{"lexical_limit":40,"vector_limit":60}} diff --git a/document/testdata/source-metadata-v1.golden.json b/document/testdata/source-metadata-v1.golden.json new file mode 100644 index 00000000..a075e5e6 --- /dev/null +++ b/document/testdata/source-metadata-v1.golden.json @@ -0,0 +1 @@ +{"contract_version":"source-metadata/v1","fields":[{"key":"created","namespace":"xmp","sensitive":false,"source_field":"CreateDate","value":{"kind":"timestamp","timestamp":{"normalized":"2024-01-02T03:04:05-07:00","offset":"-07:00","precision":"second","raw":"2024-01-02T03:04:05-07:00","timezone":"offset"}}},{"key":"creators","namespace":"pdf.info","sensitive":false,"source_field":"Author","value":{"kind":"string_list","strings":["Ada Lovelace","Grace Hopper"]}},{"key":"pdf.info.page_count","namespace":"pdf.info","sensitive":false,"source_field":"PageCount","value":{"integer":7,"kind":"integer"}},{"key":"title","namespace":"pdf.info","sensitive":false,"source_field":"Title","value":{"kind":"string","string":"Synthetic report"}}],"warnings":[{"code":"unparseable_timestamp","detail":"value retained without coercion","namespace":"pdf.info","source_field":"ModDate"}]} diff --git a/document/testdata/vector-set-v1.golden.bin b/document/testdata/vector-set-v1.golden.bin new file mode 100644 index 0000000000000000000000000000000000000000..d0716a023ba8efaa1ea5d067bd97b3448bdbd833 GIT binary patch literal 293 zcmb`B!41MN3`Ii^a9{}*sA$|cX-{+$lsGP}K&eF2o*9kNa8C>XmY(##WIcOj%G0x) zIiFXpc1o#EXrMM4qGnIQx)v=!#aII&YjVJ$HxV&9Pma*Hwy{{bfF?+#o@g%)n&vU_ qQPPsee0^}b<)M(eao84yCEZKkN&NNRR~P=FlaFLxmqJkL^!@;)I7zes literal 0 HcmV?d00001 diff --git a/document/tika/profile.go b/document/tika/profile.go new file mode 100644 index 00000000..bf3bffda --- /dev/null +++ b/document/tika/profile.go @@ -0,0 +1,317 @@ +// Package tika defines the fixed compatibility profile used when an operator +// deploys an Apache Tika adapter behind docbank-rendition/v1. +package tika + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "slices" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/bridge" +) + +const ( + // ProfileContractV1 identifies the canonical Apache Tika bridge profile. + ProfileContractV1 = "tika-bridge-profile/v1" + descriptorID = "tika.bridge.v1" + maxCredentialBindingBytes = 128 +) + +// Config supplies the only operator-specific profile values. It deliberately +// has no routes, URLs, headers, parser options, or fetch controls. +type Config struct { + DeploymentID string + RuntimeID string + CredentialBinding string +} + +// LimitsV1 fixes every finite input, output, polling, and wall-clock bound. +type LimitsV1 struct { + MaxDocumentBytes int64 `json:"max_document_bytes"` + MaxPollAttempts int `json:"max_poll_attempts"` + MaxResponseBytes int64 `json:"max_response_bytes"` + PollIntervalMillis int64 `json:"poll_interval_millis"` + RequestTimeoutMillis int64 `json:"request_timeout_millis"` + TotalTimeoutMillis int64 `json:"total_timeout_millis"` +} + +// DisclosurePolicyV1 permits only exact supplied bytes and records the +// bridge's safe-basename disclosure. Authorization binds every byte and tuple. +type DisclosurePolicyV1 struct { + DiscloseFilename bool `json:"disclose_filename"` + Source string `json:"source"` +} + +// ReferencePolicyV1 refuses both embedded and external reference fetching. +// The generic bridge cannot inspect parser internals, so compatibility requires +// an operator-pinned adapter runtime audited to enforce both refusals. +type ReferencePolicyV1 struct { + EmbeddedReferenceFetch string `json:"embedded_reference_fetch"` + EnforcementBoundary string `json:"enforcement_boundary"` + ExternalReferenceFetch string `json:"external_reference_fetch"` +} + +// EvidencePolicyV1 fixes bounded provider-neutral evidence and Markdown. +type EvidencePolicyV1 struct { + MaxProviderMarkdownBytes int `json:"max_provider_markdown_bytes"` + MaxTotalResultBytes int `json:"max_total_result_bytes"` + MaxUnits int `json:"max_units"` + SourceEvidenceContract string `json:"source_evidence_contract"` +} + +// ArtifactPolicyV1 permits only one bounded structured-evidence artifact. +type ArtifactPolicyV1 struct { + AllowedRoles []document.EvidenceArtifactRole `json:"allowed_roles"` + MaxArtifactBytes int64 `json:"max_artifact_bytes"` + MaxArtifacts int `json:"max_artifacts"` +} + +// ProfileV1 is the immutable compatibility identity expected from an +// operator-network Apache Tika bridge deployment. +type ProfileV1 struct { + ArtifactPolicy ArtifactPolicyV1 `json:"artifact_policy"` + BridgeContract string `json:"bridge_contract"` + ContractVersion string `json:"contract_version"` + CredentialBinding string `json:"credential_binding"` + DeploymentID string `json:"deployment_id"` + Disclosure DisclosurePolicyV1 `json:"disclosure"` + EvidencePolicy EvidencePolicyV1 `json:"evidence_policy"` + InputKind document.RenditionInputKind `json:"input_kind"` + Limits LimitsV1 `json:"limits"` + PolicyFingerprint string `json:"policy_fingerprint"` + ReferencePolicy ReferencePolicyV1 `json:"reference_policy"` + RuntimeID string `json:"runtime_id"` + SupportedFormats []document.RenditionFormatCapability `json:"supported_formats"` + TrustBoundary document.RenditionTrustBoundary `json:"trust_boundary"` +} + +// NewProfile returns the standard profile with operator-pinned deployment and +// runtime identity and an optional named credential binding. +func NewProfile(config Config) (ProfileV1, error) { + profile := ProfileV1{ + ArtifactPolicy: standardArtifactPolicy(), + BridgeContract: bridge.ContractVersion, ContractVersion: ProfileContractV1, + CredentialBinding: config.CredentialBinding, DeploymentID: config.DeploymentID, + Disclosure: standardDisclosurePolicy(), + EvidencePolicy: standardEvidencePolicy(), + InputKind: document.RenditionInputOriginalFile, + Limits: standardLimits(), + ReferencePolicy: standardReferencePolicy(), + RuntimeID: config.RuntimeID, SupportedFormats: standardFormats(), + TrustBoundary: document.RenditionTrustOperatorNetwork, + } + _, fingerprint, err := CanonicalProfile(profile) + if err != nil { + return ProfileV1{}, err + } + profile.PolicyFingerprint = fingerprint + return profile, nil +} + +// CanonicalProfile validates, sorts, and deterministically encodes a profile. +// A populated fingerprint must match the canonical identity. +func CanonicalProfile(profile ProfileV1) ([]byte, string, error) { + canonical := cloneProfile(profile) + slices.SortFunc(canonical.SupportedFormats, compareFormats) + slices.Sort(canonical.ArtifactPolicy.AllowedRoles) + claimed := canonical.PolicyFingerprint + canonical.PolicyFingerprint = "" + if err := validateProfile(canonical); err != nil { + return nil, "", fmt.Errorf("tika: invalid profile: %w", err) + } + identityJSON, err := json.Marshal(canonical, json.Deterministic(true)) + if err != nil { + return nil, "", fmt.Errorf("tika: encode profile identity: %w", err) + } + digest := sha256.Sum256(identityJSON) + fingerprint := hex.EncodeToString(digest[:]) + if claimed != "" && claimed != fingerprint { + return nil, "", errors.New("tika: policy fingerprint does not match canonical profile") + } + canonical.PolicyFingerprint = fingerprint + encoded, err := json.Marshal(canonical, json.Deterministic(true)) + if err != nil { + return nil, "", fmt.Errorf("tika: encode canonical profile: %w", err) + } + return encoded, fingerprint, nil +} + +// ParseProfile accepts only the exact canonical v1 representation. +func ParseProfile(raw []byte) (ProfileV1, error) { + var profile ProfileV1 + if err := json.Unmarshal(raw, &profile, json.RejectUnknownMembers(true)); err != nil { + return ProfileV1{}, fmt.Errorf("tika: decode profile: %w", err) + } + canonical, _, err := CanonicalProfile(profile) + if err != nil { + return ProfileV1{}, err + } + if !bytes.Equal(raw, canonical) { + return ProfileV1{}, errors.New("tika: profile is not canonical") + } + return cloneProfile(profile), nil +} + +// BridgeProfile projects the compatibility identity into the generic hardened +// bridge. Origin remains generic bridge configuration, not profile schema. +func BridgeProfile(profile ProfileV1, origin string) (bridge.Profile, error) { + _, fingerprint, err := CanonicalProfile(profile) + if err != nil { + return bridge.Profile{}, err + } + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: descriptorID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: fingerprint, TrustBoundary: document.RenditionTrustOperatorNetwork, + SupportedFormats: slices.Clone(profile.SupportedFormats), ReturnsMarkdown: true, + ReturnsStructured: true, + ArtifactRoles: slices.Clone(profile.ArtifactPolicy.AllowedRoles), + }) + if err != nil { + return bridge.Profile{}, fmt.Errorf("tika: construct bridge descriptor: %w", err) + } + return bridge.Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: profile.CredentialBinding, + RequestTimeout: time.Duration(profile.Limits.RequestTimeoutMillis) * time.Millisecond, + TotalTimeout: time.Duration(profile.Limits.TotalTimeoutMillis) * time.Millisecond, + PollInterval: time.Duration(profile.Limits.PollIntervalMillis) * time.Millisecond, + MaxPollAttempts: profile.Limits.MaxPollAttempts, MaxResponseBytes: profile.Limits.MaxResponseBytes, + MaxSourceBytes: profile.Limits.MaxDocumentBytes, + MaxProviderMarkdownBytes: profile.EvidencePolicy.MaxProviderMarkdownBytes, + MaxArtifactBytes: int(profile.ArtifactPolicy.MaxArtifactBytes), + MaxArtifacts: profile.ArtifactPolicy.MaxArtifacts, + MaxTotalResultBytes: profile.EvidencePolicy.MaxTotalResultBytes, + MaxEvidenceUnits: profile.EvidencePolicy.MaxUnits, + }, nil +} + +func validateProfile(profile ProfileV1) error { + if profile.ContractVersion != ProfileContractV1 || profile.BridgeContract != bridge.ContractVersion { + return errors.New("contract version is invalid") + } + if err := validateIdentity(profile.DeploymentID, "deployment ID", false); err != nil { + return err + } + if err := validateIdentity(profile.RuntimeID, "runtime ID", true); err != nil { + return err + } + if profile.CredentialBinding != "" { + if len(profile.CredentialBinding) > maxCredentialBindingBytes { + return errors.New("credential binding is invalid") + } + if err := validateIdentity(profile.CredentialBinding, "credential binding", false); err != nil { + return err + } + } + if profile.TrustBoundary != document.RenditionTrustOperatorNetwork || + profile.InputKind != document.RenditionInputOriginalFile || + profile.Disclosure != standardDisclosurePolicy() { + return errors.New("disclosure or execution boundary is invalid") + } + if profile.ReferencePolicy != standardReferencePolicy() { + return errors.New("embedded and external reference fetching must be refused") + } + if !slices.Equal(profile.SupportedFormats, standardFormats()) { + return errors.New("supported formats differ from the standard profile") + } + if !slices.Equal(profile.ArtifactPolicy.AllowedRoles, + []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}) { + return errors.New("artifact roles differ from the standard profile") + } + standardArtifacts := standardArtifactPolicy() + if profile.Limits != standardLimits() || profile.EvidencePolicy != standardEvidencePolicy() || + profile.ArtifactPolicy.MaxArtifactBytes != standardArtifacts.MaxArtifactBytes || + profile.ArtifactPolicy.MaxArtifacts != standardArtifacts.MaxArtifacts { + return errors.New("limits differ from the finite standard profile") + } + return nil +} + +func standardLimits() LimitsV1 { + return LimitsV1{ + MaxDocumentBytes: 100 << 20, MaxPollAttempts: 300, MaxResponseBytes: 128 << 20, + PollIntervalMillis: 1_000, RequestTimeoutMillis: 30_000, TotalTimeoutMillis: 600_000, + } +} + +func standardDisclosurePolicy() DisclosurePolicyV1 { + return DisclosurePolicyV1{DiscloseFilename: true, Source: "exact_supplied_bytes"} +} + +func standardReferencePolicy() ReferencePolicyV1 { + return ReferencePolicyV1{ + EmbeddedReferenceFetch: "refuse", EnforcementBoundary: "pinned_audited_adapter_runtime", + ExternalReferenceFetch: "refuse", + } +} + +func standardEvidencePolicy() EvidencePolicyV1 { + return EvidencePolicyV1{ + MaxProviderMarkdownBytes: 32 << 20, MaxTotalResultBytes: 128 << 20, + MaxUnits: 100_000, SourceEvidenceContract: document.SourceEvidenceContractV1, + } +} + +func standardArtifactPolicy() ArtifactPolicyV1 { + return ArtifactPolicyV1{ + AllowedRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + MaxArtifactBytes: 64 << 20, MaxArtifacts: 1, + } +} + +func validateIdentity(value, subject string, allowColon bool) error { + if value == "" || len(value) > 256 || !utf8.ValidString(value) || strings.TrimSpace(value) != value { + return fmt.Errorf("%s is invalid", subject) + } + for _, character := range value { + if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' || + character >= '0' && character <= '9' || strings.ContainsRune("._-", character) || + allowColon && character == ':' { + continue + } + return fmt.Errorf("%s is invalid", subject) + } + return nil +} + +func standardFormats() []document.RenditionFormatCapability { + formats := []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "text/csv", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "mail", MediaType: "message/rfc822", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "structured", MediaType: "application/xml", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/markdown", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + } + slices.SortFunc(formats, compareFormats) + return formats +} + +func compareFormats(left, right document.RenditionFormatCapability) int { + if comparison := strings.Compare(left.MediaFamily, right.MediaFamily); comparison != 0 { + return comparison + } + if comparison := strings.Compare(left.MediaType, right.MediaType); comparison != 0 { + return comparison + } + return strings.Compare(string(left.InputKind), string(right.InputKind)) +} + +func cloneProfile(profile ProfileV1) ProfileV1 { + profile.SupportedFormats = slices.Clone(profile.SupportedFormats) + profile.ArtifactPolicy.AllowedRoles = slices.Clone(profile.ArtifactPolicy.AllowedRoles) + return profile +} diff --git a/document/tika/profile_test.go b/document/tika/profile_test.go new file mode 100644 index 00000000..0f654492 --- /dev/null +++ b/document/tika/profile_test.go @@ -0,0 +1,192 @@ +package tika + +import ( + "context" + "encoding/json/v2" + "net/http" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/bridge" +) + +func TestReferenceProfileCanonicalizesCompleteCompatibilityIdentity(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-tika-primary", + RuntimeID: "sha256:" + strings.Repeat("a", 64), + CredentialBinding: "tika-api", + }) + require.NoError(t, err) + + canonical, fingerprint, err := CanonicalProfile(profile) + require.NoError(t, err) + assert.Equal(t, profile.PolicyFingerprint, fingerprint) + parsed, err := ParseProfile(canonical) + require.NoError(t, err) + assert.Equal(t, profile, parsed) + + reordered := profile + reordered.PolicyFingerprint = "" + reordered.SupportedFormats = slices.Clone(profile.SupportedFormats) + slices.Reverse(reordered.SupportedFormats) + reorderedJSON, reorderedFingerprint, err := CanonicalProfile(reordered) + require.NoError(t, err) + assert.JSONEq(t, string(canonical), string(reorderedJSON)) + assert.Equal(t, fingerprint, reorderedFingerprint) + + var wire map[string]any + require.NoError(t, json.Unmarshal(canonical, &wire)) + assert.ElementsMatch(t, []string{ + "artifact_policy", "bridge_contract", "contract_version", "credential_binding", + "deployment_id", "disclosure", "evidence_policy", "input_kind", "limits", + "policy_fingerprint", "reference_policy", "runtime_id", "supported_formats", + "trust_boundary", + }, mapKeys(wire)) + _, err = ParseProfile(append(canonical, '\n')) + require.ErrorContains(t, err, "not canonical") +} + +func TestReferenceProfileRestrictsBroadFormatsToSuppliedBytesAndRefusesReferences(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-tika-primary", + RuntimeID: "sha256:" + strings.Repeat("b", 64), + }) + require.NoError(t, err) + + assert.Equal(t, []document.RenditionFormatCapability{ + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "mail", MediaType: "message/rfc822", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "text/csv", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "structured", MediaType: "application/xml", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/markdown", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: document.RenditionInputOriginalFile}, + }, profile.SupportedFormats) + assert.Equal(t, "exact_supplied_bytes", profile.Disclosure.Source) + assert.True(t, profile.Disclosure.DiscloseFilename) + assert.Equal(t, "refuse", profile.ReferencePolicy.EmbeddedReferenceFetch) + assert.Equal(t, "refuse", profile.ReferencePolicy.ExternalReferenceFetch) + assert.Equal(t, "pinned_audited_adapter_runtime", profile.ReferencePolicy.EnforcementBoundary) + assert.Equal(t, []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + profile.ArtifactPolicy.AllowedRoles) +} + +func TestReferenceProfileBuildsStandardBridgeAndRejectsIdentityOrLimitDrift(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-tika-primary", + RuntimeID: "sha256:" + strings.Repeat("c", 64), + CredentialBinding: "tika-api", + }) + require.NoError(t, err) + + bridgeProfile, err := BridgeProfile(profile, "http://127.0.0.1:9998") + require.NoError(t, err) + assert.Equal(t, bridge.ContractVersion, profile.BridgeContract) + assert.Equal(t, profile.CredentialBinding, bridgeProfile.SecretBinding) + assert.Equal(t, document.RenditionTrustOperatorNetwork, bridgeProfile.Descriptor.TrustBoundary) + assert.Equal(t, profile.PolicyFingerprint, bridgeProfile.Descriptor.PolicyFingerprint) + assert.Equal(t, profile.SupportedFormats, bridgeProfile.Descriptor.SupportedFormats) + assert.True(t, bridgeProfile.Descriptor.ReturnsMarkdown) + assert.True(t, bridgeProfile.Descriptor.ReturnsStructured) + assert.Equal(t, profile.ArtifactPolicy.AllowedRoles, bridgeProfile.Descriptor.ArtifactRoles) + assert.Equal(t, profile.Limits.MaxDocumentBytes, bridgeProfile.MaxSourceBytes) + assert.Equal(t, profile.EvidencePolicy.MaxProviderMarkdownBytes, bridgeProfile.MaxProviderMarkdownBytes) + assert.Equal(t, int(profile.ArtifactPolicy.MaxArtifactBytes), bridgeProfile.MaxArtifactBytes) + assert.Equal(t, profile.ArtifactPolicy.MaxArtifacts, bridgeProfile.MaxArtifacts) + assert.Equal(t, profile.EvidencePolicy.MaxTotalResultBytes, bridgeProfile.MaxTotalResultBytes) + assert.Equal(t, profile.EvidencePolicy.MaxUnits, bridgeProfile.MaxEvidenceUnits) + _, err = bridge.New(bridgeProfile, staticSecretResolver{}, http.DefaultClient) + require.NoError(t, err) + + credentialless, err := NewProfile(Config{ + DeploymentID: "operator-tika-no-auth", RuntimeID: "sha256:" + strings.Repeat("e", 64), + }) + require.NoError(t, err) + credentiallessBridge, err := BridgeProfile(credentialless, "http://127.0.0.1:9998") + require.NoError(t, err) + _, err = bridge.New(credentiallessBridge, nil, http.DefaultClient) + require.NoError(t, err) + + for _, test := range []struct { + mutate func(*ProfileV1) + want string + }{ + {mutate: func(value *ProfileV1) { value.RuntimeID = "sha256:" + strings.Repeat("d", 64) }, want: "policy fingerprint does not match"}, + {mutate: func(value *ProfileV1) { value.Limits.MaxDocumentBytes++ }, want: "limits differ"}, + {mutate: func(value *ProfileV1) { value.ReferencePolicy.ExternalReferenceFetch = "allow" }, want: "reference fetching must be refused"}, + } { + drifted := profile + test.mutate(&drifted) + _, err := BridgeProfile(drifted, "http://127.0.0.1:9998") + require.ErrorContains(t, err, test.want) + } +} + +func TestReferenceProfileRejectsUnpinnedIdentityAndRecanonicalizedPolicyDrift(t *testing.T) { + tests := []struct { + name string + config Config + want string + }{ + {name: "deployment", config: Config{RuntimeID: "runtime-v1"}, want: "deployment ID"}, + {name: "runtime", config: Config{DeploymentID: "operator-tika-primary"}, want: "runtime ID"}, + {name: "credential", config: Config{DeploymentID: "operator-tika-primary", RuntimeID: "runtime-v1", CredentialBinding: "https://secret.invalid"}, want: "credential binding"}, + {name: "credential exceeds bridge bound", config: Config{DeploymentID: "operator-tika-primary", RuntimeID: "runtime-v1", CredentialBinding: strings.Repeat("a", 129)}, want: "credential binding"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := NewProfile(test.config) + require.ErrorContains(t, err, test.want) + }) + } + + profile, err := NewProfile(Config{DeploymentID: "operator-tika-primary", RuntimeID: "runtime-v1"}) + require.NoError(t, err) + profile.PolicyFingerprint = "" + profile.Limits.MaxResponseBytes++ + _, _, err = CanonicalProfile(profile) + require.ErrorContains(t, err, "limits differ") + + profile, err = NewProfile(Config{DeploymentID: "operator-tika-primary", RuntimeID: "runtime-v1"}) + require.NoError(t, err) + profile.PolicyFingerprint = "" + profile.ReferencePolicy.EmbeddedReferenceFetch = "allow" + _, _, err = CanonicalProfile(profile) + require.ErrorContains(t, err, "reference fetching must be refused") + + boundary, err := NewProfile(Config{ + DeploymentID: "operator-tika-primary", RuntimeID: "runtime-v1", + CredentialBinding: strings.Repeat("a", maxCredentialBindingBytes), + }) + require.NoError(t, err) + bridgeProfile, err := BridgeProfile(boundary, "http://127.0.0.1:9998") + require.NoError(t, err) + _, err = bridge.New(bridgeProfile, staticSecretResolver{}, http.DefaultClient) + require.NoError(t, err, "profile and bridge credential bounds must stay aligned") +} + +type staticSecretResolver struct{} + +func (staticSecretResolver) ResolveSecret(context.Context, string) (string, error) { + return "synthetic-secret", nil +} + +func mapKeys(value map[string]any) []string { + keys := make([]string, 0, len(value)) + for key := range value { + keys = append(keys, key) + } + return keys +} + +var _ bridge.SecretResolver = staticSecretResolver{} diff --git a/document/tokenizer.go b/document/tokenizer.go new file mode 100644 index 00000000..82e600a8 --- /dev/null +++ b/document/tokenizer.go @@ -0,0 +1,72 @@ +package document + +import ( + "errors" + "fmt" +) + +const ( + maxEmbeddingTokensPerGeneration = 1_000_000 + maxTokenizerIdentityBytes = 128 +) + +// ErrTokenizerLimit lets a tokenizer reject an input before allocating more +// than the caller-authorized number of token boundaries. +var ErrTokenizerLimit = errors.New("tokenizer token limit exceeded") + +// TokenizerIdentity pins the exact tokenizer vocabulary and segmentation +// revision. PrefixTokenCountsMonotonic is an explicit contract that a longer +// prefix can never have fewer tokens; fitting may use binary search only when +// it is true. Model names are deliberately not tokenizer identities. +type TokenizerIdentity struct { + Name string `json:"name"` + Revision string `json:"revision"` + PrefixTokenCountsMonotonic bool `json:"prefix_token_counts_monotonic"` +} + +// TokenBoundary is one half-open rune range. A canonical tokenization is a +// non-empty contiguous partition of the complete input string. +type TokenBoundary struct { + Start int `json:"start"` + End int `json:"end"` +} + +// Tokenizer returns exact rune boundaries and must honor limit before growing +// its result beyond that many entries. +type Tokenizer interface { + Identity() TokenizerIdentity + Tokenize(text string, limit int) ([]TokenBoundary, error) +} + +func validateTokenizerIdentity(identity TokenizerIdentity) error { + if err := validateStableToken(identity.Name, "tokenizer name", maxTokenizerIdentityBytes); err != nil { + return err + } + if err := validateStableToken(identity.Revision, "tokenizer revision", maxTokenizerIdentityBytes); err != nil { + return err + } + return nil +} + +func validateTokenBoundaries(tokens []TokenBoundary, runeCount, limit int) error { + if len(tokens) == 0 { + return errors.New("tokenizer must return at least one token for non-empty evidence") + } + if len(tokens) > limit { + return ErrTokenizerLimit + } + expectedStart := 0 + for index, token := range tokens { + if token.Start != expectedStart { + return fmt.Errorf("tokenizer boundaries must be contiguous at token %d", index) + } + if token.End <= token.Start || token.End > runeCount { + return fmt.Errorf("tokenizer boundary %d leaves input bounds", index) + } + expectedStart = token.End + } + if expectedStart != runeCount { + return errors.New("tokenizer boundaries must cover the complete input contiguously") + } + return nil +} diff --git a/document/trafilatura/doc.go b/document/trafilatura/doc.go new file mode 100644 index 00000000..4e81be01 --- /dev/null +++ b/document/trafilatura/doc.go @@ -0,0 +1,16 @@ +// Package trafilatura renders supplied HTML bytes through an operator-pinned +// local Trafilatura bridge. +// +// A nil Profile.Runner selects the production Linux runner. It executes a +// digest-verified sealed copy of the bridge inside new user, network, and PID +// namespaces, bounds stdout, and kills the PID-namespace init on cancellation +// or overflow so the kernel reaps every descendant. It fails closed when the +// host disables the required namespace controls. The native runner is not +// available on macOS or Windows; callers there must inject a separately audited +// IsolatedRunner or construction fails before any child can launch. +// +// An injected runner is an explicit trusted deployment boundary. Its immutable +// identity and exact per-run attestation are checked, but a dishonest runner can +// still lie. Deployments using one must audit and pin its platform-specific +// isolation implementation. +package trafilatura diff --git a/document/trafilatura/native_launcher_linux.go b/document/trafilatura/native_launcher_linux.go new file mode 100644 index 00000000..dd5969d3 --- /dev/null +++ b/document/trafilatura/native_launcher_linux.go @@ -0,0 +1,174 @@ +//go:build linux + +package trafilatura + +import ( + "crypto/subtle" + "encoding/hex" + "fmt" + "os" + "path/filepath" + "runtime" + "unsafe" + + "golang.org/x/sys/unix" +) + +const ( + nativeLauncherMarker = "--docbank-internal-trafilatura-launch-v2" + nativeLauncherExecutableFD = 3 + nativeLauncherControlFD = 4 + nativeLauncherStatusFD = 5 + nativeLauncherTokenBytes = 32 + nativeLauncherFailureExitCode = 125 + nativeLauncherReadyStatus = byte(1) + nativeLauncherFailureStatus = byte(2) + nativeX32SyscallBit = uint32(0x40000000) +) + +func init() { + executable, authenticated := authenticatedNativeLaunch(os.Args, nativeLauncherControlFD) + if !authenticated { + return + } + if runNativeLauncher(executable, nativeLauncherStatusFD) != nil { + _ = writeNativeLauncherStatus(nativeLauncherStatusFD, nativeLauncherFailureStatus) + os.Exit(nativeLauncherFailureExitCode) + } + _ = writeNativeLauncherStatus(nativeLauncherStatusFD, nativeLauncherFailureStatus) + os.Exit(nativeLauncherFailureExitCode) +} + +func authenticatedNativeLaunch(arguments []string, controlFD int) (string, bool) { + if len(arguments) != 4 || arguments[1] != nativeLauncherMarker || + !filepath.IsAbs(arguments[3]) || filepath.Clean(arguments[3]) != arguments[3] { + return "", false + } + want, err := hex.DecodeString(arguments[2]) + if err != nil || len(want) != nativeLauncherTokenBytes { + return "", false + } + var stat unix.Stat_t + if err := unix.Fstat(controlFD, &stat); err != nil || + stat.Mode&unix.S_IFMT != unix.S_IFREG || stat.Size != nativeLauncherTokenBytes { + return "", false + } + seals := unix.F_SEAL_WRITE | unix.F_SEAL_GROW | unix.F_SEAL_SHRINK | unix.F_SEAL_SEAL + applied, err := unix.FcntlInt(uintptr(controlFD), unix.F_GET_SEALS, 0) + if err != nil || applied&seals != seals { + return "", false + } + got := make([]byte, nativeLauncherTokenBytes) + read, err := unix.Pread(controlFD, got, 0) + if err != nil || read != len(got) || subtle.ConstantTimeCompare(got, want) != 1 { + return "", false + } + return arguments[3], true +} + +func runNativeLauncher(executable string, statusFD int) error { + if err := unix.Mount("", "/", "", unix.MS_REC|unix.MS_PRIVATE, ""); err != nil { + return fmt.Errorf("make mount propagation private: %w", err) + } + if err := unix.Mount("proc", "/proc", "proc", + unix.MS_NOSUID|unix.MS_NODEV|unix.MS_RDONLY, ""); err != nil { + return fmt.Errorf("mount private proc: %w", err) + } + if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil { + return fmt.Errorf("set no-new-privileges: %w", err) + } + if err := installNativeNetworkSeccomp(); err != nil { + return err + } + if err := unix.Close(nativeLauncherControlFD); err != nil { + return fmt.Errorf("close launcher control: %w", err) + } + unix.CloseOnExec(nativeLauncherExecutableFD) + if err := writeNativeLauncherStatus(statusFD, nativeLauncherReadyStatus); err != nil { + return err + } + unix.CloseOnExec(statusFD) + if err := unix.Exec("/proc/self/fd/3", []string{executable, "--protocol", protocolVersion}, cleanEnvironment()); err != nil { + return fmt.Errorf("execute sealed bridge: %w", err) + } + return nil +} + +func writeNativeLauncherStatus(fd int, status byte) error { + written, err := unix.Write(fd, []byte{status}) + if err != nil { + return fmt.Errorf("write launcher status: %w", err) + } + if written != 1 { + return fmt.Errorf("write launcher status: wrote %d bytes", written) + } + return nil +} + +func installNativeNetworkSeccomp() error { + architecture, ok := nativeAuditArchitecture() + if !ok { + return unix.ENOTSUP + } + filters, err := buildNativeNetworkSeccompFilters(architecture) + if err != nil { + return err + } + program := unix.SockFprog{Len: uint16(len(filters)), Filter: &filters[0]} //nolint:gosec // filter count is fixed and bounded below uint16 + //nolint:gosec // SockFprog requires the audited kernel ABI pointer for the bounded filter slice. + _, _, errno := unix.Syscall(unix.SYS_SECCOMP, unix.SECCOMP_SET_MODE_FILTER, + unix.SECCOMP_FILTER_FLAG_TSYNC, uintptr(unsafe.Pointer(&program))) + if errno != 0 { + return errno + } + return nil +} + +func buildNativeNetworkSeccompFilters(architecture uint32) ([]unix.SockFilter, error) { + if architecture != unix.AUDIT_ARCH_X86_64 && architecture != unix.AUDIT_ARCH_AARCH64 { + return nil, unix.ENOTSUP + } + filters := []unix.SockFilter{ + {Code: unix.BPF_LD | unix.BPF_W | unix.BPF_ABS, K: 4}, + {Code: unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K, Jt: 1, K: architecture}, + {Code: unix.BPF_RET | unix.BPF_K, K: unix.SECCOMP_RET_KILL_PROCESS}, + {Code: unix.BPF_LD | unix.BPF_W | unix.BPF_ABS, K: 0}, + } + if architecture == unix.AUDIT_ARCH_X86_64 { + filters = append(filters, + unix.SockFilter{Code: unix.BPF_JMP | unix.BPF_JSET | unix.BPF_K, Jf: 1, K: nativeX32SyscallBit}, + unix.SockFilter{Code: unix.BPF_RET | unix.BPF_K, K: unix.SECCOMP_RET_ERRNO | uint32(unix.EPERM)}, + ) + } + for _, syscallNumber := range nativeBlockedNetworkSyscalls() { + filters = append(filters, + unix.SockFilter{Code: unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K, Jf: 1, K: uint32(syscallNumber)}, //nolint:gosec // Linux syscall numbers are nonnegative uint32 values + unix.SockFilter{Code: unix.BPF_RET | unix.BPF_K, K: unix.SECCOMP_RET_ERRNO | uint32(unix.EPERM)}, + ) + } + filters = append(filters, unix.SockFilter{Code: unix.BPF_RET | unix.BPF_K, K: unix.SECCOMP_RET_ALLOW}) + return filters, nil +} + +func nativeAuditArchitecture() (uint32, bool) { + switch runtime.GOARCH { + case "amd64": + return unix.AUDIT_ARCH_X86_64, true + case "arm64": + return unix.AUDIT_ARCH_AARCH64, true + default: + return 0, false + } +} + +func nativeBlockedNetworkSyscalls() []uintptr { + return []uintptr{ + unix.SYS_SOCKET, unix.SYS_SOCKETPAIR, + unix.SYS_CONNECT, unix.SYS_BIND, unix.SYS_LISTEN, unix.SYS_ACCEPT, unix.SYS_ACCEPT4, + unix.SYS_SENDTO, unix.SYS_SENDMSG, unix.SYS_SENDMMSG, + unix.SYS_RECVFROM, unix.SYS_RECVMSG, unix.SYS_RECVMMSG, + unix.SYS_SHUTDOWN, unix.SYS_GETSOCKNAME, unix.SYS_GETPEERNAME, + unix.SYS_SETSOCKOPT, unix.SYS_GETSOCKOPT, + unix.SYS_IO_URING_SETUP, unix.SYS_IO_URING_ENTER, unix.SYS_IO_URING_REGISTER, + } +} diff --git a/document/trafilatura/native_runner_linux.go b/document/trafilatura/native_runner_linux.go new file mode 100644 index 00000000..67b5261d --- /dev/null +++ b/document/trafilatura/native_runner_linux.go @@ -0,0 +1,286 @@ +//go:build linux + +package trafilatura + +import ( + "bytes" + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "slices" + "sync" + "syscall" + "time" + + "golang.org/x/sys/unix" +) + +const ( + nativeRunnerIdentity = "sha256:4fb9848ec9197e12848559df3002bd39ceaa0377ae994bc8fbe9c8f082288d9e" + nativeChildDrainWindow = 250 * time.Millisecond +) + +type nativeRunner struct{} + +func newNativeRunner() (IsolatedRunner, error) { + return nativeRunner{}, nil +} + +func (nativeRunner) Identity() string { + return nativeRunnerIdentity +} + +func (runner nativeRunner) Run( + ctx context.Context, request IsolatedRunRequest, +) (IsolatedRunResult, error) { + if err := validateNativeRequest(request); err != nil { + return IsolatedRunResult{}, ErrIsolationUnavailable + } + if err := ctx.Err(); err != nil { + return IsolatedRunResult{}, errors.Join(errNativeCanceledBeforeLaunch, err) + } + executable, err := openVerifiedExecutable(ctx, request) + if err != nil { + if ctx.Err() != nil { + return IsolatedRunResult{}, errors.Join(errNativeCanceledBeforeLaunch, ctx.Err()) + } + return IsolatedRunResult{}, ErrIsolationUnavailable + } + defer func() { _ = executable.Close() }() + control, launchToken, err := openNativeLaunchControl() + if err != nil { + return IsolatedRunResult{}, ErrIsolationUnavailable + } + defer func() { _ = control.Close() }() + statusReader, statusWriter, err := os.Pipe() + if err != nil { + return IsolatedRunResult{}, ErrIsolationUnavailable + } + defer func() { _ = statusReader.Close() }() + defer func() { _ = statusWriter.Close() }() + if err := ctx.Err(); err != nil { + return IsolatedRunResult{}, errors.Join(errNativeCanceledBeforeLaunch, err) + } + + output := &nativeBoundedOutput{limit: request.MaxStdoutBytes} + command := exec.Command( //nolint:gosec // request validation fixes the executable fd and complete argument vector + "/proc/self/exe", nativeLauncherMarker, launchToken, request.Executable, + ) + command.Dir = request.Directory + command.Env = slices.Clone(request.Environment) + command.Stdin = bytes.NewReader(request.Stdin) + command.Stdout = output + command.Stderr = io.Discard + command.ExtraFiles = []*os.File{executable, control, statusWriter} + command.SysProcAttr = &syscall.SysProcAttr{ + Cloneflags: syscall.CLONE_NEWUSER | syscall.CLONE_NEWNET | + syscall.CLONE_NEWPID | syscall.CLONE_NEWNS, + UidMappings: []syscall.SysProcIDMap{{ContainerID: 0, HostID: os.Getuid(), Size: 1}}, + GidMappings: []syscall.SysProcIDMap{{ContainerID: 0, HostID: os.Getgid(), Size: 1}}, + GidMappingsEnableSetgroups: false, + Pdeathsig: syscall.SIGKILL, + } + command.WaitDelay = nativeChildDrainWindow + output.onOverflow = func() { + if command.Process != nil { + _ = command.Process.Kill() + } + } + if err := command.Start(); err != nil { + return IsolatedRunResult{}, ErrIsolationUnavailable + } + _ = statusWriter.Close() + + waited := make(chan error, 1) + go func() { waited <- command.Wait() }() + launcherStatus := make(chan bool, 1) + go func() { launcherStatus <- nativeLauncherFailed(statusReader) }() + var runErr error + select { + case runErr = <-waited: + case <-ctx.Done(): + _ = command.Process.Kill() + runErr = <-waited + } + launcherFailed := <-launcherStatus + result := IsolatedRunResult{ + Stdout: output.Bytes(), Attestation: nativeAttestation(request), + } + if output.Exceeded() { + return result, ErrChildOutputTooLarge + } + if err := ctx.Err(); err != nil { + return result, err + } + if err := classifyNativeRunError(runErr, launcherFailed); err != nil { + return result, err + } + return result, nil +} + +func nativeLauncherFailed(reader io.Reader) bool { + status, err := io.ReadAll(io.LimitReader(reader, 3)) + return err != nil || !bytes.Equal(status, []byte{nativeLauncherReadyStatus}) +} + +func classifyNativeRunError(runErr error, launcherFailed bool) error { + if launcherFailed { + return ErrIsolationUnavailable + } + if runErr != nil { + return ErrChildFailed + } + return nil +} + +func openNativeLaunchControl() (*os.File, string, error) { + token := make([]byte, nativeLauncherTokenBytes) + if _, err := rand.Read(token); err != nil { + return nil, "", fmt.Errorf("create launcher token: %w", err) + } + fd, err := unix.MemfdCreate("docbank-trafilatura-launch", unix.MFD_CLOEXEC|unix.MFD_ALLOW_SEALING) + if err != nil { + return nil, "", fmt.Errorf("create launcher control: %w", err) + } + file := os.NewFile(uintptr(fd), "docbank-trafilatura-launch") + valid := false + defer func() { + if !valid { + _ = file.Close() + } + }() + if _, err := file.Write(token); err != nil { + return nil, "", err + } + seals := unix.F_SEAL_WRITE | unix.F_SEAL_GROW | unix.F_SEAL_SHRINK | unix.F_SEAL_SEAL + if _, err := unix.FcntlInt(file.Fd(), unix.F_ADD_SEALS, seals); err != nil { + return nil, "", fmt.Errorf("seal launcher control: %w", err) + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + return nil, "", err + } + valid = true + return file, hex.EncodeToString(token), nil +} + +func validateNativeRequest(request IsolatedRunRequest) error { + stdinDigest := sha256.Sum256(request.Stdin) + if !filepath.IsAbs(request.Executable) || filepath.Clean(request.Executable) != request.Executable || + request.Directory != filepath.Dir(request.Executable) || + !slices.Equal(request.Arguments, []string{"--protocol", protocolVersion}) || + !slices.Equal(request.Environment, cleanEnvironment()) || + request.StdinSHA256 != hex.EncodeToString(stdinDigest[:]) || + request.MaxStdoutBytes <= 0 || request.MaxStdoutBytes > MaxResponseBytes || + !request.Requirements.NetworkDisabled || !request.Requirements.KillProcessTree || + !request.Requirements.VerifyExecutableSHA256 || + request.PolicyFingerprint != isolationRequestPolicyFingerprint(nativeRunnerIdentity, request) { + return errors.New("native isolation request is outside the fixed policy") + } + if err := validateSHA256(request.ExecutableSHA256, "executable SHA-256"); err != nil { + return err + } + return nil +} + +func openVerifiedExecutable(ctx context.Context, request IsolatedRunRequest) (*os.File, error) { + source, err := os.Open(request.Executable) + if err != nil { + return nil, err + } + defer func() { _ = source.Close() }() + info, err := source.Stat() + if err != nil || !info.Mode().IsRegular() || info.Size() <= 0 || info.Size() > MaxExecutableBytes { + return nil, errors.New("executable identity is outside the supported bound") + } + fd, err := unix.MemfdCreate("docbank-trafilatura", unix.MFD_CLOEXEC|unix.MFD_ALLOW_SEALING) + if err != nil { + return nil, fmt.Errorf("create sealed executable: %w", err) + } + file := os.NewFile(uintptr(fd), "docbank-trafilatura") + valid := false + defer func() { + if !valid { + _ = file.Close() + } + }() + hash := sha256.New() + written, err := io.Copy( + io.MultiWriter(file, hash), + io.LimitReader(contextReader{ctx: ctx, reader: source}, MaxExecutableBytes+1), + ) + if err != nil || written != info.Size() || hex.EncodeToString(hash.Sum(nil)) != request.ExecutableSHA256 { + return nil, errors.New("executable identity changed") + } + if err := unix.Fchmod(fd, 0o500); err != nil { + return nil, fmt.Errorf("make sealed executable runnable: %w", err) + } + seals := unix.F_SEAL_WRITE | unix.F_SEAL_GROW | unix.F_SEAL_SHRINK | unix.F_SEAL_SEAL + if _, err := unix.FcntlInt(file.Fd(), unix.F_ADD_SEALS, seals); err != nil { + return nil, fmt.Errorf("seal executable content: %w", err) + } + applied, err := unix.FcntlInt(file.Fd(), unix.F_GET_SEALS, 0) + if err != nil || applied&seals != seals { + return nil, errors.New("executable content could not be sealed") + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + return nil, err + } + valid = true + return file, nil +} + +func nativeAttestation(request IsolatedRunRequest) IsolationAttestation { + return IsolationAttestation{ + RunnerIdentity: nativeRunnerIdentity, PolicyFingerprint: request.PolicyFingerprint, + ExecutableSHA256: request.ExecutableSHA256, StdinSHA256: request.StdinSHA256, + NetworkDisabled: true, ProcessTreeContained: true, DigestVerifiedLaunch: true, + } +} + +type nativeBoundedOutput struct { + mu sync.Mutex + data []byte + limit int64 + exceeded bool + overflow sync.Once + onOverflow func() +} + +func (output *nativeBoundedOutput) Write(value []byte) (int, error) { + output.mu.Lock() + remaining := output.limit - int64(len(output.data)) + if remaining > 0 { + kept := min(int64(len(value)), remaining) + output.data = append(output.data, value[:kept]...) + } + exceeded := int64(len(value)) > remaining + if exceeded { + output.exceeded = true + } + output.mu.Unlock() + if exceeded { + output.overflow.Do(output.onOverflow) + } + return len(value), nil +} + +func (output *nativeBoundedOutput) Bytes() []byte { + output.mu.Lock() + defer output.mu.Unlock() + return slices.Clone(output.data) +} + +func (output *nativeBoundedOutput) Exceeded() bool { + output.mu.Lock() + defer output.mu.Unlock() + return output.exceeded +} + +var _ IsolatedRunner = nativeRunner{} diff --git a/document/trafilatura/native_runner_linux_test.go b/document/trafilatura/native_runner_linux_test.go new file mode 100644 index 00000000..4130aceb --- /dev/null +++ b/document/trafilatura/native_runner_linux_test.go @@ -0,0 +1,389 @@ +//go:build linux + +package trafilatura + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "net" + "os" + "os/exec" + "path/filepath" + "runtime" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" +) + +type nativeRunOutcome struct { + result IsolatedRunResult + err error +} + +func TestNativeRunnerUsesExactStdinArgumentsAndCleanEnvironment(t *testing.T) { + executable := buildIsolatedHelper(t, "echo", "") + runner, err := newNativeRunner() + require.NoError(t, err) + stdin := []byte("exact supplied bytes\x00remain data, never arguments") + request := nativeTestRequest(t, runner, executable, stdin, 1<<20) + + result, err := runner.Run(t.Context(), request) + skipUnavailableNativeIsolation(t, err) + require.NoError(t, err) + + var response struct { + Arguments []string `json:"arguments"` + Environment []string `json:"environment"` + StdinSHA256 string `json:"stdin_sha256"` + } + require.NoError(t, json.Unmarshal(result.Stdout, &response)) + digest := sha256.Sum256(stdin) + assert.Equal(t, hex.EncodeToString(digest[:]), response.StdinSHA256) + assert.Equal(t, []string{"--protocol", protocolVersion}, response.Arguments) + assert.Equal(t, cleanEnvironment(), response.Environment) + assert.Equal(t, nativeRunnerIdentity, result.Attestation.RunnerIdentity) + assert.True(t, result.Attestation.NetworkDisabled) + assert.True(t, result.Attestation.ProcessTreeContained) + assert.True(t, result.Attestation.DigestVerifiedLaunch) +} + +func TestNativeRunnerDeniesLoopbackNetworkAccess(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer func() { _ = listener.Close() }() + executable := buildIsolatedHelper(t, "network", listener.Addr().String()) + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("network probe"), 1<<20) + + result, err := runner.Run(t.Context(), request) + skipUnavailableNativeIsolation(t, err) + require.NoError(t, err) + assert.Equal(t, "denied", string(result.Stdout)) +} + +func TestNativeRunnerDeniesHostPathnameUnixSocketAccess(t *testing.T) { + socketPath := filepath.Join(t.TempDir(), "host.sock") + listener, err := net.Listen("unix", socketPath) + require.NoError(t, err) + defer func() { _ = listener.Close() }() + executable := buildIsolatedHelper(t, "unix-network", socketPath) + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("unix network probe"), 1<<20) + + result, err := runner.Run(t.Context(), request) + skipUnavailableNativeIsolation(t, err) + require.NoError(t, err) + assert.Equal(t, "denied", string(result.Stdout)) +} + +func TestNativeSeccompDeniesPathnameUnixSockets(t *testing.T) { + const helperEnvironment = "DOCBANK_TEST_TRAFILATURA_SECCOMP" + if socketPath := os.Getenv(helperEnvironment); socketPath != "" { + if err := unix.Prctl(unix.PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0); err != nil { + os.Exit(10) + } + if err := installNativeNetworkSeccomp(); err != nil { + os.Exit(11) + } + connection, err := net.DialTimeout("unix", socketPath, time.Second) + if connection != nil { + _ = connection.Close() + } + if !errors.Is(err, unix.EPERM) { + os.Exit(12) + } + os.Exit(0) + } + + socketPath := filepath.Join(t.TempDir(), "seccomp.sock") + listener, err := net.Listen("unix", socketPath) + require.NoError(t, err) + defer func() { _ = listener.Close() }() + connection, err := net.DialTimeout("unix", socketPath, time.Second) + require.NoError(t, err, "the host listener must be reachable before filtering") + require.NoError(t, connection.Close()) + command := exec.Command( //nolint:gosec // os.Args[0] is the trusted current test executable + os.Args[0], "-test.run=^TestNativeSeccompDeniesPathnameUnixSockets$", + ) + command.Env = append(os.Environ(), helperEnvironment+"="+socketPath) + output, err := command.CombinedOutput() + require.NoError(t, err, "%s", output) +} + +func TestNativeSeccompFilterDeniesX32ABI(t *testing.T) { + if runtime.GOARCH != "amd64" { + t.Skip("x32 ABI exists only on amd64") + } + filters, err := buildNativeNetworkSeccompFilters(unix.AUDIT_ARCH_X86_64) + require.NoError(t, err) + denied := unix.SECCOMP_RET_ERRNO | uint32(unix.EPERM) + for name, syscallNumber := range map[string]uint32{ + "socket": unix.SYS_SOCKET, + "connect": unix.SYS_CONNECT, + "io_uring": unix.SYS_IO_URING_SETUP, + } { + t.Run(name, func(t *testing.T) { + assert.Equal(t, denied, evaluateNativeSeccomp(t, filters, syscallNumber)) + assert.Equal(t, denied, evaluateNativeSeccomp(t, filters, syscallNumber|nativeX32SyscallBit)) + }) + } + assert.Equal(t, uint32(unix.SECCOMP_RET_ALLOW), + evaluateNativeSeccomp(t, filters, unix.SYS_GETPID)) + assert.Equal(t, denied, + evaluateNativeSeccomp(t, filters, unix.SYS_GETPID|nativeX32SyscallBit)) +} + +func TestNativeLauncherRequiresSealedMatchingInheritedControl(t *testing.T) { + control, token, err := openNativeLaunchControl() + require.NoError(t, err) + defer func() { _ = control.Close() }() + arguments := []string{"/proc/self/exe", nativeLauncherMarker, token, "/opt/trafilatura-bridge"} + + executable, authenticated := authenticatedNativeLaunch(arguments, int(control.Fd())) + assert.True(t, authenticated) + assert.Equal(t, "/opt/trafilatura-bridge", executable) + + _, authenticated = authenticatedNativeLaunch(arguments, -1) + assert.False(t, authenticated) + arguments[2] = strings.Repeat("0", nativeLauncherTokenBytes*2) + _, authenticated = authenticatedNativeLaunch(arguments, int(control.Fd())) + assert.False(t, authenticated) +} + +func TestNativeRunnerCancellationReapsDescendantProcessTree(t *testing.T) { + lockPath := filepath.Join(t.TempDir(), "descendant.lock") + require.NoError(t, os.WriteFile(lockPath, nil, 0o600)) + executable := buildIsolatedHelper(t, "descendant", lockPath) + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("descendant probe"), 1<<20) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + finished := make(chan nativeRunOutcome, 1) + started := time.Now() + go func() { + result, runErr := runner.Run(ctx, request) + finished <- nativeRunOutcome{result: result, err: runErr} + }() + + observeDescendantLockHeld(t, lockPath, finished) + cancel() + outcome := <-finished + require.ErrorIs(t, outcome.err, context.Canceled) + assert.Less(t, time.Since(started), 2*time.Second) + assert.True(t, outcome.result.Attestation.ProcessTreeContained) + require.Eventually(t, func() bool { return exclusiveLockAvailable(t, lockPath) }, + time.Second, 10*time.Millisecond, "descendant must release its inherited host lock") +} + +func TestNativeRunnerTerminatesPromptlyOnStdoutOverflow(t *testing.T) { + executable := buildIsolatedHelper(t, "overflow", "") + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("overflow probe"), 1024) + started := time.Now() + + result, err := runner.Run(t.Context(), request) + skipUnavailableNativeIsolation(t, err) + require.ErrorIs(t, err, ErrChildOutputTooLarge) + assert.Less(t, time.Since(started), 2*time.Second) + assert.LessOrEqual(t, int64(len(result.Stdout)), request.MaxStdoutBytes) +} + +func TestNativeRunnerDoesNotMisclassifyBridgeExit125AsLauncherFailure(t *testing.T) { + executable := buildIsolatedHelper(t, "exit-125", "") + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("exit probe"), 1<<20) + + _, err = runner.Run(t.Context(), request) + skipUnavailableNativeIsolation(t, err) + require.ErrorIs(t, err, ErrChildFailed) + require.NotErrorIs(t, err, ErrIsolationUnavailable) +} + +func TestNativeRunErrorClassificationUsesOutOfBandLauncherStatus(t *testing.T) { + runErr := exec.Command("/bin/sh", "-c", "exit 125").Run() + require.Error(t, runErr) + require.ErrorIs(t, classifyNativeRunError(runErr, false), ErrChildFailed) + require.ErrorIs(t, classifyNativeRunError(runErr, true), ErrIsolationUnavailable) +} + +func TestNativeRunnerHonorsCancellationBeforeExecutablePreparation(t *testing.T) { + executable := buildIsolatedHelper(t, "echo", "") + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("canceled probe"), 1<<20) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = runner.Run(ctx, request) + require.ErrorIs(t, err, context.Canceled) +} + +func TestNativeRunnerNeverLaunchesExecutableContentOutsidePinnedDigest(t *testing.T) { + executable := buildIsolatedHelper(t, "echo", "") + replacement := buildIsolatedHelper(t, "replacement", "") + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("identity probe"), 1<<20) + require.NoError(t, os.Rename(replacement, executable)) + + _, err = runner.Run(t.Context(), request) + require.ErrorIs(t, err, ErrIsolationUnavailable) +} + +func TestVerifiedExecutableContentCannotChangeAfterDigestVerification(t *testing.T) { + executable := buildIsolatedHelper(t, "echo", "") + replacement := buildIsolatedHelper(t, "replacement", "") + runner, err := newNativeRunner() + require.NoError(t, err) + request := nativeTestRequest(t, runner, executable, []byte("identity probe"), 1<<20) + want, err := os.ReadFile(executable) + require.NoError(t, err) + prepared, err := openVerifiedExecutable(t.Context(), request) + require.NoError(t, err) + defer func() { _ = prepared.Close() }() + replacementBytes, err := os.ReadFile(replacement) + require.NoError(t, err) + require.NoError(t, os.WriteFile(executable, replacementBytes, 0o700)) + + got, err := os.ReadFile("/proc/self/fd/" + strconv.FormatUint(uint64(prepared.Fd()), 10)) + require.NoError(t, err) + wantDigest := sha256.Sum256(want) + gotDigest := sha256.Sum256(got) + assert.Equal(t, wantDigest, gotDigest) +} + +func TestNewUsesNativeRunnerWhenNoneIsInjected(t *testing.T) { + profile := testProfile(t, helperExecutable(t, "complete"), time.Second, 1<<20) + profile.Runner = nil + + provider, err := New(profile) + require.NoError(t, err) + assert.Equal(t, nativeRunnerIdentity, provider.runnerIdentity) +} + +func nativeTestRequest( + t *testing.T, runner IsolatedRunner, executable string, stdin []byte, maxStdout int64, +) IsolatedRunRequest { + t.Helper() + data, err := os.ReadFile(executable) + require.NoError(t, err) + digest := sha256.Sum256(data) + provider := &Provider{ + executable: executable, executableSHA256: hex.EncodeToString(digest[:]), + runnerIdentity: runner.Identity(), environment: cleanEnvironment(), + } + return provider.isolatedRequest(stdin, maxStdout) +} + +func buildIsolatedHelper(t *testing.T, mode, networkAddress string) string { + t.Helper() + target := filepath.Join(t.TempDir(), "isolated-helper") + ldflags := strings.Join([]string{ + "-X=main.mode=" + mode, + "-X=main.networkAddress=" + networkAddress, + }, " ") + command := exec.Command("go", "build", "-trimpath", "-ldflags", ldflags, + "-o", target, "./testdata/isolatedhelper") + output, err := command.CombinedOutput() + require.NoError(t, err, "%s", output) + return target +} + +func skipUnavailableNativeIsolation(t *testing.T, err error) { + t.Helper() + if errors.Is(err, ErrIsolationUnavailable) { + t.Skipf("native Linux namespace isolation unavailable: %v", err) + } +} + +func observeDescendantLockHeld( + t *testing.T, lockPath string, finished <-chan nativeRunOutcome, +) { + t.Helper() + deadline := time.NewTimer(2 * time.Second) + defer deadline.Stop() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + select { + case outcome := <-finished: + skipUnavailableNativeIsolation(t, outcome.err) + require.FailNow(t, "isolated runner exited before descendant held its lock", "%v", outcome.err) + case <-ticker.C: + owner, err := os.ReadFile(lockPath) + require.NoError(t, err) + if !exclusiveLockAvailable(t, lockPath) && string(owner) == "descendant-ready\n" { + return + } + case <-deadline.C: + require.FailNow(t, "descendant did not acquire its host-visible lock") + } + } +} + +func exclusiveLockAvailable(t *testing.T, lockPath string) bool { + t.Helper() + file, err := os.OpenFile(lockPath, os.O_RDWR, 0) + require.NoError(t, err) + defer func() { _ = file.Close() }() + err = unix.Flock(int(file.Fd()), unix.LOCK_EX|unix.LOCK_NB) + if errors.Is(err, unix.EWOULDBLOCK) { + return false + } + require.NoError(t, err) + require.NoError(t, unix.Flock(int(file.Fd()), unix.LOCK_UN)) + return true +} + +func evaluateNativeSeccomp( + t *testing.T, filters []unix.SockFilter, syscallNumber uint32, +) uint32 { + t.Helper() + accumulator := uint32(0) + for programCounter, steps := 0, 0; programCounter < len(filters) && steps <= len(filters); steps++ { + instruction := filters[programCounter] + switch instruction.Code { + case unix.BPF_LD | unix.BPF_W | unix.BPF_ABS: + switch instruction.K { + case 0: + accumulator = syscallNumber + case 4: + accumulator = unix.AUDIT_ARCH_X86_64 + default: + require.FailNow(t, "unexpected seccomp load offset", "%d", instruction.K) + } + programCounter++ + case unix.BPF_JMP | unix.BPF_JEQ | unix.BPF_K: + if accumulator == instruction.K { + programCounter += int(instruction.Jt) + 1 + } else { + programCounter += int(instruction.Jf) + 1 + } + case unix.BPF_JMP | unix.BPF_JSET | unix.BPF_K: + if accumulator&instruction.K != 0 { + programCounter += int(instruction.Jt) + 1 + } else { + programCounter += int(instruction.Jf) + 1 + } + case unix.BPF_RET | unix.BPF_K: + return instruction.K + default: + require.FailNow(t, "unexpected seccomp instruction", "%#x", instruction.Code) + } + } + require.FailNow(t, "seccomp filter did not return") + return 0 +} diff --git a/document/trafilatura/native_runner_unsupported.go b/document/trafilatura/native_runner_unsupported.go new file mode 100644 index 00000000..356e46ad --- /dev/null +++ b/document/trafilatura/native_runner_unsupported.go @@ -0,0 +1,7 @@ +//go:build !linux + +package trafilatura + +func newNativeRunner() (IsolatedRunner, error) { + return nil, ErrIsolationUnavailable +} diff --git a/document/trafilatura/provider.go b/document/trafilatura/provider.go new file mode 100644 index 00000000..e683a6e9 --- /dev/null +++ b/document/trafilatura/provider.go @@ -0,0 +1,939 @@ +package trafilatura + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "encoding/xml" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "runtime" + "slices" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "golang.org/x/net/html" + + "go.kenn.io/docbank/document" +) + +const ( + providerID = "trafilatura.local-v1" + protocolVersion = "docbank-trafilatura/v2" + profileVersion = "docbank-trafilatura-profile/v2" + timestampForm = "2006-01-02T15:04:05.000000000Z" + + // MaxDocumentBytes is the largest supplied HTML document accepted by a profile. + MaxDocumentBytes = int64(50 << 20) + // MaxResponseBytes is the largest structured child response accepted by a profile. + MaxResponseBytes = int64(64 << 20) + // MaxUnits is the largest extracted unit sequence accepted by a profile. + MaxUnits = 100_000 + // MaxTimeout is the largest local child deadline accepted by a profile. + MaxTimeout = 30 * time.Minute + // MaxExecutableBytes bounds executable identity verification. + MaxExecutableBytes = int64(256 << 20) +) + +var ( + errInputIdentity = errors.New("authorized input identity changed") + errNativeCanceledBeforeLaunch = errors.New("native isolated runner canceled before launch") + // ErrIsolationUnavailable means the runner could not enforce the requested isolation policy. + ErrIsolationUnavailable = errors.New("isolated runner policy unavailable") + // ErrChildOutputTooLarge means the isolated child exceeded its stdout allowance. + ErrChildOutputTooLarge = errors.New("isolated child output exceeds limit") + // ErrChildFailed means the isolated child exited without a valid response. + ErrChildFailed = errors.New("isolated child failed") +) + +// IsolationRequirements are mandatory runner controls. A runner must fail +// closed rather than execute when any requested control is unavailable. +type IsolationRequirements struct { + NetworkDisabled bool + KillProcessTree bool + VerifyExecutableSHA256 bool +} + +// IsolatedRunRequest is the complete, fixed child execution authority. +type IsolatedRunRequest struct { + Executable string + ExecutableSHA256 string + Arguments []string + Environment []string + Directory string + Stdin []byte + StdinSHA256 string + MaxStdoutBytes int64 + PolicyFingerprint string + Requirements IsolationRequirements +} + +// IsolationAttestation reports the exact controls applied to a completed run. +type IsolationAttestation struct { + RunnerIdentity string + PolicyFingerprint string + ExecutableSHA256 string + StdinSHA256 string + NetworkDisabled bool + ProcessTreeContained bool + DigestVerifiedLaunch bool +} + +// IsolatedRunResult is bounded stdout plus its isolation attestation. +type IsolatedRunResult struct { + Stdout []byte + Attestation IsolationAttestation +} + +// IsolatedRunner is the trusted cross-platform process isolation boundary. +// Run must launch the digest-verified executable without a path re-open race, +// deny all network access, contain the process tree, and reap that tree on +// cancellation. It must return ErrIsolationUnavailable rather than weaken a +// requested control. +type IsolatedRunner interface { + Identity() string + Run(ctx context.Context, request IsolatedRunRequest) (IsolatedRunResult, error) +} + +// Profile fixes one executable, immutable runtime identity, and all local bounds. +type Profile struct { + Executable string + ExecutableSHA256 string + RuntimeIdentity string + Runner IsolatedRunner + MaxDocumentBytes int64 + MaxResponseBytes int64 + MaxUnits int + Timeout time.Duration +} + +// Provider renders exact, caller-supplied HTML bytes through one local bridge. +type Provider struct { + descriptor document.RenditionDescriptor + executable string + executableSHA256 string + runtimeIdentity string + runnerIdentity string + runner IsolatedRunner + environment []string + maxDocumentBytes int64 + maxResponseBytes int64 + maxUnits int + timeout time.Duration +} + +// New constructs one immutable local-process provider profile. +func New(profile Profile) (*Provider, error) { + if !filepath.IsAbs(profile.Executable) || filepath.Clean(profile.Executable) != profile.Executable { + return nil, errors.New("trafilatura: executable must be an absolute clean path") + } + if pythonInterpreter(filepath.Base(profile.Executable)) { + return nil, errors.New("trafilatura: executable must not be a Python interpreter") + } + info, err := os.Lstat(profile.Executable) + if err != nil { + return nil, errors.New("trafilatura: configured executable is unavailable") + } + if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 { + return nil, errors.New("trafilatura: executable must be a regular non-symlink file") + } + if info.Size() <= 0 || info.Size() > MaxExecutableBytes { + return nil, errors.New("trafilatura: executable size is outside the supported bound") + } + if err := validateSHA256(profile.ExecutableSHA256, "executable SHA-256"); err != nil { + return nil, err + } + executableDigest, err := hashExecutable(profile.Executable) + if err != nil || executableDigest != profile.ExecutableSHA256 { + return nil, errors.New("trafilatura: executable SHA-256 does not match configured content") + } + if err := validateImmutableIdentity(profile.RuntimeIdentity, "runtime identity"); err != nil { + return nil, err + } + runner := profile.Runner + if runner == nil { + runner, err = newNativeRunner() + if err != nil { + return nil, fmt.Errorf("trafilatura: native isolated runner: %w", err) + } + } + runnerIdentity := runner.Identity() + if err := validateImmutableIdentity(runnerIdentity, "runner identity"); err != nil { + return nil, err + } + if profile.MaxDocumentBytes <= 0 || profile.MaxDocumentBytes > MaxDocumentBytes { + return nil, fmt.Errorf("trafilatura: max document bytes must be between 1 and %d", MaxDocumentBytes) + } + if profile.MaxResponseBytes <= 0 || profile.MaxResponseBytes > MaxResponseBytes { + return nil, fmt.Errorf("trafilatura: max response bytes must be between 1 and %d", MaxResponseBytes) + } + if profile.MaxUnits <= 0 || profile.MaxUnits > MaxUnits { + return nil, fmt.Errorf("trafilatura: max units must be between 1 and %d", MaxUnits) + } + if profile.Timeout <= 0 || profile.Timeout > MaxTimeout { + return nil, fmt.Errorf("trafilatura: timeout must be between 1ns and %s", MaxTimeout) + } + environment := cleanEnvironment() + identity := strings.Join([]string{ + profileVersion, protocolVersion, profile.Executable, profile.ExecutableSHA256, + profile.RuntimeIdentity, runnerIdentity, strings.Join(environment, "\x1f"), + strconv.FormatInt(profile.MaxDocumentBytes, 10), strconv.FormatInt(profile.MaxResponseBytes, 10), + strconv.Itoa(profile.MaxUnits), strconv.FormatInt(int64(profile.Timeout), 10), + }, "\x00") + policyDigest := sha256.Sum256([]byte(identity)) + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: providerID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: hex.EncodeToString(policyDigest[:]), + TrustBoundary: document.RenditionTrustLocalProcess, + SupportedFormats: []document.RenditionFormatCapability{ + {MediaFamily: "text", MediaType: "text/html", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "application/xhtml+xml", InputKind: document.RenditionInputOriginalFile}, + }, + ReturnsStructured: true, + }) + if err != nil { + return nil, fmt.Errorf("trafilatura: construct descriptor: %w", err) + } + return &Provider{ + descriptor: cloneDescriptor(descriptor), executable: profile.Executable, + executableSHA256: profile.ExecutableSHA256, runtimeIdentity: profile.RuntimeIdentity, + runnerIdentity: runnerIdentity, runner: runner, maxDocumentBytes: profile.MaxDocumentBytes, + environment: slices.Clone(environment), maxResponseBytes: profile.MaxResponseBytes, + maxUnits: profile.MaxUnits, timeout: profile.Timeout, + }, nil +} + +// Descriptor returns the immutable provider identity fixed by the profile. +func (provider *Provider) Descriptor() document.RenditionDescriptor { + if provider == nil { + return document.RenditionDescriptor{} + } + return cloneDescriptor(provider.descriptor) +} + +// Render validates supplied HTML locally and sends its exact bytes over stdin only. +func (provider *Provider) Render( + ctx context.Context, upload document.AuthorizedUpload, + authorization document.RenditionAuthorization, +) (document.RenditionResult, error) { + if provider == nil { + return document.RenditionResult{}, errors.New("trafilatura: provider is required") + } + if _, err := document.ValidateRenditionProviderRequest(provider, upload, authorization); err != nil { + return document.RenditionResult{}, err + } + metadata := upload.Metadata() + if metadata.ByteLength > provider.maxDocumentBytes { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura input exceeds the configured byte limit", nil) + } + if !validFilename(metadata.Filename, metadata.MediaType) { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, + "Trafilatura requires a supplied HTML filename", nil) + } + expiresAt, err := time.Parse(timestampForm, authorization.ExpiresAt) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura authorization is invalid", err) + } + startedAt := time.Now().UTC() + deadline := startedAt.Add(provider.timeout) + expiryDeadline := false + if expiresAt.Before(deadline) { + deadline = expiresAt + expiryDeadline = true + } + operationCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + if err := operationCtx.Err(); err != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, err) + } + + source, err := readAuthorizedExact( + operationCtx, upload, metadata.ByteLength, provider.maxDocumentBytes, + ) + if err != nil { + if operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + if errors.Is(err, errInputIdentity) { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura input identity does not match authorization", err) + } + return document.RenditionResult{}, providerError(document.RenditionErrorTransient, + "Trafilatura input could not be read", err) + } + defer clear(source) + digest := sha256.Sum256(source) + if int64(len(source)) != metadata.ByteLength || hex.EncodeToString(digest[:]) != metadata.SHA256 { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura input identity does not match authorization", nil) + } + authority, err := inspectHTML(source, metadata.MediaType, provider.maxUnits) + if err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorUnsupportedInput, + "Trafilatura input is not locally verified HTML", err) + } + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + + if err := provider.reverifyExecutionBoundary(operationCtx); err != nil { + if operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura isolated runtime identity changed", err) + } + runnerInput := slices.Clone(source) + defer clear(runnerInput) + stdoutLimit := min(provider.maxResponseBytes, int64(authorization.MaxTotalResultBytes)) + request := provider.isolatedRequest(runnerInput, stdoutLimit) + runResult, runErr := provider.runner.Run(operationCtx, request) + defer func() { clear(runResult.Stdout) }() + if errors.Is(runErr, errNativeCanceledBeforeLaunch) && operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + if errors.Is(runErr, ErrIsolationUnavailable) { + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura isolation policy could not be enforced", runErr) + } + if err := provider.validateAttestation(request, runResult.Attestation); err != nil { + clear(runResult.Stdout) + return document.RenditionResult{}, providerError(document.RenditionErrorPolicyRejected, + "Trafilatura isolation attestation is invalid", err) + } + if operationCtx.Err() != nil { + return document.RenditionResult{}, provider.contextError(ctx, expiryDeadline, operationCtx.Err()) + } + if errors.Is(runErr, ErrChildOutputTooLarge) { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "Trafilatura output exceeds the configured byte limit", nil) + } + if runErr != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorTransient, + "Trafilatura executable failed", runErr) + } + if int64(len(runResult.Stdout)) > request.MaxStdoutBytes { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "Trafilatura output exceeds the configured byte limit", nil) + } + raw := runResult.Stdout + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + + wire, err := parseResponse(operationCtx, raw, provider.runtimeIdentity, metadata, authority, provider.maxUnits) + if err != nil { + if operationErr := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); operationErr != nil { + return document.RenditionResult{}, operationErr + } + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "Trafilatura output is malformed", err) + } + evidence := evidenceFromResponse(wire) + if err := document.ValidateSourceEvidenceV1(evidence); err != nil { + return document.RenditionResult{}, providerError(document.RenditionErrorMalformedEvidence, + "Trafilatura output is malformed", err) + } + if err := provider.postProcessError(ctx, operationCtx, deadline, expiryDeadline); err != nil { + return document.RenditionResult{}, err + } + completedAt := time.Now().UTC() + warnings := []string(nil) + if !*wire.ProvenanceComplete { + warnings = []string{"degraded_provenance"} + } + return document.RenditionResult{ + Evidence: evidence, + Receipt: document.RenditionReceipt{ + ProviderID: provider.descriptor.ID, DescriptorFingerprint: provider.descriptor.Fingerprint, + PolicyFingerprint: authorization.PolicyFingerprint, SourceSHA256: metadata.SHA256, + OperationID: "trafilatura-" + authorization.RenditionRequestFingerprint[:24], + StartedAt: startedAt.Format(timestampForm), CompletedAt: completedAt.Format(timestampForm), + Warnings: warnings, + Usage: document.RenditionUsage{Requests: 1, InputBytes: metadata.ByteLength, + OutputBytes: int64(len(raw)), Units: int64(len(evidence.Units))}, + }, + }, nil +} + +type response struct { + ContractVersion string `json:"contract_version"` + RuntimeIdentity string `json:"runtime_identity"` + SourceSHA256 string `json:"source_sha256"` + SourceBytes int64 `json:"source_bytes"` + ExtractionComplete bool `json:"extraction_complete"` + ProvenanceComplete *bool `json:"provenance_complete"` + Units []responseUnit `json:"units"` +} + +type responseUnit struct { + SourcePath string `json:"source_path,omitempty"` + Heading string `json:"heading,omitempty"` + Text string `json:"text"` +} + +type htmlAuthority struct { + visibleTokens []string + sections []htmlSectionAuthority +} + +type htmlSectionAuthority struct { + path string + heading string + tokens []string +} + +func (provider *Provider) isolatedRequest(stdin []byte, maxStdoutBytes int64) IsolatedRunRequest { + arguments := []string{"--protocol", protocolVersion} + environment := slices.Clone(provider.environment) + requirements := IsolationRequirements{ + NetworkDisabled: true, KillProcessTree: true, VerifyExecutableSHA256: true, + } + stdinDigest := sha256.Sum256(stdin) + stdinSHA256 := hex.EncodeToString(stdinDigest[:]) + request := IsolatedRunRequest{ + Executable: provider.executable, ExecutableSHA256: provider.executableSHA256, + Arguments: arguments, Environment: environment, Directory: filepath.Dir(provider.executable), + Stdin: stdin, StdinSHA256: stdinSHA256, MaxStdoutBytes: maxStdoutBytes, + Requirements: requirements, + } + request.PolicyFingerprint = isolationRequestPolicyFingerprint(provider.runnerIdentity, request) + return request +} + +func isolationRequestPolicyFingerprint(runnerIdentity string, request IsolatedRunRequest) string { + identity := strings.Join([]string{ + "docbank-isolated-run/v1", runnerIdentity, request.Executable, + request.ExecutableSHA256, strings.Join(request.Arguments, "\x1f"), strings.Join(request.Environment, "\x1f"), + request.Directory, request.StdinSHA256, strconv.FormatInt(request.MaxStdoutBytes, 10), + "network-disabled", "kill-process-tree", "digest-verified-launch", + }, "\x00") + digest := sha256.Sum256([]byte(identity)) + return hex.EncodeToString(digest[:]) +} + +func (provider *Provider) reverifyExecutionBoundary(ctx context.Context) error { + if provider.runner == nil || provider.runner.Identity() != provider.runnerIdentity { + return errors.New("isolated runner identity changed") + } + digest, err := hashExecutableContext(ctx, provider.executable) + if err != nil || digest != provider.executableSHA256 { + return errors.New("executable content changed") + } + return nil +} + +func (provider *Provider) validateAttestation( + request IsolatedRunRequest, attestation IsolationAttestation, +) error { + stdinDigest := sha256.Sum256(request.Stdin) + if attestation.RunnerIdentity != provider.runnerIdentity || + attestation.PolicyFingerprint != request.PolicyFingerprint || + attestation.ExecutableSHA256 != provider.executableSHA256 || + request.StdinSHA256 != hex.EncodeToString(stdinDigest[:]) || attestation.StdinSHA256 != request.StdinSHA256 || + !attestation.NetworkDisabled || !attestation.ProcessTreeContained || !attestation.DigestVerifiedLaunch { + return errors.New("isolated runner did not attest the exact required policy") + } + return nil +} + +func parseResponse( + ctx context.Context, raw []byte, runtimeIdentity string, + metadata document.AuthorizedUploadMetadata, authority htmlAuthority, maxUnits int, +) (response, error) { + if len(raw) == 0 { + return response{}, errors.New("empty response") + } + var wire response + if err := json.Unmarshal(raw, &wire, json.RejectUnknownMembers(true)); err != nil { + return response{}, err + } + if wire.ContractVersion != protocolVersion || wire.RuntimeIdentity != runtimeIdentity { + return response{}, errors.New("local protocol identity changed") + } + if wire.SourceSHA256 != metadata.SHA256 || wire.SourceBytes != metadata.ByteLength { + return response{}, errors.New("source identity changed") + } + if !wire.ExtractionComplete { + return response{}, errors.New("response does not attest complete extraction") + } + if wire.ProvenanceComplete == nil { + return response{}, errors.New("response does not attest provenance completeness") + } + if len(wire.Units) == 0 || len(wire.Units) > maxUnits { + return response{}, errors.New("unit count is invalid") + } + outputTokens := make([]string, 0) + for index := range wire.Units { + if err := ctx.Err(); err != nil { + return response{}, err + } + unit := &wire.Units[index] + if !validOutputText(unit.Text) { + return response{}, errors.New("unit text is invalid") + } + unit.Text = strings.TrimSpace(unit.Text) + unitTokens := strings.Fields(unit.Text) + outputTokens = append(outputTokens, unitTokens...) + } + if !slices.Equal(outputTokens, authority.visibleTokens) { + return response{}, errors.New("output does not cover exact supplied HTML text") + } + if !*wire.ProvenanceComplete { + for _, unit := range wire.Units { + if unit.SourcePath != "" || unit.Heading != "" { + return response{}, errors.New("degraded output claims unverified natural provenance") + } + } + return wire, nil + } + if len(authority.sections) == 0 || len(wire.Units) != len(authority.sections) { + return response{}, errors.New("complete section structure is not locally verified") + } + for index, unit := range wire.Units { + section := authority.sections[index] + if unit.SourcePath != section.path || strings.Join(strings.Fields(unit.Heading), " ") != section.heading || + !slices.Equal(strings.Fields(unit.Text), section.tokens) { + return response{}, errors.New("section output does not match locally verified HTML structure") + } + } + return wire, nil +} + +func evidenceFromResponse(wire response) document.SourceEvidenceV1 { + if !*wire.ProvenanceComplete { + parts := make([]string, len(wire.Units)) + for index, unit := range wire.Units { + parts[index] = unit.Text + } + return document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, + Completeness: document.EvidenceDegradedProvenance, + Family: "text", UnitKind: document.EvidenceUnitGeneric, + Omissions: []document.SourceEvidenceOmissionV1{{ + Kind: document.EvidenceOmissionField, Field: "natural_structure", + Reason: "All supplied text is extracted but natural HTML section boundaries are not proven", + }}, + Units: []document.SourceEvidenceUnitV1{{Order: 0, Text: strings.Join(parts, "\n\n"), + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorGeneric, + IndexOrigin: document.EvidenceIndexOriginNone}}}, + } + } + units := make([]document.SourceEvidenceUnitV1, len(wire.Units)) + for index, unit := range wire.Units { + units[index] = document.SourceEvidenceUnitV1{ + Order: index, ProviderID: fmt.Sprintf("trafilatura-section-%d", index+1), Text: unit.Text, + Locator: document.SourceEvidenceLocatorV1{Kind: document.EvidenceLocatorSection, + IndexOrigin: document.EvidenceIndexOriginNone, Name: unit.Heading}, + } + } + return document.SourceEvidenceV1{ + ContractVersion: document.SourceEvidenceContractV1, Completeness: document.EvidenceComplete, + Family: "text", UnitKind: document.EvidenceUnitSection, Units: units, + } +} + +func inspectHTML(source []byte, mediaType string, maxSections int) (htmlAuthority, error) { + if len(source) == 0 || !utf8.Valid(source) || bytes.IndexByte(source, 0) >= 0 { + return htmlAuthority{}, errors.New("HTML must be nonempty safe UTF-8") + } + if mediaType == "application/xhtml+xml" { + if err := validateXHTML(source); err != nil { + return htmlAuthority{}, err + } + } else if mediaType != "text/html" || !hasHTMLStructure(source) { + return htmlAuthority{}, errors.New("HTML structure is missing") + } + node, err := html.Parse(bytes.NewReader(source)) + if err != nil { + return htmlAuthority{}, fmt.Errorf("parse supplied HTML: %w", err) + } + if err := validateVisibleText(node); err != nil { + return htmlAuthority{}, err + } + visible := visibleTokens(node) + if len(visible) == 0 { + return htmlAuthority{}, errors.New("HTML has no visible supplied text") + } + sections := naturalSections(node, maxSections) + sectionTokens := make([]string, 0, len(visible)) + for _, section := range sections { + sectionTokens = append(sectionTokens, section.tokens...) + } + if !slices.Equal(sectionTokens, visible) { + sections = nil + } + return htmlAuthority{visibleTokens: visible, sections: sections}, nil +} + +func validateVisibleText(node *html.Node) error { + var walk func(*html.Node, bool) error + walk = func(current *html.Node, hidden bool) error { + if current.Type == html.ElementNode { + hidden = hidden || hiddenElement(strings.ToLower(current.Data)) + } + if current.Type == html.TextNode && !hidden { + for _, char := range current.Data { + if unicode.IsControl(char) && char != '\n' && char != '\r' && char != '\t' { + return errors.New("HTML visible text contains an unsupported control character") + } + } + } + for child := current.FirstChild; child != nil; child = child.NextSibling { + if err := walk(child, hidden); err != nil { + return err + } + } + return nil + } + return walk(node, false) +} + +func visibleTokens(node *html.Node) []string { + var tokens []string + var walk func(*html.Node, bool) + walk = func(current *html.Node, hidden bool) { + if current.Type == html.ElementNode { + hidden = hidden || hiddenElement(strings.ToLower(current.Data)) + } + if current.Type == html.TextNode && !hidden { + tokens = append(tokens, strings.Fields(current.Data)...) + } + for child := current.FirstChild; child != nil; child = child.NextSibling { + walk(child, hidden) + } + } + walk(node, false) + return tokens +} + +func hiddenElement(tag string) bool { + return tag == "head" || tag == "script" || tag == "style" || + tag == "noscript" || tag == "template" || tag == "svg" +} + +func naturalSections(root *html.Node, maxSections int) []htmlSectionAuthority { + var sections []htmlSectionAuthority + var path []string + var walk func(*html.Node) bool + walk = func(current *html.Node) bool { + isSection := current.Type == html.ElementNode && strings.EqualFold(current.Data, "section") + if isSection { + if len(sections) >= maxSections { + return false + } + sections = append(sections, htmlSectionAuthority{ + path: "/" + strings.Join(path, "/"), heading: firstSectionHeading(current), tokens: visibleTokens(current), + }) + return true + } + var ordinals map[string]int + for child := current.FirstChild; child != nil; child = child.NextSibling { + if child.Type == html.ElementNode { + if ordinals == nil { + ordinals = make(map[string]int) + } + tag := strings.ToLower(child.Data) + ordinals[tag]++ + path = append(path, tag+"["+strconv.Itoa(ordinals[tag])+"]") + } + if !walk(child) { + return false + } + if child.Type == html.ElementNode { + path = path[:len(path)-1] + } + } + return true + } + if !walk(root) { + return nil + } + return sections +} + +func firstSectionHeading(section *html.Node) string { + var find func(*html.Node) *html.Node + find = func(current *html.Node) *html.Node { + if current.Type == html.ElementNode { + tag := strings.ToLower(current.Data) + if len(tag) == 2 && tag[0] == 'h' && tag[1] >= '1' && tag[1] <= '6' { + return current + } + } + for child := current.FirstChild; child != nil; child = child.NextSibling { + if found := find(child); found != nil { + return found + } + } + return nil + } + heading := find(section) + if heading == nil { + return "" + } + return strings.Join(visibleTokens(heading), " ") +} + +func hasHTMLStructure(source []byte) bool { + tokenizer := html.NewTokenizer(bytes.NewReader(source)) + for { + switch tokenizer.Next() { + case html.ErrorToken: + return false + case html.StartTagToken, html.SelfClosingTagToken: + name, _ := tokenizer.TagName() + if strings.EqualFold(string(name), "html") || strings.EqualFold(string(name), "body") { + return true + } + default: + } + } +} + +func validateXHTML(source []byte) error { + decoder := xml.NewDecoder(bytes.NewReader(source)) + depth := 0 + roots := 0 + for { + token, err := decoder.Token() + if errors.Is(err, io.EOF) { + if roots == 1 && depth == 0 { + return nil + } + return errors.New("XHTML root is invalid") + } + if err != nil { + return errors.New("XHTML is malformed") + } + switch value := token.(type) { + case xml.StartElement: + if depth == 0 { + roots++ + if roots != 1 || value.Name.Local != "html" || + (value.Name.Space != "" && value.Name.Space != "http://www.w3.org/1999/xhtml") { + return errors.New("XHTML root is invalid") + } + } + depth++ + case xml.EndElement: + depth-- + if depth < 0 { + return errors.New("XHTML is malformed") + } + case xml.CharData: + if depth == 0 && len(bytes.TrimSpace(value)) != 0 { + return errors.New("XHTML has text outside its root") + } + } + } +} + +func readExact(ctx context.Context, reader io.Reader, expected, maximum int64) ([]byte, error) { + data := make([]byte, 0, expected) + buffer := make([]byte, 32<<10) + for { + if err := ctx.Err(); err != nil { + return nil, err + } + read, err := reader.Read(buffer) + if read > 0 { + if int64(len(data))+int64(read) > maximum { + return nil, errInputIdentity + } + data = append(data, buffer[:read]...) + } + switch { + case errors.Is(err, io.EOF): + if int64(len(data)) != expected { + return nil, errInputIdentity + } + return data, nil + case err != nil: + return nil, err + case read == 0: + return nil, io.ErrNoProgress + } + } +} + +func readAuthorizedExact( + ctx context.Context, upload document.AuthorizedUpload, expected, maximum int64, +) ([]byte, error) { + closeDone := make(chan struct{}) + stopClose := context.AfterFunc(ctx, func() { + _ = upload.Close() + close(closeDone) + }) + data, err := readExact(ctx, upload, expected, maximum) + if !stopClose() { + <-closeDone + } + return data, err +} + +func (provider *Provider) contextError(parent context.Context, expiry bool, cause error) error { + if parent.Err() != nil { + return providerError(document.RenditionErrorCanceled, "Trafilatura rendering canceled", parent.Err()) + } + if expiry { + return providerError(document.RenditionErrorPolicyRejected, + "Trafilatura authorization expired during rendering", cause) + } + return providerError(document.RenditionErrorTransient, "Trafilatura rendering timed out", cause) +} + +func (provider *Provider) postProcessError( + parent, operation context.Context, deadline time.Time, expiry bool, +) error { + if parentErr := parent.Err(); parentErr != nil { + return provider.contextError(parent, expiry, parentErr) + } + if operationErr := operation.Err(); operationErr != nil { + return provider.contextError(parent, expiry, operationErr) + } + if !time.Now().UTC().Before(deadline) { + return provider.contextError(parent, expiry, context.DeadlineExceeded) + } + return nil +} + +func providerError(code document.RenditionErrorCode, message string, cause error) error { + classified, err := document.NewRenditionProviderError(code, message, 0, cause) + if err != nil { + return fmt.Errorf("trafilatura: classify provider error: %w", err) + } + return classified +} + +func cleanEnvironment() []string { + environment := []string{ + "LANG=C.UTF-8", "LC_ALL=C.UTF-8", "TZ=UTC", "PYTHONHASHSEED=0", + "PYTHONNOUSERSITE=1", "PYTHONDONTWRITEBYTECODE=1", + } + if runtime.GOOS == "windows" { + if systemRoot := os.Getenv("SystemRoot"); systemRoot != "" { + environment = append(environment, "SystemRoot="+systemRoot) + } + } + return environment +} + +func validFilename(filename, mediaType string) bool { + if filename == "" || filename != strings.TrimSpace(filename) || strings.ContainsAny(filename, "/\\:\x00") { + return false + } + extension := strings.ToLower(filepath.Ext(filename)) + switch mediaType { + case "text/html": + return extension == ".html" || extension == ".htm" + case "application/xhtml+xml": + return extension == ".xhtml" + default: + return false + } +} + +func validateImmutableIdentity(value, subject string) error { + if !strings.HasPrefix(value, "sha256:") || len(value) != len("sha256:")+sha256.Size*2 { + return fmt.Errorf("trafilatura: %s must be an immutable sha256 identity", subject) + } + return validateSHA256(strings.TrimPrefix(value, "sha256:"), subject) +} + +func validateSHA256(value, subject string) error { + if len(value) != sha256.Size*2 { + return fmt.Errorf("trafilatura: %s must be a lowercase SHA-256 digest", subject) + } + decoded, err := hex.DecodeString(value) + if err != nil || hex.EncodeToString(decoded) != value { + return fmt.Errorf("trafilatura: %s must be a lowercase SHA-256 digest", subject) + } + return nil +} + +func hashExecutable(path string) (string, error) { + return hashExecutableContext(context.Background(), path) +} + +func hashExecutableContext(ctx context.Context, path string) (string, error) { + file, err := os.Open(path) + if err != nil { + return "", err + } + defer func() { _ = file.Close() }() + hash := sha256.New() + written, err := io.Copy(hash, io.LimitReader(contextReader{ctx: ctx, reader: file}, MaxExecutableBytes+1)) + if err != nil || written <= 0 || written > MaxExecutableBytes { + return "", errors.New("executable content could not be bounded and hashed") + } + return hex.EncodeToString(hash.Sum(nil)), nil +} + +type contextReader struct { + ctx context.Context + reader io.Reader +} + +func (reader contextReader) Read(value []byte) (int, error) { + if err := reader.ctx.Err(); err != nil { + return 0, err + } + return reader.reader.Read(value) +} + +func validOutputText(value string) bool { + if value == "" || value != strings.TrimSpace(value) || !utf8.ValidString(value) || strings.ContainsRune(value, '\x00') { + return false + } + for _, char := range value { + if unicode.IsControl(char) && char != '\n' && char != '\r' && char != '\t' { + return false + } + } + return true +} + +func pythonInterpreter(base string) bool { + name := strings.TrimSuffix(strings.ToLower(base), ".exe") + if name == "py" || name == "pyw" { + return true + } + for _, prefix := range []string{"python", "pypy"} { + if !strings.HasPrefix(name, prefix) { + continue + } + suffix := strings.TrimPrefix(name, prefix) + if prefix == "python" { + suffix = strings.TrimPrefix(suffix, "w") + } + if suffix == "" { + return true + } + for _, char := range suffix { + if (char < '0' || char > '9') && char != '.' { + return false + } + } + return true + } + return false +} + +func cloneDescriptor(value document.RenditionDescriptor) document.RenditionDescriptor { + value.SupportedFormats = slices.Clone(value.SupportedFormats) + value.ArtifactRoles = slices.Clone(value.ArtifactRoles) + return value +} + +var _ document.RenditionProvider = (*Provider)(nil) diff --git a/document/trafilatura/provider_test.go b/document/trafilatura/provider_test.go new file mode 100644 index 00000000..3d17996d --- /dev/null +++ b/document/trafilatura/provider_test.go @@ -0,0 +1,390 @@ +package trafilatura + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +const testRuntimeIdentity = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + +func TestProviderRendersExactSuppliedHTMLWithoutFetchingRemoteReferences(t *testing.T) { + t.Setenv("DOCBANK_TRAFILATURA_AMBIENT_SECRET", "must-not-reach-child") + provider := newTestProvider(t, helperExecutable(t, "complete"), time.Second, 1<<20) + source := []byte(`

Local title

Local body

link`) + upload := newTestUpload(source, "article.html", "text/html") + + result, err := document.RenderRendition(t.Context(), provider, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, "text", result.Evidence.Family) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 1) + assert.Equal(t, "Local title Local body link", result.Evidence.Units[0].Text) + assert.NotContains(t, result.Evidence.Units[0].Text, "private.example") + assert.Equal(t, int64(len(source)), result.Receipt.Usage.InputBytes) + assert.Equal(t, upload.metadata.SHA256, result.Receipt.SourceSHA256) +} + +func TestProviderDeclaresDegradedEvidenceWithoutCompleteSourceProof(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "degraded"), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + + result, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err, "cause: %v", errors.Unwrap(err)) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitGeneric, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 1) + assert.Equal(t, "Local title\n\nLocal body", result.Evidence.Units[0].Text) + assert.Equal(t, []string{"degraded_provenance"}, result.Receipt.Warnings) +} + +func TestProviderSupportsSuppliedXHTMLBytes(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "xhtml"), time.Second, 1<<20) + source := []byte(`

Local title

Local body

`) + upload := newTestUpload(source, "article.xhtml", "application/xhtml+xml") + + result, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err) + assert.Equal(t, document.EvidenceDegradedProvenance, result.Evidence.Completeness) +} + +func TestProviderRejectsIncompleteOrTruncatedExtraction(t *testing.T) { + for _, mode := range []string{"extraction-incomplete", "truncated-degraded"} { + t.Run(mode, func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, mode), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestProviderClaimsCompleteSectionsOnlyForLocallyVerifiedStructure(t *testing.T) { + source := []byte(`

Local title

Local body

`) + t.Run("verified path", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "structured"), time.Second, 1<<20) + upload := newTestUpload(source, "article.html", "text/html") + result, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err) + assert.Equal(t, document.EvidenceComplete, result.Evidence.Completeness) + assert.Equal(t, document.EvidenceUnitSection, result.Evidence.UnitKind) + require.Len(t, result.Evidence.Units, 1) + assert.Equal(t, "Local title", result.Evidence.Units[0].Locator.Name) + }) + for _, mode := range []string{"structured-path-drift", "structured-boundary-drift"} { + t.Run(mode, func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, mode), time.Second, 1<<20) + upload := newTestUpload(source, "article.html", "text/html") + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestInspectHTMLBoundsNaturalSectionsAndPreservesSiblingOrdinals(t *testing.T) { + source := []byte(` +

One

First

+
+

Two

Second

+

Three

Third

+ `) + authority, err := inspectHTML(source, "text/html", 3) + require.NoError(t, err) + require.Len(t, authority.sections, 3) + assert.Equal(t, []string{ + "/html[1]/body[1]/section[1]", + "/html[1]/body[1]/section[2]", + "/html[1]/body[1]/section[3]", + }, []string{authority.sections[0].path, authority.sections[1].path, authority.sections[2].path}) + + authority, err = inspectHTML(source, "text/html", 2) + require.NoError(t, err) + assert.Nil(t, authority.sections, "over-bound structure must use degraded provenance") +} + +func TestProviderRejectsOutputNotDerivableFromSuppliedBytes(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "fetched"), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.NotContains(t, err.Error(), "REMOTE_FETCH_TOKEN") +} + +func TestProviderRejectsMalformedPartialAndDriftedOutput(t *testing.T) { + for _, mode := range []string{ + "malformed", "unknown-field", "missing-provenance-complete", "version-drift", "runtime-drift", "source-hash-drift", + "source-size-drift", "partial", "empty", "bad-heading", "complete-inexact", + } { + t.Run(mode, func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, mode), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + }) + } +} + +func TestProviderRejectsNonHTMLMalformedXHTMLAndSubstitutedInput(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "failure"), time.Second, 1<<20) + for _, test := range []struct { + name, filename, mediaType string + content []byte + }{ + {name: "plain text", filename: "article.html", mediaType: "text/html", content: []byte("not html")}, + {name: "malformed HTML", filename: "article.html", mediaType: "text/html", content: []byte("broken\x00")}, + {name: "control character in visible text", filename: "article.html", mediaType: "text/html", content: []byte("broken\x07text")}, + {name: "vertical tab in visible text", filename: "article.html", mediaType: "text/html", content: []byte("broken\vtext")}, + {name: "form feed in visible text", filename: "article.html", mediaType: "text/html", content: []byte("broken\ftext")}, + {name: "malformed XHTML", filename: "article.xhtml", mediaType: "application/xhtml+xml", content: []byte(`broken`)}, + } { + t.Run(test.name, func(t *testing.T) { + upload := newTestUpload(test.content, test.filename, test.mediaType) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorUnsupportedInput) + }) + } + t.Run("substituted bytes", func(t *testing.T) { + upload := newTestUpload(testHTML, "article.html", "text/html") + upload.Reader = bytes.NewReader([]byte(`replacement`)) + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) +} + +func TestProviderBoundsOutputAndSanitizesProcessFailure(t *testing.T) { + t.Run("unbounded stdout stops promptly", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "unbounded-output"), 10*time.Second, 1024) + upload := newTestUpload(testHTML, "article.html", "text/html") + started := time.Now() + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorMalformedEvidence) + assert.Less(t, time.Since(started), time.Second) + }) + t.Run("private stderr", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "failure"), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorTransient) + assert.NotContains(t, err.Error(), "private-stderr-token") + assert.NotContains(t, err.Error(), provider.executable) + }) +} + +func TestProviderEnforcesTimeoutCancellationAndExpiry(t *testing.T) { + t.Run("timeout", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "wait"), 20*time.Millisecond, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + _, err := provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorTransient) + }) + t.Run("cancellation", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "wait"), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, err := provider.Render(ctx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorCanceled) + require.ErrorIs(t, err, context.Canceled) + }) + t.Run("expired after child", func(t *testing.T) { + provider := newTestProvider(t, helperExecutable(t, "slow-complete"), time.Second, 1<<20) + upload := newTestUpload(testHTML, "article.html", "text/html") + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + authorization.ExpiresAt = time.Now().UTC().Add(15 * time.Millisecond).Format(timestampForm) + _, err := provider.Render(t.Context(), upload, authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) +} + +func TestProviderCancellationClosesBlockedAuthorizedUploadBeforeStartingRunner(t *testing.T) { + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity} + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + base := newTestUpload(testHTML, "article.html", "text/html") + upload := &blockingUpload{metadata: base.metadata, started: make(chan struct{}), closed: make(chan struct{})} + ctx, cancel := context.WithCancel(t.Context()) + result := make(chan error, 1) + go func() { + _, renderErr := provider.Render(ctx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + result <- renderErr + }() + <-upload.started + cancel() + + select { + case err := <-result: + assertProviderCode(t, err, document.RenditionErrorCanceled) + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("render remained blocked in the authorized upload read after cancellation") + } + assert.Nil(t, runner.request) +} + +func TestNewPinsExecutableRuntimeIdentityAndBounds(t *testing.T) { + first := newTestProvider(t, helperExecutable(t, "complete"), time.Second, 1<<20) + second := newTestProvider(t, helperExecutable(t, "complete-copy"), time.Second, 1<<20) + descriptor := first.Descriptor() + assert.Equal(t, document.RenditionTrustLocalProcess, descriptor.TrustBoundary) + assert.Equal(t, []document.RenditionFormatCapability{ + {MediaFamily: "text", MediaType: "application/xhtml+xml", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/html", InputKind: document.RenditionInputOriginalFile}, + }, descriptor.SupportedFormats) + assert.True(t, descriptor.ReturnsStructured) + assert.NotEqual(t, descriptor.Fingerprint, second.Descriptor().Fingerprint) + + for _, name := range []string{"pyw.exe", "python3", "pythonw.exe", "pythonw3.11.exe"} { + python := filepath.Join(t.TempDir(), name) + require.NoError(t, os.WriteFile(python, []byte("synthetic"), 0o700)) + _, err := New(Profile{Executable: python, RuntimeIdentity: testRuntimeIdentity, + MaxDocumentBytes: 1, MaxResponseBytes: 1, MaxUnits: 1, Timeout: time.Second}) + require.ErrorContains(t, err, "must not be a Python interpreter") + } + + valid := testProfile(t, helperExecutable(t, "complete"), time.Second, 1024) + valid.MaxDocumentBytes = 1024 + valid.MaxUnits = 1 + for _, mutate := range []func(*Profile){ + func(profile *Profile) { profile.Executable = "renderer" }, + func(profile *Profile) { profile.RuntimeIdentity = "" }, + func(profile *Profile) { profile.MaxDocumentBytes = 0 }, + func(profile *Profile) { profile.MaxResponseBytes = 0 }, + func(profile *Profile) { profile.MaxUnits = 0 }, + func(profile *Profile) { profile.Timeout = 0 }, + } { + profile := valid + mutate(&profile) + _, err := New(profile) + require.Error(t, err) + } +} + +var testHTML = []byte(`

Local title

Local body

`) + +func newTestProvider(t *testing.T, executable string, timeout time.Duration, maxResponse int64) *Provider { + t.Helper() + profile := testProfile(t, executable, timeout, maxResponse) + provider, err := New(profile) + require.NoError(t, err) + return provider +} + +func testProfile(t *testing.T, executable string, timeout time.Duration, maxResponse int64) Profile { + t.Helper() + data, err := os.ReadFile(executable) + require.NoError(t, err) + digest := sha256.Sum256(data) + return Profile{ + Executable: executable, ExecutableSHA256: hex.EncodeToString(digest[:]), + RuntimeIdentity: testRuntimeIdentity, Runner: &recordingRunner{identity: testRunnerIdentity}, + MaxDocumentBytes: 1 << 20, MaxResponseBytes: maxResponse, MaxUnits: 10, Timeout: timeout, + } +} + +func assertProviderCode(t *testing.T, err error, want document.RenditionErrorCode) { + t.Helper() + require.Error(t, err) + providerErr, ok := errors.AsType[*document.RenditionProviderError](err) + require.True(t, ok, "%T: %v", err, err) + assert.Equal(t, want, providerErr.Code()) +} + +func helperExecutable(t *testing.T, mode string) string { + t.Helper() + target := filepath.Join(t.TempDir(), "renderer-"+mode) + require.NoError(t, os.WriteFile(target, []byte("synthetic isolated executable fixture\n"), 0o700)) + return target +} + +type testUpload struct { + *bytes.Reader + + metadata document.AuthorizedUploadMetadata +} + +func newTestUpload(data []byte, filename, mediaType string) *testUpload { + digest := sha256.Sum256(data) + return &testUpload{Reader: bytes.NewReader(data), metadata: document.AuthorizedUploadMetadata{ + Filename: filename, MediaFamily: "text", MediaType: mediaType, + ByteLength: int64(len(data)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + }} +} + +func (*testUpload) Close() error { return nil } +func (upload *testUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +func testAuthorization(descriptor document.RenditionDescriptor, metadata document.AuthorizedUploadMetadata) document.RenditionAuthorization { + started := time.Now().UTC().Add(-time.Minute) + return document.RenditionAuthorization{ + ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, RenditionRequestFingerprint: strings.Repeat("4", 64), + SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength, + CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, + ProviderMetadataChecksum: metadata.ProviderMetadataChecksum, + MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, + InputKind: metadata.InputKind, MaxTotalResultBytes: 1 << 20, + AuthorizedAt: started.Format(timestampForm), ExpiresAt: started.Add(10 * time.Minute).Format(timestampForm), + } +} + +var _ document.AuthorizedUpload = (*testUpload)(nil) +var _ io.ReadCloser = (*testUpload)(nil) + +type blockingUpload struct { + metadata document.AuthorizedUploadMetadata + started chan struct{} + closed chan struct{} + startOnce sync.Once + closeOnce sync.Once +} + +func (upload *blockingUpload) Read([]byte) (int, error) { + upload.startOnce.Do(func() { close(upload.started) }) + <-upload.closed + return 0, errors.New("synthetic read closed") +} + +func (upload *blockingUpload) Close() error { + upload.closeOnce.Do(func() { close(upload.closed) }) + return nil +} + +func (upload *blockingUpload) Metadata() document.AuthorizedUploadMetadata { return upload.metadata } + +var _ document.AuthorizedUpload = (*blockingUpload)(nil) diff --git a/document/trafilatura/runner_contract_test.go b/document/trafilatura/runner_contract_test.go new file mode 100644 index 00000000..59ce811f --- /dev/null +++ b/document/trafilatura/runner_contract_test.go @@ -0,0 +1,341 @@ +package trafilatura + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" +) + +const testRunnerIdentity = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" + +func TestNewRequiresPinnedIsolatedRunnerAndExecutableDigest(t *testing.T) { + executable := helperExecutable(t, "complete") + executableBytes, err := os.ReadFile(executable) + require.NoError(t, err) + digest := sha256.Sum256(executableBytes) + profile := Profile{ + Executable: executable, ExecutableSHA256: hex.EncodeToString(digest[:]), + RuntimeIdentity: testRuntimeIdentity, Runner: &recordingRunner{identity: testRunnerIdentity}, + MaxDocumentBytes: 1 << 20, MaxResponseBytes: 1 << 20, MaxUnits: 10, Timeout: time.Second, + } + + _, err = New(profile) + require.NoError(t, err) + + profile.Runner = nil + provider, err := New(profile) + if runtime.GOOS == "linux" { + require.NoError(t, err) + require.NotNil(t, provider.runner) + } else { + require.ErrorIs(t, err, ErrIsolationUnavailable) + } + + profile.Runner = &recordingRunner{identity: "mutable-latest"} + _, err = New(profile) + require.ErrorContains(t, err, "runner identity") + + profile.Runner = &recordingRunner{identity: testRunnerIdentity} + profile.ExecutableSHA256 = "" + _, err = New(profile) + require.ErrorContains(t, err, "executable SHA-256") + + profile.ExecutableSHA256 = hex.EncodeToString(digest[:]) + profile.RuntimeIdentity = "mutable-latest" + _, err = New(profile) + require.ErrorContains(t, err, "runtime identity") +} + +func TestProviderDelegatesOnlyAnExactFailClosedIsolationRequest(t *testing.T) { + t.Setenv("DOCBANK_TRAFILATURA_AMBIENT_SECRET", "must-not-reach-runner") + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity} + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + source := []byte(`

Local title

Local body

link`) + upload := newTestUpload(source, "article.html", "text/html") + + _, err = document.RenderRendition(t.Context(), provider, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NoError(t, err) + require.NotNil(t, runner.request, "provider bypassed the isolated runner") + assert.Equal(t, executable, runner.request.Executable) + assert.Equal(t, profile.ExecutableSHA256, runner.request.ExecutableSHA256) + assert.Equal(t, []string{"--protocol", protocolVersion}, runner.request.Arguments) + expectedEnvironment := []string{ + "LANG=C.UTF-8", "LC_ALL=C.UTF-8", "TZ=UTC", "PYTHONHASHSEED=0", + "PYTHONNOUSERSITE=1", "PYTHONDONTWRITEBYTECODE=1", + } + if runtime.GOOS == "windows" { + if systemRoot := os.Getenv("SystemRoot"); systemRoot != "" { + expectedEnvironment = append(expectedEnvironment, "SystemRoot="+systemRoot) + } + } + assert.Equal(t, expectedEnvironment, runner.request.Environment) + assert.NotContains(t, strings.Join(runner.request.Environment, "\n"), "must-not-reach-runner") + assert.Equal(t, filepath.Dir(executable), runner.request.Directory) + assert.Equal(t, source, runner.request.Stdin) + stdinDigest := sha256.Sum256(source) + assert.Equal(t, hex.EncodeToString(stdinDigest[:]), runner.request.StdinSHA256) + assert.Equal(t, int64(1<<20), runner.request.MaxStdoutBytes) + assert.Equal(t, IsolationRequirements{ + NetworkDisabled: true, KillProcessTree: true, VerifyExecutableSHA256: true, + }, runner.request.Requirements) + require.Len(t, runner.request.PolicyFingerprint, 64) +} + +func TestProviderRejectsIncompleteOrMismatchedIsolationAttestation(t *testing.T) { + for _, testCase := range []struct { + name string + mutate func(*IsolationAttestation) + }{ + {name: "runner identity", mutate: func(value *IsolationAttestation) { value.RunnerIdentity = testRuntimeIdentity }}, + {name: "policy", mutate: func(value *IsolationAttestation) { value.PolicyFingerprint = strings.Repeat("0", 64) }}, + {name: "executable", mutate: func(value *IsolationAttestation) { value.ExecutableSHA256 = strings.Repeat("0", 64) }}, + {name: "stdin", mutate: func(value *IsolationAttestation) { value.StdinSHA256 = strings.Repeat("0", 64) }}, + {name: "network", mutate: func(value *IsolationAttestation) { value.NetworkDisabled = false }}, + {name: "process tree", mutate: func(value *IsolationAttestation) { value.ProcessTreeContained = false }}, + {name: "digest launch", mutate: func(value *IsolationAttestation) { value.DigestVerifiedLaunch = false }}, + } { + t.Run(testCase.name, func(t *testing.T) { + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity} + runner.run = func(ctx context.Context, request IsolatedRunRequest) (IsolatedRunResult, error) { + result, err := defaultRun(ctx, runner.identity, request) + testCase.mutate(&result.Attestation) + return result, err + } + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + upload := newTestUpload(testHTML, "article.html", "text/html") + + _, err = provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + }) + } +} + +func TestProviderReverifiesExecutableAndRunnerIdentityBeforeEveryRun(t *testing.T) { + for _, testCase := range []struct { + name string + mutate func(t *testing.T, executable string, runner *recordingRunner) + }{ + {name: "executable replacement", mutate: func(t *testing.T, executable string, _ *recordingRunner) { + t.Helper() + require.NoError(t, os.WriteFile(executable, []byte("synthetic replacement"), 0o700)) + }}, + {name: "runner identity drift", mutate: func(_ *testing.T, _ string, runner *recordingRunner) { + runner.identity = testRuntimeIdentity + }}, + } { + t.Run(testCase.name, func(t *testing.T) { + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity} + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + testCase.mutate(t, executable, runner) + upload := newTestUpload(testHTML, "article.html", "text/html") + + _, err = provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.Nil(t, runner.request) + }) + } +} + +func TestProviderFailsClosedWhenRunnerCannotEnforceIsolation(t *testing.T) { + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity, run: func( + context.Context, IsolatedRunRequest, + ) (IsolatedRunResult, error) { + return IsolatedRunResult{}, errors.Join(ErrIsolationUnavailable, errors.New("private-runner-detail")) + }} + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + upload := newTestUpload(testHTML, "article.html", "text/html") + + _, err = provider.Render(t.Context(), upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + assert.NotContains(t, err.Error(), "private-runner-detail") +} + +func TestProviderRequiresProcessTreeCleanupAttestationAfterRunnerCancellation(t *testing.T) { + executable := helperExecutable(t, "complete") + renderCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + runner := &recordingRunner{identity: testRunnerIdentity} + runner.run = func(ctx context.Context, request IsolatedRunRequest) (IsolatedRunResult, error) { + cancel() + <-ctx.Done() + result := isolatedResult(runner.identity, request, nil) + result.Attestation.ProcessTreeContained = false + return result, ctx.Err() + } + profile := testProfile(t, executable, time.Minute, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + upload := newTestUpload(testHTML, "article.html", "text/html") + + _, err = provider.Render(renderCtx, upload, + testAuthorization(provider.Descriptor(), upload.Metadata())) + require.NotNil(t, runner.request) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) +} + +func TestProviderConstrainsRunnerStdoutToAuthorizedTotalResultBytes(t *testing.T) { + executable := helperExecutable(t, "complete") + runner := &recordingRunner{identity: testRunnerIdentity, run: func( + context.Context, IsolatedRunRequest, + ) (IsolatedRunResult, error) { + return IsolatedRunResult{}, ErrIsolationUnavailable + }} + profile := testProfile(t, executable, time.Second, 1<<20) + profile.Runner = runner + provider, err := New(profile) + require.NoError(t, err) + upload := newTestUpload(testHTML, "article.html", "text/html") + authorization := testAuthorization(provider.Descriptor(), upload.Metadata()) + authorization.MaxTotalResultBytes = 128 + + _, err = provider.Render(t.Context(), upload, authorization) + assertProviderCode(t, err, document.RenditionErrorPolicyRejected) + require.NotNil(t, runner.request) + assert.Equal(t, int64(128), runner.request.MaxStdoutBytes) +} + +type recordingRunner struct { + identity string + request *IsolatedRunRequest + run func(context.Context, IsolatedRunRequest) (IsolatedRunResult, error) +} + +func (runner *recordingRunner) Identity() string { return runner.identity } + +func (runner *recordingRunner) Run(ctx context.Context, request IsolatedRunRequest) (IsolatedRunResult, error) { + copied := request + copied.Arguments = append([]string(nil), request.Arguments...) + copied.Environment = append([]string(nil), request.Environment...) + copied.Stdin = append([]byte(nil), request.Stdin...) + runner.request = &copied + if runner.run != nil { + return runner.run(ctx, request) + } + return defaultRun(ctx, runner.identity, request) +} + +func defaultRun(ctx context.Context, runnerIdentity string, request IsolatedRunRequest) (IsolatedRunResult, error) { + mode := strings.TrimSuffix(filepath.Base(request.Executable), filepath.Ext(request.Executable)) + mode = strings.TrimPrefix(mode, "renderer-") + switch mode { + case "failure": + return isolatedResult(runnerIdentity, request, nil), + errors.Join(ErrChildFailed, errors.New("private-stderr-token")) + case "wait": + <-ctx.Done() + return isolatedResult(runnerIdentity, request, nil), ctx.Err() + case "slow-complete": + select { + case <-ctx.Done(): + return isolatedResult(runnerIdentity, request, nil), ctx.Err() + case <-time.After(40 * time.Millisecond): + } + case "unbounded-output": + return isolatedResult(runnerIdentity, request, nil), ErrChildOutputTooLarge + case "malformed": + return isolatedResult(runnerIdentity, request, []byte("{")), nil + } + digest := sha256.Sum256(request.Stdin) + value := response{ContractVersion: protocolVersion, RuntimeIdentity: testRuntimeIdentity, + SourceSHA256: hex.EncodeToString(digest[:]), SourceBytes: int64(len(request.Stdin)), + ExtractionComplete: true, ProvenanceComplete: new(false), + Units: []responseUnit{{Text: "Local title Local body"}}, + } + if strings.Contains(string(request.Stdin), "link") { + value.Units[0].Text += " link" + } + switch mode { + case "degraded": + value.Units[0].Text = "Local title\n\nLocal body" + case "fetched": + value.Units[0].Text += " REMOTE_FETCH_TOKEN" + case "version-drift": + value.ContractVersion += ".next" + case "runtime-drift": + value.RuntimeIdentity = "sha256:" + strings.Repeat("c", 64) + case "source-hash-drift": + value.SourceSHA256 = strings.Repeat("c", 64) + case "source-size-drift": + value.SourceBytes++ + case "partial": + value.Units = append(value.Units, responseUnit{}) + case "empty": + value.Units = nil + case "bad-heading": + value.Units[0].Heading = "Fetched heading" + case "complete-inexact": + value.Units[0].Text = "Local title" + case "extraction-incomplete": + value.ExtractionComplete = false + case "truncated-degraded": + value.Units[0].Text = "Local title" + case "structured": + value.ProvenanceComplete = new(true) + value.Units[0].SourcePath = "/html[1]/body[1]/section[1]" + value.Units[0].Heading = "Local title" + case "structured-path-drift": + value.ProvenanceComplete = new(true) + value.Units[0].SourcePath = "/html[1]/body[1]/section[2]" + value.Units[0].Heading = "Local title" + case "structured-boundary-drift": + value.ProvenanceComplete = new(true) + value.Units[0].SourcePath = "/html[1]/body[1]/section[1]" + value.Units[0].Heading = "Local title" + value.Units[0].Text = "Local title" + } + encoded, err := json.Marshal(value) + if err != nil { + return IsolatedRunResult{}, err + } + if mode == "unknown-field" { + encoded = append(encoded[:len(encoded)-1], []byte(`,"unexpected":true}`)...) + } + if mode == "missing-provenance-complete" { + encoded = bytes.Replace(encoded, []byte(`,"provenance_complete":false`), nil, 1) + } + return isolatedResult(runnerIdentity, request, encoded), nil +} + +func isolatedResult(runnerIdentity string, request IsolatedRunRequest, stdout []byte) IsolatedRunResult { + stdinDigest := sha256.Sum256(request.Stdin) + return IsolatedRunResult{Stdout: stdout, Attestation: IsolationAttestation{ + RunnerIdentity: runnerIdentity, PolicyFingerprint: request.PolicyFingerprint, + ExecutableSHA256: request.ExecutableSHA256, StdinSHA256: hex.EncodeToString(stdinDigest[:]), + NetworkDisabled: true, ProcessTreeContained: true, DigestVerifiedLaunch: true, + }} +} diff --git a/document/trafilatura/testdata/isolatedhelper/main.go b/document/trafilatura/testdata/isolatedhelper/main.go new file mode 100644 index 00000000..74ca3262 --- /dev/null +++ b/document/trafilatura/testdata/isolatedhelper/main.go @@ -0,0 +1,92 @@ +package main + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net" + "os" + "os/exec" + "syscall" + "time" +) + +var ( + mode = "echo" + networkAddress string +) + +type echoResponse struct { + Arguments []string `json:"arguments"` + Environment []string `json:"environment"` + StdinSHA256 string `json:"stdin_sha256"` +} + +func main() { + if len(os.Args) == 2 && os.Args[1] == "--descendant" { + lock, err := os.OpenFile(networkAddress, os.O_RDWR|os.O_TRUNC, 0) + if err != nil || syscall.Flock(int(lock.Fd()), syscall.LOCK_EX|syscall.LOCK_NB) != nil { + os.Exit(5) + } + if _, err := fmt.Fprintln(lock, "descendant-ready"); err != nil || lock.Sync() != nil { + os.Exit(6) + } + for { + time.Sleep(time.Hour) + } + } + + switch mode { + case "echo", "replacement": + data, err := io.ReadAll(os.Stdin) + if err != nil { + os.Exit(2) + } + digest := sha256.Sum256(data) + _ = json.NewEncoder(os.Stdout).Encode(echoResponse{ + Arguments: os.Args[1:], Environment: os.Environ(), + StdinSHA256: hex.EncodeToString(digest[:]), + }) + case "network": + connection, err := net.DialTimeout("tcp", networkAddress, time.Second) + if err != nil { + fmt.Print("denied") + return + } + _ = connection.Close() + fmt.Print("connected") + case "unix-network": + connection, err := net.DialTimeout("unix", networkAddress, time.Second) + if err != nil { + fmt.Print("denied") + return + } + _ = connection.Close() + fmt.Print("connected") + case "descendant": + command := exec.Command(os.Args[0], "--descendant") + command.Stdout = os.Stdout + command.Stderr = os.Stderr + if err := command.Start(); err != nil { + os.Exit(3) + } + fmt.Print("spawned") + _ = os.Stdout.Sync() + for { + time.Sleep(time.Hour) + } + case "overflow": + block := make([]byte, 32<<10) + for { + if _, err := os.Stdout.Write(block); err != nil { + return + } + } + case "exit-125": + os.Exit(125) + default: + os.Exit(4) + } +} diff --git a/document/unstructured/profile.go b/document/unstructured/profile.go new file mode 100644 index 00000000..4139d9f5 --- /dev/null +++ b/document/unstructured/profile.go @@ -0,0 +1,292 @@ +// Package unstructured defines the fixed compatibility profile used when an +// operator deploys an Unstructured adapter behind docbank-rendition/v1. +package unstructured + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "fmt" + "slices" + "strings" + "time" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/bridge" +) + +const ( + // ProfileContractV1 identifies the canonical Unstructured bridge profile. + ProfileContractV1 = "unstructured-bridge-profile/v1" + descriptorID = "unstructured.bridge.v1" +) + +// Config supplies the only operator-specific profile values. It deliberately +// has no routes, URLs, fetch controls, or provider options. +type Config struct { + DeploymentID string + RuntimeID string + CredentialBinding string +} + +// LimitsV1 fixes every finite input, output, polling, and wall-clock bound. +type LimitsV1 struct { + MaxDocumentBytes int64 `json:"max_document_bytes"` + MaxPollAttempts int `json:"max_poll_attempts"` + MaxResponseBytes int64 `json:"max_response_bytes"` + PollIntervalMillis int64 `json:"poll_interval_millis"` + RequestTimeoutMillis int64 `json:"request_timeout_millis"` + TotalTimeoutMillis int64 `json:"total_timeout_millis"` +} + +// DisclosurePolicyV1 permits only the exact supplied bytes and records the +// bridge's safe-basename disclosure. Authorization binds every byte and tuple. +type DisclosurePolicyV1 struct { + DiscloseFilename bool `json:"disclose_filename"` + Source string `json:"source"` +} + +// EvidencePolicyV1 fixes bounded provider-neutral evidence and Markdown. +type EvidencePolicyV1 struct { + MaxProviderMarkdownBytes int `json:"max_provider_markdown_bytes"` + MaxTotalResultBytes int `json:"max_total_result_bytes"` + MaxUnits int `json:"max_units"` + SourceEvidenceContract string `json:"source_evidence_contract"` +} + +// ArtifactPolicyV1 permits only one bounded structured-evidence artifact. +type ArtifactPolicyV1 struct { + AllowedRoles []document.EvidenceArtifactRole `json:"allowed_roles"` + MaxArtifactBytes int64 `json:"max_artifact_bytes"` + MaxArtifacts int `json:"max_artifacts"` +} + +// ProfileV1 is the immutable compatibility identity expected from an +// operator-network Unstructured bridge deployment. +type ProfileV1 struct { + ArtifactPolicy ArtifactPolicyV1 `json:"artifact_policy"` + BridgeContract string `json:"bridge_contract"` + ContractVersion string `json:"contract_version"` + CredentialBinding string `json:"credential_binding"` + DeploymentID string `json:"deployment_id"` + Disclosure DisclosurePolicyV1 `json:"disclosure"` + EvidencePolicy EvidencePolicyV1 `json:"evidence_policy"` + InputKind document.RenditionInputKind `json:"input_kind"` + Limits LimitsV1 `json:"limits"` + PolicyFingerprint string `json:"policy_fingerprint"` + RuntimeID string `json:"runtime_id"` + SupportedFormats []document.RenditionFormatCapability `json:"supported_formats"` + TrustBoundary document.RenditionTrustBoundary `json:"trust_boundary"` +} + +// NewProfile returns the standard profile with operator-pinned deployment and +// runtime identity and an optional named credential binding. +func NewProfile(config Config) (ProfileV1, error) { + profile := ProfileV1{ + ArtifactPolicy: standardArtifactPolicy(), + BridgeContract: bridge.ContractVersion, ContractVersion: ProfileContractV1, + CredentialBinding: config.CredentialBinding, DeploymentID: config.DeploymentID, + Disclosure: standardDisclosurePolicy(), + EvidencePolicy: standardEvidencePolicy(), + InputKind: document.RenditionInputOriginalFile, + Limits: standardLimits(), + RuntimeID: config.RuntimeID, SupportedFormats: standardFormats(), + TrustBoundary: document.RenditionTrustOperatorNetwork, + } + _, fingerprint, err := CanonicalProfile(profile) + if err != nil { + return ProfileV1{}, err + } + profile.PolicyFingerprint = fingerprint + return profile, nil +} + +// CanonicalProfile validates, sorts, and deterministically encodes a profile. +// A populated fingerprint must match the canonical identity. +func CanonicalProfile(profile ProfileV1) ([]byte, string, error) { + canonical := cloneProfile(profile) + slices.SortFunc(canonical.SupportedFormats, compareFormats) + slices.Sort(canonical.ArtifactPolicy.AllowedRoles) + claimed := canonical.PolicyFingerprint + canonical.PolicyFingerprint = "" + if err := validateProfile(canonical); err != nil { + return nil, "", fmt.Errorf("unstructured: invalid profile: %w", err) + } + identityJSON, err := json.Marshal(canonical, json.Deterministic(true)) + if err != nil { + return nil, "", fmt.Errorf("unstructured: encode profile identity: %w", err) + } + digest := sha256.Sum256(identityJSON) + fingerprint := hex.EncodeToString(digest[:]) + if claimed != "" && claimed != fingerprint { + return nil, "", errors.New("unstructured: policy fingerprint does not match canonical profile") + } + canonical.PolicyFingerprint = fingerprint + encoded, err := json.Marshal(canonical, json.Deterministic(true)) + if err != nil { + return nil, "", fmt.Errorf("unstructured: encode canonical profile: %w", err) + } + return encoded, fingerprint, nil +} + +// ParseProfile accepts only the exact canonical v1 representation. +func ParseProfile(raw []byte) (ProfileV1, error) { + var profile ProfileV1 + if err := json.Unmarshal(raw, &profile, json.RejectUnknownMembers(true)); err != nil { + return ProfileV1{}, fmt.Errorf("unstructured: decode profile: %w", err) + } + canonical, _, err := CanonicalProfile(profile) + if err != nil { + return ProfileV1{}, err + } + if !bytes.Equal(raw, canonical) { + return ProfileV1{}, errors.New("unstructured: profile is not canonical") + } + return cloneProfile(profile), nil +} + +// BridgeProfile projects the compatibility identity into the generic hardened +// bridge. Origin remains generic bridge configuration, not profile schema. +func BridgeProfile(profile ProfileV1, origin string) (bridge.Profile, error) { + _, fingerprint, err := CanonicalProfile(profile) + if err != nil { + return bridge.Profile{}, err + } + descriptor, err := document.NewRenditionDescriptor(document.RenditionDescriptor{ + ID: descriptorID, ContractVersion: document.RenditionProviderContractVersion, + PolicyFingerprint: fingerprint, TrustBoundary: document.RenditionTrustOperatorNetwork, + SupportedFormats: slices.Clone(profile.SupportedFormats), ReturnsMarkdown: true, + ReturnsStructured: true, + ArtifactRoles: slices.Clone(profile.ArtifactPolicy.AllowedRoles), + }) + if err != nil { + return bridge.Profile{}, fmt.Errorf("unstructured: construct bridge descriptor: %w", err) + } + return bridge.Profile{ + Origin: origin, Descriptor: descriptor, SecretBinding: profile.CredentialBinding, + RequestTimeout: time.Duration(profile.Limits.RequestTimeoutMillis) * time.Millisecond, + TotalTimeout: time.Duration(profile.Limits.TotalTimeoutMillis) * time.Millisecond, + PollInterval: time.Duration(profile.Limits.PollIntervalMillis) * time.Millisecond, + MaxPollAttempts: profile.Limits.MaxPollAttempts, MaxResponseBytes: profile.Limits.MaxResponseBytes, + MaxSourceBytes: profile.Limits.MaxDocumentBytes, + MaxProviderMarkdownBytes: profile.EvidencePolicy.MaxProviderMarkdownBytes, + MaxArtifactBytes: int(profile.ArtifactPolicy.MaxArtifactBytes), + MaxArtifacts: profile.ArtifactPolicy.MaxArtifacts, + MaxTotalResultBytes: profile.EvidencePolicy.MaxTotalResultBytes, + MaxEvidenceUnits: profile.EvidencePolicy.MaxUnits, + }, nil +} + +func validateProfile(profile ProfileV1) error { + if profile.ContractVersion != ProfileContractV1 || profile.BridgeContract != bridge.ContractVersion { + return errors.New("contract version is invalid") + } + if err := validateIdentity(profile.DeploymentID, "deployment ID", false, 256); err != nil { + return err + } + if err := validateIdentity(profile.RuntimeID, "runtime ID", true, 256); err != nil { + return err + } + if profile.CredentialBinding != "" { + if err := validateIdentity(profile.CredentialBinding, "credential binding", false, 128); err != nil { + return err + } + } + if profile.TrustBoundary != document.RenditionTrustOperatorNetwork || + profile.InputKind != document.RenditionInputOriginalFile || + profile.Disclosure != standardDisclosurePolicy() { + return errors.New("disclosure or execution boundary is invalid") + } + if !slices.Equal(profile.SupportedFormats, standardFormats()) { + return errors.New("supported formats differ from the standard profile") + } + if !slices.Equal(profile.ArtifactPolicy.AllowedRoles, + []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}) { + return errors.New("artifact roles differ from the standard profile") + } + standardArtifacts := standardArtifactPolicy() + if profile.Limits != standardLimits() || profile.EvidencePolicy != standardEvidencePolicy() || + profile.ArtifactPolicy.MaxArtifactBytes != standardArtifacts.MaxArtifactBytes || + profile.ArtifactPolicy.MaxArtifacts != standardArtifacts.MaxArtifacts { + return errors.New("limits differ from the finite standard profile") + } + return nil +} + +func standardLimits() LimitsV1 { + return LimitsV1{ + MaxDocumentBytes: 100 << 20, MaxPollAttempts: 300, MaxResponseBytes: 128 << 20, + PollIntervalMillis: 1_000, RequestTimeoutMillis: 30_000, TotalTimeoutMillis: 600_000, + } +} + +func standardDisclosurePolicy() DisclosurePolicyV1 { + return DisclosurePolicyV1{DiscloseFilename: true, Source: "exact_supplied_bytes"} +} + +func standardEvidencePolicy() EvidencePolicyV1 { + return EvidencePolicyV1{ + MaxProviderMarkdownBytes: 32 << 20, MaxTotalResultBytes: 128 << 20, + MaxUnits: 100_000, SourceEvidenceContract: document.SourceEvidenceContractV1, + } +} + +func standardArtifactPolicy() ArtifactPolicyV1 { + return ArtifactPolicyV1{ + AllowedRoles: []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + MaxArtifactBytes: 64 << 20, MaxArtifacts: 1, + } +} + +func validateIdentity(value, subject string, allowColon bool, maximum int) error { + if value == "" || len(value) > maximum || !utf8.ValidString(value) || strings.TrimSpace(value) != value { + return fmt.Errorf("%s is invalid", subject) + } + for _, character := range value { + if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' || + character >= '0' && character <= '9' || strings.ContainsRune("._-", character) || + allowColon && character == ':' { + continue + } + return fmt.Errorf("%s is invalid", subject) + } + return nil +} + +func standardFormats() []document.RenditionFormatCapability { + formats := []document.RenditionFormatCapability{ + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "text/csv", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "mail", MediaType: "message/rfc822", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "structured", MediaType: "application/xml", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/markdown", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + } + slices.SortFunc(formats, compareFormats) + return formats +} + +func compareFormats(left, right document.RenditionFormatCapability) int { + if comparison := strings.Compare(left.MediaFamily, right.MediaFamily); comparison != 0 { + return comparison + } + if comparison := strings.Compare(left.MediaType, right.MediaType); comparison != 0 { + return comparison + } + return strings.Compare(string(left.InputKind), string(right.InputKind)) +} + +func cloneProfile(profile ProfileV1) ProfileV1 { + profile.SupportedFormats = slices.Clone(profile.SupportedFormats) + profile.ArtifactPolicy.AllowedRoles = slices.Clone(profile.ArtifactPolicy.AllowedRoles) + return profile +} diff --git a/document/unstructured/profile_test.go b/document/unstructured/profile_test.go new file mode 100644 index 00000000..0b9ad4de --- /dev/null +++ b/document/unstructured/profile_test.go @@ -0,0 +1,161 @@ +package unstructured + +import ( + "context" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/bridge" +) + +func TestReferenceProfileCanonicalizationPinsCompatibilityIdentity(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-unstructured-primary", + RuntimeID: "sha256:" + strings.Repeat("a", 64), + CredentialBinding: "unstructured-api", + }) + require.NoError(t, err) + + canonical, fingerprint, err := CanonicalProfile(profile) + require.NoError(t, err) + assert.Equal(t, fingerprint, profile.PolicyFingerprint) + assert.JSONEq(t, `{ + "artifact_policy":{"allowed_roles":["structured_evidence"],"max_artifact_bytes":67108864,"max_artifacts":1}, + "bridge_contract":"docbank-rendition/v1", + "contract_version":"unstructured-bridge-profile/v1", + "credential_binding":"unstructured-api", + "deployment_id":"operator-unstructured-primary", + "disclosure":{"disclose_filename":true,"source":"exact_supplied_bytes"}, + "evidence_policy":{"max_provider_markdown_bytes":33554432,"max_total_result_bytes":134217728,"max_units":100000,"source_evidence_contract":"source-evidence/v1"}, + "input_kind":"original_file", + "limits":{"max_document_bytes":104857600,"max_poll_attempts":300,"max_response_bytes":134217728,"poll_interval_millis":1000,"request_timeout_millis":30000,"total_timeout_millis":600000}, + "policy_fingerprint":"`+fingerprint+`", + "runtime_id":"sha256:`+strings.Repeat("a", 64)+`", + "supported_formats":[ + {"media_family":"ebook","media_type":"application/epub+zip","input_kind":"original_file"}, + {"media_family":"image","media_type":"image/jpeg","input_kind":"original_file"}, + {"media_family":"image","media_type":"image/png","input_kind":"original_file"}, + {"media_family":"mail","media_type":"message/rfc822","input_kind":"original_file"}, + {"media_family":"pdf","media_type":"application/pdf","input_kind":"original_file"}, + {"media_family":"presentation","media_type":"application/vnd.openxmlformats-officedocument.presentationml.presentation","input_kind":"original_file"}, + {"media_family":"spreadsheet","media_type":"application/vnd.oasis.opendocument.spreadsheet","input_kind":"original_file"}, + {"media_family":"spreadsheet","media_type":"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet","input_kind":"original_file"}, + {"media_family":"spreadsheet","media_type":"text/csv","input_kind":"original_file"}, + {"media_family":"structured","media_type":"application/xml","input_kind":"original_file"}, + {"media_family":"text","media_type":"text/markdown","input_kind":"original_file"}, + {"media_family":"text","media_type":"text/plain","input_kind":"original_file"} + ], + "trust_boundary":"operator_network" + }`, string(canonical)) + + parsed, err := ParseProfile(canonical) + require.NoError(t, err) + assert.Equal(t, profile, parsed) + _, err = ParseProfile(append(canonical, '\n')) + require.ErrorContains(t, err, "not canonical") +} + +func TestReferenceProfileAdvertisesOnlyEnforceableBroadSuppliedByteFormats(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-unstructured-primary", + RuntimeID: "sha256:" + strings.Repeat("b", 64), + }) + require.NoError(t, err) + + assert.Equal(t, []document.RenditionFormatCapability{ + {MediaFamily: "ebook", MediaType: "application/epub+zip", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/jpeg", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "image", MediaType: "image/png", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "mail", MediaType: "message/rfc822", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "pdf", MediaType: "application/pdf", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "presentation", MediaType: "application/vnd.openxmlformats-officedocument.presentationml.presentation", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.oasis.opendocument.spreadsheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "spreadsheet", MediaType: "text/csv", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "structured", MediaType: "application/xml", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/markdown", InputKind: document.RenditionInputOriginalFile}, + {MediaFamily: "text", MediaType: "text/plain", InputKind: document.RenditionInputOriginalFile}, + }, profile.SupportedFormats) + assert.Equal(t, "exact_supplied_bytes", profile.Disclosure.Source) + assert.True(t, profile.Disclosure.DiscloseFilename) + assert.Equal(t, []document.EvidenceArtifactRole{document.EvidenceArtifactStructured}, + profile.ArtifactPolicy.AllowedRoles) +} + +func TestReferenceProfileBuildsStandardBridgeAndRejectsDrift(t *testing.T) { + profile, err := NewProfile(Config{ + DeploymentID: "operator-unstructured-primary", + RuntimeID: "sha256:" + strings.Repeat("c", 64), + CredentialBinding: "unstructured-api", + }) + require.NoError(t, err) + + bridgeProfile, err := BridgeProfile(profile, "http://127.0.0.1:8421") + require.NoError(t, err) + assert.Equal(t, bridge.ContractVersion, profile.BridgeContract) + assert.Equal(t, profile.CredentialBinding, bridgeProfile.SecretBinding) + assert.Equal(t, document.RenditionTrustOperatorNetwork, bridgeProfile.Descriptor.TrustBoundary) + assert.Equal(t, profile.PolicyFingerprint, bridgeProfile.Descriptor.PolicyFingerprint) + assert.Equal(t, profile.SupportedFormats, bridgeProfile.Descriptor.SupportedFormats) + assert.Equal(t, profile.Limits.MaxDocumentBytes, bridgeProfile.MaxSourceBytes) + assert.Equal(t, profile.EvidencePolicy.MaxProviderMarkdownBytes, + bridgeProfile.MaxProviderMarkdownBytes) + assert.Equal(t, int(profile.ArtifactPolicy.MaxArtifactBytes), bridgeProfile.MaxArtifactBytes) + assert.Equal(t, profile.ArtifactPolicy.MaxArtifacts, bridgeProfile.MaxArtifacts) + assert.Equal(t, profile.EvidencePolicy.MaxTotalResultBytes, bridgeProfile.MaxTotalResultBytes) + assert.Equal(t, profile.EvidencePolicy.MaxUnits, bridgeProfile.MaxEvidenceUnits) + _, err = bridge.New(bridgeProfile, staticSecretResolver{}, http.DefaultClient) + require.NoError(t, err) + + for _, test := range []struct { + mutate func(*ProfileV1) + want string + }{ + {mutate: func(value *ProfileV1) { value.RuntimeID = "sha256:" + strings.Repeat("d", 64) }, want: "policy fingerprint does not match"}, + {mutate: func(value *ProfileV1) { value.Limits.MaxDocumentBytes++ }, want: "limits differ"}, + } { + drifted := profile + test.mutate(&drifted) + _, err := BridgeProfile(drifted, "http://127.0.0.1:8421") + require.ErrorContains(t, err, test.want) + } +} + +func TestReferenceProfileRejectsUnpinnedIdentityAndUnboundedLimits(t *testing.T) { + tests := []struct { + name string + config Config + want string + }{ + {name: "deployment", config: Config{RuntimeID: "sha256:" + strings.Repeat("e", 64)}, want: "deployment ID"}, + {name: "runtime", config: Config{DeploymentID: "operator-unstructured-primary"}, want: "runtime ID"}, + {name: "credential", config: Config{DeploymentID: "operator-unstructured-primary", RuntimeID: "runtime-v1", CredentialBinding: "https://secret.invalid"}, want: "credential binding"}, + {name: "credential length", config: Config{DeploymentID: "operator-unstructured-primary", RuntimeID: "runtime-v1", CredentialBinding: strings.Repeat("c", 129)}, want: "credential binding"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := NewProfile(test.config) + require.ErrorContains(t, err, test.want) + }) + } + + profile, err := NewProfile(Config{DeploymentID: "operator-unstructured-primary", RuntimeID: "runtime-v1"}) + require.NoError(t, err) + profile.Limits.MaxResponseBytes++ + profile.PolicyFingerprint = "" + _, _, err = CanonicalProfile(profile) + require.ErrorContains(t, err, "limits") +} + +type staticSecretResolver struct{} + +func (staticSecretResolver) ResolveSecret(context.Context, string) (string, error) { + return "synthetic-secret", nil +} + +var _ bridge.SecretResolver = staticSecretResolver{} diff --git a/document/upload/doc.go b/document/upload/doc.go new file mode 100644 index 00000000..500e11aa --- /dev/null +++ b/document/upload/doc.go @@ -0,0 +1,5 @@ +// Package upload turns a locally inspected capability into a one-shot +// provider upload. It copies exact source bytes into a private spool, syncs +// them, independently reopens and verifies them, removes their pathname, and +// owns cleanup until the reader closes. +package upload diff --git a/document/upload/spool.go b/document/upload/spool.go new file mode 100644 index 00000000..3b01ef6f --- /dev/null +++ b/document/upload/spool.go @@ -0,0 +1,384 @@ +package upload + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "reflect" + "sync" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" +) + +const ( + spoolDirectoryPrefix = ".docbank-upload-" + spoolFilename = "source" + maxProviderMetadataSize = 1 << 20 +) + +type authorizeStage uint8 + +const ( + authorizeStageWritten authorizeStage = iota + 1 + authorizeStageReaderOpened + authorizeStageValidated +) + +// Source is one authoritative stream and the application-owned directory in +// which a private, descriptor-held spool may be created. Authorize always +// closes Reader. +type Source struct { + Reader io.ReadCloser + Directory string + + testHook func(authorizeStage, string) error +} + +// UploadMetadata contains bounded provider-facing metadata not derived from +// the inspected bytes. +type UploadMetadata struct { + Filename string + ProviderMetadata []byte +} + +type authorizedUpload struct { + mu sync.Mutex + reader *os.File + metadata document.AuthorizedUploadMetadata + capability media.CapabilityRecord + cleanup func() error + stop func() bool + closed bool +} + +func (upload *authorizedUpload) Read(buffer []byte) (int, error) { + upload.mu.Lock() + defer upload.mu.Unlock() + if upload.closed || upload.reader == nil { + return 0, os.ErrClosed + } + return upload.reader.Read(buffer) +} + +func (upload *authorizedUpload) Close() error { + if upload == nil { + return nil + } + upload.mu.Lock() + defer upload.mu.Unlock() + if upload.closed { + return nil + } + upload.closed = true + var result error + if upload.stop != nil { + upload.stop() + upload.stop = nil + } + if upload.reader != nil { + result = upload.reader.Close() + upload.reader = nil + } + if upload.cleanup != nil { + result = errors.Join(result, upload.cleanup()) + upload.cleanup = nil + } + return result +} + +func (upload *authorizedUpload) Metadata() document.AuthorizedUploadMetadata { + if upload == nil { + return document.AuthorizedUploadMetadata{} + } + return upload.metadata +} + +func (upload *authorizedUpload) CapabilityRecord() media.CapabilityRecord { + if upload == nil { + return media.CapabilityRecord{} + } + return upload.capability +} + +var _ document.AuthorizedUpload = (*authorizedUpload)(nil) + +// Authorize copies, syncs, independently reopens, rehashes, and reinspects an +// exact source before returning the core-owned one-shot reader. The named file +// is unlinked or delete-pended before the adapter can receive it. +func Authorize( + ctx context.Context, source Source, capability media.CapabilityRecord, + metadata UploadMetadata, +) (authorized document.AuthorizedUpload, retErr error) { + if source.Reader == nil { + return nil, errors.New("upload: source reader is required") + } + var sourceCloseOnce sync.Once + var sourceCloseErr error + closeSource := func() error { + sourceCloseOnce.Do(func() { sourceCloseErr = source.Reader.Close() }) + return sourceCloseErr + } + defer func() { + if err := closeSource(); err != nil { + retErr = errors.Join(retErr, fmt.Errorf("upload: close source: %w", err)) + } + if retErr != nil && authorized != nil { + retErr = errors.Join(retErr, authorized.Close()) + authorized = nil + } + }() + if err := ctx.Err(); err != nil { + return nil, err + } + if err := media.ValidateCapabilityRecord(capability); err != nil { + return nil, fmt.Errorf("upload: invalid capability record: %w", err) + } + if !capability.Eligible { + return nil, fmt.Errorf("upload: capability record is ineligible: %s", capability.Reason) + } + policy, local := capability.InspectionPolicy() + if !local { + return nil, errors.New("upload: capability record lacks local inspection authority") + } + if metadata.Filename != policy.Filename { + return nil, errors.New("upload: filename does not match capability record") + } + if len(metadata.ProviderMetadata) > maxProviderMetadataSize { + return nil, errors.New("upload: provider metadata exceeds byte limit") + } + if source.Directory == "" || !filepath.IsAbs(source.Directory) { + return nil, errors.New("upload: spool directory must be an absolute path") + } + + directory, err := openSpoolDirectory(source.Directory) + if err != nil { + return nil, fmt.Errorf("upload: create private spool: %w", err) + } + cleanupDirectory := true + defer func() { + if cleanupDirectory { + retErr = errors.Join(retErr, directory.cleanup()) + } + }() + writer, err := directory.create(spoolFilename) + if err != nil { + return nil, fmt.Errorf("upload: create exclusive spool: %w", err) + } + writerOpen := true + defer func() { + if writerOpen { + retErr = errors.Join(retErr, writer.Close()) + } + }() + + stopClose := context.AfterFunc(ctx, func() { _ = closeSource() }) + hasher := sha256.New() + written, err := copyContext(ctx, io.MultiWriter(writer, hasher), + io.LimitReader(source.Reader, capability.SourceBytes+1)) + stopClose() + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, contextErr + } + return nil, fmt.Errorf("upload: copy source: %w", err) + } + if err := ctx.Err(); err != nil { + return nil, err + } + if written != capability.SourceBytes || hex.EncodeToString(hasher.Sum(nil)) != capability.SourceSHA256 { + return nil, errors.New("upload: copied source does not match capability record") + } + if err := writer.Sync(); err != nil { + return nil, fmt.Errorf("upload: sync spool: %w", err) + } + writtenInfo, err := writer.Stat() + if err != nil { + return nil, fmt.Errorf("upload: stat writer: %w", err) + } + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("upload: close writer: %w", err) + } + writerOpen = false + if err := directory.sync(); err != nil { + return nil, fmt.Errorf("upload: sync spool directory: %w", err) + } + if err := callTestHook(source.testHook, authorizeStageWritten, directory.path(spoolFilename)); err != nil { + return nil, err + } + + reader, err := directory.openReader(spoolFilename, writtenInfo) + if err != nil { + return nil, fmt.Errorf("upload: open independent spool reader: %w", err) + } + readerOpen := true + defer func() { + if readerOpen { + retErr = errors.Join(retErr, reader.Close()) + } + }() + if err := callTestHook(source.testHook, authorizeStageReaderOpened, directory.path(spoolFilename)); err != nil { + return nil, err + } + if err := directory.unlink(spoolFilename); err != nil { + return nil, fmt.Errorf("upload: unlink spool before independent validation: %w", err) + } + if err := directory.sync(); err != nil { + return nil, fmt.Errorf("upload: sync spool unlink: %w", err) + } + secondHasher := sha256.New() + secondSize, err := io.Copy(secondHasher, reader) + if err != nil { + return nil, fmt.Errorf("upload: independently hash spool: %w", err) + } + if secondSize != capability.SourceBytes || + hex.EncodeToString(secondHasher.Sum(nil)) != capability.SourceSHA256 { + return nil, errors.New("upload: independent spool hash does not match authorized source") + } + if _, err := reader.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("upload: rewind spool for capability validation: %w", err) + } + reinspected, err := media.InspectCapability(reader, policy) + if err != nil { + return nil, fmt.Errorf("upload: inspect sealed spool: %w", err) + } + if !reflect.DeepEqual(reinspected, capability) { + return nil, errors.New("upload: sealed spool capability does not match authorization") + } + if _, err := reader.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("upload: rewind authorized spool: %w", err) + } + if err := callTestHook(source.testHook, authorizeStageValidated, directory.path(spoolFilename)); err != nil { + return nil, err + } + providerDigest := sha256.Sum256(metadata.ProviderMetadata) + result := &authorizedUpload{reader: reader, capability: capability, metadata: document.AuthorizedUploadMetadata{ + Filename: metadata.Filename, MediaFamily: capability.MediaFamily, + MediaType: capability.MediaType, ByteLength: capability.SourceBytes, + SHA256: capability.SourceSHA256, CapabilityRecordChecksum: capability.Checksum, + ProviderMetadataChecksum: hex.EncodeToString(providerDigest[:]), InputKind: capability.InputKind, + }} + result.cleanup = directory.cleanup + readerOpen = false + cleanupDirectory = false + result.mu.Lock() + result.stop = context.AfterFunc(ctx, func() { _ = result.Close() }) + contextErr := ctx.Err() + result.mu.Unlock() + if contextErr != nil { + _ = result.Close() + return nil, contextErr + } + authorized = result + return authorized, nil +} + +// RecoverStale removes only package-owned private spool directories. It is +// safe to call at startup after prior process crashes. +func RecoverStale(ctx context.Context, base string) (int, error) { + if base == "" || !filepath.IsAbs(base) { + return 0, errors.New("upload: stale spool root must be an absolute path") + } + root, err := openStableRoot(base) + if err != nil { + return 0, fmt.Errorf("upload: open stale spool root: %w", err) + } + defer func() { _ = root.Close() }() + entries, err := fs.ReadDir(root.FS(), ".") + if err != nil { + return 0, fmt.Errorf("upload: read stale spool root: %w", err) + } + removed := 0 + for _, entry := range entries { + if err := ctx.Err(); err != nil { + return removed, err + } + if !stringsHasPrefix(entry.Name(), spoolDirectoryPrefix) { + continue + } + if entry.Type()&os.ModeSymlink != 0 || !entry.IsDir() { + continue + } + if err := root.RemoveAll(entry.Name()); err != nil { + return removed, fmt.Errorf("upload: remove stale spool %s: %w", entry.Name(), err) + } + removed++ + } + return removed, nil +} + +func openStableRoot(path string) (*os.Root, error) { + before, err := os.Lstat(path) + if err != nil { + return nil, err + } + if before.Mode()&os.ModeSymlink != 0 || !before.IsDir() { + return nil, errors.New("upload: spool root is a symlink, reparse point, or non-directory") + } + root, err := os.OpenRoot(path) + if err != nil { + return nil, err + } + after, err := root.Stat(".") + if err != nil || !after.IsDir() || !os.SameFile(before, after) { + _ = root.Close() + if err != nil { + return nil, err + } + return nil, errors.New("upload: spool root identity changed while opening") + } + return root, nil +} + +func copyContext(ctx context.Context, writer io.Writer, reader io.Reader) (int64, error) { + buffer := make([]byte, 128<<10) + var total int64 + for { + if err := ctx.Err(); err != nil { + return total, err + } + count, readErr := reader.Read(buffer) + if count > 0 { + written, writeErr := writer.Write(buffer[:count]) + total += int64(written) + if writeErr != nil { + return total, writeErr + } + if written != count { + return total, io.ErrShortWrite + } + } + if readErr != nil { + if errors.Is(readErr, io.EOF) { + return total, nil + } + return total, readErr + } + } +} + +func callTestHook(hook func(authorizeStage, string) error, stage authorizeStage, path string) error { + if hook == nil { + return nil + } + return hook(stage, path) +} + +func randomSpoolName() (string, error) { + var value [16]byte + if _, err := rand.Read(value[:]); err != nil { + return "", fmt.Errorf("read spool name entropy: %w", err) + } + return spoolDirectoryPrefix + hex.EncodeToString(value[:]), nil +} + +func stringsHasPrefix(value, prefix string) bool { + return len(value) >= len(prefix) && value[:len(prefix)] == prefix +} diff --git a/document/upload/spool_test.go b/document/upload/spool_test.go new file mode 100644 index 00000000..e89f5df3 --- /dev/null +++ b/document/upload/spool_test.go @@ -0,0 +1,445 @@ +package upload + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "io" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/media" +) + +func TestAuthorizeReturnsOnlyValidatedUnlinkedExactBytes(t *testing.T) { + data := []byte("alpha\nbeta\n") + record := inspectCapability(t, data) + directory := t.TempDir() + var validated bool + upload, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: directory, + testHook: func(stage authorizeStage, _ string) error { + if stage == authorizeStageValidated { + validated = true + } + return nil + }, + }, record, UploadMetadata{Filename: "notes.txt", ProviderMetadata: []byte(`{"mode":"synthetic"}`)}) + require.NoError(t, err) + require.True(t, validated, "the adapter-facing reader cannot exist before validation completes") + metadata := upload.Metadata() + assert.Equal(t, record.SourceSHA256, metadata.SHA256) + assert.Equal(t, record.Checksum, metadata.CapabilityRecordChecksum) + assert.Equal(t, sha256Hex([]byte(`{"mode":"synthetic"}`)), metadata.ProviderMetadataChecksum) + assert.Equal(t, "text", metadata.MediaFamily) + + got, err := io.ReadAll(upload) + require.NoError(t, err) + assert.Equal(t, data, got) + require.NoError(t, upload.Close()) + _, err = upload.Read(make([]byte, 1)) + require.Error(t, err) + assert.Empty(t, spoolEntries(t, directory), "the named spool must be gone before handoff") +} + +// TestAuthorizeCarriesExactLocalCapabilityValue catches an authorized upload +// that drops the locally inspected capability authority or exposes mutable +// shared state to a provider adapter. +func TestAuthorizeCarriesExactLocalCapabilityValue(t *testing.T) { + data := []byte("alpha\nbeta\n") + record := inspectCapability(t, data) + upload, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: t.TempDir(), + }, record, UploadMetadata{Filename: "notes.txt"}) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, upload.Close()) }) + + carrier, ok := upload.(interface { + CapabilityRecord() media.CapabilityRecord + }) + require.True(t, ok) + carried := carrier.CapabilityRecord() + assert.Equal(t, record, carried) + require.NoError(t, media.ValidateCapabilityRecord(carried)) + _, local := carried.InspectionPolicy() + assert.True(t, local) + + carried.Eligible = false + assert.Equal(t, record, carrier.CapabilityRecord(), "the carrier must return a value copy") +} + +func TestAuthorizeRejectsPathReplacementSymlinkAndWriterMutation(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows reparse-point cases run in the platform suite") + } + data := []byte("authoritative bytes\n") + record := inspectCapability(t, data) + tests := []struct { + name string + hook func(string) error + }{ + { + name: "path replacement", + hook: func(path string) error { + if err := os.Rename(path, path+".original"); err != nil { + return err + } + return os.WriteFile(path, data, 0o600) + }, + }, + { + name: "symlink replacement", + hook: func(path string) error { + if err := os.Rename(path, path+".original"); err != nil { + return err + } + return os.Symlink(filepath.Base(path)+".original", path) + }, + }, + { + name: "writer mutation and second hash mismatch", + hook: func(path string) error { + return os.WriteFile(path, []byte("mutated bytes differ\n"), 0o600) + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + directory := t.TempDir() + _, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: directory, + testHook: func(stage authorizeStage, path string) error { + if stage == authorizeStageWritten { + return tt.hook(path) + } + return nil + }, + }, record, UploadMetadata{Filename: "notes.txt"}) + require.Error(t, err) + assert.Empty(t, spoolEntries(t, directory)) + }) + } +} + +func TestAuthorizeCleanupRemainsBoundToOriginalParentDirectory(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("directory renames with open handles are covered by Windows CI") + } + data := []byte("authoritative bytes\n") + record := inspectCapability(t, data) + root := t.TempDir() + base := filepath.Join(root, "spool") + relocated := filepath.Join(root, "relocated") + require.NoError(t, os.Mkdir(base, 0o700)) + var replacementSpool string + + upload, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: base, + testHook: func(stage authorizeStage, path string) error { + if stage != authorizeStageWritten { + return nil + } + if err := os.Rename(base, relocated); err != nil { + return err + } + if err := os.Mkdir(base, 0o700); err != nil { + return err + } + replacementSpool = filepath.Join(base, filepath.Base(filepath.Dir(path))) + return os.Mkdir(replacementSpool, 0o700) + }, + }, record, UploadMetadata{Filename: "notes.txt"}) + require.NoError(t, err) + require.NoError(t, upload.Close()) + _, err = os.Stat(replacementSpool) + require.NoError(t, err, "cleanup must not remove a same-named directory under a replacement parent") + _, err = os.Stat(filepath.Join(relocated, filepath.Base(replacementSpool))) + require.ErrorIs(t, err, os.ErrNotExist, "cleanup must remove the descriptor-held original spool") +} + +func TestAuthorizeRejectsMutatedCapabilityAndSource(t *testing.T) { + data := []byte("alpha\n") + record := inspectCapability(t, data) + record.Measurements.TextLines++ + _, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: t.TempDir(), + }, record, UploadMetadata{Filename: "notes.txt"}) + require.ErrorContains(t, err, "capability") + + record = inspectCapability(t, data) + _, err = Authorize(t.Context(), Source{ + Reader: io.NopCloser(strings.NewReader("omega\n")), Directory: t.TempDir(), + }, record, UploadMetadata{Filename: "notes.txt"}) + require.ErrorContains(t, err, "source") +} + +func TestAuthorizeRejectsSerializedCapabilityBeforeReadingSource(t *testing.T) { + data := []byte("alpha\n") + record := inspectCapability(t, data) + encoded, err := json.Marshal(record) + require.NoError(t, err) + var decoded media.CapabilityRecord + require.NoError(t, json.Unmarshal(encoded, &decoded)) + reader := &countingReadCloser{Reader: bytes.NewReader(data)} + + _, err = Authorize(t.Context(), Source{Reader: reader, Directory: t.TempDir()}, decoded, + UploadMetadata{Filename: "notes.txt"}) + require.ErrorContains(t, err, "local inspection authority") + assert.Zero(t, reader.reads) + assert.True(t, reader.didClose) +} + +func TestAuthorizeCapsUntrustedSourceAtAuthorizedBytesPlusOne(t *testing.T) { + authorized := []byte("alpha\n") + record := inspectCapability(t, authorized) + reader := &countingReadCloser{Reader: bytes.NewReader(bytes.Repeat([]byte("x"), 1<<20))} + _, err := Authorize(t.Context(), Source{Reader: reader, Directory: t.TempDir()}, record, + UploadMetadata{Filename: "notes.txt"}) + require.ErrorContains(t, err, "source") + assert.LessOrEqual(t, reader.bytesRead, int(record.SourceBytes)+1) +} + +func TestAuthorizeSourceCloseFailureDoesNotLeakAuthorizedUpload(t *testing.T) { + data := []byte("alpha\n") + record := inspectCapability(t, data) + directory := t.TempDir() + reader := &countingReadCloser{Reader: bytes.NewReader(data), closeErr: errors.New("synthetic close failure")} + + upload, err := Authorize(t.Context(), Source{Reader: reader, Directory: directory}, record, + UploadMetadata{Filename: "notes.txt"}) + require.ErrorContains(t, err, "close source") + assert.Nil(t, upload) + assert.Equal(t, 1, reader.closes) + assert.Empty(t, spoolEntries(t, directory)) +} + +func TestAuthorizeCancellationClosesSourceAndCleansSpool(t *testing.T) { + record := inspectCapability(t, []byte("alpha\n")) + directory := t.TempDir() + reader := newBlockingReadCloser() + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { + _, err := Authorize(ctx, Source{Reader: reader, Directory: directory}, record, + UploadMetadata{Filename: "notes.txt"}) + done <- err + }() + reader.waitStarted(t) + cancel() + require.ErrorIs(t, <-done, context.Canceled) + assert.True(t, reader.closed()) + assert.Empty(t, spoolEntries(t, directory)) +} + +func TestAuthorizeCancellationClosesReturnedUpload(t *testing.T) { + data := []byte("alpha\n") + record := inspectCapability(t, data) + directory := t.TempDir() + ctx, cancel := context.WithCancel(t.Context()) + upload, err := Authorize(ctx, Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: directory, + }, record, UploadMetadata{Filename: "notes.txt"}) + require.NoError(t, err) + cancel() + require.Eventually(t, func() bool { + _, readErr := upload.Read(make([]byte, 1)) + return readErr != nil + }, 5*time.Second, 10*time.Millisecond) + require.NoError(t, upload.Close()) + assert.Empty(t, spoolEntries(t, directory)) +} + +func TestAuthorizeCancellationAtValidatedHandoffCleansWithoutRacing(t *testing.T) { + data := []byte("alpha\n") + record := inspectCapability(t, data) + directory := t.TempDir() + ctx, cancel := context.WithCancel(t.Context()) + reader := &countingReadCloser{Reader: bytes.NewReader(data)} + + _, err := Authorize(ctx, Source{ + Reader: reader, Directory: directory, + testHook: func(stage authorizeStage, _ string) error { + if stage == authorizeStageValidated { + cancel() + } + return nil + }, + }, record, UploadMetadata{Filename: "notes.txt"}) + require.ErrorIs(t, err, context.Canceled) + assert.True(t, reader.didClose) + assert.Empty(t, spoolEntries(t, directory)) +} + +func TestRecoverStaleRemovesOnlyOwnedSpoolDirectories(t *testing.T) { + base := t.TempDir() + stale := filepath.Join(base, spoolDirectoryPrefix+"stale") + require.NoError(t, os.Mkdir(stale, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(stale, "source"), []byte("stale"), 0o600)) + unrelated := filepath.Join(base, "keep") + require.NoError(t, os.Mkdir(unrelated, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(unrelated, "source"), []byte("keep"), 0o600)) + spoofedFile := filepath.Join(base, spoolDirectoryPrefix+"file") + require.NoError(t, os.WriteFile(spoofedFile, []byte("keep"), 0o600)) + spoofedLink := filepath.Join(base, spoolDirectoryPrefix+"link") + if runtime.GOOS != "windows" { + require.NoError(t, os.Symlink("keep", spoofedLink)) + } + + recovered, err := RecoverStale(t.Context(), base) + require.NoError(t, err) + assert.Equal(t, 1, recovered) + _, err = os.Stat(stale) + require.ErrorIs(t, err, os.ErrNotExist) + _, err = os.Stat(unrelated) + require.NoError(t, err) + _, err = os.Lstat(spoofedFile) + require.NoError(t, err) + if runtime.GOOS != "windows" { + _, err = os.Lstat(spoofedLink) + require.NoError(t, err) + } +} + +func TestRecoverStaleRejectsSymlinkRoot(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("reparse-root coverage runs in the Windows platform suite") + } + target := t.TempDir() + stale := filepath.Join(target, spoolDirectoryPrefix+"stale") + require.NoError(t, os.Mkdir(stale, 0o700)) + link := filepath.Join(t.TempDir(), "spool-link") + require.NoError(t, os.Symlink(target, link)) + + _, err := RecoverStale(t.Context(), link) + require.Error(t, err) + _, err = os.Stat(stale) + require.NoError(t, err) +} + +func TestAuthorizeRejectsIneligibleRecordBeforeReadingSource(t *testing.T) { + data := []byte("%PDF-1.7\n") + policy := inspectionPolicy(data, "report.pdf", "application/pdf") + record, err := media.InspectCapability(bytes.NewReader(data), policy) + require.NoError(t, err) + require.False(t, record.Eligible) + reader := &countingReadCloser{Reader: bytes.NewReader(data)} + _, err = Authorize(t.Context(), Source{Reader: reader, Directory: t.TempDir()}, record, + UploadMetadata{Filename: "report.pdf"}) + require.ErrorContains(t, err, "ineligible") + assert.Zero(t, reader.reads) + assert.True(t, reader.didClose) +} + +func inspectCapability(t *testing.T, data []byte) media.CapabilityRecord { + t.Helper() + record, err := media.InspectCapability(bytes.NewReader(data), inspectionPolicy(data, "notes.txt", "text/plain")) + require.NoError(t, err) + require.True(t, record.Eligible) + return record +} + +func inspectionPolicy(data []byte, filename, mediaType string) media.InspectionPolicy { + return media.InspectionPolicy{ + Filename: filename, DeclaredMediaType: mediaType, + ExpectedBytes: int64(len(data)), ExpectedSHA256: sha256Hex(data), + DescriptorFingerprint: strings.Repeat("a", 64), ProfileFingerprint: strings.Repeat("b", 64), + DisclosureFingerprint: strings.Repeat("c", 64), InputKind: document.RenditionInputOriginalFile, + MaxSourceBytes: 1 << 20, MaxExpandedBytes: 1 << 20, MaxEntryBytes: 1 << 20, + MaxEntries: 100, MaxNestingDepth: 1, MaxTextLines: 1_000, MaxCharacters: 1 << 20, + MaxPages: 100, MaxSlides: 100, MaxSheets: 100, MaxCells: 10_000, MaxSpineItems: 1_000, MaxResources: 10_000, + } +} + +func spoolEntries(t *testing.T, base string) []string { + t.Helper() + entries, err := os.ReadDir(base) + require.NoError(t, err) + result := make([]string, 0, len(entries)) + for _, entry := range entries { + result = append(result, entry.Name()) + } + return result +} + +func sha256Hex(data []byte) string { + digest := sha256.Sum256(data) + return hex.EncodeToString(digest[:]) +} + +type countingReadCloser struct { + io.Reader + + reads int + bytesRead int + didClose bool + closes int + closeErr error +} + +func (reader *countingReadCloser) Read(buffer []byte) (int, error) { + reader.reads++ + count, err := reader.Reader.Read(buffer) + reader.bytesRead += count + return count, err +} + +func (reader *countingReadCloser) Close() error { + reader.didClose = true + reader.closes++ + return reader.closeErr +} + +type blockingReadCloser struct { + started chan struct{} + once sync.Once + closedC chan struct{} +} + +func newBlockingReadCloser() *blockingReadCloser { + return &blockingReadCloser{started: make(chan struct{}), closedC: make(chan struct{})} +} + +func (reader *blockingReadCloser) Read([]byte) (int, error) { + reader.once.Do(func() { close(reader.started) }) + <-reader.closedC + return 0, errors.New("closed") +} + +func (reader *blockingReadCloser) Close() error { + select { + case <-reader.closedC: + default: + close(reader.closedC) + } + return nil +} + +func (reader *blockingReadCloser) waitStarted(t *testing.T) { + t.Helper() + select { + case <-reader.started: + case <-time.After(5 * time.Second): + t.Fatal("source read did not start") + } +} + +func (reader *blockingReadCloser) closed() bool { + select { + case <-reader.closedC: + return true + default: + return false + } +} diff --git a/document/upload/spool_unix.go b/document/upload/spool_unix.go new file mode 100644 index 00000000..3dda5c55 --- /dev/null +++ b/document/upload/spool_unix.go @@ -0,0 +1,116 @@ +//go:build !windows + +package upload + +import ( + "errors" + "fmt" + "os" + "path/filepath" + + "golang.org/x/sys/unix" +) + +type spoolDirectory struct { + base string + name string + root *os.File + dir *os.File +} + +func openSpoolDirectory(base string) (*spoolDirectory, error) { + name, err := randomSpoolName() + if err != nil { + return nil, fmt.Errorf("create no-follow spool: %w", err) + } + rootFD, err := unix.Open(base, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return nil, fmt.Errorf("open no-follow spool root: %w", err) + } + root := os.NewFile(uintptr(rootFD), base) + if err := unix.Mkdirat(rootFD, name, 0o700); err != nil { + _ = root.Close() + return nil, fmt.Errorf("create private spool directory: %w", err) + } + dirFD, err := unix.Openat(rootFD, name, + unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + _ = unix.Unlinkat(rootFD, name, unix.AT_REMOVEDIR) + _ = root.Close() + return nil, fmt.Errorf("open private spool directory: %w", err) + } + dir := os.NewFile(uintptr(dirFD), name) + return &spoolDirectory{base: base, name: name, root: root, dir: dir}, nil +} + +func (directory *spoolDirectory) create(name string) (*os.File, error) { + fd, err := unix.Openat(int(directory.dir.Fd()), name, + unix.O_WRONLY|unix.O_CREAT|unix.O_EXCL|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0o600) + if err != nil { + return nil, fmt.Errorf("create no-follow spool: %w", err) + } + return os.NewFile(uintptr(fd), name), nil +} + +func (directory *spoolDirectory) openReader(name string, written os.FileInfo) (*os.File, error) { + fd, err := unix.Openat(int(directory.dir.Fd()), name, + unix.O_RDONLY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return nil, fmt.Errorf("open no-follow spool: %w", err) + } + reader := os.NewFile(uintptr(fd), name) + info, err := reader.Stat() + if err != nil { + _ = reader.Close() + return nil, fmt.Errorf("stat no-follow spool reader: %w", err) + } + if !info.Mode().IsRegular() || !os.SameFile(written, info) || info.Size() != written.Size() { + _ = reader.Close() + return nil, errors.New("spool reader does not name the synced writer identity") + } + return reader, nil +} + +func (directory *spoolDirectory) unlink(name string) error { + if err := unix.Unlinkat(int(directory.dir.Fd()), name, 0); err != nil { + return fmt.Errorf("unlink descriptor-relative spool: %w", err) + } + return nil +} + +func (directory *spoolDirectory) sync() error { + return errors.Join(directory.dir.Sync(), directory.root.Sync()) +} + +func (directory *spoolDirectory) path(name string) string { + return filepath.Join(directory.base, directory.name, name) +} + +func (directory *spoolDirectory) cleanup() error { + if directory == nil { + return nil + } + if directory.dir != nil { + _ = unix.Unlinkat(int(directory.dir.Fd()), spoolFilename, 0) + _, _ = directory.dir.Seek(0, 0) + if names, err := directory.dir.Readdirnames(-1); err == nil { + for _, name := range names { + _ = unix.Unlinkat(int(directory.dir.Fd()), name, 0) + } + } + } + var result error + if directory.dir != nil { + result = directory.dir.Close() + directory.dir = nil + } + if directory.root != nil { + if err := unix.Unlinkat(int(directory.root.Fd()), directory.name, unix.AT_REMOVEDIR); err != nil && + !errors.Is(err, os.ErrNotExist) { + result = errors.Join(result, fmt.Errorf("remove private spool directory: %w", err)) + } + result = errors.Join(result, directory.root.Sync(), directory.root.Close()) + directory.root = nil + } + return result +} diff --git a/document/upload/spool_windows.go b/document/upload/spool_windows.go new file mode 100644 index 00000000..fcf3f1d7 --- /dev/null +++ b/document/upload/spool_windows.go @@ -0,0 +1,106 @@ +//go:build windows + +package upload + +import ( + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" +) + +type spoolDirectory struct { + base string + name string + baseRoot *os.Root + root *os.Root +} + +func openSpoolDirectory(base string) (*spoolDirectory, error) { + name, err := randomSpoolName() + if err != nil { + return nil, err + } + baseRoot, err := openStableRoot(base) + if err != nil { + return nil, err + } + if err := baseRoot.Mkdir(name, 0o700); err != nil { + _ = baseRoot.Close() + return nil, err + } + root, err := baseRoot.OpenRoot(name) + if err != nil { + _ = baseRoot.Remove(name) + _ = baseRoot.Close() + return nil, err + } + return &spoolDirectory{base: base, name: name, baseRoot: baseRoot, root: root}, nil +} + +func (directory *spoolDirectory) create(name string) (*os.File, error) { + return directory.root.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) +} + +func (directory *spoolDirectory) openReader(name string, written os.FileInfo) (*os.File, error) { + linkInfo, err := directory.root.Lstat(name) + if err != nil { + return nil, err + } + if linkInfo.Mode()&os.ModeSymlink != 0 || !linkInfo.Mode().IsRegular() { + return nil, errors.New("spool reader path is a reparse point or non-regular file") + } + reader, err := directory.root.Open(name) + if err != nil { + return nil, err + } + info, err := reader.Stat() + if err != nil { + _ = reader.Close() + return nil, err + } + if !os.SameFile(written, info) || info.Size() != written.Size() { + _ = reader.Close() + return nil, errors.New("spool reader does not name the synced writer identity") + } + return reader, nil +} + +func (directory *spoolDirectory) unlink(name string) error { + return directory.root.Remove(name) +} + +func (directory *spoolDirectory) sync() error { return nil } + +func (directory *spoolDirectory) path(name string) string { + return filepath.Join(directory.base, directory.name, name) +} + +func (directory *spoolDirectory) cleanup() error { + if directory == nil { + return nil + } + var result error + if directory.root != nil { + entries, err := fs.ReadDir(directory.root.FS(), ".") + if err != nil && !errors.Is(err, os.ErrNotExist) { + result = err + } + for _, entry := range entries { + if err := directory.root.RemoveAll(entry.Name()); err != nil && !errors.Is(err, os.ErrNotExist) { + result = errors.Join(result, err) + } + } + result = errors.Join(result, directory.root.Close()) + directory.root = nil + } + if directory.baseRoot != nil { + if err := directory.baseRoot.Remove(directory.name); err != nil && !errors.Is(err, os.ErrNotExist) { + result = errors.Join(result, fmt.Errorf("remove private spool directory: %w", err)) + } + result = errors.Join(result, directory.baseRoot.Close()) + directory.baseRoot = nil + } + return result +} diff --git a/document/upload/spool_windows_test.go b/document/upload/spool_windows_test.go new file mode 100644 index 00000000..2ad82afd --- /dev/null +++ b/document/upload/spool_windows_test.go @@ -0,0 +1,47 @@ +//go:build windows + +package upload + +import ( + "bytes" + "io" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAuthorizeWindowsRejectsReparseReplacement(t *testing.T) { + data := []byte("authoritative bytes\n") + record := inspectCapability(t, data) + directory := t.TempDir() + _, err := Authorize(t.Context(), Source{ + Reader: io.NopCloser(bytes.NewReader(data)), Directory: directory, + testHook: func(stage authorizeStage, path string) error { + if stage != authorizeStageWritten { + return nil + } + if err := os.Rename(path, path+".original"); err != nil { + return err + } + return os.Symlink(filepath.Base(path)+".original", path) + }, + }, record, UploadMetadata{Filename: "notes.txt"}) + require.Error(t, err) + assert.Empty(t, spoolEntries(t, directory)) +} + +func TestRecoverStaleWindowsRejectsReparseRoot(t *testing.T) { + target := t.TempDir() + stale := filepath.Join(target, spoolDirectoryPrefix+"stale") + require.NoError(t, os.Mkdir(stale, 0o700)) + link := filepath.Join(t.TempDir(), "spool-link") + require.NoError(t, os.Symlink(target, link)) + + _, err := RecoverStale(t.Context(), link) + require.Error(t, err) + _, err = os.Stat(stale) + require.NoError(t, err) +} diff --git a/document/vector_set.go b/document/vector_set.go new file mode 100644 index 00000000..d8cb40c0 --- /dev/null +++ b/document/vector_set.go @@ -0,0 +1,411 @@ +package document + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "errors" + "fmt" + "io" + "math" + "unicode/utf8" +) + +const ( + vectorSetV1Magic = "vector-set/v1\x00" + vectorSetV1Domain = "docbank-vector-set/v1\x00" + vectorSetV1Version = uint32(1) + maxVectorSetString = 1 << 16 + maxVectorSetScalars = 10_000_000 +) + +const ( + VectorNormalizationNone = "none" + VectorNormalizationUnitLength = "unit_length" +) + +// VectorSetV1 is the immutable vector payload apart from its storage record. +type VectorSetV1 struct { + VectorSpaceFingerprint string `json:"vector_space_fingerprint"` + Metric string `json:"metric"` + Normalization string `json:"normalization"` + Dimension int `json:"dimension"` + InputKeys []string `json:"input_keys"` + InputChecksums []string `json:"input_checksums"` + Vectors [][]float32 `json:"vectors"` +} + +// VectorSetV1Input accepts provider precision before it crosses the durable +// float32 boundary. Conversion follows IEEE-754 round-to-nearest-even. +type VectorSetV1Input struct { + VectorSpaceFingerprint string + Metric string + Normalization string + Dimension int + InputKeys []string + InputChecksums []string + Values [][]float64 +} + +// NewVectorSetV1 constructs canonical float32 rows from provider float64 +// values. It rejects values that are non-finite before or after quantization. +func NewVectorSetV1(input VectorSetV1Input) (VectorSetV1, error) { + if err := preflightVectorSetInput(input); err != nil { + return VectorSetV1{}, err + } + set := VectorSetV1{VectorSpaceFingerprint: input.VectorSpaceFingerprint, Metric: input.Metric, Normalization: input.Normalization, Dimension: input.Dimension, InputKeys: append([]string(nil), input.InputKeys...), InputChecksums: append([]string(nil), input.InputChecksums...), Vectors: make([][]float32, len(input.Values))} + for row, values := range input.Values { + set.Vectors[row] = make([]float32, len(values)) + for column, value := range values { + if math.IsNaN(value) || math.IsInf(value, 0) { + return VectorSetV1{}, errors.New("vector input contains non-finite scalar") + } + quantized := float32(value) + if math.IsInf(float64(quantized), 0) { + return VectorSetV1{}, errors.New("vector input exceeds finite float32 range") + } + if quantized == 0 { + quantized = 0 + } + set.Vectors[row][column] = quantized + } + } + if err := validateVectorSetV1(set); err != nil { + return VectorSetV1{}, err + } + return set, nil +} + +func preflightVectorSetInput(input VectorSetV1Input) error { + if err := validateFingerprint(input.VectorSpaceFingerprint, "vector-space fingerprint"); err != nil { + return err + } + if !validVectorMetric(input.Metric) { + return errors.New("vector metric is invalid") + } + if !validVectorNormalization(input.Normalization) { + return errors.New("vector normalization is invalid") + } + if input.Dimension < 1 || input.Dimension > maxEmbeddingDimensions || len(input.Values) == 0 || len(input.Values) > maxEmbeddingBatchItems || len(input.InputKeys) != len(input.Values) || len(input.InputChecksums) != len(input.Values) || uint64(len(input.Values))*uint64(input.Dimension) > maxVectorSetScalars { + return errors.New("vector input allocation exceeds bounds") + } + seen := make(map[string]struct{}, len(input.InputKeys)) + for index, row := range input.Values { + if len(row) != input.Dimension { + return errors.New("vector input row dimension does not match header") + } + if err := validateVectorSetKey(input.InputKeys[index]); err != nil { + return err + } + if _, exists := seen[input.InputKeys[index]]; exists { + return errors.New("vector input keys must be unique") + } + seen[input.InputKeys[index]] = struct{}{} + if err := validateFingerprint(input.InputChecksums[index], "vector input checksum"); err != nil { + return err + } + for _, value := range row { + if math.IsNaN(value) || math.IsInf(value, 0) { + return errors.New("vector input contains non-finite scalar") + } + if math.IsInf(float64(float32(value)), 0) { + return errors.New("vector input exceeds finite float32 range") + } + } + } + return nil +} + +// VectorBounds limits untrusted durable payloads before their rows or vectors +// are allocated. All fields must be positive for a decode operation. +type VectorBounds struct { + MaxRows int + MaxDimension int + MaxBytes int +} + +// EncodeVectorSetV1 emits the platform-independent vector-set/v1 framing and +// a SHA-256 over its domain prefix followed by the exact emitted bytes. +func EncodeVectorSetV1(set VectorSetV1) ([]byte, string, error) { + if err := validateVectorSetV1(set); err != nil { + return nil, "", err + } + var buffer bytes.Buffer + buffer.WriteString(vectorSetV1Magic) + writeUint32(&buffer, vectorSetV1Version) + for _, value := range []string{set.VectorSpaceFingerprint, set.Metric, set.Normalization} { + if err := writeString(&buffer, value); err != nil { + return nil, "", err + } + } + if err := writeBoundedUint32(&buffer, len(set.Vectors)); err != nil { + return nil, "", err + } + if err := writeBoundedUint32(&buffer, set.Dimension); err != nil { + return nil, "", err + } + for index := range set.Vectors { + if err := writeString(&buffer, set.InputKeys[index]); err != nil { + return nil, "", err + } + if err := writeString(&buffer, set.InputChecksums[index]); err != nil { + return nil, "", err + } + } + for _, vector := range set.Vectors { + for _, value := range vector { + if value == 0 { + value = 0 + } + writeUint32(&buffer, math.Float32bits(value)) + } + } + encoded := buffer.Bytes() + checksum := vectorSetChecksum(encoded) + return encoded, checksum, nil +} + +// DecodeVectorSetV1 validates and reads one exactly framed vector-set/v1 payload. +func DecodeVectorSetV1(encoded []byte, bounds VectorBounds) (VectorSetV1, error) { + if bounds.MaxRows < 1 || bounds.MaxDimension < 1 || bounds.MaxBytes < 1 { + return VectorSetV1{}, errors.New("vector bounds must permit positive rows, dimension, and bytes") + } + if len(encoded) > bounds.MaxBytes { + return VectorSetV1{}, errors.New("vector payload exceeds byte bounds") + } + reader := bytes.NewReader(encoded) + magic := make([]byte, len(vectorSetV1Magic)) + if _, err := io.ReadFull(reader, magic); err != nil || string(magic) != vectorSetV1Magic { + return VectorSetV1{}, errors.New("vector payload has an invalid header") + } + if version, err := readUint32(reader); err != nil || version != vectorSetV1Version { + return VectorSetV1{}, errors.New("vector payload has an unsupported version") + } + space, err := readString(reader) + if err != nil { + return VectorSetV1{}, err + } + metric, err := readString(reader) + if err != nil { + return VectorSetV1{}, err + } + normalization, err := readString(reader) + if err != nil { + return VectorSetV1{}, err + } + if err := validateFingerprint(space, "vector-space fingerprint"); err != nil { + return VectorSetV1{}, err + } + if !validVectorMetric(metric) || !validVectorNormalization(normalization) { + return VectorSetV1{}, errors.New("vector payload has an invalid header") + } + rows, err := readUint32(reader) + if err != nil { + return VectorSetV1{}, err + } + dimension, err := readUint32(reader) + if err != nil { + return VectorSetV1{}, err + } + if rows == 0 || uint64(rows) > uint64(bounds.MaxRows) || rows > maxEmbeddingBatchItems { + return VectorSetV1{}, errors.New("vector payload rows exceed bounds") + } + if dimension == 0 || uint64(dimension) > uint64(bounds.MaxDimension) || dimension > maxEmbeddingDimensions { + return VectorSetV1{}, errors.New("vector payload dimension exceeds bounds") + } + scalars := uint64(rows) * uint64(dimension) + if scalars > maxVectorSetScalars { + return VectorSetV1{}, errors.New("vector payload allocation exceeds bounds") + } + metadata := *reader + if err := preflightVectorMetadata(&metadata, rows, scalars); err != nil { + return VectorSetV1{}, err + } + set := VectorSetV1{VectorSpaceFingerprint: space, Metric: metric, Normalization: normalization, Dimension: int(dimension), InputKeys: make([]string, int(rows)), InputChecksums: make([]string, int(rows)), Vectors: make([][]float32, int(rows))} + for index := range set.Vectors { + key, err := readString(reader) + if err != nil { + return VectorSetV1{}, err + } + checksum, err := readString(reader) + if err != nil { + return VectorSetV1{}, err + } + if err := validateVectorSetKey(key); err != nil { + return VectorSetV1{}, err + } + if err := validateFingerprint(checksum, "vector input checksum"); err != nil { + return VectorSetV1{}, err + } + set.InputKeys[index], set.InputChecksums[index] = key, checksum + } + for index := range set.Vectors { + set.Vectors[index] = make([]float32, int(dimension)) + for scalar := range set.Vectors[index] { + bits, err := readUint32(reader) + if err != nil { + return VectorSetV1{}, err + } + value := math.Float32frombits(bits) + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return VectorSetV1{}, errors.New("vector payload contains non-finite scalar") + } + if value == 0 && bits != 0 { + return VectorSetV1{}, errors.New("vector payload contains non-canonical negative zero") + } + set.Vectors[index][scalar] = value + } + } + if reader.Len() != 0 { + return VectorSetV1{}, errors.New("vector payload has trailing bytes") + } + if err := validateVectorSetV1(set); err != nil { + return VectorSetV1{}, err + } + return set, nil +} + +// preflightVectorMetadata proves the entire remaining frame is canonical before +// DecodeVectorSetV1 allocates the caller-selected rows and scalar buffers. +func preflightVectorMetadata(reader *bytes.Reader, rows uint32, scalars uint64) error { + seen := make(map[string]struct{}, rows) + for range rows { + key, err := readString(reader) + if err != nil { + return err + } + checksum, err := readString(reader) + if err != nil { + return err + } + if err := validateVectorSetKey(key); err != nil { + return err + } + if _, exists := seen[key]; exists { + return errors.New("vector input keys must be unique") + } + seen[key] = struct{}{} + if err := validateFingerprint(checksum, "vector input checksum"); err != nil { + return err + } + } + for range scalars { + bits, err := readUint32(reader) + if err != nil { + return err + } + value := math.Float32frombits(bits) + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("vector payload contains non-finite scalar") + } + if value == 0 && bits != 0 { + return errors.New("vector payload contains non-canonical negative zero") + } + } + if reader.Len() != 0 { + return errors.New("vector payload scalar bytes do not match header") + } + return nil +} + +func validateVectorSetV1(set VectorSetV1) error { + if err := validateFingerprint(set.VectorSpaceFingerprint, "vector-space fingerprint"); err != nil { + return err + } + if !validVectorMetric(set.Metric) { + return errors.New("vector metric is invalid") + } + if !validVectorNormalization(set.Normalization) { + return errors.New("vector normalization is invalid") + } + if set.Dimension < 1 || set.Dimension > maxEmbeddingDimensions { + return fmt.Errorf("vector dimension must be between 1 and %d", maxEmbeddingDimensions) + } + if len(set.Vectors) == 0 || len(set.Vectors) > maxEmbeddingBatchItems || len(set.InputKeys) != len(set.Vectors) || len(set.InputChecksums) != len(set.Vectors) { + return errors.New("vector rows, keys, and checksums must have equal non-zero length") + } + if uint64(len(set.Vectors))*uint64(set.Dimension) > maxVectorSetScalars { + return errors.New("vector scalar allocation exceeds bounds") + } + seen := make(map[string]struct{}, len(set.InputKeys)) + for index, vector := range set.Vectors { + if err := validateVectorSetKey(set.InputKeys[index]); err != nil { + return err + } + if _, exists := seen[set.InputKeys[index]]; exists { + return errors.New("vector input keys must be unique") + } + seen[set.InputKeys[index]] = struct{}{} + if err := validateFingerprint(set.InputChecksums[index], "vector input checksum"); err != nil { + return err + } + if len(vector) != set.Dimension { + return errors.New("vector row dimension does not match header") + } + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("vector contains non-finite scalar") + } + } + } + return nil +} + +func validVectorNormalization(value string) bool { + return value == VectorNormalizationNone || value == VectorNormalizationUnitLength +} + +func validateVectorSetKey(value string) error { + if value == "" || len(value) > maxVectorSetString || !utf8.ValidString(value) { + return errors.New("vector input key must be bounded non-empty UTF-8") + } + return nil +} +func vectorSetChecksum(encoded []byte) string { + digest := sha256.Sum256(append([]byte(vectorSetV1Domain), encoded...)) + return hex.EncodeToString(digest[:]) +} +func writeUint32(buffer *bytes.Buffer, value uint32) { + var data [4]byte + binary.LittleEndian.PutUint32(data[:], value) + buffer.Write(data[:]) +} +func writeString(buffer *bytes.Buffer, value string) error { + if err := writeBoundedUint32(buffer, len(value)); err != nil { + return err + } + buffer.WriteString(value) + return nil +} +func writeBoundedUint32(buffer *bytes.Buffer, value int) error { + if value < 0 || uint64(value) > math.MaxUint32 { + return errors.New("vector value cannot fit uint32 framing") + } + writeUint32(buffer, uint32(value)) + return nil +} +func readUint32(reader *bytes.Reader) (uint32, error) { + var data [4]byte + if _, err := io.ReadFull(reader, data[:]); err != nil { + return 0, errors.New("vector payload is truncated") + } + return binary.LittleEndian.Uint32(data[:]), nil +} +func readString(reader *bytes.Reader) (string, error) { + length, err := readUint32(reader) + if err != nil { + return "", err + } + if length == 0 || length > maxVectorSetString || int64(length) > int64(reader.Len()) { + return "", errors.New("vector payload has invalid length-prefixed string") + } + data := make([]byte, int(length)) + if _, err := io.ReadFull(reader, data); err != nil { + return "", errors.New("vector payload is truncated") + } + if !utf8.Valid(data) { + return "", errors.New("vector payload string is not valid UTF-8") + } + return string(data), nil +} diff --git a/document/vector_set_test.go b/document/vector_set_test.go new file mode 100644 index 00000000..57dbd039 --- /dev/null +++ b/document/vector_set_test.go @@ -0,0 +1,125 @@ +package document_test + +import ( + "bytes" + "encoding/binary" + "math" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +// This test fails if a future codec change alters the durable bytes, including +// its header, framed identities, float byte order, or negative-zero handling. +func TestVectorSetV1EncodesCanonicalGoldenBytes(t *testing.T) { + set, err := document.NewVectorSetV1(document.VectorSetV1Input{ + VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, Dimension: 2, + InputKeys: []string{"chunk-a", "chunk-b"}, + InputChecksums: []string{testFingerprint(), testFingerprint()}, + Values: [][]float64{{math.Copysign(0, -1), 1 + math.Exp2(-24)}, {1 + 3*math.Exp2(-24), -2.5}}, + }) + require.NoError(t, err) + got, checksum, err := document.EncodeVectorSetV1(set) + require.NoError(t, err) + want, err := os.ReadFile("testdata/vector-set-v1.golden.bin") + require.NoError(t, err) + assert.Equal(t, want, got) + assert.Equal(t, "6aa367d8ac92d6893dd8768980ea06cc2024edf0822e58d34de2bafcdf65263f", checksum) + assert.Equal(t, uint32(0), math.Float32bits(set.Vectors[0][0])) + assert.Equal(t, uint32(0x3f800000), math.Float32bits(set.Vectors[0][1])) + assert.Equal(t, uint32(0x3f800002), math.Float32bits(set.Vectors[1][0])) + decoded, err := document.DecodeVectorSetV1(got, document.VectorBounds{MaxRows: 2, MaxDimension: 2, MaxBytes: len(got)}) + require.NoError(t, err) + assert.False(t, math.Signbit(float64(decoded.Vectors[0][0]))) +} + +// This test fails if untrusted vector bytes can allocate beyond caller policy, +// retain a non-canonical negative zero, or ignore trailing bytes. +func TestVectorSetV1RejectsMalformedOrUnboundedBytes(t *testing.T) { + set := document.VectorSetV1{ + VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationNone, Dimension: 1, + InputKeys: []string{"chunk-a"}, InputChecksums: []string{testFingerprint()}, Vectors: [][]float32{{1}}, + } + encoded, _, err := document.EncodeVectorSetV1(set) + require.NoError(t, err) + + _, err = document.DecodeVectorSetV1(encoded, document.VectorBounds{MaxRows: 0, MaxDimension: 1, MaxBytes: len(encoded)}) + require.ErrorContains(t, err, "rows") + _, err = document.DecodeVectorSetV1(append(encoded, 0), document.VectorBounds{MaxRows: 1, MaxDimension: 1, MaxBytes: len(encoded) + 1}) + require.ErrorContains(t, err, "scalar bytes") +} + +// This test fails if vector-set framing cannot retain the same normalization +// identifier that a canonical embedding binding uses. +func TestVectorSetV1PreservesUnitLengthNormalization(t *testing.T) { + set, err := document.NewVectorSetV1(document.VectorSetV1Input{ + VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, Dimension: 1, + InputKeys: []string{"chunk-a"}, InputChecksums: []string{testFingerprint()}, Values: [][]float64{{1}}, + }) + require.NoError(t, err) + assert.Equal(t, document.VectorNormalizationUnitLength, set.Normalization) +} + +func TestVectorSetV1RejectsL2NormalizationAlias(t *testing.T) { + _, err := document.NewVectorSetV1(document.VectorSetV1Input{ + VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, + Normalization: "l2", Dimension: 1, + InputKeys: []string{"chunk-a"}, InputChecksums: []string{testFingerprint()}, Values: [][]float64{{1}}, + }) + require.ErrorContains(t, err, "normalization") +} + +func TestVectorSetV1RejectsEveryTruncatedFixedWidthField(t *testing.T) { + set, err := document.NewVectorSetV1(document.VectorSetV1Input{VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationNone, Dimension: 1, InputKeys: []string{"chunk-a"}, InputChecksums: []string{testFingerprint()}, Values: [][]float64{{1}}}) + require.NoError(t, err) + encoded, _, err := document.EncodeVectorSetV1(set) + require.NoError(t, err) + for _, payload := range [][]byte{encoded[:3], encoded[:len(encoded)-1]} { + _, err := document.DecodeVectorSetV1(payload, document.VectorBounds{MaxRows: 1, MaxDimension: 1, MaxBytes: len(encoded)}) + require.Error(t, err) + } +} + +func TestVectorSetV1PreflightsUntrustedFrameBeforeAllocatingRows(t *testing.T) { + set, err := document.NewVectorSetV1(document.VectorSetV1Input{VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationNone, Dimension: 1, InputKeys: []string{"chunk-a"}, InputChecksums: []string{testFingerprint()}, Values: [][]float64{{1}}}) + require.NoError(t, err) + encoded, _, err := document.EncodeVectorSetV1(set) + require.NoError(t, err) + + // Header framing is fixed through the two uint32 row/dimension fields. + rowsOffset := 14 + 4 + 4 + 64 + 4 + len(document.VectorMetricCosine) + 4 + len(document.VectorNormalizationNone) + oversizedRows := append([]byte(nil), encoded...) + binary.LittleEndian.PutUint32(oversizedRows[rowsOffset:], ^uint32(0)) + _, err = document.DecodeVectorSetV1(oversizedRows, document.VectorBounds{MaxRows: int(^uint32(0)), MaxDimension: 1, MaxBytes: len(oversizedRows)}) + require.ErrorContains(t, err, "rows exceed bounds") + + invalidHeader := append([]byte(nil), encoded...) + invalidHeader[14+4+4] = 'x' + _, err = document.DecodeVectorSetV1(invalidHeader, document.VectorBounds{MaxRows: 1, MaxDimension: 1, MaxBytes: len(invalidHeader)}) + require.ErrorContains(t, err, "fingerprint") + + _, err = document.NewVectorSetV1(document.VectorSetV1Input{VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationNone, Dimension: 1, InputKeys: make([]string, 10_001), InputChecksums: make([]string, 10_001), Values: make([][]float64, 10_001)}) + require.ErrorContains(t, err, "allocation exceeds bounds") +} + +// This test fails if duplicate metadata can proceed to scalar validation or +// row allocation instead of being rejected during metadata preflight. +func TestVectorSetV1RejectsDuplicateMetadataKeysBeforeScalarAllocation(t *testing.T) { + set, err := document.NewVectorSetV1(document.VectorSetV1Input{VectorSpaceFingerprint: testFingerprint(), Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationNone, Dimension: 1, InputKeys: []string{"chunk-a", "chunk-b"}, InputChecksums: []string{testFingerprint(), testFingerprint()}, Values: [][]float64{{1}, {2}}}) + require.NoError(t, err) + encoded, _, err := document.EncodeVectorSetV1(set) + require.NoError(t, err) + duplicate := append([]byte(nil), encoded...) + secondKey := bytes.LastIndex(duplicate, []byte("chunk-b")) + require.GreaterOrEqual(t, secondKey, 0) + duplicate[secondKey+len("chunk-")] = 'a' + binary.LittleEndian.PutUint32(duplicate[len(duplicate)-4:], math.Float32bits(float32(math.NaN()))) + _, err = document.DecodeVectorSetV1(duplicate, document.VectorBounds{MaxRows: 2, MaxDimension: 1, MaxBytes: len(duplicate)}) + require.ErrorContains(t, err, "keys must be unique") +} diff --git a/document/voyage/client.go b/document/voyage/client.go index 88ad0ffc..e1a5c77b 100644 --- a/document/voyage/client.go +++ b/document/voyage/client.go @@ -16,6 +16,7 @@ import ( "go.kenn.io/docbank/document/internal/manifestjson" "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/providerhttp" ) const ( @@ -112,7 +113,7 @@ func NewClient(policy Policy, config ClientConfig) (*Client, error) { clone := *config.HTTPClient httpClient = &clone } - httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } + httpClient.CheckRedirect = providerhttp.RefuseRedirects return &Client{ policy: policy, apiKey: config.APIKey, timeout: config.Timeout, maxRetries: config.MaxRetries, retryBaseDelay: config.RetryBaseDelay, http: httpClient, now: time.Now, diff --git a/document/voyage/embedding.go b/document/voyage/embedding.go new file mode 100644 index 00000000..7042bb46 --- /dev/null +++ b/document/voyage/embedding.go @@ -0,0 +1,740 @@ +package voyage + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + json "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "net/http" + "net/netip" + "net/url" + "reflect" + "slices" + "strconv" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/internal/manifestjson" + "go.kenn.io/docbank/document/media" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + EmbeddingProviderID = "voyage.embeddings-v1" + EmbeddingDocumentFormatterV1 = "voyage/document/v1" + EmbeddingQueryFormatterV1 = "voyage/query/v1" + EmbeddingScalarFloat32 = "float32" + TextModel = "voyage-4" + ContextualModel = "voyage-context-4" + HostedAliasRevision = "mutable-alias-export-only" + + textEmbeddingsPath = "/embeddings" + contextualEmbeddingsPath = "/contextualizedembeddings" + embeddingAdapterContract = "docbank-voyage-embeddings/v1" + + defaultEmbeddingMaxBatchItems = 128 + defaultEmbeddingMaxInputBytes = int64(1 << 20) + defaultEmbeddingMaxRequestBytes = int64(2 << 20) + defaultEmbeddingMaxResponseBytes = int64(32 << 20) + maxEmbeddingSecretBytes = 64 << 10 + unitLengthTolerance = 1e-4 +) + +type EmbeddingMode string + +const ( + EmbeddingModeText EmbeddingMode = "text" + EmbeddingModeContextual EmbeddingMode = "contextual" + EmbeddingModeDirectFile EmbeddingMode = "direct_file" +) + +type SecretResolver interface { + ResolveSecret(ctx context.Context, name string) (string, error) +} + +type EmbeddingProfile struct { + Mode EmbeddingMode + Endpoint string + EgressPolicy providerhttp.EgressPolicy + Descriptor document.EmbeddingDescriptor + ModelInput document.ModelInputContract + SecretBinding string + ChunkerVersion string + RequestTimeout time.Duration + MaxRetries int + RetryBaseDelay time.Duration + MaxBatchItems int + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + Policy Policy + CapabilityManifest CapabilityManifest +} + +type embeddingPolicyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Endpoint string `json:"endpoint"` + Route string `json:"route"` + Mode EmbeddingMode `json:"mode"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + ModelInput document.ModelInputContract `json:"model_input"` + CredentialBinding string `json:"credential_binding"` + Egress embeddingEgressIdentity `json:"egress"` + ChunkerVersion string `json:"chunker_version,omitempty"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxRetries int `json:"max_retries"` + RetryBaseDelay int64 `json:"retry_base_delay_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + CapabilityPolicy string `json:"capability_policy,omitempty"` +} + +type EmbeddingClient struct { + profile EmbeddingProfile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client + now func() time.Time +} + +var _ document.EmbeddingProvider = (*EmbeddingClient)(nil) + +type embeddingWireRequest struct { + Input []string `json:"input,omitempty"` + Inputs [][]string `json:"inputs,omitempty"` + Model string `json:"model"` + InputType string `json:"input_type"` + Truncation bool `json:"truncation"` + OutputDimension int `json:"output_dimension"` + OutputDType string `json:"output_dtype"` +} + +type contextualWireRequest struct { + Inputs [][]string `json:"inputs"` + Model string `json:"model"` + InputType string `json:"input_type"` + OutputDimension int `json:"output_dimension"` + OutputDType string `json:"output_dtype"` +} + +type embeddingWireItem struct { + Object string `json:"object"` + Embedding []float32 `json:"embedding"` + Index *int `json:"index"` +} + +type embeddingWireResponse struct { + Object string `json:"object"` + Data []embeddingWireItem `json:"data"` + Model string `json:"model"` + Usage *struct { + TotalTokens int64 `json:"total_tokens"` + } `json:"usage,omitempty"` +} + +type contextualWireGroup struct { + Data []contextualWireItem `json:"data"` + Index *int `json:"index"` +} + +type contextualWireItem struct { + Embedding []float32 `json:"embedding"` + Index *int `json:"index"` + Text string `json:"text"` +} + +type contextualWireResponse struct { + Data []contextualWireGroup `json:"data"` + Model string `json:"model"` + ChunkerVersion string `json:"chunker_version"` + Usage *struct { + TotalTokens int64 `json:"total_tokens"` + } `json:"usage,omitempty"` +} + +type embeddingEgressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +func EmbeddingPolicyFingerprint(profile EmbeddingProfile) (string, error) { + normalized, descriptorIdentity, route, err := normalizeEmbeddingProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(embeddingPolicyIdentity{ + AdapterContract: embeddingAdapterContract, Endpoint: normalized.Endpoint, Route: route, + Mode: normalized.Mode, Descriptor: descriptorIdentity, ModelInput: normalized.ModelInput, + CredentialBinding: normalized.SecretBinding, Egress: embeddingEgressPolicyIdentity(normalized.EgressPolicy), ChunkerVersion: normalized.ChunkerVersion, + RequestTimeout: int64(normalized.RequestTimeout), MaxRetries: normalized.MaxRetries, + RetryBaseDelay: int64(normalized.RetryBaseDelay), MaxBatchItems: normalized.MaxBatchItems, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, + CapabilityPolicy: directCapabilityFingerprint(normalized), + }, json.Deterministic(true)) + if err != nil { + return "", fmt.Errorf("voyage embedding: encode policy identity: %w", err) + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +func NewEmbeddingProvider(profile EmbeddingProfile, secrets SecretResolver, resolver providerhttp.Resolver) (*EmbeddingClient, error) { + normalized, _, _, err := normalizeEmbeddingProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + if err == nil { + err = errors.New("descriptor is not canonical") + } + return nil, fmt.Errorf("voyage embedding: invalid descriptor: %w", err) + } + fingerprint, err := EmbeddingPolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("voyage embedding: descriptor policy fingerprint does not match profile") + } + if descriptor.SupportsTextQuery { + return nil, errors.New("voyage embedding: hosted mutable alias is export-only") + } + if normalized.SecretBinding == "" || nilEmbeddingValue(secrets) { + return nil, errors.New("voyage embedding: named secret binding and resolver are required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, fmt.Errorf("voyage embedding: invalid sealed egress policy: %w", err) + } + normalized.Descriptor = cloneEmbeddingDescriptor(descriptor) + return &EmbeddingClient{profile: normalized, descriptor: cloneEmbeddingDescriptor(descriptor), secrets: secrets, http: &http.Client{Transport: transport, CheckRedirect: providerhttp.RefuseRedirects}, now: time.Now}, nil +} + +func (client *EmbeddingClient) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneEmbeddingDescriptor(client.descriptor) +} + +func (client *EmbeddingClient) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + if client == nil { + return document.EmbeddingResult{}, errors.New("voyage embedding: client is required") + } + if ctx == nil { + return document.EmbeddingResult{}, errors.New("voyage embedding: context is required") + } + if err := document.ValidateEmbeddingProviderRequest(client, inputs, authorization); err != nil { + return document.EmbeddingResult{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || authorization.MaxInputBytes > client.profile.MaxInputBytes || authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return document.EmbeddingResult{}, errors.New("voyage embedding: authorization exceeds profile capacity") + } + if client.profile.Mode == EmbeddingModeDirectFile { + return client.embedDirectFiles(ctx, inputs, authorization) + } + result := document.EmbeddingResult{Vectors: make([]document.EmbeddingVector, len(inputs))} + for _, role := range []document.EmbeddingRole{document.EmbeddingRoleDocument, document.EmbeddingRoleQuery} { + positions := make([]int, 0, len(inputs)) + rendered := make([]string, 0, len(inputs)) + for index, input := range inputs { + if input.Role != role { + continue + } + positions = append(positions, index) + if role == document.EmbeddingRoleDocument { + rendered = append(rendered, client.profile.ModelInput.EncodeDocument(input.Text)) + } else { + rendered = append(rendered, client.profile.ModelInput.EncodeQuery(input.Text)) + } + } + if len(positions) == 0 { + continue + } + vectors, err := client.embedTextRole(ctx, rendered, string(role)) + if err != nil { + return document.EmbeddingResult{}, err + } + for local, global := range positions { + result.Vectors[global] = document.EmbeddingVector{Key: inputs[global].Key, Values: vectors[local]} + } + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, err + } + return result, nil +} + +func (client *EmbeddingClient) embedDirectFiles(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + secret, err := client.secrets.ResolveSecret(ctx, client.profile.SecretBinding) + if err != nil || !validEmbeddingSecret(secret) { + if contextErr := ctx.Err(); contextErr != nil { + return document.EmbeddingResult{}, fmt.Errorf("voyage embedding: credential resolution canceled: %w", contextErr) + } + return document.EmbeddingResult{}, errors.New("voyage embedding: credential is unavailable") + } + legacy, err := NewClient(client.profile.Policy, ClientConfig{ + APIKey: secret, Timeout: client.profile.RequestTimeout, MaxRetries: client.profile.MaxRetries, + RetryBaseDelay: client.profile.RetryBaseDelay, HTTPClient: client.http, + }) + if err != nil { + return document.EmbeddingResult{}, errors.New("voyage embedding: direct-file client configuration failed") + } + authorities, err := client.profile.Policy.AuthorizeAll(client.profile.CapabilityManifest) + if err != nil { + return document.EmbeddingResult{}, errors.New("voyage embedding: direct-file capability authority is invalid") + } + direct := make([]Input, len(inputs)) + for index, input := range inputs { + metadata := input.Source.Metadata() + data, readErr := io.ReadAll(io.LimitReader(input.Source, metadata.ByteLength+1)) + closeErr := input.Source.Close() + if readErr != nil || closeErr != nil || int64(len(data)) != metadata.ByteLength { + clear(data) + return document.EmbeddingResult{}, errors.New("voyage embedding: direct-file source could not be read exactly") + } + digest := sha256.Sum256(data) + if hex.EncodeToString(digest[:]) != metadata.SHA256 { + clear(data) + return document.EmbeddingResult{}, errors.New("voyage embedding: direct-file source identity changed") + } + detected, detectErr := media.DetectBytes(data, metadata.MediaType) + if detectErr != nil || string(detected.Kind) != metadata.MediaFamily || detected.MediaType != metadata.MediaType { + clear(data) + return document.EmbeddingResult{}, errors.New("voyage embedding: direct-file media identity could not be verified") + } + direct[index] = Input{Parts: []Part{{Media: &Media{Metadata: detected, Bytes: data}}}} + } + defer func() { + for index := range direct { + clear(direct[index].Parts[0].Media.Bytes) + } + }() + providerResult, err := legacy.EmbedDocuments(ctx, direct, authorities) + if err != nil { + return document.EmbeddingResult{}, fmt.Errorf("voyage embedding: direct-file provider request failed: %w", err) + } + result := document.EmbeddingResult{Vectors: make([]document.EmbeddingVector, len(inputs))} + for index, vector := range providerResult.Vectors { + result.Vectors[index] = document.EmbeddingVector{Key: inputs[index].Key, Values: slices.Clone(vector)} + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return document.EmbeddingResult{}, err + } + return result, nil +} + +func (client *EmbeddingClient) embedTextRole(ctx context.Context, rendered []string, inputType string) ([][]float32, error) { + route := textEmbeddingsPath + var payload []byte + var err error + if client.profile.Mode == EmbeddingModeContextual { + route = contextualEmbeddingsPath + payload, err = json.Marshal(contextualWireRequest{Inputs: [][]string{rendered}, Model: client.descriptor.Model, InputType: inputType, OutputDimension: client.descriptor.Dimension, OutputDType: "float"}) + } else { + payload, err = json.Marshal(embeddingWireRequest{Input: rendered, Model: client.descriptor.Model, InputType: inputType, Truncation: false, OutputDimension: client.descriptor.Dimension, OutputDType: "float"}) + } + if err != nil { + return nil, errors.New("voyage embedding: could not encode request") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + return nil, &ProviderError{Kind: ErrBatchTooLarge} + } + started := time.Now() + metrics := RequestMetrics{} + for attempt := 1; ; attempt++ { + vectors, retryAfter, retry, requested, err := client.embeddingAttempt(ctx, route, payload, rendered) + if requested { + metrics.Requests++ + } + if err == nil { + return vectors, nil + } + if !retry || attempt >= client.profile.MaxRetries { + if providerErr, ok := errors.AsType[*ProviderError](err); ok { + providerErr.Metrics = metrics + providerErr.Metrics.Latency = time.Since(started) + } + return nil, err + } + metrics.Retries++ + delay := retryBackoffDelay(client.profile.RetryBaseDelay, attempt) + if retryAfter >= 0 { + delay = retryAfter + } + if waitErr := sleepContext(ctx, delay); waitErr != nil { + return nil, &ProviderError{Kind: waitErr, cause: waitErr, Metrics: RequestMetrics{Requests: metrics.Requests, Retries: metrics.Retries, Latency: time.Since(started)}} + } + } +} + +func (client *EmbeddingClient) embeddingAttempt(ctx context.Context, route string, payload []byte, rendered []string) ([][]float32, time.Duration, bool, bool, error) { + attemptCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + secret, err := client.secrets.ResolveSecret(attemptCtx, client.profile.SecretBinding) + if err != nil || !validEmbeddingSecret(secret) { + if contextErr := attemptCtx.Err(); contextErr != nil { + return nil, -1, false, false, &ProviderError{Kind: contextErr, cause: contextErr} + } + return nil, -1, false, false, &ProviderError{Kind: ErrPermanentResponse} + } + request, err := http.NewRequestWithContext(attemptCtx, http.MethodPost, client.profile.Endpoint+route, bytes.NewReader(payload)) + if err != nil { + return nil, -1, false, false, &ProviderError{Kind: ErrPermanentResponse} + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+secret) + response, err := client.http.Do(request) + if err != nil { + if contextErr := attemptCtx.Err(); contextErr != nil { + return nil, -1, false, true, &ProviderError{Kind: contextErr, cause: contextErr} + } + return nil, -1, true, true, &ProviderError{Kind: ErrTransientResponse, cause: err} + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode >= 300 && response.StatusCode < 400 { + return nil, -1, false, true, &ProviderError{Kind: ErrPermanentResponse, StatusCode: response.StatusCode} + } + if response.StatusCode == http.StatusUnauthorized || response.StatusCode == http.StatusForbidden { + return nil, -1, false, true, &ProviderError{Kind: ErrUnauthorized, StatusCode: response.StatusCode} + } + if response.StatusCode == http.StatusRequestEntityTooLarge { + return nil, -1, false, true, &ProviderError{Kind: ErrBatchTooLarge, StatusCode: response.StatusCode} + } + if response.StatusCode == http.StatusTooManyRequests || response.StatusCode >= 500 { + delay, set := parseRetryAfter(response.Header.Get("Retry-After"), client.now()) + if !set { + delay = -1 + } + providerErr := &ProviderError{Kind: ErrTransientResponse, StatusCode: response.StatusCode, RetryAfter: delay, RetrySet: delay >= 0} + return nil, delay, true, true, providerErr + } + if response.StatusCode != http.StatusOK { + return nil, -1, false, true, &ProviderError{Kind: ErrPermanentResponse, StatusCode: response.StatusCode} + } + if err := validateEmbeddingContentType(response.Header.Get("Content-Type")); err != nil { + return nil, -1, false, true, &ProviderError{Kind: ErrMalformedResponse} + } + body, err := readEmbeddingBody(attemptCtx, response.Body, client.profile.MaxResponseBytes) + if err != nil { + return nil, -1, false, true, &ProviderError{Kind: ErrMalformedResponse, cause: err} + } + if err := manifestjson.RejectDuplicateKeys(body, "voyage embedding response"); err != nil { + return nil, -1, false, true, &ProviderError{Kind: ErrMalformedResponse} + } + if client.profile.Mode == EmbeddingModeContextual { + vectors, decodeErr := client.decodeContextual(body, rendered) + return vectors, -1, false, true, decodeErr + } + vectors, decodeErr := client.decodeText(body, len(rendered)) + return vectors, -1, false, true, decodeErr +} + +func (client *EmbeddingClient) decodeText(body []byte, want int) ([][]float32, error) { + var response embeddingWireResponse + if err := json.Unmarshal(body, &response, json.RejectUnknownMembers(true)); err != nil { + return nil, &ProviderError{Kind: ErrMalformedResponse} + } + vectors, err := client.orderItems(response.Object, response.Model, response.Data, want) + if err != nil { + err = &ProviderError{Kind: ErrMalformedResponse, cause: err} + } + return vectors, err +} + +func (client *EmbeddingClient) decodeContextual(body []byte, rendered []string) ([][]float32, error) { + var response contextualWireResponse + if err := json.Unmarshal(body, &response, json.RejectUnknownMembers(true)); err != nil || response.Model != client.descriptor.Model || response.ChunkerVersion != client.profile.ChunkerVersion { + return nil, &ProviderError{Kind: ErrMalformedResponse} + } + if len(response.Data) != 1 || response.Data[0].Index == nil || *response.Data[0].Index != 0 || len(response.Data[0].Data) != len(rendered) { + return nil, &ProviderError{Kind: ErrMalformedResponse} + } + vectors := make([][]float32, len(rendered)) + seen := make([]bool, len(rendered)) + for _, item := range response.Data[0].Data { + if item.Index == nil || *item.Index < 0 || *item.Index >= len(rendered) || seen[*item.Index] || item.Text != rendered[*item.Index] || client.validateEmbeddingVector(item.Embedding) != nil { + return nil, &ProviderError{Kind: ErrMalformedResponse} + } + seen[*item.Index] = true + vectors[*item.Index] = slices.Clone(item.Embedding) + } + if slices.Contains(seen, false) { + return nil, &ProviderError{Kind: ErrMalformedResponse} + } + return vectors, nil +} + +func (client *EmbeddingClient) orderItems(object, model string, items []embeddingWireItem, want int) ([][]float32, error) { + if object != "list" || model != client.descriptor.Model { + return nil, errors.New("voyage embedding: provider model or response contract drifted") + } + if len(items) != want { + return nil, errors.New("voyage embedding: provider response has a missing vector") + } + vectors := make([][]float32, want) + seen := make([]bool, want) + for _, item := range items { + if item.Object != "embedding" || item.Index == nil || *item.Index < 0 || *item.Index >= want || seen[*item.Index] { + return nil, errors.New("voyage embedding: provider response index contract drifted") + } + if err := client.validateEmbeddingVector(item.Embedding); err != nil { + return nil, err + } + seen[*item.Index] = true + vectors[*item.Index] = slices.Clone(item.Embedding) + } + if slices.Contains(seen, false) { + return nil, errors.New("voyage embedding: provider response has a missing vector index") + } + return vectors, nil +} + +func (client *EmbeddingClient) validateEmbeddingVector(vector []float32) error { + if len(vector) != client.descriptor.Dimension { + return errors.New("voyage embedding: provider vector dimension does not match profile") + } + var norm float64 + for _, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return errors.New("voyage embedding: provider vector contains a non-finite value") + } + norm += float64(value) * float64(value) + } + if norm == 0 { + return errors.New("voyage embedding: provider returned a zero vector") + } + if client.descriptor.Normalization == document.VectorNormalizationUnitLength && math.Abs(norm-1) > unitLengthTolerance { + return errors.New("voyage embedding: provider vector normalization does not match profile") + } + return nil +} + +func normalizeEmbeddingProfile(profile EmbeddingProfile) (EmbeddingProfile, document.EmbeddingDescriptor, string, error) { + if profile.Endpoint == "" { + profile.Endpoint = DefaultEndpoint + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = DefaultTimeout + } + if profile.MaxRetries == 0 { + profile.MaxRetries = DefaultMaxRetries + } + if profile.RetryBaseDelay == 0 { + profile.RetryBaseDelay = defaultRetryBaseDelay + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultEmbeddingMaxBatchItems + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultEmbeddingMaxInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultEmbeddingMaxRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultEmbeddingMaxResponseBytes + } + if profile.RequestTimeout <= 0 || profile.RequestTimeout > MaxTimeout || profile.MaxRetries < 1 || profile.MaxRetries > MaxRetries || profile.RetryBaseDelay < 0 || profile.RetryBaseDelay > maxRetryAfter || profile.MaxBatchItems < 1 || profile.MaxBatchItems > 1000 || profile.MaxInputBytes < 1 || profile.MaxRequestBytes < 1 || profile.MaxResponseBytes < 1 { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: execution bounds are invalid") + } + if !validEmbeddingToken(profile.SecretBinding) { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: binding is invalid") + } + if err := normalizeAndValidateEmbeddingEgress(&profile); err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", err + } + descriptorIdentity := profile.Descriptor + descriptorIdentity.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptorIdentity.Fingerprint = "" + var err error + descriptorIdentity, err = document.NewEmbeddingDescriptor(descriptorIdentity) + if err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", fmt.Errorf("voyage embedding: invalid descriptor identity: %w", err) + } + descriptorIdentity.PolicyFingerprint = "" + descriptorIdentity.Fingerprint = "" + if descriptorIdentity.ID != EmbeddingProviderID || descriptorIdentity.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptorIdentity.ScalarEncoding != EmbeddingScalarFloat32 || descriptorIdentity.DocumentFormatter != EmbeddingDocumentFormatterV1 || descriptorIdentity.QueryFormatter != EmbeddingQueryFormatterV1 || descriptorIdentity.ModelInput != profile.ModelInput || descriptorIdentity.CompatibilityID != profile.ModelInput.CompatibilityID { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: descriptor does not match adapter contract") + } + if profile.Mode != EmbeddingModeDirectFile && descriptorIdentity.SupportsTextQuery { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: hosted mutable alias is export-only") + } + if profile.Mode != EmbeddingModeDirectFile && (!slices.Equal(descriptorIdentity.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}) || profile.ModelInput.Document.Mode != document.ModelInputModeDocument || profile.ModelInput.Query.Mode != document.ModelInputModeQuery) { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: descriptor native role modes do not match document/query behavior") + } + route := textEmbeddingsPath + switch profile.Mode { + case EmbeddingModeText: + if !slices.Contains([]string{"voyage-4-large", TextModel, "voyage-4-lite"}, descriptorIdentity.Model) || descriptorIdentity.ModelRevision != HostedAliasRevision || descriptorIdentity.SupportsTextQuery || !slices.Contains([]int{2048, 1024, 512, 256}, descriptorIdentity.Dimension) || !slices.Equal(descriptorIdentity.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || !slices.Equal(descriptorIdentity.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}) || profile.ModelInput.Document.Mode != document.ModelInputModeDocument || profile.ModelInput.Query.Mode != document.ModelInputModeQuery { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: descriptor is not a pinned Voyage 4 text profile") + } + case EmbeddingModeContextual: + route = contextualEmbeddingsPath + if !validEmbeddingToken(profile.ChunkerVersion) || descriptorIdentity.Model != ContextualModel || descriptorIdentity.ModelRevision != HostedAliasRevision || descriptorIdentity.SupportsTextQuery || !slices.Contains([]int{2048, 1024, 512, 256}, descriptorIdentity.Dimension) || !slices.Equal(descriptorIdentity.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || !slices.Equal(descriptorIdentity.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}) || profile.ModelInput.Document.Mode != document.ModelInputModeDocument || profile.ModelInput.Query.Mode != document.ModelInputModeQuery { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: descriptor is not a pinned contextual profile") + } + case EmbeddingModeDirectFile: + route = embeddingsPath + if !profile.Policy.valid() { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: direct-file policy is invalid") + } + if profile.Endpoint != profile.Policy.values.Endpoint { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: direct-file endpoint differs from capability policy") + } + if err := profile.CapabilityManifest.ValidateComplete(); err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", fmt.Errorf("voyage embedding: direct-file capability evidence: %w", err) + } + capabilityFingerprint, err := profile.Policy.Fingerprint(profile.CapabilityManifest) + if err != nil { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", fmt.Errorf("voyage embedding: direct-file capability policy: %w", err) + } + expectedRevision := "capability-" + capabilityFingerprint[:32] + if descriptorIdentity.Model != profile.Policy.values.Model || descriptorIdentity.Dimension != profile.Policy.values.Dimension || descriptorIdentity.ModelRevision != expectedRevision || descriptorIdentity.SupportsTextQuery || !slices.Equal(descriptorIdentity.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile}) || !slices.Contains(descriptorIdentity.SupportedRequestModes, document.ModelInputModeDocument) { + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: direct-file descriptor is not the exact capability-attested export profile") + } + default: + return EmbeddingProfile{}, document.EmbeddingDescriptor{}, "", errors.New("voyage embedding: mode is invalid") + } + return profile, descriptorIdentity, route, nil +} + +func normalizeAndValidateEmbeddingEgress(profile *EmbeddingProfile) error { + if profile.EgressPolicy.ConnectTimeout == 0 { + profile.EgressPolicy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if profile.EgressPolicy.KeepAlive == 0 { + profile.EgressPolicy.KeepAlive = providerhttp.DefaultKeepAlive + } + if profile.EgressPolicy.TLSHandshakeTimeout == 0 { + profile.EgressPolicy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if profile.EgressPolicy.ProxyMode == "" { + profile.EgressPolicy.ProxyMode = providerhttp.ProxyDisabled + } + if profile.EgressPolicy.TLS.RootCAs != nil { + return errors.New("voyage embedding: custom egress roots cannot enter canonical identity") + } + parsed, err := url.Parse(profile.Endpoint) + if err != nil || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" || strings.TrimSuffix(parsed.Path, "/") != "/v1" { + return errors.New("voyage embedding: endpoint must be an exact /v1 provider root") + } + port := parsed.Port() + if port == "" { + switch parsed.Scheme { + case "https": + port = "443" + case "http": + port = "80" + } + } + if parsed.Scheme != profile.EgressPolicy.Scheme || !strings.EqualFold(parsed.Hostname(), profile.EgressPolicy.Host) || port != strconv.FormatUint(uint64(profile.EgressPolicy.Port), 10) { + return errors.New("voyage embedding: endpoint and egress authority differ") + } + profile.Endpoint = strings.TrimSuffix(profile.Endpoint, "/") + slices.SortFunc(profile.EgressPolicy.AllowedCIDRs, func(a, b netip.Prefix) int { return strings.Compare(a.Masked().String(), b.Masked().String()) }) + slices.Sort(profile.EgressPolicy.TLS.SPKISHA256) + return nil +} + +func embeddingEgressPolicyIdentity(policy providerhttp.EgressPolicy) embeddingEgressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.Masked().String() + } + return embeddingEgressIdentity{Scheme: policy.Scheme, Host: strings.ToLower(policy.Host), Port: policy.Port, AllowedCIDRs: cidrs, ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), KeepAlive: int64(policy.KeepAlive), TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func DirectFileModelRevision(policy Policy, manifest CapabilityManifest) (string, error) { + fingerprint, err := policy.Fingerprint(manifest) + if err != nil { + return "", err + } + return "capability-" + fingerprint[:32], nil +} + +func directCapabilityFingerprint(profile EmbeddingProfile) string { + if profile.Mode != EmbeddingModeDirectFile { + return "" + } + fingerprint, _ := profile.Policy.Fingerprint(profile.CapabilityManifest) + return fingerprint +} + +func readEmbeddingBody(ctx context.Context, reader io.Reader, maximum int64) ([]byte, error) { + body, err := io.ReadAll(io.LimitReader(reader, maximum+1)) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, fmt.Errorf("voyage embedding: response read canceled: %w", contextErr) + } + return nil, errors.New("voyage embedding: could not read provider response") + } + if int64(len(body)) > maximum { + return nil, errors.New("voyage embedding: provider response byte limit exceeded") + } + return body, nil +} + +func validateEmbeddingContentType(value string) error { + mediaType, parameters, err := mime.ParseMediaType(value) + if err != nil || mediaType != "application/json" || (len(parameters) != 0 && (len(parameters) != 1 || !strings.EqualFold(parameters["charset"], "utf-8"))) { + return errors.New("voyage embedding: provider response content type is invalid") + } + return nil +} + +func validEmbeddingSecret(secret string) bool { + if secret == "" || len(secret) > maxEmbeddingSecretBytes { + return false + } + for _, character := range secret { + if unicode.IsControl(character) || unicode.IsSpace(character) { + return false + } + } + return true +} + +func validEmbeddingToken(value string) bool { + return value != "" && len(value) <= 128 && value == strings.TrimSpace(value) && utf8.ValidString(value) && strings.IndexFunc(value, unicode.IsControl) < 0 +} + +func cloneEmbeddingDescriptor(value document.EmbeddingDescriptor) document.EmbeddingDescriptor { + value.InputKinds = slices.Clone(value.InputKinds) + value.SupportedRequestModes = slices.Clone(value.SupportedRequestModes) + return value +} + +func nilEmbeddingValue(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/voyage/embedding_review_test.go b/document/voyage/embedding_review_test.go new file mode 100644 index 00000000..0f649f08 --- /dev/null +++ b/document/voyage/embedding_review_test.go @@ -0,0 +1,151 @@ +package voyage_test + +import ( + "context" + json "encoding/json/v2" + "net" + "net/http" + "net/http/httptest" + "net/netip" + "net/url" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" + "go.kenn.io/docbank/document/voyage" +) + +type embeddingResolver struct{ address netip.Addr } + +func (resolver embeddingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return []netip.Addr{resolver.address}, nil +} + +func voyageFixture(t *testing.T, handler http.Handler) (string, providerhttp.EgressPolicy, providerhttp.Resolver) { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + parsed, err := url.Parse(server.URL) + require.NoError(t, err) + _, portText, err := net.SplitHostPort(parsed.Host) + require.NoError(t, err) + port, err := strconv.ParseUint(portText, 10, 16) + require.NoError(t, err) + return "http://voyage.invalid:" + portText + "/v1", providerhttp.EgressPolicy{ + Scheme: "http", Host: "voyage.invalid", Port: uint16(port), + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, + }, embeddingResolver{address: netip.MustParseAddr("127.0.0.1")} +} + +func TestContextualEmbeddingUsesDocumentedWireShape(t *testing.T) { + var request map[string]any + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, incoming *http.Request) { + assert.NoError(t, json.UnmarshalRead(incoming.Body, &request)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(contextualOfficialBody(t, []string{"document envelope: first", "document envelope: second"}, []int{1, 0})) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeContextual) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "synthetic-secret"}, resolver) + require.NoError(t, err) + inputs := []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "first"}, + {Key: "second", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "second"}, + } + result, err := provider.Embed(t.Context(), inputs, voyageAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.NotContains(t, request, "truncation") + assert.Equal(t, []any{[]any{"document envelope: first", "document envelope: second"}}, request["inputs"]) + assert.InDelta(t, float32(1), result.Vectors[0].Values[1], 0) + assert.InDelta(t, float32(1), result.Vectors[1].Values[2], 0) +} + +func TestContextualEmbeddingStrictlyRejectsDocumentedShapeDrift(t *testing.T) { + texts := []string{"document envelope: first", "document envelope: second"} + valid := contextualOfficialBody(t, texts, []int{0, 1}) + tests := []struct { + name string + body []byte + }{ + {"unknown", []byte(strings.Replace(string(valid), `"chunker_version":`, `"unknown":"PRIVATE_RAW_BODY","chunker_version":`, 1))}, + {"duplicate", []byte(strings.Replace(string(valid), `"model":`, `"model":"duplicate","model":`, 1))}, + {"model drift", []byte(strings.Replace(string(valid), voyage.ContextualModel, "voyage-context-drift", 1))}, + {"chunker drift", []byte(strings.Replace(string(valid), `"1.0.0"`, `"2.0.0"`, 1))}, + {"text drift", []byte(strings.Replace(string(valid), texts[0], "PRIVATE_RETURNED_TEXT", 1))}, + {"partial indices", contextualOfficialBody(t, texts[:1], []int{0})}, + {"duplicate index", contextualOfficialBody(t, []string{texts[0], texts[0]}, []int{0, 0})}, + } + inputs := []document.EmbeddingInput{ + {Key: "first", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "first"}, + {Key: "second", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "second"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(test.body) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeContextual) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "PRIVATE_SECRET"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), inputs, voyageAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, voyage.ErrMalformedResponse) + assert.NotContains(t, err.Error(), "PRIVATE") + }) + } +} + +func TestHostedVoyageAliasIsAlwaysExportOnly(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.NotFoundHandler()) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile.Descriptor.SupportsTextQuery = true + _, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "synthetic-secret"}, resolver) + require.ErrorContains(t, err, "export-only") +} + +func TestVoyageRoleContractMustMatchNativeModes(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.NotFoundHandler()) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile.Descriptor.SupportedRequestModes = []document.ModelInputMode{document.ModelInputModeDocument} + _, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "synthetic-secret"}, resolver) + require.ErrorContains(t, err, "native role") +} + +func contextualOfficialBody(t *testing.T, texts []string, indices []int) []byte { + t.Helper() + items := make([]map[string]any, len(indices)) + for position, index := range indices { + items[position] = map[string]any{"embedding": unitEmbedding(index + 1), "index": index, "text": texts[index]} + } + body, err := json.Marshal(map[string]any{ + "data": []any{map[string]any{"data": items, "index": 0}}, "model": voyage.ContextualModel, + "usage": map[string]any{"total_tokens": 2}, "chunker_version": "1.0.0", + }) + require.NoError(t, err) + return body +} + +func refingerprintVoyageProfile(t *testing.T, profile voyage.EmbeddingProfile) voyage.EmbeddingProfile { + t.Helper() + profile.Descriptor.PolicyFingerprint = "0000000000000000000000000000000000000000000000000000000000000000" + profile.Descriptor.Fingerprint = "" + profile.Descriptor, _ = document.NewEmbeddingDescriptor(profile.Descriptor) + fingerprint, err := voyage.EmbeddingPolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + profile.Descriptor.Fingerprint = "" + profile.Descriptor, err = document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return profile +} diff --git a/document/voyage/embedding_test.go b/document/voyage/embedding_test.go new file mode 100644 index 00000000..1527abc9 --- /dev/null +++ b/document/voyage/embedding_test.go @@ -0,0 +1,350 @@ +package voyage_test + +import ( + "context" + json "encoding/json/v2" + "errors" + "io" + "net/http" + "net/netip" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" + "go.kenn.io/docbank/document/voyage" + "go.kenn.io/docbank/document/voyage/voyagetest" +) + +type embeddingSecrets map[string]string + +func (secrets embeddingSecrets) ResolveSecret(_ context.Context, name string) (string, error) { + return secrets[name], nil +} + +type textEmbeddingRequest struct { + Input []string `json:"input"` + Inputs [][]string `json:"inputs"` + Model string `json:"model"` + InputType string `json:"input_type"` + Truncation *bool `json:"truncation"` + OutputDimension int `json:"output_dimension"` + OutputDType string `json:"output_dtype"` +} + +func TestEmbeddingProviderSendsVoyageDocumentRoleAndEnvelope(t *testing.T) { + var decoded textEmbeddingRequest + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + assert.Equal(t, "/v1/embeddings", request.URL.Path) + assert.Equal(t, "Bearer synthetic-secret", request.Header.Get("Authorization")) + assert.NoError(t, json.UnmarshalRead(request.Body, &decoded, json.RejectUnknownMembers(true))) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(voyageTextBody(t, voyage.TextModel, []int{0}, []int{1})) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "synthetic-secret"}, resolver) + require.NoError(t, err) + inputs := []document.EmbeddingInput{{Key: "doc", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "passage"}} + result, err := document.ExecuteEmbedding(t.Context(), provider, inputs, voyageAuthorization(profile.Descriptor)) + require.NoError(t, err) + assert.Equal(t, "document", decoded.InputType) + assert.Equal(t, []string{"document envelope: passage"}, decoded.Input) + require.NotNil(t, decoded.Truncation) + assert.False(t, *decoded.Truncation) + assert.Equal(t, "doc", result.Vectors[0].Key) +} + +func TestVoyageEmbeddingPolicyFingerprintCoversDisclosureAndRoleContract(t *testing.T) { + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + baseline, err := voyage.EmbeddingPolicyFingerprint(profile) + require.NoError(t, err) + + changedEnvelope := profile + changedEnvelope.ModelInput = customVoyageInput(t, "changed: {{content}}", "query envelope: {{content}}") + changedEnvelope.Descriptor.ModelInput = changedEnvelope.ModelInput + changedEnvelope.Descriptor.CompatibilityID = changedEnvelope.ModelInput.CompatibilityID + changedEnvelope = refingerprintVoyageProfile(t, changedEnvelope) + assert.NotEqual(t, baseline, changedEnvelope.Descriptor.PolicyFingerprint) + + changedEgress := profile + changedEgress.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24")} + changedEgress = refingerprintVoyageProfile(t, changedEgress) + assert.NotEqual(t, baseline, changedEgress.Descriptor.PolicyFingerprint) +} + +func TestEmbeddingProviderRejectsMalformedIndexedResponsesPrivately(t *testing.T) { + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + one, two := unitEmbedding(0), unitEmbedding(1) + valid := func(indices []int, vectors [][]float32) []byte { + items := make([]map[string]any, len(indices)) + for index := range indices { + items[index] = map[string]any{"object": "embedding", "embedding": vectors[index], "index": indices[index]} + } + body, err := json.Marshal(map[string]any{"object": "list", "data": items, "model": profile.Descriptor.Model, "usage": map[string]any{"total_tokens": 1}}) + require.NoError(t, err) + return body + } + tests := []struct { + name string + body []byte + }{ + {"malformed", []byte(`{"object":`)}, + {"unknown", []byte(`{"object":"list","unknown":"PRIVATE_RAW_BODY","data":[],"model":"voyage-4"}`)}, + {"duplicate key", []byte(`{"object":"list","model":"voyage-4","model":"PRIVATE_RAW_BODY","data":[]}`)}, + {"partial indices", valid([]int{0}, [][]float32{one})}, + {"duplicate index", valid([]int{0, 0}, [][]float32{one, two})}, + {"out of range", valid([]int{0, 2}, [][]float32{one, two})}, + {"wrong dimension", valid([]int{0, 1}, [][]float32{one[:255], two})}, + {"zero vector", valid([]int{0, 1}, [][]float32{make([]float32, 256), two})}, + {"non finite", []byte(`{"object":"list","data":[{"object":"embedding","embedding":[1e999],"index":0}],"model":"voyage-4"}`)}, + } + inputs := []document.EmbeddingInput{ + {Key: "a", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "PRIVATE_INPUT"}, + {Key: "b", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "other"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write(test.body) + })) + local := profile + local.Endpoint, local.EgressPolicy = endpoint, egress + local = refingerprintVoyageProfile(t, local) + provider, err := voyage.NewEmbeddingProvider(local, embeddingSecrets{"credential:voyage": "PRIVATE_SECRET"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), inputs, voyageAuthorization(local.Descriptor)) + require.ErrorIs(t, err, voyage.ErrMalformedResponse) + assert.NotContains(t, err.Error(), "PRIVATE") + }) + } +} + +func TestEmbeddingProviderClassifiesCapacityAndExhaustedTransient(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + for _, test := range []struct { + name string + status int + kind error + }{ + {"auth", http.StatusUnauthorized, voyage.ErrUnauthorized}, + {"capacity", http.StatusRequestEntityTooLarge, voyage.ErrBatchTooLarge}, + {"permanent", http.StatusBadRequest, voyage.ErrPermanentResponse}, + {"rate limit", http.StatusTooManyRequests, voyage.ErrTransientResponse}, + {"transient", http.StatusServiceUnavailable, voyage.ErrTransientResponse}, + } { + t.Run(test.name, func(t *testing.T) { + var calls atomic.Int32 + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + calls.Add(1) + writer.WriteHeader(test.status) + _, _ = writer.Write([]byte("PRIVATE_BODY")) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy, profile.MaxRetries, profile.RetryBaseDelay = endpoint, egress, 2, time.Millisecond + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, voyageAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, test.kind) + assert.NotContains(t, err.Error(), "PRIVATE_BODY") + metrics := voyage.MetricsFromError(err) + assert.Equal(t, int(calls.Load()), metrics.Requests) + if errors.Is(test.kind, voyage.ErrTransientResponse) { + assert.Equal(t, 2, metrics.Requests) + assert.Equal(t, 1, metrics.Retries) + } + }) + } +} + +func TestVoyageEmbeddingRejectsIdentityDrift(t *testing.T) { + for _, test := range []struct { + name string + mutate func(*voyage.EmbeddingProfile) + }{ + {"model", func(profile *voyage.EmbeddingProfile) { profile.Descriptor.Model = "voyage-4-drift" }}, + {"revision", func(profile *voyage.EmbeddingProfile) { profile.Descriptor.ModelRevision = "invented-revision" }}, + {"compatibility", func(profile *voyage.EmbeddingProfile) { profile.Descriptor.CompatibilityID = "voyage/other-space/v1" }}, + } { + t.Run(test.name, func(t *testing.T) { + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + test.mutate(&profile) + _, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, nil) + require.Error(t, err) + }) + } +} + +func TestVoyageEmbeddingResponseLimitCancellationAndTransportFailure(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "PRIVATE_INPUT"}} + t.Run("response limit", func(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(strings.Repeat("x", 4097))) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy, profile.MaxResponseBytes = endpoint, egress, 4096 + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, resolver) + require.NoError(t, err) + authorization := voyageAuthorization(profile.Descriptor) + authorization.MaxResponseBytes = 4096 + _, err = provider.Embed(t.Context(), input, authorization) + require.ErrorIs(t, err, voyage.ErrMalformedResponse) + }) + t.Run("cancellation", func(t *testing.T) { + endpoint, egress, resolver := voyageFixture(t, http.NotFoundHandler()) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, resolver) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + _, err = provider.Embed(ctx, input, voyageAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, context.Canceled) + assert.NotContains(t, err.Error(), "PRIVATE_INPUT") + }) + t.Run("transport exhaustion", func(t *testing.T) { + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.MaxRetries, profile.RetryBaseDelay = 2, time.Millisecond + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, failingEmbeddingResolver{}) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, voyageAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, voyage.ErrTransientResponse) + assert.Equal(t, 2, voyage.MetricsFromError(err).Requests) + }) +} + +type failingEmbeddingResolver struct{} + +func (failingEmbeddingResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return nil, errors.New("synthetic DNS failure") +} + +func TestEmbeddingProviderBoundsRequestBytesAndRefusesRedirect(t *testing.T) { + input := []document.EmbeddingInput{{Key: "one", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "text"}} + endpoint, egress, resolver := voyageFixture(t, http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Location", "http://elsewhere.invalid/private") + writer.WriteHeader(http.StatusTemporaryRedirect) + })) + profile := voyageTextProfile(t, voyage.EmbeddingModeText) + profile.Endpoint, profile.EgressPolicy = endpoint, egress + profile = refingerprintVoyageProfile(t, profile) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, voyageAuthorization(profile.Descriptor)) + require.ErrorIs(t, err, voyage.ErrPermanentResponse) + + bounded := profile + bounded.MaxRequestBytes = 1 + bounded = refingerprintVoyageProfile(t, bounded) + provider, err = voyage.NewEmbeddingProvider(bounded, embeddingSecrets{"credential:voyage": "secret"}, resolver) + require.NoError(t, err) + _, err = provider.Embed(t.Context(), input, voyageAuthorization(bounded.Descriptor)) + require.ErrorIs(t, err, voyage.ErrBatchTooLarge) +} + +func TestDirectFileProfileRemainsCapabilityAttestedExportOnly(t *testing.T) { + policy := testPolicy(t) + manifest, err := voyagetest.SyntheticManifest(policy) + require.NoError(t, err) + profile := voyageDirectFileProfile(t, policy, manifest) + provider, err := voyage.NewEmbeddingProvider(profile, embeddingSecrets{"credential:voyage": "secret"}, embeddingResolver{address: netip.MustParseAddr("192.0.2.1")}) + require.NoError(t, err) + assert.False(t, provider.Descriptor().SupportsTextQuery) +} + +func unitEmbedding(hot int) []float32 { + vector := make([]float32, 256) + vector[hot] = 1 + return vector +} + +func voyageDirectFileProfile(t *testing.T, policy voyage.Policy, manifest voyage.CapabilityManifest) voyage.EmbeddingProfile { + t.Helper() + modelInput := customVoyageInput(t, "document envelope: {{content}}", "query envelope: {{content}}") + revision, err := voyage.DirectFileModelRevision(policy, manifest) + require.NoError(t, err) + profile := voyage.EmbeddingProfile{ + Mode: voyage.EmbeddingModeDirectFile, Endpoint: voyage.DefaultEndpoint, EgressPolicy: productionVoyageEgress(), ModelInput: modelInput, + SecretBinding: "credential:voyage", MaxBatchItems: 8, MaxInputBytes: 1 << 20, MaxRequestBytes: voyage.MaxRequestBytes, + MaxResponseBytes: 1 << 20, Policy: policy, CapabilityManifest: manifest, + Descriptor: document.EmbeddingDescriptor{ + ID: voyage.EmbeddingProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("0", 64), TrustBoundary: document.EmbeddingTrustHostedProvider, + Model: voyage.DefaultModel, ModelRevision: revision, Dimension: voyage.DefaultDimension, + Metric: document.VectorMetricCosine, Normalization: document.VectorNormalizationUnitLength, + ScalarEncoding: voyage.EmbeddingScalarFloat32, DocumentFormatter: voyage.EmbeddingDocumentFormatterV1, + QueryFormatter: voyage.EmbeddingQueryFormatterV1, InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputOriginalFile}, + CompatibilityID: modelInput.CompatibilityID, ModelInput: modelInput, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeDocument}, + }, + } + return refingerprintVoyageProfile(t, profile) +} + +func voyageTextProfile(t *testing.T, mode voyage.EmbeddingMode) voyage.EmbeddingProfile { + t.Helper() + model := voyage.TextModel + if mode == voyage.EmbeddingModeContextual { + model = voyage.ContextualModel + } + modelInput := customVoyageInput(t, "document envelope: {{content}}", "query envelope: {{content}}") + profile := voyage.EmbeddingProfile{ + Mode: mode, Endpoint: voyage.DefaultEndpoint, EgressPolicy: productionVoyageEgress(), ModelInput: modelInput, + SecretBinding: "credential:voyage", ChunkerVersion: "1.0.0", MaxBatchItems: 8, MaxInputBytes: 4096, + MaxRequestBytes: 8192, MaxResponseBytes: 1 << 20, + Descriptor: document.EmbeddingDescriptor{ + ID: voyage.EmbeddingProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("0", 64), TrustBoundary: document.EmbeddingTrustHostedProvider, + Model: model, ModelRevision: voyage.HostedAliasRevision, Dimension: 256, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationUnitLength, ScalarEncoding: voyage.EmbeddingScalarFloat32, + DocumentFormatter: voyage.EmbeddingDocumentFormatterV1, QueryFormatter: voyage.EmbeddingQueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, CompatibilityID: modelInput.CompatibilityID, + ModelInput: modelInput, SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}, + }, + } + return refingerprintVoyageProfile(t, profile) +} + +func productionVoyageEgress() providerhttp.EgressPolicy { + return providerhttp.EgressPolicy{Scheme: "https", Host: "api.voyageai.com", Port: 443, AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("0.0.0.0/0"), netip.MustParsePrefix("::/0")}, ProxyMode: providerhttp.ProxyDisabled} +} + +func customVoyageInput(t *testing.T, documentTemplate, queryTemplate string) document.ModelInputContract { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "voyage/test-space/v1", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: documentTemplate}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: queryTemplate}, + }) + require.NoError(t, err) + return contract +} + +func voyageAuthorization(descriptor document.EmbeddingDescriptor) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: 8, MaxInputBytes: 4096, MaxResponseBytes: 1 << 20} +} + +func voyageTextBody(t *testing.T, model string, indices, hot []int) []byte { + t.Helper() + items := make([]map[string]any, len(indices)) + for position, index := range indices { + items[position] = map[string]any{"object": "embedding", "embedding": unitEmbedding(hot[position]), "index": index} + } + body, err := json.Marshal(map[string]any{"object": "list", "data": items, "model": model, "usage": map[string]any{"total_tokens": 2}}) + require.NoError(t, err) + return body +} + +var _ io.Reader diff --git a/document/zeroentropyembed/client.go b/document/zeroentropyembed/client.go new file mode 100644 index 00000000..7499e2bb --- /dev/null +++ b/document/zeroentropyembed/client.go @@ -0,0 +1,308 @@ +package zeroentropyembed + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/binary" + jsonv1 "encoding/json" + "encoding/json/v2" + "errors" + "fmt" + "io" + "math" + "mime" + "net/http" + "slices" + "time" + + "go.kenn.io/docbank/document" +) + +const ( + maximumSecretBytes = 64 << 10 + maximumUsage = int64(1 << 50) + providerPayloadMax = int64(5_000_000) +) + +var _ document.EmbeddingProvider = (*Client)(nil) + +type wireRequest struct { + Model string `json:"model"` + InputType string `json:"input_type"` + Input []string `json:"input"` + Dimensions int `json:"dimensions"` + EncodingFormat EncodingFormat `json:"encoding_format"` + Latency Latency `json:"latency,omitempty"` +} + +type wireResponse struct { + Results []wireResult `json:"results"` + Usage *wireUsage `json:"usage"` +} + +type wireResult struct { + Embedding jsonv1.RawMessage `json:"embedding"` +} + +type wireUsage struct { + TotalBytes *int64 `json:"total_bytes"` + TotalTokens *int64 `json:"total_tokens"` +} + +type preparedRequest struct { + positions []int + payload []byte +} + +type Receipt struct { + ProviderID string + DescriptorFingerprint string + PolicyFingerprint string + Model string + ModelRevision string + EncodingFormat EncodingFormat + RequestedLatency Latency + RequestCount int + TotalBytes int64 + TotalTokens int64 +} + +type Execution struct { + Result document.EmbeddingResult + Receipt Receipt +} + +func (client *Client) Embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (document.EmbeddingResult, error) { + execution, err := client.embed(ctx, inputs, authorization, false) + return execution.Result, err +} + +func (client *Client) EmbedWithReceipt(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization) (Execution, error) { + return client.embed(ctx, inputs, authorization, true) +} + +func (client *Client) embed(ctx context.Context, inputs []document.EmbeddingInput, authorization document.EmbeddingAuthorization, includeReceipt bool) (Execution, error) { + if client == nil || ctx == nil { + return Execution{}, errors.New("zeroentropy embed: client and context are required") + } + requestCtx, cancel := context.WithTimeout(ctx, client.profile.RequestTimeout) + defer cancel() + if err := document.ValidateEmbeddingProviderRequest(client, inputs, authorization); err != nil { + return Execution{}, err + } + if authorization.MaxBatchItems > client.profile.MaxBatchItems || authorization.MaxInputBytes > client.profile.MaxInputBytes || + authorization.MaxResponseBytes > client.profile.MaxResponseBytes { + return Execution{}, &ProviderError{Kind: ErrCapacityResponse} + } + prepared, err := client.prepareRequests(inputs) + if err != nil { + return Execution{}, err + } + defer func() { + for index := range prepared { + clear(prepared[index].payload) + } + }() + secret, err := client.secrets.ResolveSecret(requestCtx, client.profile.SecretBinding) + if err != nil || !validSecret(secret) { + if contextErr := requestCtx.Err(); contextErr != nil { + return Execution{}, fmt.Errorf("zeroentropy embed: credential resolution canceled: %w", contextErr) + } + return Execution{}, errors.New("zeroentropy embed: API-key resolution failed") + } + result := document.EmbeddingResult{Vectors: make([]document.EmbeddingVector, len(inputs))} + receipt := Receipt{ProviderID: ProviderID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: client.descriptor.PolicyFingerprint, Model: Model, ModelRevision: client.descriptor.ModelRevision, + EncodingFormat: client.profile.EncodingFormat, RequestedLatency: client.profile.Latency} + for _, request := range prepared { + vectors, usage, executeErr := client.execute(requestCtx, request.payload, len(request.positions), secret) + if executeErr != nil { + return Execution{}, executeErr + } + for local, global := range request.positions { + result.Vectors[global] = document.EmbeddingVector{Key: inputs[global].Key, Values: vectors[local]} + } + receipt.RequestCount++ + if receipt.TotalBytes > maximumUsage-*usage.TotalBytes || receipt.TotalTokens > maximumUsage-*usage.TotalTokens { + return Execution{}, &ProviderError{Kind: ErrPermanentResponse} + } + receipt.TotalBytes += *usage.TotalBytes + receipt.TotalTokens += *usage.TotalTokens + } + if err := document.ValidateEmbeddingProviderResult(client.descriptor, inputs, authorization, result); err != nil { + return Execution{}, &ProviderError{Kind: ErrPermanentResponse} + } + execution := Execution{Result: result} + if includeReceipt { + execution.Receipt = receipt + } + return execution, nil +} + +func (client *Client) prepareRequests(inputs []document.EmbeddingInput) ([]preparedRequest, error) { + documentPositions, queryPositions := []int{}, []int{} + documents, queries := []string{}, []string{} + for index, input := range inputs { + var rendered string + var target *[]string + var positions *[]int + switch input.Role { + case document.EmbeddingRoleDocument: + rendered, target, positions = client.descriptor.ModelInput.EncodeDocument(input.Text), &documents, &documentPositions + case document.EmbeddingRoleQuery: + rendered, target, positions = client.descriptor.ModelInput.EncodeQuery(input.Text), &queries, &queryPositions + default: + return nil, &ProviderError{Kind: ErrPermanentResponse} + } + if int64(len(rendered)) > client.profile.MaxInputItemBytes { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + *target = append(*target, rendered) + *positions = append(*positions, index) + } + requests := make([]preparedRequest, 0, 2) + for _, group := range []struct { + positions []int + values []string + inputType string + }{{documentPositions, documents, "document"}, {queryPositions, queries, "query"}} { + if len(group.positions) == 0 { + continue + } + var providerBytes int64 + for _, value := range group.values { + if int64(len(value))+150 > providerPayloadMax-providerBytes { + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + providerBytes += int64(len(value)) + 150 + } + latency := client.profile.Latency + if latency == LatencyAuto { + latency = "" + } + payload, err := json.Marshal(wireRequest{Model: Model, InputType: group.inputType, Input: group.values, + Dimensions: client.descriptor.Dimension, EncodingFormat: client.profile.EncodingFormat, Latency: latency}) + if err != nil { + return nil, errors.New("zeroentropy embed: request encoding failed") + } + if int64(len(payload)) > client.profile.MaxRequestBytes { + clear(payload) + return nil, &ProviderError{Kind: ErrCapacityResponse} + } + requests = append(requests, preparedRequest{positions: slices.Clone(group.positions), payload: payload}) + } + return requests, nil +} + +func (client *Client) execute(ctx context.Context, payload []byte, expected int, secret string) ([][]float32, wireUsage, error) { + request, err := http.NewRequestWithContext(ctx, http.MethodPost, origin+embedPath, bytes.NewReader(payload)) + if err != nil { + return nil, wireUsage{}, errors.New("zeroentropy embed: request construction failed") + } + request.Header.Set("Accept", "application/json") + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+secret) + response, err := client.http.Do(request) + if err != nil { + if contextErr := ctx.Err(); contextErr != nil { + return nil, wireUsage{}, fmt.Errorf("zeroentropy embed: request canceled: %w", contextErr) + } + return nil, wireUsage{}, &ProviderError{Kind: ErrTransientResponse} + } + defer func() { _ = response.Body.Close() }() + if response.StatusCode != http.StatusOK { + return nil, wireUsage{}, statusError(response.StatusCode, response.Header.Get("Retry-After"), time.Now().UTC()) + } + if !isJSONContentType(response.Header.Get("Content-Type")) { + return nil, wireUsage{}, &ProviderError{Kind: ErrPermanentResponse} + } + body, readErr := readBounded(response.Body, client.profile.MaxResponseBytes) + defer clear(body) + if contextErr := ctx.Err(); contextErr != nil { + return nil, wireUsage{}, fmt.Errorf("zeroentropy embed: response read canceled: %w", contextErr) + } + if readErr != nil { + if errors.Is(readErr, errCapacity) { + return nil, wireUsage{}, &ProviderError{Kind: ErrCapacityResponse} + } + return nil, wireUsage{}, &ProviderError{Kind: ErrTransientResponse} + } + var decoded wireResponse + if err := json.Unmarshal(body, &decoded, json.RejectUnknownMembers(true)); err != nil || + len(decoded.Results) != expected || !validUsage(decoded.Usage) { + return nil, wireUsage{}, &ProviderError{Kind: ErrPermanentResponse} + } + vectors := make([][]float32, expected) + for index, result := range decoded.Results { + vector, decodeErr := client.decodeVector(result.Embedding) + if decodeErr != nil { + return nil, wireUsage{}, &ProviderError{Kind: ErrPermanentResponse} + } + vectors[index] = vector + } + return vectors, *decoded.Usage, nil +} + +func (client *Client) decodeVector(raw jsonv1.RawMessage) ([]float32, error) { + var values []float32 + if client.profile.EncodingFormat == EncodingFloat { + if err := json.Unmarshal(raw, &values, json.RejectUnknownMembers(true)); err != nil { + return nil, err + } + } else { + var encoded string + if err := json.Unmarshal(raw, &encoded); err != nil || base64.StdEncoding.DecodedLen(len(encoded)) < client.descriptor.Dimension*4 { + return nil, errors.New("invalid base64 vector") + } + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil || len(decoded) != client.descriptor.Dimension*4 { + clear(decoded) + return nil, errors.New("invalid base64 vector") + } + defer clear(decoded) + values = make([]float32, client.descriptor.Dimension) + for index := range values { + values[index] = math.Float32frombits(binary.LittleEndian.Uint32(decoded[index*4:])) + } + } + if len(values) != client.descriptor.Dimension { + return nil, errors.New("invalid vector dimension") + } + for _, value := range values { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return nil, errors.New("non-finite vector") + } + } + return values, nil +} + +func validUsage(usage *wireUsage) bool { + return usage != nil && usage.TotalBytes != nil && usage.TotalTokens != nil && + *usage.TotalBytes >= 0 && *usage.TotalBytes <= maximumUsage && + *usage.TotalTokens >= 0 && *usage.TotalTokens <= maximumUsage +} + +var errCapacity = errors.New("response capacity exceeded") + +func readBounded(reader io.Reader, maximum int64) ([]byte, error) { + limited := io.LimitReader(reader, maximum+1) + body, err := io.ReadAll(limited) + if err != nil { + return nil, err + } + if int64(len(body)) > maximum { + clear(body) + return nil, errCapacity + } + return body, nil +} + +func isJSONContentType(value string) bool { + mediaType, _, err := mime.ParseMediaType(value) + return err == nil && (mediaType == "application/json" || len(mediaType) > 5 && mediaType[len(mediaType)-5:] == "+json") +} + +func validSecret(value string) bool { + return len(value) <= maximumSecretBytes && validToken(value) +} diff --git a/document/zeroentropyembed/client_test.go b/document/zeroentropyembed/client_test.go new file mode 100644 index 00000000..d1ad7b10 --- /dev/null +++ b/document/zeroentropyembed/client_test.go @@ -0,0 +1,218 @@ +package zeroentropyembed + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/binary" + "encoding/json/v2" + "io" + "math" + "net/http" + "net/netip" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" +) + +func TestEmbedSeparatesInputTypesAndPreservesInputOrder(t *testing.T) { + profile := testProfile(t, 40, EncodingFloat, LatencyFast) + var calls atomic.Int32 + client := testClient(t, profile, testSecrets{"secret:zeroentropy": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + assert.Equal(t, "https://api.zeroentropy.dev/v1/models/embed", request.URL.String()) + assert.Equal(t, "Bearer synthetic-key", request.Header.Get("Authorization")) + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + var payload struct { + Model string `json:"model"` + InputType string `json:"input_type"` + Input []string `json:"input"` + Dimensions int `json:"dimensions"` + EncodingFormat EncodingFormat `json:"encoding_format"` + Latency Latency `json:"latency"` + } + require.NoError(t, json.Unmarshal(body, &payload, json.RejectUnknownMembers(true))) + assert.Equal(t, Model, payload.Model) + assert.Equal(t, 40, payload.Dimensions) + assert.Equal(t, EncodingFloat, payload.EncodingFormat) + assert.Equal(t, LatencyFast, payload.Latency) + value := float64(calls.Add(1)) + vector := make([]float64, 40) + vector[0] = value + response, err := json.Marshal(map[string]any{ + "results": []any{map[string]any{"embedding": vector}}, + "usage": map[string]any{"total_bytes": 155, "total_tokens": 2}, + }) + require.NoError(t, err) + if payload.InputType == "document" { + assert.Equal(t, []string{"first passage"}, payload.Input) + } else { + assert.Equal(t, "query", payload.InputType) + assert.Equal(t, []string{"search words"}, payload.Input) + } + return jsonResponse(request, http.StatusOK, response), nil + })) + inputs := []document.EmbeddingInput{ + {Key: "document", Role: document.EmbeddingRoleDocument, Kind: document.EmbeddingInputRenditionChunk, Text: "first passage"}, + {Key: "query", Role: document.EmbeddingRoleQuery, Kind: document.EmbeddingInputQueryText, Text: "search words"}, + } + + execution, err := client.EmbedWithReceipt(context.Background(), inputs, authorization(client.descriptor, 2)) + require.NoError(t, err) + require.Len(t, execution.Result.Vectors, 2) + assert.Equal(t, "document", execution.Result.Vectors[0].Key) + assert.InDelta(t, 1, execution.Result.Vectors[0].Values[0], 0) + assert.Equal(t, "query", execution.Result.Vectors[1].Key) + assert.InDelta(t, 2, execution.Result.Vectors[1].Values[0], 0) + assert.Equal(t, Receipt{ProviderID: ProviderID, DescriptorFingerprint: client.descriptor.Fingerprint, + PolicyFingerprint: client.descriptor.PolicyFingerprint, Model: Model, ModelRevision: profile.Descriptor.ModelRevision, + EncodingFormat: EncodingFloat, RequestedLatency: LatencyFast, RequestCount: 2, + TotalBytes: 310, TotalTokens: 4}, execution.Receipt) +} + +func TestEmbedDecodesExactLittleEndianBase64(t *testing.T) { + profile := testProfile(t, 40, EncodingBase64, LatencyAuto) + raw := make([]byte, 40*4) + for index := range 40 { + binary.LittleEndian.PutUint32(raw[index*4:], math.Float32bits(float32(index)+0.25)) + } + encoded := base64.StdEncoding.EncodeToString(raw) + client := testClient(t, profile, testSecrets{"secret:zeroentropy": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + body, err := io.ReadAll(request.Body) + require.NoError(t, err) + assert.NotContains(t, string(body), `"latency"`) + return jsonResponse(request, http.StatusOK, []byte(`{"results":[{"embedding":"`+encoded+`"}],"usage":{"total_bytes":155,"total_tokens":2}}`)), nil + })) + + result, err := client.Embed(context.Background(), oneInput(), authorization(client.descriptor, 1)) + require.NoError(t, err) + require.Len(t, result.Vectors, 1) + assert.InDelta(t, 0.25, result.Vectors[0].Values[0], 0) + assert.InDelta(t, 39.25, result.Vectors[0].Values[39], 0) +} + +func TestEmbedRejectsResponseDriftAndInvalidVectors(t *testing.T) { + validVector := strings.TrimSuffix(strings.Repeat("0,", 40), ",") + valid := `{"results":[{"embedding":[` + validVector + `]}],"usage":{"total_bytes":155,"total_tokens":2}}` + tests := map[string]string{ + "missing result": `{"results":[],"usage":{"total_bytes":155,"total_tokens":2}}`, + "wrong dimension": strings.Replace(valid, validVector, "0,0", 1), + "non finite": strings.Replace(valid, "0,", "1e1000,", 1), + "missing usage": strings.Replace(valid, `,"usage":{"total_bytes":155,"total_tokens":2}`, "", 1), + "negative usage": strings.Replace(valid, `"total_tokens":2`, `"total_tokens":-1`, 1), + "unknown field": strings.Replace(valid, `"usage":`, `"private":true,"usage":`, 1), + } + for name, body := range tests { + t.Run(name, func(t *testing.T) { + profile := testProfile(t, 40, EncodingFloat, LatencySlow) + client := testClient(t, profile, testSecrets{"secret:zeroentropy": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + return jsonResponse(request, http.StatusOK, []byte(body)), nil + })) + _, err := client.Embed(context.Background(), oneInput(), authorization(client.descriptor, 1)) + require.ErrorIs(t, err, ErrPermanentResponse) + assert.NotContains(t, err.Error(), "private") + }) + } +} + +func TestEmbedRejectsProviderPayloadLimitBeforeSecretsOrEgress(t *testing.T) { + profile := testProfile(t, 40, EncodingFloat, LatencyFast) + profile.MaxInputItemBytes = maximumInputBytes + profile.MaxInputBytes = maximumInputBytes + profile.Descriptor = descriptorFor(t, profile) + secrets := &countingSecrets{value: "synthetic-key"} + var requests atomic.Int32 + client := testClient(t, profile, secrets, roundTripFunc(func(*http.Request) (*http.Response, error) { + requests.Add(1) + return nil, assert.AnError + })) + input := oneInput() + input[0].Text = strings.Repeat("a", int(maximumInputBytes-100)) + authorization := authorization(client.descriptor, 1) + authorization.MaxInputBytes = maximumInputBytes + + _, err := client.Embed(context.Background(), input, authorization) + require.ErrorIs(t, err, ErrCapacityResponse) + assert.Zero(t, secrets.calls.Load()) + assert.Zero(t, requests.Load()) +} + +func TestEmbedPreservesCancellationDuringResponseRead(t *testing.T) { + started := make(chan struct{}) + profile := testProfile(t, 40, EncodingFloat, LatencyFast) + client := testClient(t, profile, testSecrets{"secret:zeroentropy": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) { + close(started) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &contextBody{ctx: request.Context()}, Request: request}, nil + })) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := client.Embed(ctx, oneInput(), authorization(client.descriptor, 1)) + done <- err + }() + <-started + cancel() + assert.ErrorIs(t, <-done, context.Canceled) +} + +func TestRetryAfterClampsBeforeDurationConversion(t *testing.T) { + err := statusError(http.StatusTooManyRequests, "9223372036854775807", time.Now().UTC()) + delay, ok := RetryAfter(err) + assert.True(t, ok) + assert.Equal(t, time.Hour, delay) +} + +func oneInput() []document.EmbeddingInput { + return []document.EmbeddingInput{{Key: "document", Role: document.EmbeddingRoleDocument, + Kind: document.EmbeddingInputRenditionChunk, Text: "alpha"}} +} + +func authorization(descriptor document.EmbeddingDescriptor, batch int) document.EmbeddingAuthorization { + return document.EmbeddingAuthorization{ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, + PolicyFingerprint: descriptor.PolicyFingerprint, MaxBatchItems: batch, MaxInputBytes: 4096, + MaxResponseBytes: int64(batch*descriptor.Dimension*4 + 1024)} +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return function(request) +} + +func testClient(t *testing.T, profile Profile, secrets SecretResolver, transport http.RoundTripper) *Client { + t.Helper() + client, err := New(profile, secrets, testResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + client.http.Transport = transport + return client +} + +func jsonResponse(request *http.Request, status int, body []byte) *http.Response { + return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader(body)), Request: request} +} + +type countingSecrets struct { + value string + calls atomic.Int32 +} + +func (resolver *countingSecrets) ResolveSecret(context.Context, string) (string, error) { + resolver.calls.Add(1) + return resolver.value, nil +} + +type contextBody struct{ ctx context.Context } + +func (body *contextBody) Read([]byte) (int, error) { + <-body.ctx.Done() + return 0, body.ctx.Err() +} + +func (*contextBody) Close() error { return nil } diff --git a/document/zeroentropyembed/errors.go b/document/zeroentropyembed/errors.go new file mode 100644 index 00000000..164e7ee4 --- /dev/null +++ b/document/zeroentropyembed/errors.go @@ -0,0 +1,70 @@ +package zeroentropyembed + +import ( + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" +) + +var ( + ErrTransientResponse = errors.New("zeroentropy embed: transient provider response") + ErrCapacityResponse = errors.New("zeroentropy embed: provider capacity exceeded") + ErrPermanentResponse = errors.New("zeroentropy embed: permanent provider response") +) + +type ProviderError struct { + Kind error + StatusCode int + RetryDelay time.Duration + RetrySet bool +} + +func (failure *ProviderError) Error() string { + if failure.StatusCode != 0 { + return fmt.Sprintf("zeroentropy embed: HTTP %d: %v", failure.StatusCode, failure.Kind) + } + return failure.Kind.Error() +} + +func (failure *ProviderError) Unwrap() error { return failure.Kind } + +func RetryAfter(err error) (time.Duration, bool) { + failure, ok := errors.AsType[*ProviderError](err) + if !ok || !failure.RetrySet { + return 0, false + } + return failure.RetryDelay, true +} + +func statusError(status int, retryAfter string, now time.Time) error { + kind := ErrPermanentResponse + switch { + case status == http.StatusRequestEntityTooLarge: + kind = ErrCapacityResponse + case status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500 && status <= 599: + kind = ErrTransientResponse + } + delay, set := parseRetryAfter(retryAfter, now) + if !errors.Is(kind, ErrTransientResponse) { + delay, set = 0, false + } + return &ProviderError{Kind: kind, StatusCode: status, RetryDelay: delay, RetrySet: set} +} + +func parseRetryAfter(value string, now time.Time) (time.Duration, bool) { + value = strings.TrimSpace(value) + if seconds, err := strconv.ParseInt(value, 10, 64); err == nil && seconds >= 0 { + if seconds >= int64(time.Hour/time.Second) { + return time.Hour, true + } + return time.Duration(seconds) * time.Second, true + } + when, err := http.ParseTime(value) + if err != nil { + return 0, false + } + return min(max(when.Sub(now), 0), time.Hour), true +} diff --git a/document/zeroentropyembed/profile.go b/document/zeroentropyembed/profile.go new file mode 100644 index 00000000..b50604f3 --- /dev/null +++ b/document/zeroentropyembed/profile.go @@ -0,0 +1,360 @@ +// Package zeroentropyembed implements the fixed hosted ZeroEntropy zembed-1 contract. +package zeroentropyembed + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json/v2" + "errors" + "net/http" + "net/netip" + "reflect" + "slices" + "strings" + "time" + "unicode" + "unicode/utf8" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +const ( + ProviderID = "zeroentropy.hosted.zembed-1-v1" + Model = "zembed-1" + DocumentFormatterV1 = "zeroentropy-zembed-1/document/v1" + QueryFormatterV1 = "zeroentropy-zembed-1/query/v1" + ScalarEncodingFloat32 = "float32" + TransformNone = "none" + + host = "api.zeroentropy.dev" + origin = "https://api.zeroentropy.dev" + embedPath = "/v1/models/embed" + adapterContract = "docbank-zeroentropy-zembed-1/v1" + compatibilityID = "zeroentropy/zembed-1/retrieval/v1" + defaultTimeout = 30 * time.Second + maximumTimeout = 5 * time.Minute + defaultBatch = 128 + maximumBatch = 2048 + defaultItemBytes = int64(1 << 20) + defaultInputBytes = int64(4 << 20) + maximumInputBytes = int64(5_000_000) + defaultRequestBytes = int64(8 << 20) + maximumRequestBytes = int64(16 << 20) + defaultResponseBytes = int64(32 << 20) + maximumResponseBytes = int64(128 << 20) + maximumTokenBytes = 128 +) + +var supportedDimensions = []int{40, 80, 160, 320, 640, 1280, 2560} + +type EncodingFormat string + +const ( + EncodingFloat EncodingFormat = "float" + EncodingBase64 EncodingFormat = "base64" +) + +type Latency string + +const ( + LatencyAuto Latency = "auto" + LatencyFast Latency = "fast" + LatencySlow Latency = "slow" +) + +type SecretResolver interface { + ResolveSecret(ctx context.Context, binding string) (string, error) +} + +type Profile struct { + Descriptor document.EmbeddingDescriptor + CompatibilityEpoch string + SecretBinding string + EncodingFormat EncodingFormat + Latency Latency + ClientTransform string + RequestTimeout time.Duration + MaxBatchItems int + MaxInputItemBytes int64 + MaxInputBytes int64 + MaxRequestBytes int64 + MaxResponseBytes int64 + EgressPolicy providerhttp.EgressPolicy +} + +type Client struct { + profile Profile + descriptor document.EmbeddingDescriptor + secrets SecretResolver + http *http.Client +} + +type policyIdentity struct { + AdapterContract string `json:"adapter_contract"` + Origin string `json:"origin"` + Route string `json:"route"` + Descriptor document.EmbeddingDescriptor `json:"descriptor"` + CompatibilityEpoch string `json:"compatibility_epoch"` + SecretBinding string `json:"secret_binding"` + EncodingFormat EncodingFormat `json:"encoding_format"` + Latency Latency `json:"latency"` + ClientTransform string `json:"client_transform"` + RequestTimeout int64 `json:"request_timeout_nanos"` + MaxBatchItems int `json:"max_batch_items"` + MaxInputItemBytes int64 `json:"max_input_item_bytes"` + MaxInputBytes int64 `json:"max_input_bytes"` + MaxRequestBytes int64 `json:"max_request_bytes"` + MaxResponseBytes int64 `json:"max_response_bytes"` + Egress egressIdentity `json:"egress"` +} + +type egressIdentity struct { + Scheme string `json:"scheme"` + Host string `json:"host"` + Port uint16 `json:"port"` + AllowedCIDRs []string `json:"allowed_cidrs"` + ProxyMode string `json:"proxy_mode"` + ConnectTimeout int64 `json:"connect_timeout_nanos"` + KeepAlive int64 `json:"keep_alive_nanos"` + TLSHandshakeTimeout int64 `json:"tls_handshake_timeout_nanos"` + SPKISHA256 []string `json:"spki_sha256,omitempty"` +} + +func PolicyFingerprint(profile Profile) (string, error) { + normalized, descriptor, err := normalizeProfile(profile) + if err != nil { + return "", err + } + encoded, err := json.Marshal(policyIdentity{ + AdapterContract: adapterContract, Origin: origin, Route: embedPath, Descriptor: descriptor, + CompatibilityEpoch: normalized.CompatibilityEpoch, SecretBinding: normalized.SecretBinding, + EncodingFormat: normalized.EncodingFormat, Latency: normalized.Latency, + ClientTransform: normalized.ClientTransform, RequestTimeout: int64(normalized.RequestTimeout), + MaxBatchItems: normalized.MaxBatchItems, MaxInputItemBytes: normalized.MaxInputItemBytes, + MaxInputBytes: normalized.MaxInputBytes, MaxRequestBytes: normalized.MaxRequestBytes, + MaxResponseBytes: normalized.MaxResponseBytes, Egress: profileEgressIdentity(normalized.EgressPolicy), + }, json.Deterministic(true)) + if err != nil { + return "", errors.New("zeroentropy embed: policy identity encoding failed") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +func New(profile Profile, secrets SecretResolver, resolver providerhttp.Resolver, supplied *http.Client) (*Client, error) { + if supplied == nil { + return nil, errors.New("zeroentropy embed: HTTP client settings source is required") + } + normalized, _, err := normalizeProfile(profile) + if err != nil { + return nil, err + } + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + if err != nil || !reflect.DeepEqual(descriptor, profile.Descriptor) { + return nil, errors.New("zeroentropy embed: descriptor is not canonical") + } + fingerprint, err := PolicyFingerprint(profile) + if err != nil { + return nil, err + } + if descriptor.PolicyFingerprint != fingerprint { + return nil, errors.New("zeroentropy embed: descriptor policy fingerprint does not match profile") + } + if nilInterface(secrets) { + return nil, errors.New("zeroentropy embed: named API-key resolver is required") + } + transport, err := providerhttp.NewTransport(normalized.EgressPolicy, resolver) + if err != nil { + return nil, errors.New("zeroentropy embed: sealed egress policy is invalid") + } + isolated := *supplied + isolated.Transport = transport + isolated.CheckRedirect = providerhttp.RefuseRedirects + isolated.Jar = nil + isolated.Timeout = 0 + normalized.Descriptor = cloneDescriptor(descriptor) + return &Client{profile: normalized, descriptor: cloneDescriptor(descriptor), secrets: secrets, http: &isolated}, nil +} + +func (client *Client) Descriptor() document.EmbeddingDescriptor { + if client == nil { + return document.EmbeddingDescriptor{} + } + return cloneDescriptor(client.descriptor) +} + +func normalizeProfile(profile Profile) (Profile, document.EmbeddingDescriptor, error) { + profile.EgressPolicy.AllowedCIDRs = slices.Clone(profile.EgressPolicy.AllowedCIDRs) + profile.EgressPolicy.TLS.SPKISHA256 = slices.Clone(profile.EgressPolicy.TLS.SPKISHA256) + if profile.EncodingFormat == "" { + profile.EncodingFormat = EncodingFloat + } + if profile.Latency == "" { + profile.Latency = LatencyAuto + } + if profile.ClientTransform == "" { + profile.ClientTransform = TransformNone + } + if profile.RequestTimeout == 0 { + profile.RequestTimeout = defaultTimeout + } + if profile.MaxBatchItems == 0 { + profile.MaxBatchItems = defaultBatch + } + if profile.MaxInputItemBytes == 0 { + profile.MaxInputItemBytes = defaultItemBytes + } + if profile.MaxInputBytes == 0 { + profile.MaxInputBytes = defaultInputBytes + } + if profile.MaxRequestBytes == 0 { + profile.MaxRequestBytes = defaultRequestBytes + } + if profile.MaxResponseBytes == 0 { + profile.MaxResponseBytes = defaultResponseBytes + } + if !slices.Contains(supportedDimensions, profile.Descriptor.Dimension) || + (profile.EncodingFormat != EncodingFloat && profile.EncodingFormat != EncodingBase64) || + (profile.Latency != LatencyAuto && profile.Latency != LatencyFast && profile.Latency != LatencySlow) || + profile.ClientTransform != TransformNone || profile.RequestTimeout <= 0 || profile.RequestTimeout > maximumTimeout || + profile.MaxBatchItems < 1 || profile.MaxBatchItems > maximumBatch || + profile.MaxInputItemBytes < 1 || profile.MaxInputItemBytes > maximumInputBytes || + profile.MaxInputBytes < profile.MaxInputItemBytes || profile.MaxInputBytes > maximumInputBytes || + profile.MaxRequestBytes < 1 || profile.MaxRequestBytes > maximumRequestBytes || + profile.MaxResponseBytes < 1 || profile.MaxResponseBytes > maximumResponseBytes { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("zeroentropy embed: profile bounds or execution policy are invalid") + } + if !validToken(profile.CompatibilityEpoch) || profile.Descriptor.ModelRevision != profile.CompatibilityEpoch { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("zeroentropy embed: compatibility epoch must match descriptor revision") + } + if !validToken(profile.SecretBinding) { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("zeroentropy embed: named API-key binding is required") + } + if err := normalizeEgress(&profile.EgressPolicy); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + descriptor := cloneDescriptor(profile.Descriptor) + descriptor.PolicyFingerprint = strings.Repeat("0", sha256.Size*2) + descriptor.Fingerprint = "" + var err error + descriptor, err = document.NewEmbeddingDescriptor(descriptor) + if err != nil { + return Profile{}, document.EmbeddingDescriptor{}, errors.New("zeroentropy embed: descriptor identity is invalid") + } + descriptor.PolicyFingerprint, descriptor.Fingerprint = "", "" + if err := validateDescriptor(descriptor, profile.CompatibilityEpoch); err != nil { + return Profile{}, document.EmbeddingDescriptor{}, err + } + return profile, descriptor, nil +} + +func validateDescriptor(descriptor document.EmbeddingDescriptor, epoch string) error { + expected, err := modelInputContract() + if err != nil { + return errors.New("zeroentropy embed: fixed model-input contract is invalid") + } + if descriptor.ID != ProviderID || descriptor.ContractVersion != document.EmbeddingProviderContractVersion || + descriptor.TrustBoundary != document.EmbeddingTrustHostedProvider || descriptor.Model != Model || + descriptor.ModelRevision != epoch || !slices.Contains(supportedDimensions, descriptor.Dimension) || + descriptor.Metric != document.VectorMetricCosine || descriptor.Normalization != document.VectorNormalizationNone || + descriptor.ScalarEncoding != ScalarEncodingFloat32 || descriptor.DocumentFormatter != DocumentFormatterV1 || + descriptor.QueryFormatter != QueryFormatterV1 || !descriptor.SupportsTextQuery || + !reflect.DeepEqual(descriptor.ModelInput, expected) || descriptor.CompatibilityID != compatibilityID || + !slices.Equal(descriptor.InputKinds, []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}) || + !slices.Equal(descriptor.SupportedRequestModes, []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}) { + return errors.New("zeroentropy embed: descriptor does not match the fixed hosted zembed-1 contract") + } + return nil +} + +func modelInputContract() (document.ModelInputContract, error) { + return document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: compatibilityID, + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}, + }) +} + +func normalizeEgress(policy *providerhttp.EgressPolicy) error { + if policy.ConnectTimeout == 0 { + policy.ConnectTimeout = providerhttp.DefaultConnectTimeout + } + if policy.KeepAlive == 0 { + policy.KeepAlive = providerhttp.DefaultKeepAlive + } + if policy.TLSHandshakeTimeout == 0 { + policy.TLSHandshakeTimeout = providerhttp.DefaultTLSHandshakeTimeout + } + if policy.ProxyMode == "" { + policy.ProxyMode = providerhttp.ProxyDisabled + } + if policy.Scheme != "https" || policy.Host != host || policy.Port != 443 || + policy.ProxyMode != providerhttp.ProxyDisabled || policy.TLS.RootCAs != nil { + return errors.New("zeroentropy embed: egress authority must be exactly api.zeroentropy.dev:443") + } + for index := range policy.AllowedCIDRs { + policy.AllowedCIDRs[index] = policy.AllowedCIDRs[index].Masked() + } + slices.SortFunc(policy.AllowedCIDRs, func(left, right netip.Prefix) int { return strings.Compare(left.String(), right.String()) }) + for index := 1; index < len(policy.AllowedCIDRs); index++ { + if policy.AllowedCIDRs[index] == policy.AllowedCIDRs[index-1] { + return errors.New("zeroentropy embed: duplicate egress CIDR") + } + } + for index := range policy.TLS.SPKISHA256 { + policy.TLS.SPKISHA256[index] = strings.ToLower(policy.TLS.SPKISHA256[index]) + } + slices.Sort(policy.TLS.SPKISHA256) + for index := 1; index < len(policy.TLS.SPKISHA256); index++ { + if policy.TLS.SPKISHA256[index] == policy.TLS.SPKISHA256[index-1] { + return errors.New("zeroentropy embed: duplicate SPKI pin") + } + } + if _, err := providerhttp.NewTransport(*policy, nil); err != nil { + return errors.New("zeroentropy embed: sealed egress policy is invalid") + } + return nil +} + +func profileEgressIdentity(policy providerhttp.EgressPolicy) egressIdentity { + cidrs := make([]string, len(policy.AllowedCIDRs)) + for index, prefix := range policy.AllowedCIDRs { + cidrs[index] = prefix.String() + } + return egressIdentity{Scheme: policy.Scheme, Host: policy.Host, Port: policy.Port, AllowedCIDRs: cidrs, + ProxyMode: string(policy.ProxyMode), ConnectTimeout: int64(policy.ConnectTimeout), KeepAlive: int64(policy.KeepAlive), + TLSHandshakeTimeout: int64(policy.TLSHandshakeTimeout), SPKISHA256: slices.Clone(policy.TLS.SPKISHA256)} +} + +func cloneDescriptor(descriptor document.EmbeddingDescriptor) document.EmbeddingDescriptor { + descriptor.InputKinds = slices.Clone(descriptor.InputKinds) + descriptor.SupportedRequestModes = slices.Clone(descriptor.SupportedRequestModes) + return descriptor +} + +func validToken(value string) bool { + if value == "" || len(value) > maximumTokenBytes || !utf8.ValidString(value) || value != strings.TrimSpace(value) { + return false + } + for _, current := range value { + if unicode.IsControl(current) { + return false + } + } + return true +} + +func nilInterface(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/document/zeroentropyembed/profile_test.go b/document/zeroentropyembed/profile_test.go new file mode 100644 index 00000000..8464bd4b --- /dev/null +++ b/document/zeroentropyembed/profile_test.go @@ -0,0 +1,142 @@ +package zeroentropyembed + +import ( + "context" + "errors" + "net/http" + "net/netip" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/providerhttp" +) + +func TestNewRequiresFixedZembedProfile(t *testing.T) { + for _, dimension := range []int{2560, 1280, 640, 320, 160, 80, 40} { + for _, encoding := range []EncodingFormat{EncodingFloat, EncodingBase64} { + profile := testProfile(t, dimension, encoding, LatencyAuto) + client, err := New(profile, testSecrets{"secret:zeroentropy": "synthetic-key"}, + testResolver{netip.MustParseAddr("192.0.2.10")}, &http.Client{}) + require.NoError(t, err) + assert.Equal(t, profile.Descriptor, client.Descriptor()) + } + } + + mutations := map[string]func(*Profile){ + "model": func(value *Profile) { value.Descriptor.Model = "zembed-2" }, + "dimension": func(value *Profile) { value.Descriptor.Dimension = 768 }, + "epoch": func(value *Profile) { value.CompatibilityEpoch = "other" }, + "encoding": func(value *Profile) { value.EncodingFormat = "json" }, + "latency": func(value *Profile) { value.Latency = "instant" }, + "transform": func(value *Profile) { value.ClientTransform = "truncate" }, + "secret": func(value *Profile) { value.SecretBinding = "" }, + "host": func(value *Profile) { value.EgressPolicy.Host = "example.com" }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + profile := testProfile(t, 640, EncodingBase64, LatencyFast) + mutate(&profile) + _, err := PolicyFingerprint(profile) + require.Error(t, err) + }) + } +} + +func TestPolicyFingerprintBindsZembedExecutionIdentityWithoutMutatingProfile(t *testing.T) { + profile := testProfile(t, 640, EncodingBase64, LatencyFast) + original := slices.Clone(profile.EgressPolicy.AllowedCIDRs) + base, err := PolicyFingerprint(profile) + require.NoError(t, err) + mutations := map[string]func(*Profile){ + "dimension": func(value *Profile) { value.Descriptor.Dimension = 320 }, + "epoch": func(value *Profile) { + value.CompatibilityEpoch, value.Descriptor.ModelRevision = "deployment-2026-09", "deployment-2026-09" + }, + "encoding": func(value *Profile) { value.EncodingFormat = EncodingFloat }, + "latency": func(value *Profile) { value.Latency = LatencySlow }, + "binding": func(value *Profile) { value.SecretBinding = "secret:other" }, + "batch": func(value *Profile) { value.MaxBatchItems-- }, + "item": func(value *Profile) { value.MaxInputItemBytes-- }, + "input": func(value *Profile) { value.MaxInputBytes-- }, + "request": func(value *Profile) { value.MaxRequestBytes-- }, + "response": func(value *Profile) { value.MaxResponseBytes-- }, + "timeout": func(value *Profile) { value.RequestTimeout += time.Second }, + "egress": func(value *Profile) { + value.EgressPolicy.AllowedCIDRs = []netip.Prefix{netip.MustParsePrefix("198.51.100.0/24")} + }, + } + for name, mutate := range mutations { + t.Run(name, func(t *testing.T) { + changed := profile + mutate(&changed) + fingerprint, fingerprintErr := PolicyFingerprint(changed) + require.NoError(t, fingerprintErr) + assert.NotEqual(t, base, fingerprint) + }) + } + assert.Equal(t, original, profile.EgressPolicy.AllowedCIDRs) +} + +func testProfile(t *testing.T, dimension int, encoding EncodingFormat, latency Latency) Profile { + t.Helper() + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileCustom, CompatibilityID: "zeroentropy/zembed-1/retrieval/v1", + Document: document.ModelInputEncoder{Mode: document.ModelInputModeDocument, Template: "{{content}}"}, + Query: document.ModelInputEncoder{Mode: document.ModelInputModeQuery, Template: "{{content}}"}, + }) + require.NoError(t, err) + profile := Profile{ + Descriptor: document.EmbeddingDescriptor{ + ID: ProviderID, ContractVersion: document.EmbeddingProviderContractVersion, + TrustBoundary: document.EmbeddingTrustHostedProvider, Model: Model, + ModelRevision: "deployment-2026-08", Dimension: dimension, Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationNone, ScalarEncoding: ScalarEncodingFloat32, + DocumentFormatter: DocumentFormatterV1, QueryFormatter: QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, SupportsTextQuery: true, ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{document.ModelInputModeDocument, document.ModelInputModeQuery}, + }, + CompatibilityEpoch: "deployment-2026-08", SecretBinding: "secret:zeroentropy", + EncodingFormat: encoding, Latency: latency, ClientTransform: TransformNone, + RequestTimeout: time.Second, MaxBatchItems: 128, MaxInputItemBytes: 1 << 20, + MaxInputBytes: 4 << 20, MaxRequestBytes: 8 << 20, MaxResponseBytes: 32 << 20, + EgressPolicy: providerhttp.EgressPolicy{Scheme: "https", Host: "api.zeroentropy.dev", Port: 443, + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("192.0.2.0/24")}, + ProxyMode: providerhttp.ProxyDisabled, ConnectTimeout: time.Second, + KeepAlive: time.Second, TLSHandshakeTimeout: time.Second}, + } + profile.Descriptor = descriptorFor(t, profile) + return profile +} + +func descriptorFor(t *testing.T, profile Profile) document.EmbeddingDescriptor { + t.Helper() + profile.Descriptor.PolicyFingerprint = "" + profile.Descriptor.Fingerprint = "" + fingerprint, err := PolicyFingerprint(profile) + require.NoError(t, err) + profile.Descriptor.PolicyFingerprint = fingerprint + descriptor, err := document.NewEmbeddingDescriptor(profile.Descriptor) + require.NoError(t, err) + return descriptor +} + +type testSecrets map[string]string + +func (secrets testSecrets) ResolveSecret(_ context.Context, binding string) (string, error) { + value, ok := secrets[binding] + if !ok { + return "", errors.New("missing synthetic secret") + } + return value, nil +} + +type testResolver []netip.Addr + +func (resolver testResolver) LookupNetIP(context.Context, string, string) ([]netip.Addr, error) { + return append([]netip.Addr(nil), resolver...), nil +} diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 3fb54e03..7707ccd4 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -11,7 +11,8 @@ "@kenn-io/kit-ui": "git+https://github.com/kenn-io/kit-ui.git#b35e4ca7df66969f4650327d6e315986afef8716", "@lucide/svelte": "1.24.0", "@noble/hashes": "^2.2.0", - "svelte": "5.56.4" + "svelte": "5.56.4", + "yaml": "2.8.1" }, "devDependencies": { "@playwright/test": "1.61.1", @@ -3994,6 +3995,18 @@ "dev": true, "license": "MIT" }, + "node_modules/yaml": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.8.1.tgz", + "integrity": "sha512-lcYcMxX2PO9XMGvAJkJ3OsNMw+/7FKes7/hgerGUYWIoWu5j/+YQqcZr5JnPZWzOsEBgMbSbiSTn/dv/69Mkpw==", + "license": "ISC", + "bin": { + "yaml": "bin.mjs" + }, + "engines": { + "node": ">= 14.6" + } + }, "node_modules/zimmerframe": { "version": "1.1.4", "resolved": "https://registry.npmjs.org/zimmerframe/-/zimmerframe-1.1.4.tgz", diff --git a/frontend/package.json b/frontend/package.json index 2887235c..3fa4a2d4 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -19,7 +19,8 @@ "@kenn-io/kit-ui": "git+https://github.com/kenn-io/kit-ui.git#b35e4ca7df66969f4650327d6e315986afef8716", "@lucide/svelte": "1.24.0", "@noble/hashes": "^2.2.0", - "svelte": "5.56.4" + "svelte": "5.56.4", + "yaml": "2.8.1" }, "devDependencies": { "@playwright/test": "1.61.1", diff --git a/frontend/screenshots/README.md b/frontend/screenshots/README.md index b4d44499..e254168d 100644 --- a/frontend/screenshots/README.md +++ b/frontend/screenshots/README.md @@ -1,3 +1,7 @@ +--- +last_edited: 2026-08-28 +--- + # Web screenshots This Playwright harness captures the actual daemon-served Docbank interface @@ -9,6 +13,7 @@ Install Chromium once: ```sh cd frontend npm ci +node node_modules/@playwright/test/cli.js install-deps chromium node node_modules/@playwright/test/cli.js install chromium ``` @@ -26,7 +31,9 @@ inspection and PR attachment; the current case captures both move-to-trash and restore confirmations, the tag-definition catalog, and a completed tag assignment, current vault browsing, extracted-text search, retained-version selection, packed-storage status, and independently verified permanent-audit -evidence. Generated images are intentionally not committed. +evidence. It also captures the reviewed document-processing plan, partial +semantic coverage, and the retained sanitized Markdown rendition. Generated +images are intentionally not committed. Pass ordinary Playwright arguments after `--` to select a case: diff --git a/frontend/screenshots/processing-profile.go b/frontend/screenshots/processing-profile.go new file mode 100644 index 00000000..9cd7e82d --- /dev/null +++ b/frontend/screenshots/processing-profile.go @@ -0,0 +1,107 @@ +//go:build ignore + +// processing-profile prints the dynamic descriptor identities used by the +// real-daemon screenshot harness. The embedding identity binds the loopback +// endpoint selected by the test, so it cannot be checked in as a constant. +package main + +import ( + "encoding/json" + "fmt" + "net/netip" + "net/url" + "os" + "strconv" + "strings" + "time" + + "go.kenn.io/docbank/document" + "go.kenn.io/docbank/document/openaiembed" + "go.kenn.io/docbank/document/plaintext" + "go.kenn.io/docbank/document/providerhttp" +) + +type identities struct { + RenditionID string `json:"rendition_id"` + RenditionFingerprint string `json:"rendition_fingerprint"` + EmbeddingID string `json:"embedding_id"` + EmbeddingFingerprint string `json:"embedding_fingerprint"` + CompatibilityID string `json:"compatibility_id"` +} + +func main() { + if len(os.Args) != 2 { + panic("usage: go run processing-profile.go ") + } + origin := os.Args[1] + parsed, err := url.Parse(origin) + if err != nil || parsed.Scheme != "http" || parsed.Hostname() != "127.0.0.1" || parsed.Port() == "" { + panic("loopback origin must be an exact http://127.0.0.1: URL") + } + port, err := strconv.ParseUint(parsed.Port(), 10, 16) + if err != nil { + panic(err) + } + + renderer, err := plaintext.New(plaintext.Profile{MaxDocumentBytes: plaintext.MaxDocumentBytes}) + if err != nil { + panic(err) + } + contract, err := document.NewModelInputContract(document.ModelInputContractConfig{ + Profile: document.ModelInputProfileNomic, + }) + if err != nil { + panic(err) + } + descriptor, err := document.NewEmbeddingDescriptor(document.EmbeddingDescriptor{ + ID: openaiembed.ProviderID, + ContractVersion: document.EmbeddingProviderContractVersion, + PolicyFingerprint: strings.Repeat("0", 64), + TrustBoundary: document.EmbeddingTrustOperatorNetwork, + Model: "synthetic-model", + ModelRevision: "deployment-v1", + Dimension: 2, + Metric: document.VectorMetricCosine, + Normalization: document.VectorNormalizationNone, + ScalarEncoding: openaiembed.ScalarEncodingFloat32, + DocumentFormatter: openaiembed.DocumentFormatterV1, + QueryFormatter: openaiembed.QueryFormatterV1, + InputKinds: []document.EmbeddingInputKind{document.EmbeddingInputRenditionChunk}, + CompatibilityID: contract.CompatibilityID, + SupportsTextQuery: true, + ModelInput: contract, + SupportedRequestModes: []document.ModelInputMode{contract.Document.Mode}, + }) + if err != nil { + panic(err) + } + profile := openaiembed.Profile{ + Origin: origin, Descriptor: descriptor, ModelInput: contract, + SecretBinding: "credential:semantic", DeploymentEpoch: "deployment-v1", + RequestTimeout: time.Second, MaxBatchItems: 8, MaxInputBytes: 1 << 20, + MaxRequestBytes: 1 << 20, MaxResponseBytes: 1 << 20, + EgressPolicy: providerhttp.EgressPolicy{ + Scheme: "http", Host: "127.0.0.1", Port: uint16(port), + AllowedCIDRs: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + ProxyMode: providerhttp.ProxyDisabled, ConnectTimeout: time.Second, + KeepAlive: time.Second, TLSHandshakeTimeout: time.Second, + }, + } + policyFingerprint, err := openaiembed.PolicyFingerprint(profile) + if err != nil { + panic(err) + } + descriptor.PolicyFingerprint, descriptor.Fingerprint = policyFingerprint, "" + descriptor, err = document.NewEmbeddingDescriptor(descriptor) + if err != nil { + panic(err) + } + result := identities{ + RenditionID: renderer.Descriptor().ID, RenditionFingerprint: renderer.Descriptor().Fingerprint, + EmbeddingID: descriptor.ID, EmbeddingFingerprint: descriptor.Fingerprint, + CompatibilityID: contract.CompatibilityID, + } + if err := json.NewEncoder(os.Stdout).Encode(result); err != nil { + panic(fmt.Errorf("encode identities: %w", err)) + } +} diff --git a/frontend/screenshots/web-trash.screenshot.ts b/frontend/screenshots/web-trash.screenshot.ts index ea8a5c8e..0f70b6c5 100644 --- a/frontend/screenshots/web-trash.screenshot.ts +++ b/frontend/screenshots/web-trash.screenshot.ts @@ -1,6 +1,8 @@ import { expect, test } from "@playwright/test"; import { execFile } from "node:child_process"; import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; import { tmpdir } from "node:os"; import path from "node:path"; import { promisify } from "node:util"; @@ -76,11 +78,30 @@ const packedStorageScreenshotPath = path.join( "screenshots", "web-storage-status.png", ); +const processingPlanScreenshotPath = path.join( + repositoryRoot, + ".superpowers", + "screenshots", + "web-document-processing-plan.png", +); +const processingPartialScreenshotPath = path.join( + repositoryRoot, + ".superpowers", + "screenshots", + "web-document-processing-partial.png", +); +const renditionScreenshotPath = path.join( + repositoryRoot, + ".superpowers", + "screenshots", + "web-document-rendition.png", +); test.describe("Docbank web screenshots", () => { let workspace = ""; let vault = ""; let webURL = ""; + let embeddingServer: Server | undefined; async function runDocbank(args: string[]): Promise { const result = await execFileAsync(binary, args, { @@ -88,6 +109,7 @@ test.describe("Docbank web screenshots", () => { env: { ...process.env, DOCBANK_HOME: vault, + DOCBANK_SCREENSHOT_EMBEDDING_KEY: "synthetic-secret", }, maxBuffer: 1024 * 1024, timeout: 60_000, @@ -154,12 +176,127 @@ test.describe("Docbank web screenshots", () => { await rm(searchResultsScreenshotPath, { force: true }); await rm(retainedVersionScreenshotPath, { force: true }); await rm(packedStorageScreenshotPath, { force: true }); + await rm(processingPlanScreenshotPath, { force: true }); + await rm(processingPartialScreenshotPath, { force: true }); + await rm(renditionScreenshotPath, { force: true }); const archive = path.join(workspace, "archive-store"); await mkdir(vault, { recursive: true, mode: 0o700 }); await mkdir(archive, { recursive: true, mode: 0o700 }); + + embeddingServer = createServer((_request, response) => { + response.writeHead(503, { "Content-Type": "application/json" }); + response.end(JSON.stringify({ error: { message: "synthetic provider unavailable" } })); + }); + await new Promise((resolve, reject) => { + embeddingServer!.once("error", reject); + embeddingServer!.listen(0, "127.0.0.1", resolve); + }); + const embeddingAddress = embeddingServer.address() as AddressInfo; + const embeddingOrigin = `http://127.0.0.1:${embeddingAddress.port}`; + const identityResult = await execFileAsync( + "go", + ["run", "-tags", "fts5", path.join(here, "processing-profile.go"), embeddingOrigin], + { cwd: repositoryRoot, maxBuffer: 1024 * 1024, timeout: 60_000 }, + ); + const identities = JSON.parse(identityResult.stdout) as { + rendition_id: string; + rendition_fingerprint: string; + embedding_id: string; + embedding_fingerprint: string; + compatibility_id: string; + }; await writeFile( path.join(vault, "config.toml"), - `[store_bindings.archive]\nkind = "filesystem"\npath = ${JSON.stringify(archive)}\npriority = 20\n`, + `[store_bindings.archive] +kind = "filesystem" +path = ${JSON.stringify(archive)} +priority = 20 + +[credential_bindings.semantic] +environment_variable = "DOCBANK_SCREENSHOT_EMBEDDING_KEY" + +[rendition_profiles.plaintext] +adapter_contract = "docbank-plaintext-rendition/v1" +authorization_fingerprint = "1111111111111111111111111111111111111111111111111111111111111111" +credential_binding = "credential:none" +deployment_fingerprint = "2222222222222222222222222222222222222222222222222222222222222222" +descriptor_id = ${JSON.stringify(identities.rendition_id)} +descriptor_fingerprint = ${JSON.stringify(identities.rendition_fingerprint)} +disclose_filename = false +disclosure_fingerprint = "3333333333333333333333333333333333333333333333333333333333333333" +max_document_bytes = 16777216 +max_response_bytes = 16777216 +max_units = 1 +requested_artifacts = ["structured_evidence"] +trust_boundary = "local_process" +upload_options_fingerprint = "4444444444444444444444444444444444444444444444444444444444444444" + +[embedding_profiles.semantic] +activation = "optional" +authorization_fingerprint = "5555555555555555555555555555555555555555555555555555555555555555" +compatibility_id = ${JSON.stringify(identities.compatibility_id)} +credential_binding = "credential:semantic" +descriptor_id = ${JSON.stringify(identities.embedding_id)} +descriptor_fingerprint = ${JSON.stringify(identities.embedding_fingerprint)} +dimensions = 2 +disclosure_fingerprint = "6666666666666666666666666666666666666666666666666666666666666666" +document_formatter = "openai-compatible/document/v1" +input_kind = "rendition_chunk" +max_batch_items = 8 +max_input_bytes = 1048576 +max_response_bytes = 1048576 +metric = "cosine" +model = "synthetic-model" +normalization = "none" +query_formatter = "openai-compatible/query/v1" +scalar_encoding = "float32" +trust_boundary = "operator_network" + +[embedding_profiles.semantic.chunk] +context_fingerprint = "7777777777777777777777777777777777777777777777777777777777777777" +formatter = "rendition-chunk/v1" +max_tokens = 128 +overlap_tokens = 8 +tokenizer = "unicode-runes@v1" +truncation_policy = "reject_indivisible" + +[embedding_profiles.semantic.model_input] +profile = "nomic/v1" + +[embedding_profiles.semantic.runtime] +adapter_contract = "docbank-openai-compatible-embeddings/v1" +endpoint = ${JSON.stringify(embeddingOrigin)} +model_revision = "deployment-v1" +deployment_epoch = "deployment-v1" +request_timeout = "1s" +max_request_bytes = 1048576 +max_retries = 1 +allowed_cidrs = ["127.0.0.0/8"] +proxy_mode = "disabled" +connect_timeout = "1s" +keep_alive = "1s" +tls_handshake_timeout = "1s" + +[retrieval_profiles.hybrid] +lexical_limit = 20 +vector_limit = 20 + +[processing_profiles.private_text] +rendition = "plaintext" +embeddings = ["semantic"] +retrieval = "hybrid" +attachment_policy_fingerprint = "8888888888888888888888888888888888888888888888888888888888888888" +completeness_fingerprint = "9999999999999999999999999999999999999999999999999999999999999999" +consent_fingerprint = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" +lexical_segmenter_fingerprint = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" +max_segment_runes = 2000 +max_unit_runes = 100000 +normalizer_fingerprint = "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc" +sanitizer_fingerprint = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd" +retain_sanitized_markdown = true +retain_typed_artifacts = true +trust_boundary = "local_process" +`, { mode: 0o600 }, ); const reports = path.join(workspace, "synthetic", "Reports"); @@ -326,6 +463,11 @@ test.describe("Docbank web screenshots", () => { } } if (workspace) await rm(workspace, { recursive: true, force: true }); + if (embeddingServer) { + await new Promise((resolve, reject) => { + embeddingServer!.close((error) => (error ? reject(error) : resolve())); + }); + } }); test("trash confirmation", async ({ page }) => { @@ -359,6 +501,36 @@ test.describe("Docbank web screenshots", () => { animations: "disabled", }); + await page.getByRole("button", { name: "Process and retrieve" }).click(); + const processing = page.getByRole("dialog", { + name: "Document processing and coverage", + }); + await expect(processing).toContainText("Private network"); + await expect(processing).toContainText("sanitized_markdown"); + await page.screenshot({ + path: processingPlanScreenshotPath, + fullPage: true, + animations: "disabled", + }); + await processing.getByRole("button", { name: "Consent and run" }).click(); + await expect(processing.getByRole("button", { name: "Read sanitized Markdown" })).toBeVisible({ timeout: 30_000 }); + await expect(processing).toContainText(/semantic.*unavailable/i); + await page.screenshot({ + path: processingPartialScreenshotPath, + fullPage: true, + animations: "disabled", + }); + await processing.getByRole("button", { name: "Read sanitized Markdown" }).click(); + const rendition = page.getByRole("dialog", { name: "Sanitized Markdown rendition" }); + await expect(rendition).toContainText("Synthetic quarterly tax report with reviewed totals"); + await expect(rendition).toContainText("degraded_provenance"); + await page.screenshot({ + path: renditionScreenshotPath, + fullPage: true, + animations: "disabled", + }); + await rendition.getByRole("button", { name: "Close sanitized Markdown" }).click(); + await page.getByRole("button", { name: "Version history" }).click(); const versions = page.getByRole("dialog", { name: "Immutable version history for /Reports/quarterly-tax-report.txt", diff --git a/frontend/src/App.processing.test.ts b/frontend/src/App.processing.test.ts new file mode 100644 index 00000000..b1734b5c --- /dev/null +++ b/frontend/src/App.processing.test.ts @@ -0,0 +1,40 @@ +import { afterEach, expect, it, vi } from "vitest"; +import { cleanup, fireEvent, render, screen } from "@testing-library/svelte"; +import App from "./App.svelte"; + +afterEach(() => { + cleanup(); + history.replaceState(null, "", "/"); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +it("opens processing for the selected exact document version", async () => { + history.replaceState(null, "", "/#web_session=short-lived&web_upload_secret=proof"); + vi.stubGlobal("ResizeObserver", class { observe() {} unobserve() {} disconnect() {} }); + Object.defineProperty(Element.prototype, "scrollIntoView", { configurable: true, value: vi.fn() }); + const versionID = "11111111-1111-4111-8111-111111111111"; + const fingerprint = "a".repeat(64); + const root = { id: 1, name: "", kind: "dir", size: 0, revision: 1, created_at: "2026-08-28T00:00:00Z", modified_at: "2026-08-28T00:00:00Z", path: "/" }; + const file = { id: 2, parent_id: 1, name: "report.pdf", kind: "file", current_version_id: versionID, blob_hash: "b".repeat(64), size: 2048, mime_type: "application/pdf", revision: 1, created_at: "2026-08-28T00:00:00Z", modified_at: "2026-08-28T00:00:00Z" }; + const json = (value: unknown) => Response.json(value); + vi.spyOn(globalThis, "fetch").mockImplementation(async (input) => { + const url = String(input); + if (url === "/api/v1/path?path=%2F") return json(root); + if (url === "/api/v1/nodes/1/children?limit=1000&offset=0") return json({ directory: root, items: [file], total: 1, limit: 1000, offset: 0 }); + if (url === "/api/v1/tags?limit=1000&offset=0") return json({ items: [], total: 0, limit: 1000, offset: 0 }); + if (url === "/api/v1/audit/status?node_id=2") return json({ enabled: false, scopes: [] }); + if (url === "/api/v1/nodes/2/tags?limit=1000&offset=0") return json({ items: [], total: 0, limit: 1000, offset: 0 }); + if (url === "/api/v1/processing/profiles") return json([{ name: "private", fingerprint, rendition: true, embedding_bindings: [] }]); + if (url === "/api/v1/processing/plans") return json({ fingerprint, vault_uid: versionID, selector: { node_id: 2, content_version_id: versionID, profile: "private" }, profile_fingerprint: fingerprint, flow: [], disclosed_classes: [], retained_classes: ["sanitized_markdown"], estimate: { source_bytes: 2048, provider_calls: 1, vector_spaces: 0 }, consent_required: false, consent_state: "active", backup_consequence: "retained derivatives enter future backups" }); + if (url.startsWith("/api/v1/coverage?")) return json({ vault_uid: versionID, profile_fingerprint: fingerprint, state: "missing", renditions: { name: "rendition", required: true, state: "missing", complete: 0, unavailable: 0, stale: 0, ineligible: 0, total: 1 }, embeddings: [] }); + throw new Error(`unexpected request: ${url}`); + }); + + render(App); + await fireEvent.click(await screen.findByRole("cell", { name: "report.pdf" })); + await fireEvent.click(screen.getByRole("button", { name: "Process and retrieve" })); + expect(await screen.findByRole("dialog", { name: "Document processing and coverage" })).toBeTruthy(); + expect(screen.getByText(`Exact version ${versionID}`)).toBeTruthy(); + expect(await screen.findByText("Consent active")).toBeTruthy(); +}); diff --git a/frontend/src/App.svelte b/frontend/src/App.svelte index 4930a30a..2e12044b 100644 --- a/frontend/src/App.svelte +++ b/frontend/src/App.svelte @@ -40,7 +40,9 @@ import DownloadButton from "./DownloadButton.svelte"; import JobsDrawer from "./JobsDrawer.svelte"; import ManageTagsModal from "./ManageTagsModal.svelte"; + import ProcessingDrawer from "./ProcessingDrawer.svelte"; import ProvenanceDrawer from "./ProvenanceDrawer.svelte"; + import RenditionDrawer from "./RenditionDrawer.svelte"; import StorageDrawer from "./StorageDrawer.svelte"; import TagCatalogModal, { type TagDefinitionChange, @@ -123,6 +125,8 @@ let historyOpen = $state(false); let versionsOpen = $state(false); let provenanceOpen = $state(false); + let processingTarget = $state(null); + let renditionTarget = $state<{ attachmentID: string; path: string } | null>(null); let jobsOpen = $state(false); let auditEvidenceOpen = $state(false); let storageOpen = $state(false); @@ -186,6 +190,8 @@ historyOpen = false; versionsOpen = false; provenanceOpen = false; + processingTarget = null; + renditionTarget = null; jobsOpen = false; auditEvidenceOpen = false; storageOpen = false; @@ -440,6 +446,8 @@ historyOpen = false; versionsOpen = false; provenanceOpen = false; + processingTarget = null; + renditionTarget = null; } selectedID = nodeID; selectedAudit = null; @@ -771,6 +779,8 @@ historyOpen = false; versionsOpen = false; provenanceOpen = false; + processingTarget = null; + renditionTarget = null; jobsOpen = false; auditEvidenceOpen = false; storageOpen = false; @@ -1276,6 +1286,28 @@