diff --git a/document/mistral/policy_authorization_test.go b/document/mistral/policy_authorization_test.go new file mode 100644 index 00000000..f70f1534 --- /dev/null +++ b/document/mistral/policy_authorization_test.go @@ -0,0 +1,20 @@ +package mistral + +import "testing" + +func TestPolicyDoesNotAuthorizeUnprovedNonPDFFormats(t *testing.T) { + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + + for _, formatID := range []string{"docx", "pptx", "xlsx", "odt", "epub", "txt"} { + t.Run(formatID, func(t *testing.T) { + if expectedUnitBound(formatID) == UnitBoundLocalExact || localUnitCounters[formatID] != nil { + t.Fatalf("unproved format %q is registered for local authority", formatID) + } + _, err := policy.Authorize(manifest, formatID) + if err == nil { + t.Fatalf("Policy.Authorize(%q) succeeded without provider-authentic unit evidence", formatID) + } + }) + } +} diff --git a/document/mistral/rendition.go b/document/mistral/rendition.go index 8d7e9bb2..3aa8ace6 100644 --- a/document/mistral/rendition.go +++ b/document/mistral/rendition.go @@ -156,7 +156,7 @@ func (client *RenditionClient) Render( if err != nil { return document.RenditionResult{}, err } - providerResult, err := client.process(operation, source, metadata, candidate, + providerResult, err := client.process(operation, source, metadata, candidate, localUnits, min(client.policy.values.MaxResponseBytes, int64(authorization.MaxTotalResultBytes))) if err != nil { return document.RenditionResult{}, err @@ -199,8 +199,8 @@ func (client *RenditionClient) Render( return document.RenditionResult{Evidence: evidence, ProviderMarkdown: markdown, Receipt: receipt}, nil } -// verifySource re-detects the exact format and proves the PDF page count -// before any byte leaves the process. +// verifySource re-detects the exact format and proves any registered local +// unit count before any byte leaves the process. func (client *RenditionClient) verifySource( source []byte, metadata document.AuthorizedUploadMetadata, ) (CandidateFormat, int64, error) { @@ -213,24 +213,35 @@ func (client *RenditionClient) verifySource( return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorPolicyRejected, "Mistral input identity does not match authorization", nil) } - if candidate.ID != formatIDPDF { + var localUnits int64 + countFailure := "Mistral OCR local unit count could not be verified" + limitFailure := "Mistral OCR exceeds the complete unit limit" + switch { + case candidate.ID == formatIDPDF: + localUnits, err = formatdetect.CountPDFPages(source) + countFailure = "Mistral PDF page count could not be verified" + limitFailure = "Mistral PDF exceeds the complete unit limit" + case expectedUnitBound(candidate.ID) == UnitBoundLocalExact: + var units int + units, err = countLocalUnits(candidate, bytes.NewReader(source), int64(len(source))) + localUnits = int64(units) + default: return candidate, 0, nil } - localUnits, err := formatdetect.CountPDFPages(source) if err != nil { return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorUnsupportedInput, - "Mistral PDF page count could not be verified", err) + countFailure, err) } if localUnits <= 0 || localUnits > int64(client.policy.values.MaxUnits) { return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorPolicyRejected, - "Mistral PDF exceeds the complete unit limit", nil) + limitFailure, nil) } return candidate, localUnits, nil } func (client *RenditionClient) process( operation *providerutil.Operation, source []byte, metadata document.AuthorizedUploadMetadata, - candidate CandidateFormat, maxResponseBytes int64, + candidate CandidateFormat, localUnits int64, maxResponseBytes int64, ) (Result, error) { formatAuthorization, err := client.policy.Authorize(client.manifest, candidate.ID) if err != nil { @@ -239,7 +250,7 @@ func (client *RenditionClient) process( } snapshot := preparedSnapshot{ size: int64(len(source)), sha256: metadata.SHA256, format: candidate, - mediaType: metadata.MediaType, + mediaType: metadata.MediaType, localUnits: int(localUnits), } snapshotForAttempt := func() (preparedSnapshot, error) { digest := sha256.Sum256(source) diff --git a/document/mistral/rendition_test.go b/document/mistral/rendition_test.go index 50678ee7..327c6a17 100644 --- a/document/mistral/rendition_test.go +++ b/document/mistral/rendition_test.go @@ -114,6 +114,84 @@ func TestRenditionClientMapsExactMistralOCRResponse(t *testing.T) { assert.NotContains(t, fmt.Sprintf("%+v", result.Receipt), "synthetic-key") } +func TestRenditionClientUsesLocalExactUnitsForAuthorizedNonPDF(t *testing.T) { + // Keep this test and its subtests sequential: the synthetic authority below + // replaces package globals until cleanup. It does not enable DOCX in production. + policy := testPolicy(t, 1<<20, 10) + manifest := syntheticManifest(t, policy, true) + + previousMethod, hadMethod := expectedUnitBounds["docx"] + previousCounter, hadCounter := localUnitCounters["docx"] + t.Cleanup(func() { + if hadMethod { + expectedUnitBounds["docx"] = previousMethod + } else { + delete(expectedUnitBounds, "docx") + } + if hadCounter { + localUnitCounters["docx"] = previousCounter + } else { + delete(localUnitCounters, "docx") + } + }) + expectedUnitBounds["docx"] = UnitBoundLocalExact + localUnitCounters["docx"] = func(io.ReaderAt, int64) (int, error) { return 2, nil } + for index := range manifest.Results { + if manifest.Results[index].FormatID == "docx" { + manifest.Results[index].UnitCount = 2 + manifest.Results[index].UnitsProcessed = 2 + manifest.Results[index].LocalUnits = 2 + manifest.Results[index].UnitBoundMethod = UnitBoundLocalExact + } + } + + descriptor := renditionDescriptor(t, policy, manifest, "docx") + source := documentZIP(t, map[string]string{ + ooxmlContentTypesName: docxContentTypes("application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"), + "word/document.xml": "", + }) + response := fmt.Sprintf( + `{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"first"},{"index":1,"markdown":"second"}],"usage_info":{"pages_processed":2,"doc_size_bytes":%d}}`, + len(source), + ) + digest := sha256.Sum256(source) + metadata := document.AuthorizedUploadMetadata{ + Filename: "document.docx", MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]), + CapabilityRecordChecksum: strings.Repeat("2", 64), ProviderMetadataChecksum: strings.Repeat("3", 64), + InputKind: document.RenditionInputOriginalFile, + } + 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: time.Now().UTC().Add(-time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + ExpiresAt: time.Now().UTC().Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"), + } + 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(response)), Request: request, + }, nil + })) + + result, err := client.Render(t.Context(), &renditionUpload{Reader: bytes.NewReader(source), metadata: metadata}, authorization) + require.NoError(t, err) + assert.Equal(t, int64(2), result.Receipt.Usage.Units) + assert.Len(t, result.Evidence.Units, 2) + + response = fmt.Sprintf( + `{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"first"},{"index":1,"markdown":"second"},{"index":2,"markdown":"third"}],"usage_info":{"pages_processed":3,"doc_size_bytes":%d}}`, + len(source), + ) + _, err = client.Render(t.Context(), &renditionUpload{Reader: bytes.NewReader(source), metadata: metadata}, authorization) + assertRenditionCode(t, err, document.RenditionErrorPolicyRejected) + assert.ErrorIs(t, err, ErrCapabilityContract) +} + func TestRenditionClientClassifiesHTTPAndModelFailures(t *testing.T) { for _, testCase := range []struct { name string diff --git a/document/mistral/spool.go b/document/mistral/spool.go index 71e8f02a..c2a7ed09 100644 --- a/document/mistral/spool.go +++ b/document/mistral/spool.go @@ -329,16 +329,6 @@ func wrapSpoolIOError(operation string, err error) error { return fmt.Errorf("%w: %w", ErrSpoolUnavailable, operationError) } -func countLocalUnits(format CandidateFormat, reader io.ReaderAt, size int64) (int, error) { - counter := localUnitCounters[format.ID] - if counter == nil { - return 0, nil - } - return counter(reader, size) -} - -var localUnitCounters = map[string]func(io.ReaderAt, int64) (int, error){} - // ScavengeSpoolDirectory removes stale package-created regular files. It // leaves unrelated regular files and fails closed on unsafe entries. func ScavengeSpoolDirectory(directory string, staleBefore time.Time) (int, error) { diff --git a/document/mistral/units.go b/document/mistral/units.go new file mode 100644 index 00000000..2f6d44bf --- /dev/null +++ b/document/mistral/units.go @@ -0,0 +1,17 @@ +package mistral + +import "io" + +type localUnitCounter func(io.ReaderAt, int64) (int, error) + +// localUnitCounters is the Mistral-owned authority registry. Only formats +// with provider-authentic unit evidence may be added here. +var localUnitCounters = map[string]localUnitCounter{} + +func countLocalUnits(format CandidateFormat, reader io.ReaderAt, size int64) (int, error) { + counter := localUnitCounters[format.ID] + if counter == nil { + return 0, nil + } + return counter(reader, size) +} diff --git a/document/mistral/units_test.go b/document/mistral/units_test.go new file mode 100644 index 00000000..d47894c1 --- /dev/null +++ b/document/mistral/units_test.go @@ -0,0 +1,37 @@ +package mistral + +import ( + "bytes" + "testing" +) + +func TestLocalUnitCounterRegistryIsMistralOwnedAndBounded(t *testing.T) { + // This guards future registrations; no local counters are authorized today. + for formatID, counter := range localUnitCounters { + if counter == nil { + t.Fatalf("localUnitCounters[%q] is nil", formatID) + } + if _, ok := CandidateFormatByID(formatID); !ok { + t.Fatalf("localUnitCounters[%q] has no candidate", formatID) + } + } + for formatID, method := range expectedUnitBounds { + if method == UnitBoundLocalExact && localUnitCounters[formatID] == nil { + t.Fatalf("local exact format %q has no local counter", formatID) + } + } +} + +func TestCountLocalUnitsLeavesUnprovedFormatsUnbounded(t *testing.T) { + docx, ok := CandidateFormatByID("docx") + if !ok { + t.Fatal("docx candidate is missing") + } + units, err := countLocalUnits(docx, bytes.NewReader([]byte("synthetic")), 9) + if err != nil { + t.Fatal(err) + } + if units != 0 { + t.Fatalf("countLocalUnits(docx) = %d, want 0 for unproved format", units) + } +}