From 22dacbed01f5f3c1148b5e20f92a0c06092b5338 Mon Sep 17 00:00:00 2001 From: ConnorYoh <40631091+ConnorYoh@users.noreply.github.com> Date: Thu, 4 Jun 2026 16:07:58 +0100 Subject: [PATCH] PAYG B-2: shadow-mode filter + interceptor (engine activation) (#6519) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What this PR does Wires the **B-1 shadow charging engine** into real HTTP request flow. After this lands, flipping an internal team to ``PAYG_SHADOW`` via SQL begins populating ``payg_shadow_charge`` automatically — with **zero impact** on the legacy credit deduction path. **This is the load-bearing PR for shadow mode.** Without it, B-1's engine sits idle — nothing in the codebase calls ``JobChargeService.openProcess()`` from a real HTTP request. Stacks on top of #6477 (PR B-1). ## Components | Class | Role | |---|---| | ``PaygResponseBodyWrapperFilter`` | Servlet filter, installs tee'ing response wrapper. Defers wrapper close to ``AsyncListener`` for ``DeferredResult`` / ``CompletableFuture`` controllers so the lifetime spans the async window. | | ``PaygResponseBodyWrapper`` | ``HttpServletResponseWrapper`` — in-memory ``ByteArrayOutputStream`` up to 10 MiB; spills to ``TempFile`` above. ``materialisedPath()`` always returns a uniform ``Path`` interface. | | ``PaygChargeInterceptor`` | ``AsyncHandlerInterceptor`` mirroring ``UnifiedCreditInterceptor`` shape. ``preHandle`` gates on ``@AutoJobPostMapping``, materialises multipart inputs, calls ``JobChargeService.openProcess``. ``afterCompletion`` branches on HTTP status. | | ``PaygOutputExtractor`` | Pulls PDFs out of the response body. Direct ``application/pdf`` returns body verbatim; ``application/zip`` iterates entries and keeps each ``.pdf`` entry whose first bytes match the ``%PDF-`` magic. | | ``PaygWebMvcConfig`` | Registers filter at end of Spring filter chain (after security); interceptor after ``UnifiedCreditInterceptor``. | | ``PaygFilterProperties`` | ``payg.filter.enabled`` master switch + in-memory threshold + optional max-bytes ceiling. | ## Status branching in afterCompletion | HTTP status | Action | |---|---| | **2xx** | Append OK step; extract PDFs from response; ``JobService.recordOutput`` per PDF | | **4xx** | Append FAILED step with ``errorCode``. No refund — customer paid for the attempt. No OUTPUT recording. | | **5xx + OPENED** (first-step) | ``JobChargeService.markFirstStepFailed`` → shadow row flipped to ``REFUNDED``, process CLOSED. Refund counter incremented. | | **5xx + JOINED** (mid-chain) | ``JobChargeService.decrementStepCount`` — step slot returned without resetting ``lastStepAt`` (workflow window stays active for retry). | ## New ``JobChargeService`` methods - **``markFirstStepFailed(jobId, reason)``** — flips shadow row to ``REFUNDED`` with ``refundedAt`` + ``refundReason``, closes the process. Idempotent. Mimics the eventual Stripe ``meter_event_adjustment(cancel)`` flow that real-mode will invoke at the same callsite. **Refund implies close** so a same-input retry can't lineage-join into a refunded chain for free work. - **``decrementStepCount(jobId)``** — defensive floor at 1; never drives count negative. ## Schema - Backend: ``V13__payg_shadow_charge_status.sql`` adds ``status`` (``CHARGED`` | ``REFUNDED``) + ``refunded_at`` + ``refund_reason``. ``DEFAULT 'CHARGED'`` so existing B-1 rows stay correct without backfill. - Supabase: matching migration in [Stirling-PDF-SaaS#payg-shadow-charge-status](https://github.com/Stirling-Tools/Stirling-PDF-SaaS/tree/payg-shadow-charge-status) ## Fail-open semantics in shadow Any ``RuntimeException`` in ``preHandle`` / ``afterCompletion`` is logged at WARN, increments ``payg.filter.errors``, and lets the customer's tool call proceed unbilled. This **reverses to fail-closed** when ``wallet_policy.engine = PAYG`` (real charging) — that reversal lives inside ``JobChargeService`` and ships with the cap evaluator PR (PR-C1 in PAYG_DESIGN.md). ## Observability Micrometer metrics: - ``payg.filter.errors`` Counter — internal failures (preHandle + afterCompletion). Alert source. - ``payg.filter.calls`` Counter, tagged ``disposition`` (``OPENED`` | ``JOINED`` | ``SHORT_CIRCUIT``) - ``payg.filter.refunds`` Counter — first-step 5xx refunds - ``payg.filter.duration`` Timer — preHandle + afterCompletion wall-clock per request ## Test coverage (38 tests across 4 classes) - **PaygResponseBodyWrapperTest** (12 tests) — in-memory, spill, threshold crossing mid-chunk, writer vs outputStream exclusivity, ``resetBuffer`` with and without spill, close idempotency, single-byte writes across threshold. - **PaygOutputExtractorTest** (7 tests) — direct PDF, parametrised content type, ZIP with mixed entries + magic-byte gate, corrupt ZIP fail-open, empty ZIP. - **PaygChargeInterceptorTest** (13 tests) — all preHandle short-circuits, OPENED disposition stash, fail-open on chargeService exception, 2xx recordOutputs path, 5xx OPENED → markFirstStepFailed, 5xx JOINED → decrementStepCount, 4xx FAILED step append, max-bytes ceiling skip, PIPELINE header detection. - **JobChargeServiceTest extended** (+6 tests) — markFirstStepFailed happy path, idempotency, missing-shadow-row case, long-reason trim; decrementStepCount happy path, floor-at-1 defence, missing-job no-op. ## What's NOT in this PR (deliberate) - **No SpringBootTest layer.** The saas module doesn't have bootstrap test infrastructure (Supabase JWT config + H2 schema harness). Integration confidence comes from the manual staging deploy + SQL-flip of an internal team. Bootstrap-test infra is a focused follow-up if needed. - **No saas-mode Behave / docker-compose.** Per design §17 — deferred. Existing ``testing/cucumber/`` infrastructure doesn't yet have a saas-profile compose target; that's its own PR when warranted. - **No CreditService wire-in** (per design §13 decision). Per-row comparison data moves to the reconciliation report PR (PR-S2). ``legacy_credits_charged`` + ``diff_pct`` columns stay at 0 in shadow rows. - **No reconciliation report endpoint.** Direct SQL queries against ``payg_shadow_charge`` cover the data-access need until patterns emerge. ## Rollback levers | Symptom | Lever | |---|---| | Some / all tool calls breaking due to filter | ``payg.filter.enabled=false`` + restart (~20s) | | Shadow rows look wrong for a specific team | ``UPDATE wallet_policy SET engine = 'LEGACY' WHERE team_id = ?`` | | Mass shadow weirdness | ``UPDATE wallet_policy SET engine = 'LEGACY'`` | | Memory exhaustion from response tee | Lower ``payg.filter.response.in-memory-threshold-bytes`` | ## Test plan - [ ] CI green (build + tests) - [ ] Aikido / Snyk / SonarCloud clean - [ ] Manual: deploy to staging - [ ] Manual: flip one internal team via ``UPDATE wallet_policy SET engine = 'PAYG_SHADOW' WHERE team_id = ?`` - [ ] Manual: hit ``/api/v1/security/add-password`` with that team's JWT; verify a ``payg_shadow_charge`` row appears with ``status='CHARGED'`` - [ ] Manual: trigger a 503 (e.g. via temporary backend kill mid-request); verify the resulting row is ``status='REFUNDED'`` + the process is ``CLOSED`` - [ ] Manual: hit ``/api/v1/general/split`` with a multi-page PDF; verify one OUTPUT signature per inner PDF appears in ``job_artifact_hash`` - [ ] Manual: chain ``add-password`` → ``compress`` on the output; verify the second call JOINS the first process (no new shadow row) and the inner output OUTPUT signature is what drove the lineage join ## Stacks on / references - Stacks on: #6477 (B-1 — shadow charging engine) - Schema mirror: Stirling-PDF-SaaS#payg-shadow-charge-status branch - Design doc: ``notes/PAYG_FILTER_DESIGN.md`` (all 19 decisions DECIDED) --- .../saas/payg/charge/JobChargeService.java | 97 +++- .../payg/docs/DefaultDocumentClassifier.java | 60 ++- .../saas/payg/docs/DocumentClassifier.java | 28 ++ .../payg/filter/PaygChargeInterceptor.java | 420 ++++++++++++++++++ .../payg/filter/PaygFilterProperties.java | 44 ++ .../saas/payg/filter/PaygOutputExtractor.java | 168 +++++++ .../payg/filter/PaygResponseBodyWrapper.java | 289 ++++++++++++ .../filter/PaygResponseBodyWrapperFilter.java | 138 ++++++ .../saas/payg/filter/PaygWebMvcConfig.java | 60 +++ .../software/saas/payg/job/JobService.java | 41 +- .../lineage/DefaultHashLineageDetector.java | 32 +- .../payg/lineage/HashLineageDetector.java | 22 + .../saas/payg/model/ShadowChargeStatus.java | 18 + .../payg/policy/PricingPolicyService.java | 17 +- .../PaygShadowChargeRepository.java | 10 + .../saas/payg/shadow/PaygShadowCharge.java | 20 + .../saas/V13__payg_shadow_charge_status.sql | 20 + .../payg/charge/JobChargeServiceTest.java | 126 +++++- .../filter/PaygChargeInterceptorTest.java | 420 ++++++++++++++++++ .../payg/filter/PaygOutputExtractorTest.java | 120 +++++ .../filter/PaygResponseBodyWrapperTest.java | 210 +++++++++ .../saas/payg/job/JobServiceTest.java | 48 +- .../payg/policy/PricingPolicyServiceTest.java | 27 ++ 23 files changed, 2384 insertions(+), 51 deletions(-) create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygChargeInterceptor.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygFilterProperties.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygOutputExtractor.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapper.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperFilter.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/filter/PaygWebMvcConfig.java create mode 100644 app/saas/src/main/java/stirling/software/saas/payg/model/ShadowChargeStatus.java create mode 100644 app/saas/src/main/resources/db/migration/saas/V13__payg_shadow_charge_status.sql create mode 100644 app/saas/src/test/java/stirling/software/saas/payg/filter/PaygChargeInterceptorTest.java create mode 100644 app/saas/src/test/java/stirling/software/saas/payg/filter/PaygOutputExtractorTest.java create mode 100644 app/saas/src/test/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperTest.java diff --git a/app/saas/src/main/java/stirling/software/saas/payg/charge/JobChargeService.java b/app/saas/src/main/java/stirling/software/saas/payg/charge/JobChargeService.java index ee342f4e9..e5832bfd6 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/charge/JobChargeService.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/charge/JobChargeService.java @@ -2,8 +2,11 @@ package stirling.software.saas.payg.charge; import java.io.IOException; import java.nio.file.Path; +import java.time.LocalDateTime; import java.util.List; import java.util.Objects; +import java.util.Optional; +import java.util.UUID; import org.springframework.context.annotation.Profile; import org.springframework.stereotype.Service; @@ -17,10 +20,14 @@ import stirling.software.saas.payg.docs.DocumentMetrics; import stirling.software.saas.payg.job.JobContext; import stirling.software.saas.payg.job.JobService; import stirling.software.saas.payg.job.JoinOrOpenResult; +import stirling.software.saas.payg.job.ProcessingJob; import stirling.software.saas.payg.model.JobSource; +import stirling.software.saas.payg.model.JobStatus; +import stirling.software.saas.payg.model.ShadowChargeStatus; import stirling.software.saas.payg.policy.PricingPolicy; import stirling.software.saas.payg.policy.PricingPolicyService; import stirling.software.saas.payg.repository.PaygShadowChargeRepository; +import stirling.software.saas.payg.repository.ProcessingJobRepository; import stirling.software.saas.payg.shadow.PaygShadowCharge; /** @@ -47,16 +54,19 @@ public class JobChargeService { private final PricingPolicyService policyService; private final DocumentClassifier classifier; private final PaygShadowChargeRepository shadowRepository; + private final ProcessingJobRepository jobRepository; public JobChargeService( JobService jobService, PricingPolicyService policyService, DocumentClassifier classifier, - PaygShadowChargeRepository shadowRepository) { + PaygShadowChargeRepository shadowRepository, + ProcessingJobRepository jobRepository) { this.jobService = Objects.requireNonNull(jobService, "jobService"); this.policyService = Objects.requireNonNull(policyService, "policyService"); this.classifier = Objects.requireNonNull(classifier, "classifier"); this.shadowRepository = Objects.requireNonNull(shadowRepository, "shadowRepository"); + this.jobRepository = Objects.requireNonNull(jobRepository, "jobRepository"); } /** @@ -117,10 +127,14 @@ public class JobChargeService { private int computeUnits(List inputs, PricingPolicy policy) { List multiparts = inputs.stream().map(JobInput::multipart).toList(); + // Reuse the temp file the caller already wrote (in PaygChargeInterceptor.preHandle for + // lineage hashing) instead of materialising the same bytes a second time inside the + // classifier. Saves one write + one read per PDF input. + List paths = inputs.stream().map(JobInput::path).toList(); DocumentMetrics metrics = multiparts.size() == 1 - ? classifier.classify(multiparts.get(0), policy) - : classifier.classify(multiparts, policy); + ? classifier.classify(multiparts.get(0), paths.get(0), policy) + : classifier.classify(multiparts, paths, policy); // Apply the policy-level minChargeUnits floor per design § 3.4. The classifier returns // raw docUnits with a "non-empty input → ≥1" floor; the charge formula's // max(min_charge_units, docUnits) layers on top. @@ -138,6 +152,83 @@ public class JobChargeService { // CreditService in the follow-up PR. Until then, diff stays at 0. row.setLegacyCreditsCharged(0); row.setDiffPct(0); + row.setStatus(ShadowChargeStatus.CHARGED); shadowRepository.save(row); } + + /** + * First-step failure on a freshly-opened process: mimic a successful Stripe + * meter_event_adjustment(cancel) by flipping the shadow row to {@link + * ShadowChargeStatus#REFUNDED}, and close the process so a same-input retry can't lineage-join + * into a refunded chain for free work. + * + *

Idempotent: re-invoking on an already-REFUNDED row or already-CLOSED process is a silent + * no-op. + */ + @Transactional + public void markFirstStepFailed(UUID jobId, String refundReason) { + Objects.requireNonNull(jobId, "jobId"); + LocalDateTime now = LocalDateTime.now(); + + Optional rowOpt = shadowRepository.findFirstByJobIdOrderByIdAsc(jobId); + if (rowOpt.isEmpty()) { + log.debug( + "markFirstStepFailed: no shadow row for job {} (PAYG not active for team?)", + jobId); + } else { + PaygShadowCharge row = rowOpt.get(); + if (row.getStatus() != ShadowChargeStatus.REFUNDED) { + row.setStatus(ShadowChargeStatus.REFUNDED); + row.setRefundedAt(now); + row.setRefundReason(trimReason(refundReason)); + shadowRepository.save(row); + } + } + + ProcessingJob job = jobRepository.findById(jobId).orElse(null); + if (job == null) { + log.warn("markFirstStepFailed: no ProcessingJob with id {}", jobId); + return; + } + if (job.getStatus() == JobStatus.OPEN) { + job.setStatus(JobStatus.CLOSED); + job.setClosedAt(now); + jobRepository.save(job); + } + } + + /** + * Mid-chain 5xx on a JOINED step: return the step slot. The {@code lastStepAt} timestamp stays + * advanced (workflow window intentionally remains active for the next retry). No shadow-row + * change — only OPENED-disposition calls wrote a row. + * + *

Defensive lower bound: never drives {@code stepCount} below 1 (which would imply we'd + * decremented an already-decremented slot, or were called on a fresh process). + */ + @Transactional + public void decrementStepCount(UUID jobId) { + Objects.requireNonNull(jobId, "jobId"); + ProcessingJob job = jobRepository.findById(jobId).orElse(null); + if (job == null) { + log.warn("decrementStepCount: no ProcessingJob with id {}", jobId); + return; + } + int current = job.getStepCount() == null ? 0 : job.getStepCount(); + if (current <= 1) { + log.debug( + "decrementStepCount: stepCount already at {} for job {}; no-op", + current, + jobId); + return; + } + job.setStepCount(current - 1); + jobRepository.save(job); + } + + private static String trimReason(String reason) { + if (reason == null) { + return null; + } + return reason.length() > 128 ? reason.substring(0, 128) : reason; + } } diff --git a/app/saas/src/main/java/stirling/software/saas/payg/docs/DefaultDocumentClassifier.java b/app/saas/src/main/java/stirling/software/saas/payg/docs/DefaultDocumentClassifier.java index a46d97587..603f490a4 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/docs/DefaultDocumentClassifier.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/docs/DefaultDocumentClassifier.java @@ -4,6 +4,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.nio.file.Files; +import java.nio.file.Path; import java.util.List; import java.util.Objects; @@ -48,10 +49,16 @@ public class DefaultDocumentClassifier implements DocumentClassifier { @Override public DocumentMetrics classify(MultipartFile file, PricingPolicy policy) { + return classify(file, null, policy); + } + + @Override + public DocumentMetrics classify( + MultipartFile file, Path materialisedPath, PricingPolicy policy) { Objects.requireNonNull(file, "file"); Objects.requireNonNull(policy, "policy"); - FileFacts facts = inspect(file); + FileFacts facts = inspect(file, materialisedPath); long rawUnits = computeRawUnits(facts.pages, facts.bytes, policy); // toIntExact: fail loud on overflow rather than silently wrapping a billing number. int units = @@ -64,19 +71,35 @@ public class DefaultDocumentClassifier implements DocumentClassifier { @Override public DocumentMetrics classify(List files, PricingPolicy policy) { + return classify(files, null, policy); + } + + @Override + public DocumentMetrics classify( + List files, List materialisedPaths, PricingPolicy policy) { Objects.requireNonNull(files, "files"); Objects.requireNonNull(policy, "policy"); if (files.isEmpty()) { throw new IllegalArgumentException("files must not be empty"); } + if (materialisedPaths != null && materialisedPaths.size() != files.size()) { + throw new IllegalArgumentException( + "materialisedPaths size (" + + materialisedPaths.size() + + ") must equal files size (" + + files.size() + + ")"); + } int totalPages = 0; long totalBytes = 0; long rawUnitsSum = 0; String firstContentType = null; - for (MultipartFile file : files) { - FileFacts facts = inspect(file); + for (int i = 0; i < files.size(); i++) { + MultipartFile file = files.get(i); + Path path = materialisedPaths == null ? null : materialisedPaths.get(i); + FileFacts facts = inspect(file, path); // Sum the *raw* (unclamped) per-file units so the group cap below can actually bind. // Per-file clamping in this loop would make the group cap a no-op. rawUnitsSum = @@ -103,11 +126,17 @@ public class DefaultDocumentClassifier implements DocumentClassifier { totalUnits); } - private FileFacts inspect(MultipartFile file) { + private FileFacts inspect(MultipartFile file, Path materialisedPath) { long bytes = file.getSize(); String contentType = file.getContentType() != null ? file.getContentType() : DEFAULT_CONTENT_TYPE; - int pages = isPdf(contentType, file.getOriginalFilename()) ? readPageCount(file) : 0; + int pages = 0; + if (isPdf(contentType, file.getOriginalFilename())) { + pages = + materialisedPath != null + ? readPageCountFromPath(materialisedPath, file.getOriginalFilename()) + : readPageCount(file); + } return new FileFacts(pages, bytes, contentType); } @@ -141,9 +170,7 @@ public class DefaultDocumentClassifier implements DocumentClassifier { OutputStream out = Files.newOutputStream(temp.getPath())) { in.transferTo(out); } - try (PdfDocument doc = PdfDocument.open(temp.getPath())) { - return doc.pageCount(); - } + return readPageCountFromPath(temp.getPath(), file.getOriginalFilename()); } catch (IOException | RuntimeException e) { log.debug( "Could not read PDF page count for {} ({}); falling back to bytes-only units", @@ -153,6 +180,23 @@ public class DefaultDocumentClassifier implements DocumentClassifier { } } + /** + * Page-count read against an already-materialised file. Used by callers that already wrote the + * bytes to disk (the PAYG interceptor materialises every input for the lineage hash) so we + * avoid a second copy. + */ + private int readPageCountFromPath(Path path, String displayName) { + try (PdfDocument doc = PdfDocument.open(path)) { + return doc.pageCount(); + } catch (RuntimeException e) { + log.debug( + "Could not read PDF page count for {} ({}); falling back to bytes-only units", + displayName, + e.getClass().getSimpleName()); + return 0; + } + } + private static int saturatedAdd(int a, int b) { long sum = (long) a + b; if (sum > Integer.MAX_VALUE) { diff --git a/app/saas/src/main/java/stirling/software/saas/payg/docs/DocumentClassifier.java b/app/saas/src/main/java/stirling/software/saas/payg/docs/DocumentClassifier.java index c7af9deaf..9c4966fde 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/docs/DocumentClassifier.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/docs/DocumentClassifier.java @@ -1,5 +1,6 @@ package stirling.software.saas.payg.docs; +import java.nio.file.Path; import java.util.List; import org.springframework.web.multipart.MultipartFile; @@ -11,6 +12,17 @@ import stirling.software.saas.payg.policy.PricingPolicy; * *

Returns {@code docUnits} with an absolute floor of 1 for non-empty input. {@code * policy.minChargeUnits} is applied at charge time, not here. + * + *

Each overload comes in two flavours: + * + *

    + *
  • {@code classify(MultipartFile, ...)} — classifier reads bytes via {@code getInputStream()} + * and writes its own temp file to feed jpdfium (for PDFs). Use when no on-disk copy exists. + *
  • {@code classify(MultipartFile, Path, ...)} — caller has already materialised the bytes to + * {@code Path}; classifier reads page count directly from there without re-writing. Hot-path + * callers (the PAYG interceptor) should use this form since they materialise inputs anyway + * for the lineage hash. + *
*/ public interface DocumentClassifier { @@ -22,4 +34,20 @@ public interface DocumentClassifier { * units, capped at {@code fileUnitCap × files.size()} and floored at 1. */ DocumentMetrics classify(List files, PricingPolicy policy); + + /** + * Same as {@link #classify(MultipartFile, PricingPolicy)} but uses {@code materialisedPath} for + * PDF page-count extraction, avoiding a second copy of the upload bytes to disk. Callers that + * hold an on-disk copy (e.g. the PAYG interceptor materialises every input for the lineage + * hash) should prefer this form. + */ + DocumentMetrics classify(MultipartFile file, Path materialisedPath, PricingPolicy policy); + + /** + * Multi-file variant of {@link #classify(MultipartFile, Path, PricingPolicy)}. {@code + * materialisedPaths} must align positionally with {@code files} — entry {@code i} is the + * on-disk copy of {@code files.get(i)}. + */ + DocumentMetrics classify( + List files, List materialisedPaths, PricingPolicy policy); } diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygChargeInterceptor.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygChargeInterceptor.java new file mode 100644 index 000000000..febcbeb78 --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygChargeInterceptor.java @@ -0,0 +1,420 @@ +package stirling.software.saas.payg.filter; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.UUID; + +import org.springframework.context.annotation.Profile; +import org.springframework.security.core.Authentication; +import org.springframework.security.core.context.SecurityContextHolder; +import org.springframework.stereotype.Component; +import org.springframework.util.MultiValueMap; +import org.springframework.web.method.HandlerMethod; +import org.springframework.web.multipart.MultipartFile; +import org.springframework.web.multipart.MultipartHttpServletRequest; +import org.springframework.web.servlet.AsyncHandlerInterceptor; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.Timer; + +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.annotations.AutoJobPostMapping; +import stirling.software.common.util.TempFile; +import stirling.software.common.util.TempFileManager; +import stirling.software.proprietary.security.database.repository.UserRepository; +import stirling.software.proprietary.security.model.ApiKeyAuthenticationToken; +import stirling.software.proprietary.security.model.User; +import stirling.software.saas.payg.charge.ChargeContext; +import stirling.software.saas.payg.charge.ChargeOutcome; +import stirling.software.saas.payg.charge.JobChargeService; +import stirling.software.saas.payg.charge.JobInput; +import stirling.software.saas.payg.job.JobService; +import stirling.software.saas.payg.model.JobSource; +import stirling.software.saas.payg.model.JobStepStatus; +import stirling.software.saas.payg.model.ProcessType; +import stirling.software.saas.util.AuthenticationUtils; + +/** + * The hot-path PAYG interceptor. Mirrors the {@code UnifiedCreditInterceptor} shape: registered + * after it in {@code PaygWebMvcConfig} so legacy credit-rejection short-circuits before we waste + * work hashing inputs. + * + *

{@code preHandle}: gates on {@code @AutoJobPostMapping}, reads the parsed multipart parts, + * materialises each input to a {@code TempFile}, and asks {@link JobChargeService#openProcess} to + * open (or join) a process. The resulting {@link ChargeOutcome} plus input temp-files are stashed + * as request attributes for {@code afterCompletion}. + * + *

{@code afterCompletion}: branches on HTTP status — 2xx hashes the response body for OUTPUT + * lineage; 4xx records a step append for audit; 5xx triggers refund-and-close (OPENED) or + * step-quota return (JOINED). Closes all input temp files and the response wrapper at the end. + * + *

Fail-open everywhere: any unexpected {@link RuntimeException} is swallowed, logged at WARN, + * and counted on {@code payg.filter.errors}. The customer's tool call always proceeds. + * + *

Async controllers ({@code DeferredResult}, {@code CompletableFuture}) are handled + * transparently — {@link AsyncHandlerInterceptor#afterConcurrentHandlingStarted} is a no-op; the + * normal {@code afterCompletion} fires when the async work resolves. + */ +@Slf4j +@Component +@Profile("saas") +public class PaygChargeInterceptor implements AsyncHandlerInterceptor { + + static final String ATTR_JOB_ID = PaygChargeInterceptor.class.getName() + ".JOB_ID"; + static final String ATTR_DISPOSITION = PaygChargeInterceptor.class.getName() + ".DISPOSITION"; + static final String ATTR_INPUT_TEMP_FILES = + PaygChargeInterceptor.class.getName() + ".INPUT_TEMP_FILES"; + static final String ATTR_INPUT_BYTES = PaygChargeInterceptor.class.getName() + ".INPUT_BYTES"; + static final String ATTR_FAILED = PaygChargeInterceptor.class.getName() + ".FAILED"; + static final String ATTR_TOOL_ID = PaygChargeInterceptor.class.getName() + ".TOOL_ID"; + + private static final String AUTOMATION_HEADER = "X-Stirling-Automation"; + + private final JobChargeService chargeService; + private final JobService jobService; + private final UserRepository userRepository; + private final TempFileManager tempFileManager; + private final PaygOutputExtractor outputExtractor; + private final PaygFilterProperties properties; + + private final Counter errorsCounter; + private final Counter callsOpened; + private final Counter callsJoined; + private final Counter callsShortCircuit; + private final Counter refundsCounter; + private final Timer durationTimer; + + public PaygChargeInterceptor( + JobChargeService chargeService, + JobService jobService, + UserRepository userRepository, + TempFileManager tempFileManager, + PaygOutputExtractor outputExtractor, + PaygFilterProperties properties, + MeterRegistry meterRegistry) { + this.chargeService = chargeService; + this.jobService = jobService; + this.userRepository = userRepository; + this.tempFileManager = tempFileManager; + this.outputExtractor = outputExtractor; + this.properties = properties; + + this.errorsCounter = + Counter.builder("payg.filter.errors") + .description("PAYG interceptor / filter internal failures") + .register(meterRegistry); + this.callsOpened = + Counter.builder("payg.filter.calls") + .tag("disposition", "OPENED") + .register(meterRegistry); + this.callsJoined = + Counter.builder("payg.filter.calls") + .tag("disposition", "JOINED") + .register(meterRegistry); + this.callsShortCircuit = + Counter.builder("payg.filter.calls") + .tag("disposition", "SHORT_CIRCUIT") + .register(meterRegistry); + this.refundsCounter = + Counter.builder("payg.filter.refunds") + .description("First-step 5xx refunds applied to shadow rows") + .register(meterRegistry); + this.durationTimer = + Timer.builder("payg.filter.duration") + .description("preHandle + afterCompletion wall-clock per request") + .register(meterRegistry); + } + + @Override + public boolean preHandle( + HttpServletRequest request, HttpServletResponse response, Object handler) { + Timer.Sample sample = Timer.start(); + try { + if (!properties.isEnabled()) { + return true; + } + if (!(handler instanceof HandlerMethod hm) + || hm.getMethodAnnotation(AutoJobPostMapping.class) == null) { + callsShortCircuit.increment(); + return true; + } + try { + doPreHandle(request); + } catch (RuntimeException e) { + log.warn("PAYG preHandle failed; passing through unbilled", e); + errorsCounter.increment(); + request.setAttribute(ATTR_FAILED, Boolean.TRUE); + cleanupInputs(request); + } + return true; + } finally { + sample.stop(durationTimer); + } + } + + private void doPreHandle(HttpServletRequest request) { + Authentication auth = SecurityContextHolder.getContext().getAuthentication(); + User currentUser = resolveUser(auth); + if (currentUser == null) { + callsShortCircuit.increment(); + return; + } + + if (!(request instanceof MultipartHttpServletRequest mreq)) { + callsShortCircuit.increment(); + return; + } + + MultiValueMap map = mreq.getMultiFileMap(); + List nonEmpty = new ArrayList<>(); + for (List bucket : map.values()) { + for (MultipartFile mp : bucket) { + if (mp.getSize() > 0) { + nonEmpty.add(mp); + } + } + } + if (nonEmpty.isEmpty()) { + callsShortCircuit.increment(); + return; + } + + List tempFiles = new ArrayList<>(nonEmpty.size()); + List inputs = new ArrayList<>(nonEmpty.size()); + long totalInputBytes = 0L; + try { + for (MultipartFile mp : nonEmpty) { + TempFile tf = tempFileManager.createManagedTempFile(".upload"); + tempFiles.add(tf); + try (InputStream in = mp.getInputStream(); + OutputStream out = Files.newOutputStream(tf.getPath())) { + in.transferTo(out); + } + inputs.add(new JobInput(mp, tf.getPath())); + totalInputBytes += mp.getSize(); + } + } catch (IOException e) { + for (TempFile tf : tempFiles) { + tf.close(); + } + throw new RuntimeException("Failed to materialise multipart input", e); + } + // Publish as an immutable list. preHandle fully populates `tempFiles` before this + // setAttribute, and cleanupInputs / afterCompletion are read-only consumers — exposing + // an unmodifiable view (a) makes the safe-publication guarantee from the container's + // attribute-map synchronization unambiguous to readers on different threads (sync vs + // async-dispatch), and (b) hard-fails any future caller that tries to mutate it. + request.setAttribute(ATTR_INPUT_TEMP_FILES, Collections.unmodifiableList(tempFiles)); + request.setAttribute(ATTR_INPUT_BYTES, totalInputBytes); + request.setAttribute(ATTR_TOOL_ID, request.getRequestURI()); + + ChargeContext ctx = + new ChargeContext( + currentUser.getId(), + currentUser.getTeam() == null ? null : currentUser.getTeam().getId(), + determineSource(request, auth), + ProcessType.SINGLE_TOOL); + + ChargeOutcome outcome; + try { + outcome = chargeService.openProcess(ctx, inputs); + } catch (IOException e) { + throw new RuntimeException("openProcess IO failure", e); + } + request.setAttribute(ATTR_JOB_ID, outcome.processId()); + request.setAttribute(ATTR_DISPOSITION, outcome.disposition()); + + if (outcome.disposition() == ChargeOutcome.Disposition.OPENED) { + callsOpened.increment(); + } else { + callsJoined.increment(); + } + } + + @Override + public void afterCompletion( + HttpServletRequest request, + HttpServletResponse response, + Object handler, + Exception ex) { + Timer.Sample sample = Timer.start(); + try { + if (Boolean.TRUE.equals(request.getAttribute(ATTR_FAILED))) { + cleanupInputs(request); + closeWrapper(request); + return; + } + UUID jobId = (UUID) request.getAttribute(ATTR_JOB_ID); + if (jobId == null) { + closeWrapper(request); + return; + } + try { + doAfterCompletion(request, response, jobId); + } catch (RuntimeException e) { + log.warn( + "PAYG afterCompletion failed for job {}; lineage may be incomplete", + jobId, + e); + errorsCounter.increment(); + } finally { + cleanupInputs(request); + closeWrapper(request); + } + } finally { + sample.stop(durationTimer); + } + } + + private void doAfterCompletion( + HttpServletRequest request, HttpServletResponse response, UUID jobId) { + int status = response.getStatus(); + ChargeOutcome.Disposition disposition = + (ChargeOutcome.Disposition) request.getAttribute(ATTR_DISPOSITION); + String toolId = + Optional.ofNullable((String) request.getAttribute(ATTR_TOOL_ID)).orElse("unknown"); + Long inputBytes = (Long) request.getAttribute(ATTR_INPUT_BYTES); + + // Step audit row — appended for every disposition + every outcome class. Done first so a + // refund-and-close still has the failure recorded against the now-CLOSED process. + JobStepStatus stepStatus = status < 400 ? JobStepStatus.OK : JobStepStatus.FAILED; + String errorCode = status >= 400 ? String.valueOf(status) : null; + try { + jobService.appendStep(jobId, toolId, stepStatus, null, inputBytes, errorCode); + } catch (RuntimeException e) { + log.debug("appendStep failed for job {}: {}", jobId, e.getMessage()); + } + + if (status >= 500) { + if (disposition == ChargeOutcome.Disposition.OPENED) { + chargeService.markFirstStepFailed(jobId, "first-step-5xx:" + status); + refundsCounter.increment(); + } else { + chargeService.decrementStepCount(jobId); + } + return; + } + if (status >= 400) { + // 4xx: customer paid for the attempt. No OUTPUT recording, no refund. + return; + } + + recordOutputs(request, response, jobId); + } + + private void recordOutputs( + HttpServletRequest request, HttpServletResponse response, UUID jobId) { + PaygResponseBodyWrapper wrapper = + (PaygResponseBodyWrapper) + request.getAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE); + if (wrapper == null) { + return; + } + Long maxBytes = properties.getResponse().getMaxBytes(); + if (maxBytes != null && wrapper.bytesWritten() > maxBytes) { + log.debug( + "Response size {} exceeds payg.filter.response.max-bytes={}; skipping OUTPUT recording", + wrapper.bytesWritten(), + maxBytes); + return; + } + Path bodyPath; + try { + bodyPath = wrapper.materialisedPath(); + } catch (IOException e) { + log.debug("materialisedPath failed for job {}: {}", jobId, e.getMessage()); + return; + } + if (bodyPath == null) { + return; + } + List pdfs = + outputExtractor.extract(response.getContentType(), bodyPath); + try { + for (PaygOutputExtractor.ExtractedPdf pdf : pdfs) { + try { + jobService.recordOutput(jobId, pdf.path()); + } catch (IOException e) { + log.debug( + "recordOutput failed for job {} path {}: {}", + jobId, + pdf.path(), + e.getMessage()); + } + } + } finally { + for (PaygOutputExtractor.ExtractedPdf pdf : pdfs) { + pdf.close(); + } + } + } + + @Override + public void afterConcurrentHandlingStarted( + HttpServletRequest request, HttpServletResponse response, Object handler) { + // Async handoff: don't touch state. afterCompletion will fire when the async work resolves. + } + + @SuppressWarnings("unchecked") + private void cleanupInputs(HttpServletRequest request) { + Object raw = request.getAttribute(ATTR_INPUT_TEMP_FILES); + if (raw instanceof List) { + for (Object o : (List) raw) { + if (o instanceof TempFile tf) { + tf.close(); + } + } + } + request.removeAttribute(ATTR_INPUT_TEMP_FILES); + } + + private void closeWrapper(HttpServletRequest request) { + Object wrapper = request.getAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE); + if (wrapper instanceof PaygResponseBodyWrapper w) { + w.close(); + } + } + + private User resolveUser(Authentication auth) { + if (auth == null || !auth.isAuthenticated()) { + return null; + } + if (auth instanceof ApiKeyAuthenticationToken && auth.getPrincipal() instanceof User u) { + return u; + } + try { + String supabaseId = AuthenticationUtils.extractSupabaseId(auth); + if (supabaseId == null) { + return null; + } + UUID supabaseUuid = UUID.fromString(supabaseId); + return userRepository.findBySupabaseId(supabaseUuid).orElse(null); + } catch (RuntimeException e) { + log.debug("PAYG resolveUser failed: {}", e.getMessage()); + return null; + } + } + + private static JobSource determineSource(HttpServletRequest request, Authentication auth) { + String automationHeader = request.getHeader(AUTOMATION_HEADER); + if (automationHeader != null && "true".equalsIgnoreCase(automationHeader.trim())) { + return JobSource.PIPELINE; + } + if (auth instanceof ApiKeyAuthenticationToken) { + return JobSource.API; + } + return JobSource.WEB; + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygFilterProperties.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygFilterProperties.java new file mode 100644 index 000000000..0df2bf01b --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygFilterProperties.java @@ -0,0 +1,44 @@ +package stirling.software.saas.payg.filter; + +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.context.annotation.Profile; +import org.springframework.stereotype.Component; + +import lombok.Getter; +import lombok.Setter; + +/** + * Configuration knobs for the PAYG filter + interceptor stack. See {@code PAYG_FILTER_DESIGN.md} + * §16 for the rationale on each default. + * + *
    + *
  • {@code payg.filter.enabled} — master kill switch. Restart-required; no + * {@code @RefreshScope}. When {@code false}, the wrapper filter passes through and the + * interceptor short-circuits in {@code preHandle}. + *
  • {@code payg.filter.response.in-memory-threshold-bytes} — the wrapper buffers below this in + * a {@link java.io.ByteArrayOutputStream}; above it spills to a {@code TempFile}. + *
  • {@code payg.filter.response.max-bytes} — optional ceiling. Responses exceeding this skip + * OUTPUT recording in {@code afterCompletion}. {@code null} = unbounded. + *
+ */ +@Component +@Profile("saas") +@ConfigurationProperties(prefix = "payg.filter") +@Getter +@Setter +public class PaygFilterProperties { + + private boolean enabled = true; + + private final Response response = new Response(); + + @Getter + @Setter + public static class Response { + /** 10 MiB. Tiny responses stay in RAM; large responses spill. */ + private long inMemoryThresholdBytes = 10L * 1024L * 1024L; + + /** Optional ceiling: when set, OUTPUT recording is skipped past this size. */ + private Long maxBytes; + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygOutputExtractor.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygOutputExtractor.java new file mode 100644 index 000000000..a19a78e0a --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygOutputExtractor.java @@ -0,0 +1,168 @@ +package stirling.software.saas.payg.filter; + +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; + +import org.springframework.context.annotation.Profile; +import org.springframework.stereotype.Component; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.util.TempFile; +import stirling.software.common.util.TempFileManager; + +/** + * Extracts the PDF artefacts from a controller's response body for lineage OUTPUT recording. Two + * content-type paths: + * + *
    + *
  • {@code application/pdf} — return the response body path verbatim (the wrapper already + * materialised it to disk). + *
  • {@code application/zip} — iterate entries; each that has a {@code .pdf} extension AND + * starts with the {@code %PDF-} magic-byte sequence becomes one extracted {@link Path}. + *
+ * + *

Anything else yields an empty list — non-PDF responses don't drive lineage matching. Encrypted + * ZIPs, malformed entries, and IO failures log at debug and yield an empty list (per design §15 + * fail-open). Nested ZIP-of-ZIPs is intentionally out of scope: only outer-level PDFs are + * extracted. + * + *

Each extracted ZIP entry is written to its own {@link TempFile} returned in the result list. + * Callers are responsible for closing those temp files — typically the interceptor closes them in + * the same {@code afterCompletion} that calls this extractor. The whole-response body path is owned + * by the wrapper and is NOT in the returned list when extraction takes the direct-PDF path. + */ +@Slf4j +@Component +@Profile("saas") +public class PaygOutputExtractor { + + /** {@code %PDF-} in ASCII. */ + private static final byte[] PDF_MAGIC = {0x25, 0x50, 0x44, 0x46, 0x2D}; + + private static final String ZIP_CONTENT_TYPE = "application/zip"; + private static final String PDF_CONTENT_TYPE = "application/pdf"; + + private final TempFileManager tempFileManager; + + public PaygOutputExtractor(TempFileManager tempFileManager) { + this.tempFileManager = Objects.requireNonNull(tempFileManager, "tempFileManager"); + } + + /** + * Extract PDF paths from {@code bodyPath} according to {@code contentType}. The returned list + * carries {@link ExtractedPdf} records — each wraps a {@link Path} and an indicator of whether + * the caller owns the temp file lifecycle. + * + * @param contentType the response Content-Type header (may be null / parametrised — only the + * base media type is inspected) + * @param bodyPath the on-disk full response body (from the wrapper's {@code materialisedPath}) + * @return zero-or-more PDFs extracted from the body. Empty when content type is not PDF/ZIP, + * when extraction failed, when no entries matched, or when {@code bodyPath} is null. + */ + public List extract(String contentType, Path bodyPath) { + if (bodyPath == null) { + return List.of(); + } + String mediaType = stripParameters(contentType); + if (PDF_CONTENT_TYPE.equalsIgnoreCase(mediaType)) { + // Wrapper-owned path. Don't claim ownership. + return List.of(new ExtractedPdf(bodyPath, null)); + } + if (ZIP_CONTENT_TYPE.equalsIgnoreCase(mediaType)) { + return extractZip(bodyPath); + } + return List.of(); + } + + private List extractZip(Path bodyPath) { + List results = new ArrayList<>(); + try (InputStream rawIn = Files.newInputStream(bodyPath); + ZipInputStream zin = new ZipInputStream(rawIn)) { + ZipEntry entry; + while ((entry = zin.getNextEntry()) != null) { + try { + if (entry.isDirectory()) { + continue; + } + String name = entry.getName(); + if (name == null || !name.toLowerCase(Locale.ROOT).endsWith(".pdf")) { + continue; + } + TempFile temp = tempFileManager.createManagedTempFile(".pdf"); + Files.copy(zin, temp.getPath(), StandardCopyOption.REPLACE_EXISTING); + if (isPdfMagic(temp.getPath())) { + results.add(new ExtractedPdf(temp.getPath(), temp)); + } else { + temp.close(); + } + } finally { + zin.closeEntry(); + } + } + } catch (IOException | IllegalArgumentException e) { + // ZipException extends IOException; IllegalArgumentException covers ZipEntry + // "MALFORMED" surfaces from zlib. Fail-open: caller still serves the response. + log.debug( + "ZIP unpack failed for response body {} ({}); skipping per-PDF OUTPUT recording", + bodyPath, + e.getClass().getSimpleName()); + // Close anything we already opened. + for (ExtractedPdf p : results) { + p.close(); + } + return List.of(); + } + return results; + } + + private boolean isPdfMagic(Path path) { + byte[] head = new byte[PDF_MAGIC.length]; + try (InputStream in = Files.newInputStream(path)) { + int read = in.read(head); + if (read != PDF_MAGIC.length) { + return false; + } + for (int i = 0; i < PDF_MAGIC.length; i++) { + if (head[i] != PDF_MAGIC[i]) { + return false; + } + } + return true; + } catch (IOException e) { + log.debug("PDF magic-byte check failed for {}", path, e); + return false; + } + } + + private static String stripParameters(String contentType) { + if (contentType == null) { + return null; + } + int semi = contentType.indexOf(';'); + return (semi < 0 ? contentType : contentType.substring(0, semi)).trim(); + } + + /** + * One PDF extracted from the response body. {@link #ownedTempFile} is non-null only when this + * extractor created the temp file (ZIP entries); the {@link #path} for direct-PDF responses + * remains owned by the wrapper. + */ + public record ExtractedPdf(Path path, TempFile ownedTempFile) implements AutoCloseable { + @Override + public void close() { + if (ownedTempFile != null) { + ownedTempFile.close(); + } + } + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapper.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapper.java new file mode 100644 index 000000000..2a3c38336 --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapper.java @@ -0,0 +1,289 @@ +package stirling.software.saas.payg.filter; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.io.OutputStreamWriter; +import java.io.PrintWriter; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Objects; + +import jakarta.servlet.ServletOutputStream; +import jakarta.servlet.WriteListener; +import jakarta.servlet.http.HttpServletResponse; +import jakarta.servlet.http.HttpServletResponseWrapper; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.util.TempFile; +import stirling.software.common.util.TempFileManager; + +/** + * Tees the controller's response body so the PAYG interceptor can hash it for OUTPUT lineage + * recording. Writes still flow through to the real client output stream unmodified — the wrapper + * just keeps a parallel copy. + * + *

Memory model: bytes accumulate in an in-memory {@link ByteArrayOutputStream} until the + * configured {@code in-memory-threshold-bytes} is crossed; after that the wrapper spills the + * existing buffer plus all subsequent writes to a {@link TempFile} owned by {@link + * TempFileManager}. Tiny responses (error JSONs, small PDFs) stay entirely in RAM with zero disk + * IO; large responses (split-to-ZIP, big compressed PDFs) spill cleanly. See design doc §8. + * + *

The interceptor calls {@link #materialisedPath()} from {@code afterCompletion} to get a {@code + * Path} suitable for the lineage detector + ZIP unpack. The wrapper guarantees a Path even if the + * response stayed in memory — it materialises the buffer to a {@link TempFile} on demand so callers + * have a uniform file-based interface. + * + *

{@link #close()} closes any {@link TempFile} the wrapper created. Callers MUST invoke close in + * a finally — typically the interceptor's {@code afterCompletion} after it's done hashing. + * + *

Thread safety: all mutating methods (record paths, materialisedPath, close, + * resetBuffer) synchronize on the wrapper instance. The Servlet spec serialises controller writes + * onto a single dispatch thread, but the {@link jakarta.servlet.AsyncListener#onComplete} callback + * that closes the wrapper for async controllers runs on a container thread distinct from the + * dispatch thread that produced the body. The synchronization makes that handoff safe and also + * guards against future callers (e.g. tests) that might invoke {@link #materialisedPath} or {@link + * #close} from a non-dispatch thread. + */ +@Slf4j +public class PaygResponseBodyWrapper extends HttpServletResponseWrapper implements AutoCloseable { + + private final TempFileManager tempFileManager; + private final long inMemoryThresholdBytes; + + /** Lazily-created on the first getOutputStream() / getWriter() call. */ + private TeeingServletOutputStream teeOut; + + private PrintWriter writer; + + /** + * In-memory accumulator until the threshold is crossed. Becomes null after spill (helps GC of + * potentially large buffers). + */ + private ByteArrayOutputStream memoryBuffer = new ByteArrayOutputStream(); + + /** Non-null once we've spilled. Owns the {@link TempFile} below. */ + private OutputStream spillStream; + + /** Non-null once we've spilled, OR once {@link #materialisedPath()} forced materialisation. */ + private TempFile spillFile; + + private long bytesWritten; + private boolean spilled; + + public PaygResponseBodyWrapper( + HttpServletResponse response, + TempFileManager tempFileManager, + long inMemoryThresholdBytes) { + super(response); + this.tempFileManager = Objects.requireNonNull(tempFileManager, "tempFileManager"); + if (inMemoryThresholdBytes < 0) { + throw new IllegalArgumentException( + "inMemoryThresholdBytes must be >= 0, got " + inMemoryThresholdBytes); + } + this.inMemoryThresholdBytes = inMemoryThresholdBytes; + } + + @Override + public ServletOutputStream getOutputStream() throws IOException { + if (writer != null) { + // Servlet spec: getOutputStream() and getWriter() are mutually exclusive per request. + throw new IllegalStateException( + "getWriter() was already called on this response; cannot switch to getOutputStream()"); + } + if (teeOut == null) { + teeOut = new TeeingServletOutputStream(super.getOutputStream()); + } + return teeOut; + } + + @Override + public PrintWriter getWriter() throws IOException { + if (teeOut != null) { + throw new IllegalStateException( + "getOutputStream() was already called on this response; cannot switch to getWriter()"); + } + if (writer == null) { + String encoding = getCharacterEncoding() != null ? getCharacterEncoding() : "UTF-8"; + // Wrap super.getOutputStream() directly with our TeeingServletOutputStream so writes + // through the Writer path also get tee'd. The Writer just adds character→byte encoding. + teeOut = new TeeingServletOutputStream(super.getOutputStream()); + writer = new PrintWriter(new OutputStreamWriter(teeOut, encoding)); + } + return writer; + } + + @Override + public synchronized void resetBuffer() { + super.resetBuffer(); + if (memoryBuffer != null) { + memoryBuffer.reset(); + } + // If we'd already spilled, the only safe move is to abandon the spill file: the client + // hasn't seen the body yet (otherwise resetBuffer would be illegal), but our tee captured + // bytes we now want to forget. + if (spilled) { + closeSpillQuietly(); + spillFile = null; + spillStream = null; + spilled = false; + memoryBuffer = new ByteArrayOutputStream(); + } + bytesWritten = 0; + } + + /** + * Returns a {@link Path} containing the full response body, or {@code null} if no bytes were + * written. The returned path is owned by this wrapper — do NOT delete or modify it. Use {@link + * #close()} to release. + * + *

If the response stayed under the threshold, this materialises the in-memory buffer to a + * {@link TempFile} on demand so the caller always gets a file-based handle (uniform with the + * spilled path). + */ + public synchronized Path materialisedPath() throws IOException { + if (bytesWritten == 0) { + return null; + } + // Flush the writer so any character data lands in the underlying byte stream first. + if (writer != null) { + writer.flush(); + } + if (spilled) { + spillStream.flush(); + return spillFile.getPath(); + } + // Stayed in memory — materialise on demand for a uniform Path-based interface. + if (spillFile == null) { + spillFile = tempFileManager.createManagedTempFile(".body"); + Files.write(spillFile.getPath(), memoryBuffer.toByteArray()); + } + return spillFile.getPath(); + } + + public synchronized long bytesWritten() { + return bytesWritten; + } + + @Override + public synchronized void close() { + closeSpillQuietly(); + } + + private void closeSpillQuietly() { + if (spillStream != null) { + try { + spillStream.close(); + } catch (IOException e) { + log.debug("Ignoring close error on spill stream: {}", e.getMessage()); + } + spillStream = null; + } + if (spillFile != null) { + spillFile.close(); // TempFile.close() deletes the file + spillFile = null; + } + } + + /** + * Routes bytes both to the real client output stream AND to our buffer (in-memory or spilled). + * Writes are not re-batched — each call to the delegate corresponds exactly to one call here. + */ + private final class TeeingServletOutputStream extends ServletOutputStream { + + private final ServletOutputStream delegate; + + TeeingServletOutputStream(ServletOutputStream delegate) { + this.delegate = delegate; + } + + @Override + public void write(int b) throws IOException { + delegate.write(b); + recordSingleByte((byte) b); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + delegate.write(b, off, len); + recordRange(b, off, len); + } + + @Override + public void flush() throws IOException { + delegate.flush(); + // Read spilled / spillStream under the outer monitor so we observe the publication + // written by spillToDisk() on another thread. The writers (recordSingleByte, + // recordRange, spillToDisk) already mutate these fields under + // synchronized(PaygResponseBodyWrapper.this); without matching locking here the + // JMM permits stale reads (false `spilled`, null `spillStream`) and Aikido AI + // flagged that gap. + synchronized (PaygResponseBodyWrapper.this) { + if (spilled && spillStream != null) { + spillStream.flush(); + } + } + } + + @Override + public void close() throws IOException { + delegate.close(); + synchronized (PaygResponseBodyWrapper.this) { + if (spilled && spillStream != null) { + spillStream.flush(); + } + } + } + + @Override + public boolean isReady() { + return delegate.isReady(); + } + + @Override + public void setWriteListener(WriteListener writeListener) { + delegate.setWriteListener(writeListener); + } + } + + private synchronized void recordSingleByte(byte b) throws IOException { + if (spilled) { + spillStream.write(b & 0xFF); + } else if (bytesWritten + 1 > inMemoryThresholdBytes) { + spillToDisk(); + spillStream.write(b & 0xFF); + } else { + memoryBuffer.write(b & 0xFF); + } + bytesWritten++; + } + + private synchronized void recordRange(byte[] b, int off, int len) throws IOException { + if (spilled) { + spillStream.write(b, off, len); + } else if (bytesWritten + len > inMemoryThresholdBytes) { + // This write crosses the threshold. Spill the existing in-memory buffer, then write + // this entire chunk to disk too — we don't bother splitting it for the sake of staying + // exactly at the threshold. Going over by at most one chunk is fine. + spillToDisk(); + spillStream.write(b, off, len); + } else { + memoryBuffer.write(b, off, len); + } + bytesWritten += len; + } + + private void spillToDisk() throws IOException { + spillFile = tempFileManager.createManagedTempFile(".body"); + spillStream = Files.newOutputStream(spillFile.getPath()); + memoryBuffer.writeTo(spillStream); + memoryBuffer = null; // help GC of potentially large buffer + spilled = true; + log.debug( + "PaygResponseBodyWrapper spilled to {} after {} bytes (threshold {})", + spillFile.getPath(), + bytesWritten, + inMemoryThresholdBytes); + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperFilter.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperFilter.java new file mode 100644 index 000000000..8c6876fb6 --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperFilter.java @@ -0,0 +1,138 @@ +package stirling.software.saas.payg.filter; + +import java.io.IOException; + +import org.springframework.context.annotation.Profile; +import org.springframework.stereotype.Component; +import org.springframework.web.filter.OncePerRequestFilter; + +import jakarta.servlet.AsyncEvent; +import jakarta.servlet.AsyncListener; +import jakarta.servlet.FilterChain; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.util.TempFileManager; + +/** + * Wraps the {@link HttpServletResponse} with a {@link PaygResponseBodyWrapper} for every request, + * making the response body retrievable by the downstream {@code PaygChargeInterceptor} in {@code + * afterCompletion}. The wrapper is stashed as a request attribute under {@link #REQUEST_ATTRIBUTE} + * so the interceptor can find it. + * + *

Pure plumbing — no business logic. When {@code payg.filter.enabled=false} the filter passes + * through unchanged. + * + *

Lifecycle: the wrapper is closed in a {@code finally} after the chain returns for sync + * requests; for async controllers ({@code DeferredResult}, {@code CompletableFuture}), close is + * deferred to an {@link AsyncListener} so the wrapper survives the async window. Close is + * idempotent so a defensive call by the interceptor's {@code afterCompletion} is harmless. + */ +@Slf4j +@Component +@Profile("saas") +public class PaygResponseBodyWrapperFilter extends OncePerRequestFilter { + + /** Request-attribute key under which the wrapper is exposed to the interceptor. */ + public static final String REQUEST_ATTRIBUTE = + PaygResponseBodyWrapperFilter.class.getName() + ".WRAPPER"; + + private final TempFileManager tempFileManager; + private final PaygFilterProperties properties; + + public PaygResponseBodyWrapperFilter( + TempFileManager tempFileManager, PaygFilterProperties properties) { + this.tempFileManager = tempFileManager; + this.properties = properties; + } + + @Override + protected void doFilterInternal( + HttpServletRequest request, HttpServletResponse response, FilterChain chain) + throws ServletException, IOException { + + if (!properties.isEnabled()) { + chain.doFilter(request, response); + return; + } + + PaygResponseBodyWrapper wrapper; + try { + wrapper = + new PaygResponseBodyWrapper( + response, + tempFileManager, + properties.getResponse().getInMemoryThresholdBytes()); + } catch (RuntimeException e) { + // Wrapper construction failure: fail-open. Pass through unwrapped — OUTPUT recording + // is lost but the customer's tool call still runs. + log.warn("PaygResponseBodyWrapper construction failed; passing through unwrapped", e); + chain.doFilter(request, response); + return; + } + + request.setAttribute(REQUEST_ATTRIBUTE, wrapper); + boolean asyncStarted = false; + try { + chain.doFilter(request, wrapper); + asyncStarted = request.isAsyncStarted(); + if (asyncStarted) { + // Async controller: defer attribute removal + close to async dispatch completion. + // The interceptor's afterCompletion fires on the async dispatch and needs the + // wrapper attribute still present at that point. close() is idempotent so a + // defensive call by the interceptor is harmless. + request.getAsyncContext().addListener(new ReleaseOnAsyncComplete(request, wrapper)); + } + } finally { + if (!asyncStarted) { + // Sync path: interceptor.afterCompletion has already run inside chain.doFilter. + request.removeAttribute(REQUEST_ATTRIBUTE); + wrapper.close(); + } + } + } + + /** + * For async dispatches, removes the wrapper attribute and closes the wrapper after the async + * dispatch completes. The Servlet container fires exactly one of {@code onComplete} / {@code + * onError} / {@code onTimeout} per async context lifecycle. + */ + private static final class ReleaseOnAsyncComplete implements AsyncListener { + + private final HttpServletRequest request; + private final PaygResponseBodyWrapper wrapper; + + ReleaseOnAsyncComplete(HttpServletRequest request, PaygResponseBodyWrapper wrapper) { + this.request = request; + this.wrapper = wrapper; + } + + private void release() { + request.removeAttribute(REQUEST_ATTRIBUTE); + wrapper.close(); + } + + @Override + public void onComplete(AsyncEvent event) { + release(); + } + + @Override + public void onTimeout(AsyncEvent event) { + release(); + } + + @Override + public void onError(AsyncEvent event) { + release(); + } + + @Override + public void onStartAsync(AsyncEvent event) { + // re-dispatch retains the listener — no-op + } + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygWebMvcConfig.java b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygWebMvcConfig.java new file mode 100644 index 000000000..8b29874bc --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/filter/PaygWebMvcConfig.java @@ -0,0 +1,60 @@ +package stirling.software.saas.payg.filter; + +import org.springframework.boot.web.servlet.FilterRegistrationBean; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.Profile; +import org.springframework.web.servlet.config.annotation.InterceptorRegistry; +import org.springframework.web.servlet.config.annotation.WebMvcConfigurer; + +import lombok.RequiredArgsConstructor; + +/** + * Wires the PAYG filter + interceptor into Spring MVC. Two registrations: + * + *

    + *
  • {@link PaygResponseBodyWrapperFilter} as a Servlet filter — registered with no explicit + * order so it sits at the end of the Spring filter chain (after all security filters). Pure + * response-wrapping plumbing. + *
  • {@link PaygChargeInterceptor} as a Spring MVC interceptor — registered AFTER {@code + * UnifiedCreditInterceptor} so legacy credit rejections short-circuit before we hash inputs. + * Both intercept {@code /api/**} with the same admin/info/health exclusions as the legacy + * config. + *
+ */ +@Configuration +@Profile("saas") +@RequiredArgsConstructor +public class PaygWebMvcConfig implements WebMvcConfigurer { + + private final PaygChargeInterceptor paygChargeInterceptor; + + @Bean + public FilterRegistrationBean + paygResponseBodyWrapperFilterRegistration(PaygResponseBodyWrapperFilter filter) { + FilterRegistrationBean reg = + new FilterRegistrationBean<>(filter); + reg.addUrlPatterns("/api/*"); + return reg; + } + + /** + * Interceptor ordering: the legacy {@code UnifiedCreditInterceptor} (registered with default + * order = 0 in {@code CreditInterceptorConfig}) must run BEFORE this one so credit rejections + * short-circuit before we hash inputs. Explicit positive order guarantees this regardless of + * {@code WebMvcConfigurer} bean discovery order. + */ + public static final int INTERCEPTOR_ORDER = 1000; + + @Override + public void addInterceptors(InterceptorRegistry registry) { + registry.addInterceptor(paygChargeInterceptor) + .addPathPatterns("/api/**") + .excludePathPatterns( + "/api/v1/credits/**", + "/api/v1/config/**", + "/api/v1/info/**", + "/api/v1/admin/**") + .order(INTERCEPTOR_ORDER); + } +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/job/JobService.java b/app/saas/src/main/java/stirling/software/saas/payg/job/JobService.java index 005833c70..e6a4f04a8 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/job/JobService.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/job/JobService.java @@ -6,9 +6,12 @@ import java.time.Duration; import java.time.LocalDateTime; import java.util.ArrayList; import java.util.Comparator; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; import org.springframework.beans.factory.annotation.Value; @@ -20,6 +23,7 @@ import lombok.extern.slf4j.Slf4j; import stirling.software.saas.payg.lineage.HashLineageDetector; import stirling.software.saas.payg.lineage.LineageMatch; +import stirling.software.saas.payg.lineage.LineageSignature; import stirling.software.saas.payg.model.ArtifactKind; import stirling.software.saas.payg.model.JobStatus; import stirling.software.saas.payg.model.JobStepStatus; @@ -86,7 +90,15 @@ public class JobService { throw new IllegalArgumentException("inputs must not be empty"); } - Optional bestMatch = findBestMatch(ctx.ownerUserId(), inputs); + // Extract signatures ONCE per input, then reuse for both the lineage lookup and the + // post-decision record() call. Avoids hashing every input twice on the hot path. + Map> signaturesByInput = new HashMap<>(inputs.size()); + for (Path input : inputs) { + signaturesByInput.put(input, detector.extractSignatures(input)); + } + + Optional bestMatch = + findBestMatch(ctx.ownerUserId(), inputs, signaturesByInput); if (bestMatch.isPresent()) { ProcessingJob existing = @@ -100,7 +112,7 @@ public class JobService { + " but no such ProcessingJob row" + " exists (stale signature?)")); if (existing.getStepCount() < ctx.stepLimit()) { - return joinExisting(existing, inputs); + return joinExisting(existing, signaturesByInput); } // Step-limit hit: spawn a fresh job. The new job will share input signatures with // the existing chain so future tool calls still lineage-match into the workflow, @@ -111,7 +123,7 @@ public class JobService { ctx.stepLimit()); } - return openFresh(ctx, inputs); + return openFresh(ctx, signaturesByInput); } /** @@ -199,25 +211,26 @@ public class JobService { return stale.size(); } - private Optional findBestMatch(Long userId, List inputs) - throws IOException { + private Optional findBestMatch( + Long userId, List inputs, Map> signaturesByInput) { List matches = new ArrayList<>(inputs.size()); for (Path input : inputs) { - detector.detect(userId, input).ifPresent(matches::add); + detector.detect(userId, signaturesByInput.get(input)).ifPresent(matches::add); } return matches.stream().max(Comparator.comparing(LineageMatch::jobLastStepAt)); } - private JoinOrOpenResult joinExisting(ProcessingJob existing, List inputs) - throws IOException { + private JoinOrOpenResult joinExisting( + ProcessingJob existing, Map> signaturesByInput) { existing.setStepCount(existing.getStepCount() + 1); existing.setLastStepAt(LocalDateTime.now()); ProcessingJob saved = jobRepository.save(existing); - recordAllInputs(saved.getId(), inputs); + recordAllInputs(saved.getId(), signaturesByInput); return new JoinOrOpenResult(saved, JoinOrOpenResult.Disposition.JOINED); } - private JoinOrOpenResult openFresh(JobContext ctx, List inputs) throws IOException { + private JoinOrOpenResult openFresh( + JobContext ctx, Map> signaturesByInput) { ProcessingJob fresh = new ProcessingJob(); fresh.setId(UUID.randomUUID()); fresh.setOwnerUserId(ctx.ownerUserId()); @@ -231,13 +244,13 @@ public class JobService { fresh.setLastStepAt(now); fresh.setStatus(JobStatus.OPEN); ProcessingJob saved = jobRepository.save(fresh); - recordAllInputs(saved.getId(), inputs); + recordAllInputs(saved.getId(), signaturesByInput); return new JoinOrOpenResult(saved, JoinOrOpenResult.Disposition.OPENED); } - private void recordAllInputs(UUID jobId, List inputs) throws IOException { - for (Path input : inputs) { - detector.record(jobId, input, ArtifactKind.INPUT); + private void recordAllInputs(UUID jobId, Map> signaturesByInput) { + for (Set signatures : signaturesByInput.values()) { + detector.record(jobId, signatures, ArtifactKind.INPUT); } } } diff --git a/app/saas/src/main/java/stirling/software/saas/payg/lineage/DefaultHashLineageDetector.java b/app/saas/src/main/java/stirling/software/saas/payg/lineage/DefaultHashLineageDetector.java index 2d261ee02..082054a09 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/lineage/DefaultHashLineageDetector.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/lineage/DefaultHashLineageDetector.java @@ -58,37 +58,45 @@ public class DefaultHashLineageDetector implements HashLineageDetector { @Override public Optional detect(Long userId, Path inputFile) throws IOException { - Objects.requireNonNull(userId, "userId"); Objects.requireNonNull(inputFile, "inputFile"); + return detect(userId, extractSignatures(inputFile)); + } - Set signatures = extractAll(inputFile); + @Override + public Optional detect(Long userId, Set signatures) { + Objects.requireNonNull(userId, "userId"); + Objects.requireNonNull(signatures, "signatures"); if (signatures.isEmpty()) { // No extractor produced anything for this content. Treat as no-match. - log.debug("No signatures extracted from {}; lineage check returns empty.", inputFile); return Optional.empty(); } - return store.findOpenJobForSignatures(userId, signatures, workflowWindow); } @Override public void record(UUID jobId, Path file, ArtifactKind kind) throws IOException { - Objects.requireNonNull(jobId, "jobId"); Objects.requireNonNull(file, "file"); - Objects.requireNonNull(kind, "kind"); + record(jobId, extractSignatures(file), kind); + } - Set signatures = extractAll(file); + @Override + public void record(UUID jobId, Set signatures, ArtifactKind kind) { + Objects.requireNonNull(jobId, "jobId"); + Objects.requireNonNull(signatures, "signatures"); + Objects.requireNonNull(kind, "kind"); if (signatures.isEmpty()) { - log.debug( - "No signatures extracted from {} for job {} ({}); nothing recorded.", - file, - jobId, - kind); + log.debug("No signatures to record for job {} ({}); skipping.", jobId, kind); return; } store.record(jobId, signatures, kind); } + @Override + public Set extractSignatures(Path file) { + Objects.requireNonNull(file, "file"); + return extractAll(file); + } + private Set extractAll(Path file) { Set union = new HashSet<>(); for (LineageSignatureExtractor extractor : extractors) { diff --git a/app/saas/src/main/java/stirling/software/saas/payg/lineage/HashLineageDetector.java b/app/saas/src/main/java/stirling/software/saas/payg/lineage/HashLineageDetector.java index 60aa589bf..66559eaf4 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/lineage/HashLineageDetector.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/lineage/HashLineageDetector.java @@ -3,6 +3,7 @@ package stirling.software.saas.payg.lineage; import java.io.IOException; import java.nio.file.Path; import java.util.Optional; +import java.util.Set; import java.util.UUID; import stirling.software.saas.payg.model.ArtifactKind; @@ -33,4 +34,25 @@ public interface HashLineageDetector { /** Records the file's signatures against the given job as either INPUT or OUTPUT. */ void record(UUID jobId, Path file, ArtifactKind kind) throws IOException; + + /** + * Pre-computes the signature set for {@code file} so callers can avoid hashing the same bytes + * twice when they need both {@link #detect} and {@link #record} for the same file. Returned set + * may be empty (no extractor recognised the content) — treat as "no signatures to match or + * record" by both consumers. + */ + Set extractSignatures(Path file); + + /** + * Same as {@link #detect(Long, Path)} but operating on pre-computed signatures. Useful when the + * caller has already extracted them (e.g. via {@link #extractSignatures}) and wants to avoid + * re-hashing. Empty {@code signatures} short-circuits to {@link Optional#empty}. + */ + Optional detect(Long userId, Set signatures); + + /** + * Same as {@link #record(UUID, Path, ArtifactKind)} but operating on pre-computed signatures. + * Empty {@code signatures} is a no-op. + */ + void record(UUID jobId, Set signatures, ArtifactKind kind); } diff --git a/app/saas/src/main/java/stirling/software/saas/payg/model/ShadowChargeStatus.java b/app/saas/src/main/java/stirling/software/saas/payg/model/ShadowChargeStatus.java new file mode 100644 index 000000000..74653712e --- /dev/null +++ b/app/saas/src/main/java/stirling/software/saas/payg/model/ShadowChargeStatus.java @@ -0,0 +1,18 @@ +package stirling.software.saas.payg.model; + +/** + * Lifecycle state for a {@code payg_shadow_charge} row. Mimics the Stripe meter_event_adjustment + * (type=cancel) mechanism that real-mode will invoke when a freshly-opened process fails on its + * first step. The reconciliation report's "true net" query is {@code SUM(payg_units) WHERE status = + * 'CHARGED'}. + */ +public enum ShadowChargeStatus { + /** Default at process open; the would-be charge stands. */ + CHARGED, + /** + * Set when a 5xx first-step failure was observed in the same request's {@code afterCompletion}. + * Real-mode equivalent is a successful Stripe {@code meter_event_adjustment} posted in the same + * transaction-after-commit hook. + */ + REFUNDED +} diff --git a/app/saas/src/main/java/stirling/software/saas/payg/policy/PricingPolicyService.java b/app/saas/src/main/java/stirling/software/saas/payg/policy/PricingPolicyService.java index e273a391d..0fd7bb45d 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/policy/PricingPolicyService.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/policy/PricingPolicyService.java @@ -79,19 +79,28 @@ public class PricingPolicyService { * {@code is_default = TRUE}. Throws {@link IllegalStateException} if no default exists — the * seed migration is expected to put one there. * + *

A {@code null} {@code teamId} (admin-created user, deleted team, pre-team-migration + * account) is valid and returns the default policy directly. Throwing here would silently + * exclude the team-less cohort from PAYG via the filter's fail-open path, biasing the + * shadow-vs-legacy reconciliation. + * *

{@link Transactional}({@code readOnly = true}) so the eager-loaded {@code stepLimits} and * {@code stripePriceIds} collections initialize inside the same session. */ @Transactional(readOnly = true) public PricingPolicy getEffectivePolicy(Long teamId) { - Objects.requireNonNull(teamId, "teamId"); + if (teamId == null) { + return loadDefaultPolicy(); + } return byTeamCache.get(teamId, this::loadEffectivePolicy); } /** Bypasses the cache. Useful for admin endpoints that want a fresh read after a mutation. */ @Transactional(readOnly = true) public PricingPolicy getEffectivePolicyUncached(Long teamId) { - Objects.requireNonNull(teamId, "teamId"); + if (teamId == null) { + return loadDefaultPolicy(); + } return loadEffectivePolicy(teamId); } @@ -243,6 +252,10 @@ public class PricingPolicyService { teamId, id); } + return loadDefaultPolicy(); + } + + private PricingPolicy loadDefaultPolicy() { return policyRepository .findFirstByIsDefaultTrue() .orElseThrow( diff --git a/app/saas/src/main/java/stirling/software/saas/payg/repository/PaygShadowChargeRepository.java b/app/saas/src/main/java/stirling/software/saas/payg/repository/PaygShadowChargeRepository.java index ec1a51963..e953285ee 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/repository/PaygShadowChargeRepository.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/repository/PaygShadowChargeRepository.java @@ -2,6 +2,8 @@ package stirling.software.saas.payg.repository; import java.time.LocalDateTime; import java.util.List; +import java.util.Optional; +import java.util.UUID; import org.springframework.data.jpa.repository.JpaRepository; import org.springframework.data.jpa.repository.Query; @@ -19,4 +21,12 @@ public interface PaygShadowChargeRepository extends JpaRepository findInWindow( @Param("from") LocalDateTime from, @Param("to") LocalDateTime to); + + /** + * The shadow row written when the given process was opened. At most one row per {@code jobId} + * exists by construction — {@code openProcess} writes exactly one row on OPENED and zero on + * JOINED — so callers can treat the result as a single optional. Returns the first row by id + * defensively if a duplicate ever appears. + */ + Optional findFirstByJobIdOrderByIdAsc(UUID jobId); } diff --git a/app/saas/src/main/java/stirling/software/saas/payg/shadow/PaygShadowCharge.java b/app/saas/src/main/java/stirling/software/saas/payg/shadow/PaygShadowCharge.java index 495930c08..b2b8ed0c2 100644 --- a/app/saas/src/main/java/stirling/software/saas/payg/shadow/PaygShadowCharge.java +++ b/app/saas/src/main/java/stirling/software/saas/payg/shadow/PaygShadowCharge.java @@ -8,6 +8,8 @@ import org.hibernate.annotations.CreationTimestamp; import jakarta.persistence.Column; import jakarta.persistence.Entity; +import jakarta.persistence.EnumType; +import jakarta.persistence.Enumerated; import jakarta.persistence.GeneratedValue; import jakarta.persistence.GenerationType; import jakarta.persistence.Id; @@ -17,10 +19,16 @@ import lombok.Getter; import lombok.NoArgsConstructor; import lombok.Setter; +import stirling.software.saas.payg.model.ShadowChargeStatus; + /** * Per-job comparison row written while a team is in {@code PAYG_SHADOW} mode: what the legacy * engine actually charged vs. what the PAYG engine would have charged. Aggregated daily by the * shadow-reconciliation report; deletable after promotion. + * + *

The {@link #status} column mirrors the eventual Stripe meter_event_adjustment(cancel) flow: + * rows are written {@code CHARGED}; a freshly-opened process whose first step 5xx's is flipped to + * {@code REFUNDED} in the same request's {@code afterCompletion}. */ @Entity @Table(name = "payg_shadow_charge") @@ -55,6 +63,18 @@ public class PaygShadowCharge implements Serializable { @Column(name = "diff_pct", nullable = false) private Integer diffPct; + @Enumerated(EnumType.STRING) + @Column(name = "status", nullable = false, length = 16) + private ShadowChargeStatus status = ShadowChargeStatus.CHARGED; + + /** Set when {@link #status} flips to {@link ShadowChargeStatus#REFUNDED}. */ + @Column(name = "refunded_at") + private LocalDateTime refundedAt; + + /** Free-form reason, e.g. {@code "first-step-5xx:503"}. */ + @Column(name = "refund_reason", length = 128) + private String refundReason; + @CreationTimestamp @Column(name = "occurred_at", nullable = false, updatable = false) private LocalDateTime occurredAt; diff --git a/app/saas/src/main/resources/db/migration/saas/V13__payg_shadow_charge_status.sql b/app/saas/src/main/resources/db/migration/saas/V13__payg_shadow_charge_status.sql new file mode 100644 index 000000000..66da92a10 --- /dev/null +++ b/app/saas/src/main/resources/db/migration/saas/V13__payg_shadow_charge_status.sql @@ -0,0 +1,20 @@ +-- Refund tracking on shadow rows. Shadow models the eventual Stripe meter_event_adjustment(cancel) +-- by flipping status from CHARGED to REFUNDED in the same request's afterCompletion when a +-- freshly-opened process fails with 5xx on its first step. +-- +-- Reconciliation report selects SUM(payg_units) WHERE status = 'CHARGED' to get the true net +-- Stripe would bill. + +ALTER TABLE payg_shadow_charge + ADD COLUMN IF NOT EXISTS status VARCHAR(16) NOT NULL DEFAULT 'CHARGED', + ADD COLUMN IF NOT EXISTS refunded_at TIMESTAMP, + ADD COLUMN IF NOT EXISTS refund_reason VARCHAR(128); + +CREATE INDEX IF NOT EXISTS idx_payg_shadow_status_time + ON payg_shadow_charge (status, occurred_at); + +-- Hot-path index for findFirstByJobIdOrderByIdAsc: hit on every 5xx-first-step refund to flip +-- the row to REFUNDED. UNIQUE because at most one shadow row exists per processing_job by +-- construction (openProcess writes exactly one on OPENED, zero on JOINED). +CREATE UNIQUE INDEX IF NOT EXISTS uq_payg_shadow_job_id + ON payg_shadow_charge (job_id); diff --git a/app/saas/src/test/java/stirling/software/saas/payg/charge/JobChargeServiceTest.java b/app/saas/src/test/java/stirling/software/saas/payg/charge/JobChargeServiceTest.java index 00e584316..29e7041bf 100644 --- a/app/saas/src/test/java/stirling/software/saas/payg/charge/JobChargeServiceTest.java +++ b/app/saas/src/test/java/stirling/software/saas/payg/charge/JobChargeServiceTest.java @@ -36,9 +36,11 @@ import stirling.software.saas.payg.job.ProcessingJob; import stirling.software.saas.payg.model.JobSource; import stirling.software.saas.payg.model.JobStatus; import stirling.software.saas.payg.model.ProcessType; +import stirling.software.saas.payg.model.ShadowChargeStatus; import stirling.software.saas.payg.policy.PricingPolicy; import stirling.software.saas.payg.policy.PricingPolicyService; import stirling.software.saas.payg.repository.PaygShadowChargeRepository; +import stirling.software.saas.payg.repository.ProcessingJobRepository; import stirling.software.saas.payg.shadow.PaygShadowCharge; /** @@ -51,6 +53,7 @@ class JobChargeServiceTest { private PricingPolicyService policyService; private DocumentClassifier classifier; private PaygShadowChargeRepository shadowRepo; + private ProcessingJobRepository jobRepo; private JobChargeService service; @BeforeEach @@ -59,7 +62,8 @@ class JobChargeServiceTest { policyService = Mockito.mock(PricingPolicyService.class); classifier = Mockito.mock(DocumentClassifier.class); shadowRepo = Mockito.mock(PaygShadowChargeRepository.class); - service = new JobChargeService(jobService, policyService, classifier, shadowRepo); + jobRepo = Mockito.mock(ProcessingJobRepository.class); + service = new JobChargeService(jobService, policyService, classifier, shadowRepo, jobRepo); } @Test @@ -97,7 +101,7 @@ class JobChargeServiceTest { .thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED)); JobInput in = jobInput(tmp, "in.pdf", "application/pdf"); - when(classifier.classify(any(MultipartFile.class), eq(policy))) + when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy))) .thenReturn(new DocumentMetrics(50, 1024L, "application/pdf", 4)); ChargeOutcome out = @@ -107,9 +111,10 @@ class JobChargeServiceTest { assertThat(out.disposition()).isEqualTo(ChargeOutcome.Disposition.OPENED); assertThat(out.units()).isEqualTo(4); - // Single-file path called single-file classifier overload, not the list one. - verify(classifier, times(1)).classify(any(MultipartFile.class), eq(policy)); - verify(classifier, never()).classify(anyList(), eq(policy)); + // Single-file path called single-file classifier overload (with Path), not the list one. + verify(classifier, times(1)) + .classify(any(MultipartFile.class), any(Path.class), eq(policy)); + verify(classifier, never()).classify(anyList(), anyList(), eq(policy)); ArgumentCaptor captor = ArgumentCaptor.forClass(PaygShadowCharge.class); verify(shadowRepo).save(captor.capture()); @@ -136,7 +141,7 @@ class JobChargeServiceTest { JobInput a = jobInput(tmp, "a.pdf", "application/pdf"); JobInput b = jobInput(tmp, "b.pdf", "application/pdf"); - when(classifier.classify(anyList(), eq(policy))) + when(classifier.classify(anyList(), anyList(), eq(policy))) .thenReturn(new DocumentMetrics(100, 2048L, "application/pdf", 7)); ChargeOutcome out = @@ -145,8 +150,8 @@ class JobChargeServiceTest { List.of(a, b)); assertThat(out.units()).isEqualTo(7); - verify(classifier, never()).classify(any(MultipartFile.class), any()); - verify(classifier, times(1)).classify(anyList(), eq(policy)); + verify(classifier, never()).classify(any(MultipartFile.class), any(Path.class), any()); + verify(classifier, times(1)).classify(anyList(), anyList(), eq(policy)); } @Test @@ -157,7 +162,7 @@ class JobChargeServiceTest { ProcessingJob newJob = openJob(UUID.randomUUID()); when(jobService.joinOrOpen(any(JobContext.class), anyList())) .thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED)); - when(classifier.classify(any(MultipartFile.class), eq(policy))) + when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy))) .thenReturn(new DocumentMetrics(10, 1024L, "application/pdf", 2)); ChargeOutcome out = @@ -178,7 +183,7 @@ class JobChargeServiceTest { ProcessingJob newJob = openJob(UUID.randomUUID()); when(jobService.joinOrOpen(any(JobContext.class), anyList())) .thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED)); - when(classifier.classify(any(MultipartFile.class), eq(policy))) + when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy))) .thenReturn(new DocumentMetrics(1, 100L, "application/pdf", 1)); service.openProcess( @@ -202,7 +207,7 @@ class JobChargeServiceTest { ProcessingJob newJob = openJob(UUID.randomUUID()); when(jobService.joinOrOpen(any(JobContext.class), anyList())) .thenReturn(new JoinOrOpenResult(newJob, JoinOrOpenResult.Disposition.OPENED)); - when(classifier.classify(any(MultipartFile.class), eq(policy))) + when(classifier.classify(any(MultipartFile.class), any(Path.class), eq(policy))) .thenReturn(new DocumentMetrics(1, 100L, "application/pdf", 1)); service.openProcess( @@ -226,6 +231,105 @@ class JobChargeServiceTest { .hasMessageContaining("inputs must not be empty"); } + @Test + void markFirstStepFailed_flipsShadowRowAndClosesProcess() { + UUID jobId = UUID.randomUUID(); + PaygShadowCharge row = new PaygShadowCharge(); + row.setJobId(jobId); + row.setStatus(ShadowChargeStatus.CHARGED); + when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row)); + ProcessingJob job = openJob(jobId); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job)); + + service.markFirstStepFailed(jobId, "first-step-5xx:503"); + + assertThat(row.getStatus()).isEqualTo(ShadowChargeStatus.REFUNDED); + assertThat(row.getRefundedAt()).isNotNull(); + assertThat(row.getRefundReason()).isEqualTo("first-step-5xx:503"); + assertThat(job.getStatus()).isEqualTo(JobStatus.CLOSED); + assertThat(job.getClosedAt()).isNotNull(); + verify(shadowRepo).save(row); + verify(jobRepo).save(job); + } + + @Test + void markFirstStepFailed_alreadyRefunded_isNoOp() { + UUID jobId = UUID.randomUUID(); + PaygShadowCharge row = new PaygShadowCharge(); + row.setJobId(jobId); + row.setStatus(ShadowChargeStatus.REFUNDED); + when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row)); + ProcessingJob job = openJob(jobId); + job.setStatus(JobStatus.CLOSED); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job)); + + service.markFirstStepFailed(jobId, "first-step-5xx:500"); + + verify(shadowRepo, never()).save(any()); + verify(jobRepo, never()).save(any()); + } + + @Test + void markFirstStepFailed_noShadowRow_stillClosesProcess() { + UUID jobId = UUID.randomUUID(); + when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.empty()); + ProcessingJob job = openJob(jobId); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job)); + + service.markFirstStepFailed(jobId, "first-step-5xx:503"); + + assertThat(job.getStatus()).isEqualTo(JobStatus.CLOSED); + verify(jobRepo).save(job); + } + + @Test + void markFirstStepFailed_trimsLongRefundReason() { + UUID jobId = UUID.randomUUID(); + PaygShadowCharge row = new PaygShadowCharge(); + row.setStatus(ShadowChargeStatus.CHARGED); + when(shadowRepo.findFirstByJobIdOrderByIdAsc(jobId)).thenReturn(java.util.Optional.of(row)); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.empty()); + + String oversized = "x".repeat(200); + service.markFirstStepFailed(jobId, oversized); + + assertThat(row.getRefundReason()).hasSize(128); + } + + @Test + void decrementStepCount_decrementsByOne() { + UUID jobId = UUID.randomUUID(); + ProcessingJob job = openJob(jobId); + job.setStepCount(3); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job)); + + service.decrementStepCount(jobId); + + assertThat(job.getStepCount()).isEqualTo(2); + verify(jobRepo).save(job); + } + + @Test + void decrementStepCount_floorAtOne_neverDrivesNegative() { + UUID jobId = UUID.randomUUID(); + ProcessingJob job = openJob(jobId); + job.setStepCount(1); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.of(job)); + + service.decrementStepCount(jobId); + + assertThat(job.getStepCount()).isEqualTo(1); + verify(jobRepo, never()).save(any()); + } + + @Test + void decrementStepCount_missingJob_isNoOp() { + UUID jobId = UUID.randomUUID(); + when(jobRepo.findById(jobId)).thenReturn(java.util.Optional.empty()); + service.decrementStepCount(jobId); // must not throw + verify(jobRepo, never()).save(any()); + } + // --- helpers -------------------------------------------------------------------------------- private static PricingPolicy stubPolicy(int minCharge, Map stepLimits) { diff --git a/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygChargeInterceptorTest.java b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygChargeInterceptorTest.java new file mode 100644 index 000000000..318dc72ff --- /dev/null +++ b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygChargeInterceptorTest.java @@ -0,0 +1,420 @@ +package stirling.software.saas.payg.filter; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Optional; +import java.util.UUID; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.mock.web.MockHttpServletRequest; +import org.springframework.mock.web.MockHttpServletResponse; +import org.springframework.mock.web.MockMultipartFile; +import org.springframework.mock.web.MockMultipartHttpServletRequest; +import org.springframework.security.authentication.UsernamePasswordAuthenticationToken; +import org.springframework.security.core.authority.SimpleGrantedAuthority; +import org.springframework.security.core.context.SecurityContextHolder; +import org.springframework.web.method.HandlerMethod; + +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; + +import stirling.software.common.annotations.AutoJobPostMapping; +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.util.TempFileManager; +import stirling.software.common.util.TempFileRegistry; +import stirling.software.proprietary.security.database.repository.UserRepository; +import stirling.software.proprietary.security.model.User; +import stirling.software.saas.payg.charge.ChargeOutcome; +import stirling.software.saas.payg.charge.JobChargeService; +import stirling.software.saas.payg.job.JobService; +import stirling.software.saas.payg.model.JobStepStatus; + +/** + * Pure-Mockito tests for {@link PaygChargeInterceptor}. Real {@link TempFileManager} but mocked + * downstream charge/job services so the test runs without a Spring context. + */ +class PaygChargeInterceptorTest { + + private JobChargeService chargeService; + private JobService jobService; + private UserRepository userRepository; + private PaygOutputExtractor outputExtractor; + private PaygFilterProperties properties; + private MeterRegistry meterRegistry; + private TempFileManager tempFileManager; + private PaygChargeInterceptor interceptor; + + @BeforeEach + void setUp() { + chargeService = org.mockito.Mockito.mock(JobChargeService.class); + jobService = org.mockito.Mockito.mock(JobService.class); + userRepository = org.mockito.Mockito.mock(UserRepository.class); + outputExtractor = org.mockito.Mockito.mock(PaygOutputExtractor.class); + properties = new PaygFilterProperties(); + meterRegistry = new SimpleMeterRegistry(); + tempFileManager = new TempFileManager(new TempFileRegistry(), new ApplicationProperties()); + interceptor = + new PaygChargeInterceptor( + chargeService, + jobService, + userRepository, + tempFileManager, + outputExtractor, + properties, + meterRegistry); + SecurityContextHolder.clearContext(); + } + + @Test + void preHandle_filterDisabled_isShortCircuitNoop() throws Exception { + properties.setEnabled(false); + MockHttpServletRequest req = new MockHttpServletRequest(); + MockHttpServletResponse res = new MockHttpServletResponse(); + + boolean cont = interceptor.preHandle(req, res, handlerMethodForFakeController()); + + assertThat(cont).isTrue(); + verifyNoInteractions(chargeService); + } + + @Test + void preHandle_handlerNotAnnotated_isShortCircuitNoop() throws Exception { + MockHttpServletRequest req = new MockHttpServletRequest(); + MockHttpServletResponse res = new MockHttpServletResponse(); + + boolean cont = interceptor.preHandle(req, res, handlerMethodForPlain()); + + assertThat(cont).isTrue(); + verifyNoInteractions(chargeService); + } + + @Test + void preHandle_noAuth_isShortCircuitNoop() throws Exception { + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "hi".getBytes())); + + boolean cont = + interceptor.preHandle( + req, new MockHttpServletResponse(), handlerMethodForFakeController()); + + assertThat(cont).isTrue(); + verifyNoInteractions(chargeService); + } + + @Test + void preHandle_noMultipartParts_isShortCircuitNoop() throws Exception { + // Authenticated but no file parts. + authenticateWithUser(makeUser(1L, null)); + MockHttpServletRequest req = new MockHttpServletRequest(); + + boolean cont = + interceptor.preHandle( + req, new MockHttpServletResponse(), handlerMethodForFakeController()); + + assertThat(cont).isTrue(); + verifyNoInteractions(chargeService); + } + + @Test + void preHandle_openedDisposition_stashesJobIdAndDisposition() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 4, ChargeOutcome.Disposition.OPENED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + + boolean cont = + interceptor.preHandle( + req, new MockHttpServletResponse(), handlerMethodForFakeController()); + + assertThat(cont).isTrue(); + assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_JOB_ID)).isEqualTo(jobId); + assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_DISPOSITION)) + .isEqualTo(ChargeOutcome.Disposition.OPENED); + assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_INPUT_TEMP_FILES)).isNotNull(); + assertThat(meterRegistry.counter("payg.filter.calls", "disposition", "OPENED").count()) + .isEqualTo(1.0); + } + + @Test + void preHandle_chargeServiceThrows_failsOpenAndIncrementsErrorCounter() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + when(chargeService.openProcess(any(), anyList())).thenThrow(new RuntimeException("boom")); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + + boolean cont = + interceptor.preHandle( + req, new MockHttpServletResponse(), handlerMethodForFakeController()); + + assertThat(cont).isTrue(); // fail-open + assertThat(req.getAttribute(PaygChargeInterceptor.ATTR_FAILED)).isEqualTo(Boolean.TRUE); + assertThat(meterRegistry.counter("payg.filter.errors").count()).isEqualTo(1.0); + } + + @Test + void afterCompletion_2xx_appendsStepAndRecordsOutputs() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + MockHttpServletResponse res = new MockHttpServletResponse(); + res.setStatus(200); + res.setContentType("application/pdf"); + + // Install a wrapper as the filter would. Pre-populate it with some bytes so + // materialisedPath returns non-null. + PaygResponseBodyWrapper wrapper = new PaygResponseBodyWrapper(res, tempFileManager, 1024); + wrapper.getOutputStream().write("body".getBytes(StandardCharsets.UTF_8)); + req.setAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE, wrapper); + + interceptor.preHandle(req, res, handlerMethodForFakeController()); + + when(outputExtractor.extract(eq("application/pdf"), any())) + .thenReturn( + List.of( + new PaygOutputExtractor.ExtractedPdf( + wrapper.materialisedPath(), null))); + + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + ArgumentCaptor status = ArgumentCaptor.forClass(JobStepStatus.class); + verify(jobService).appendStep(eq(jobId), any(), status.capture(), any(), any(), any()); + assertThat(status.getValue()).isEqualTo(JobStepStatus.OK); + verify(jobService).recordOutput(eq(jobId), any()); + verify(chargeService, never()).markFirstStepFailed(any(), any()); + verify(chargeService, never()).decrementStepCount(any()); + } + + @Test + void afterCompletion_5xx_opened_callsMarkFirstStepFailed() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + MockHttpServletResponse res = new MockHttpServletResponse(); + res.setStatus(503); + + interceptor.preHandle(req, res, handlerMethodForFakeController()); + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + verify(chargeService).markFirstStepFailed(eq(jobId), eq("first-step-5xx:503")); + verify(chargeService, never()).decrementStepCount(any()); + verify(jobService, never()).recordOutput(any(), any()); + verify(jobService) + .appendStep(eq(jobId), any(), eq(JobStepStatus.FAILED), any(), any(), eq("503")); + assertThat(meterRegistry.counter("payg.filter.refunds").count()).isEqualTo(1.0); + } + + @Test + void afterCompletion_5xx_joined_callsDecrementStepCount() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 0, ChargeOutcome.Disposition.JOINED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + MockHttpServletResponse res = new MockHttpServletResponse(); + res.setStatus(500); + + interceptor.preHandle(req, res, handlerMethodForFakeController()); + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + verify(chargeService).decrementStepCount(jobId); + verify(chargeService, never()).markFirstStepFailed(any(), any()); + } + + @Test + void afterCompletion_4xx_appendsFailedStepNoRefundNoOutputs() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + MockHttpServletResponse res = new MockHttpServletResponse(); + res.setStatus(422); + + interceptor.preHandle(req, res, handlerMethodForFakeController()); + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + verify(chargeService, never()).markFirstStepFailed(any(), any()); + verify(chargeService, never()).decrementStepCount(any()); + verify(jobService, never()).recordOutput(any(), any()); + verify(jobService) + .appendStep(eq(jobId), any(), eq(JobStepStatus.FAILED), any(), any(), eq("422")); + } + + @Test + void afterCompletion_preHandleFailedFlag_cleansUpWithoutCharging() throws Exception { + MockMultipartHttpServletRequest req = newMultipart(); + req.setAttribute(PaygChargeInterceptor.ATTR_FAILED, Boolean.TRUE); + MockHttpServletResponse res = new MockHttpServletResponse(); + + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + verifyNoInteractions(chargeService); + verifyNoInteractions(jobService); + } + + @Test + void afterCompletion_noJobId_isNoop() throws Exception { + MockMultipartHttpServletRequest req = newMultipart(); + MockHttpServletResponse res = new MockHttpServletResponse(); + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + verifyNoInteractions(chargeService); + verifyNoInteractions(jobService); + } + + @Test + void afterCompletion_maxBytesExceeded_skipsOutputRecording() throws Exception { + properties.getResponse().setMaxBytes(2L); + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED)); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + MockHttpServletResponse res = new MockHttpServletResponse(); + res.setStatus(200); + res.setContentType("application/pdf"); + + PaygResponseBodyWrapper wrapper = new PaygResponseBodyWrapper(res, tempFileManager, 1024); + wrapper.getOutputStream().write("1234567890".getBytes(StandardCharsets.UTF_8)); + req.setAttribute(PaygResponseBodyWrapperFilter.REQUEST_ATTRIBUTE, wrapper); + + interceptor.preHandle(req, res, handlerMethodForFakeController()); + interceptor.afterCompletion(req, res, handlerMethodForFakeController(), null); + + verify(outputExtractor, never()).extract(any(), any()); + verify(jobService, never()).recordOutput(any(), any()); + } + + @Test + void preHandle_pipelineHeader_setsJobSourcePipeline() throws Exception { + authenticateWithUser(makeUser(7L, 42L)); + UUID jobId = UUID.randomUUID(); + when(chargeService.openProcess(any(), anyList())) + .thenReturn(new ChargeOutcome(jobId, 1, ChargeOutcome.Disposition.OPENED)); + org.mockito.ArgumentCaptor ctxCaptor = + org.mockito.ArgumentCaptor.forClass( + stirling.software.saas.payg.charge.ChargeContext.class); + + MockMultipartHttpServletRequest req = newMultipart(); + req.addFile(new MockMultipartFile("file", "x.pdf", "application/pdf", "abc".getBytes())); + req.addHeader("X-Stirling-Automation", "true"); + + interceptor.preHandle(req, new MockHttpServletResponse(), handlerMethodForFakeController()); + + verify(chargeService).openProcess(ctxCaptor.capture(), anyList()); + assertThat(ctxCaptor.getValue().source()) + .isEqualTo(stirling.software.saas.payg.model.JobSource.PIPELINE); + } + + // --- helpers -------------------------------------------------------------------------------- + + private MockMultipartHttpServletRequest newMultipart() { + MockMultipartHttpServletRequest r = new MockMultipartHttpServletRequest(); + r.setRequestURI("/api/v1/security/test-tool"); + return r; + } + + private void authenticateWithUser(User user) { + UsernamePasswordAuthenticationToken token = + new UsernamePasswordAuthenticationToken( + "supabase-id-here", null, List.of(new SimpleGrantedAuthority("ROLE_USER"))); + SecurityContextHolder.getContext().setAuthentication(token); + // resolveUser does UUID.fromString on the name, so use a real UUID string. + String supabaseId = UUID.randomUUID().toString(); + UsernamePasswordAuthenticationToken realToken = + new UsernamePasswordAuthenticationToken( + supabaseId, null, List.of(new SimpleGrantedAuthority("ROLE_USER"))); + SecurityContextHolder.getContext().setAuthentication(realToken); + when(userRepository.findBySupabaseId(UUID.fromString(supabaseId))) + .thenReturn(Optional.of(user)); + } + + private static User makeUser(Long id, Long teamId) { + User user = new User(); + try { + // User entity doesn't expose Lombok setters for all fields. Set via reflection. + java.lang.reflect.Field idField = User.class.getDeclaredField("id"); + idField.setAccessible(true); + idField.set(user, id); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + if (teamId != null) { + stirling.software.proprietary.model.Team team = + new stirling.software.proprietary.model.Team(); + team.setId(teamId); + user.setTeam(team); + } + return user; + } + + private static HandlerMethod handlerMethodForFakeController() { + try { + Method m = FakeController.class.getDeclaredMethod("handleAuto"); + return new HandlerMethod(new FakeController(), m); + } catch (NoSuchMethodException e) { + throw new RuntimeException(e); + } + } + + private static HandlerMethod handlerMethodForPlain() { + try { + Method m = FakeController.class.getDeclaredMethod("handlePlain"); + return new HandlerMethod(new FakeController(), m); + } catch (NoSuchMethodException e) { + throw new RuntimeException(e); + } + } + + static class FakeController { + @AutoJobPostMapping(value = "/x", resourceWeight = 1) + public void handleAuto() {} + + public void handlePlain() {} + } + + /** Placeholder so AutoCloseable resources flow in some helper methods. */ + @SuppressWarnings("unused") + private static void closeQuietly(AutoCloseable c) { + try { + c.close(); + } catch (Exception ignored) { + // ignored + } + } + + @SuppressWarnings("unused") + private static IOException unused() { + return null; + } +} diff --git a/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygOutputExtractorTest.java b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygOutputExtractorTest.java new file mode 100644 index 000000000..201a0d68e --- /dev/null +++ b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygOutputExtractorTest.java @@ -0,0 +1,120 @@ +package stirling.software.saas.payg.filter; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.zip.ZipEntry; +import java.util.zip.ZipOutputStream; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.util.TempFileManager; +import stirling.software.common.util.TempFileRegistry; + +class PaygOutputExtractorTest { + + private final TempFileManager tempFileManager = + new TempFileManager(new TempFileRegistry(), new ApplicationProperties()); + private final PaygOutputExtractor extractor = new PaygOutputExtractor(tempFileManager); + + @Test + void pdfContentType_returnsBodyPathVerbatim(@TempDir Path tmp) throws IOException { + Path body = tmp.resolve("body.pdf"); + Files.write(body, pdfBytes("hello")); + List out = extractor.extract("application/pdf", body); + assertThat(out).hasSize(1); + assertThat(out.get(0).path()).isEqualTo(body); + assertThat(out.get(0).ownedTempFile()).isNull(); + } + + @Test + void pdfContentType_withParameters_stillMatches(@TempDir Path tmp) throws IOException { + Path body = tmp.resolve("body.pdf"); + Files.write(body, pdfBytes("x")); + List out = + extractor.extract("application/pdf; charset=binary", body); + assertThat(out).hasSize(1); + } + + @Test + void zipContentType_extractsOnlyPdfEntriesWithValidMagicBytes(@TempDir Path tmp) + throws IOException { + Path zip = tmp.resolve("out.zip"); + try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zip))) { + writeEntry(zos, "doc1.pdf", pdfBytes("one")); + writeEntry(zos, "doc2.pdf", pdfBytes("two")); + writeEntry(zos, "fake.pdf", "not a pdf at all".getBytes(StandardCharsets.UTF_8)); + writeEntry(zos, "notes.txt", "ignored".getBytes(StandardCharsets.UTF_8)); + } + + List out = extractor.extract("application/zip", zip); + try { + assertThat(out).hasSize(2); + for (PaygOutputExtractor.ExtractedPdf p : out) { + byte[] head = new byte[5]; + Files.newInputStream(p.path()).read(head); + assertThat(head).startsWith("%PDF-".getBytes(StandardCharsets.UTF_8)); + assertThat(p.ownedTempFile()).isNotNull(); + } + } finally { + for (PaygOutputExtractor.ExtractedPdf p : out) { + p.close(); + } + } + } + + @Test + void nonPdfNonZipContentType_returnsEmpty(@TempDir Path tmp) throws IOException { + Path body = tmp.resolve("body.json"); + Files.writeString(body, "{\"error\":\"bad request\"}"); + assertThat(extractor.extract("application/json", body)).isEmpty(); + assertThat(extractor.extract("text/plain", body)).isEmpty(); + assertThat(extractor.extract(null, body)).isEmpty(); + } + + @Test + void nullBodyPath_returnsEmpty() { + assertThat(extractor.extract("application/pdf", null)).isEmpty(); + } + + @Test + void corruptZip_failsClosedAndReturnsEmpty(@TempDir Path tmp) throws IOException { + Path corrupt = tmp.resolve("garbage.zip"); + Files.write(corrupt, "not a zip file at all".getBytes(StandardCharsets.UTF_8)); + // Should NOT throw — fail-open: empty list, response still serves normally. + List out = extractor.extract("application/zip", corrupt); + assertThat(out).isEmpty(); + } + + @Test + void zipContentType_emptyZip_returnsEmpty(@TempDir Path tmp) throws IOException { + Path zip = tmp.resolve("empty.zip"); + try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zip))) { + // no entries + } + assertThat(extractor.extract("application/zip", zip)).isEmpty(); + } + + private static void writeEntry(ZipOutputStream zos, String name, byte[] data) + throws IOException { + zos.putNextEntry(new ZipEntry(name)); + zos.write(data); + zos.closeEntry(); + } + + private static byte[] pdfBytes(String payload) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + // Minimal "looks like a PDF" — just the magic-byte prefix + filler. The extractor only + // checks magic bytes, not full PDF validity. + out.write("%PDF-1.4\n".getBytes(StandardCharsets.UTF_8)); + out.write(payload.getBytes(StandardCharsets.UTF_8)); + return out.toByteArray(); + } +} diff --git a/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperTest.java b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperTest.java new file mode 100644 index 000000000..64e6ccb57 --- /dev/null +++ b/app/saas/src/test/java/stirling/software/saas/payg/filter/PaygResponseBodyWrapperTest.java @@ -0,0 +1,210 @@ +package stirling.software.saas.payg.filter; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.IOException; +import java.io.PrintWriter; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; + +import org.junit.jupiter.api.Test; +import org.springframework.mock.web.MockHttpServletResponse; + +import stirling.software.common.model.ApplicationProperties; +import stirling.software.common.util.TempFileManager; +import stirling.software.common.util.TempFileRegistry; + +class PaygResponseBodyWrapperTest { + + private final TempFileManager tempFileManager = + new TempFileManager(new TempFileRegistry(), new ApplicationProperties()); + + @Test + void inMemory_smallWrites_clientReceivesAndPathContainsSameBytes() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + wrapper.getOutputStream().write("hello world".getBytes(StandardCharsets.UTF_8)); + wrapper.getOutputStream().flush(); + + assertThat(downstream.getContentAsString()).isEqualTo("hello world"); + assertThat(wrapper.bytesWritten()).isEqualTo(11); + + Path materialised = wrapper.materialisedPath(); + assertThat(materialised).isNotNull(); + assertThat(Files.readString(materialised, StandardCharsets.UTF_8)) + .isEqualTo("hello world"); + } + } + + @Test + void spill_crossThresholdMidChunk_capturesEverything() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 8)) { + // First write: 5 bytes — fits in memory. + wrapper.getOutputStream().write("HELLO".getBytes(StandardCharsets.UTF_8)); + // Second write: 6 bytes — crosses the 8-byte threshold mid-chunk; whole chunk + // ends up on disk per the design (we don't split the crossing chunk). + wrapper.getOutputStream().write(" WORLD".getBytes(StandardCharsets.UTF_8)); + wrapper.getOutputStream().flush(); + + assertThat(downstream.getContentAsString()).isEqualTo("HELLO WORLD"); + assertThat(wrapper.bytesWritten()).isEqualTo(11); + + Path path = wrapper.materialisedPath(); + assertThat(Files.readString(path, StandardCharsets.UTF_8)).isEqualTo("HELLO WORLD"); + } + } + + @Test + void spill_largeSingleWrite_capturedToDisk() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + byte[] payload = new byte[64 * 1024]; // 64 KiB + for (int i = 0; i < payload.length; i++) { + payload[i] = (byte) (i % 251); + } + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + wrapper.getOutputStream().write(payload); + wrapper.getOutputStream().flush(); + + assertThat(downstream.getContentAsByteArray()).isEqualTo(payload); + assertThat(wrapper.bytesWritten()).isEqualTo(payload.length); + + byte[] disk = Files.readAllBytes(wrapper.materialisedPath()); + assertThat(disk).isEqualTo(payload); + } + } + + @Test + void writer_pathTeesThroughToBuffer() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + downstream.setCharacterEncoding("UTF-8"); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + PrintWriter writer = wrapper.getWriter(); + writer.write("a-string-from-writer"); + writer.flush(); + + assertThat(downstream.getContentAsString()).isEqualTo("a-string-from-writer"); + Path path = wrapper.materialisedPath(); + assertThat(Files.readString(path, StandardCharsets.UTF_8)) + .isEqualTo("a-string-from-writer"); + } + } + + @Test + void mixingOutputStreamAndWriter_throwsPerServletSpec() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + wrapper.getOutputStream(); + assertThatThrownBy(wrapper::getWriter).isInstanceOf(IllegalStateException.class); + } + + MockHttpServletResponse downstream2 = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper2 = + new PaygResponseBodyWrapper(downstream2, tempFileManager, 1024)) { + wrapper2.getWriter(); + assertThatThrownBy(wrapper2::getOutputStream).isInstanceOf(IllegalStateException.class); + } + } + + @Test + void noBytesWritten_materialisedPathIsNull() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + // Don't touch the output stream. + assertThat(wrapper.materialisedPath()).isNull(); + assertThat(wrapper.bytesWritten()).isZero(); + } + } + + @Test + void resetBuffer_inMemory_clearsAccumulatedBytes() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + wrapper.getOutputStream().write("draft".getBytes(StandardCharsets.UTF_8)); + wrapper.resetBuffer(); + assertThat(wrapper.bytesWritten()).isZero(); + assertThat(wrapper.materialisedPath()).isNull(); + + wrapper.getOutputStream().write("final".getBytes(StandardCharsets.UTF_8)); + wrapper.getOutputStream().flush(); + assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8)) + .isEqualTo("final"); + } + } + + @Test + void resetBuffer_afterSpill_dropsSpillFileAndStartsFresh() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 4)) { + wrapper.getOutputStream().write("draftover".getBytes(StandardCharsets.UTF_8)); // spill + assertThat(wrapper.bytesWritten()).isEqualTo(9); + wrapper.resetBuffer(); + assertThat(wrapper.bytesWritten()).isZero(); + assertThat(wrapper.materialisedPath()).isNull(); + + // Write again — should land in fresh memory, not the dropped spill. + wrapper.getOutputStream().write("ok".getBytes(StandardCharsets.UTF_8)); + wrapper.getOutputStream().flush(); + assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8)) + .isEqualTo("ok"); + } + } + + @Test + void close_isIdempotent() throws IOException { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 8); + wrapper.getOutputStream().write("more-than-eight-bytes".getBytes(StandardCharsets.UTF_8)); + wrapper.materialisedPath(); // forces flush + wrapper.close(); + wrapper.close(); // second call must not throw + } + + @Test + void materialisedPath_isStableAcrossCalls_whenInMemory() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 1024)) { + wrapper.getOutputStream().write("abc".getBytes(StandardCharsets.UTF_8)); + Path first = wrapper.materialisedPath(); + Path second = wrapper.materialisedPath(); + assertThat(first).isEqualTo(second); + } + } + + @Test + void singleByteWrites_areCorrectlyAccountedAcrossThreshold() throws Exception { + MockHttpServletResponse downstream = new MockHttpServletResponse(); + try (PaygResponseBodyWrapper wrapper = + new PaygResponseBodyWrapper(downstream, tempFileManager, 3)) { + for (int b : "ABCDE".getBytes(StandardCharsets.UTF_8)) { + wrapper.getOutputStream().write(b); + } + wrapper.getOutputStream().flush(); + assertThat(downstream.getContentAsString()).isEqualTo("ABCDE"); + assertThat(wrapper.bytesWritten()).isEqualTo(5); + assertThat(Files.readString(wrapper.materialisedPath(), StandardCharsets.UTF_8)) + .isEqualTo("ABCDE"); + } + } + + @Test + void negativeThreshold_rejected() { + assertThatThrownBy( + () -> + new PaygResponseBodyWrapper( + new MockHttpServletResponse(), tempFileManager, -1)) + .isInstanceOf(IllegalArgumentException.class); + } +} diff --git a/app/saas/src/test/java/stirling/software/saas/payg/job/JobServiceTest.java b/app/saas/src/test/java/stirling/software/saas/payg/job/JobServiceTest.java index 864860e8b..ca7033405 100644 --- a/app/saas/src/test/java/stirling/software/saas/payg/job/JobServiceTest.java +++ b/app/saas/src/test/java/stirling/software/saas/payg/job/JobServiceTest.java @@ -322,10 +322,17 @@ class JobServiceTest { return p; } - /** Test double: programmable detector + records observations. */ + /** + * Test double: programmable detector + records observations. Synthesises one signature per + * path; every path seen by extractSignatures is reverse-mapped so the post-dedupe flow + * (extractSignatures → detect/record by signature) keeps the path-based test assertions working + * unchanged. + */ private static class FakeDetector implements HashLineageDetector { private final Map matches = new HashMap<>(); private final Set observations = new java.util.HashSet<>(); + private final Map + pathBySignature = new HashMap<>(); void willMatch(Path input, LineageMatch match) { matches.put(input, match); @@ -335,14 +342,53 @@ class JobServiceTest { return observations.contains(jobId + "|" + file + "|" + kind); } + private stirling.software.saas.payg.lineage.LineageSignature sigFor(Path file) { + stirling.software.saas.payg.lineage.LineageSignature sig = + new stirling.software.saas.payg.lineage.LineageSignature( + "test", Integer.toHexString(file.toString().hashCode())); + pathBySignature.put(sig, file); + return sig; + } + @Override public Optional detect(Long userId, Path inputFile) { return Optional.ofNullable(matches.get(inputFile)); } + @Override + public Optional detect( + Long userId, Set signatures) { + for (stirling.software.saas.payg.lineage.LineageSignature sig : signatures) { + Path p = pathBySignature.get(sig); + if (p != null && matches.containsKey(p)) { + return Optional.of(matches.get(p)); + } + } + return Optional.empty(); + } + @Override public void record(UUID jobId, Path file, ArtifactKind kind) { observations.add(jobId + "|" + file + "|" + kind); } + + @Override + public void record( + UUID jobId, + Set signatures, + ArtifactKind kind) { + for (stirling.software.saas.payg.lineage.LineageSignature sig : signatures) { + Path p = pathBySignature.get(sig); + if (p != null) { + observations.add(jobId + "|" + p + "|" + kind); + } + } + } + + @Override + public Set extractSignatures( + Path file) { + return Set.of(sigFor(file)); + } } } diff --git a/app/saas/src/test/java/stirling/software/saas/payg/policy/PricingPolicyServiceTest.java b/app/saas/src/test/java/stirling/software/saas/payg/policy/PricingPolicyServiceTest.java index dace97d36..3c552d205 100644 --- a/app/saas/src/test/java/stirling/software/saas/payg/policy/PricingPolicyServiceTest.java +++ b/app/saas/src/test/java/stirling/software/saas/payg/policy/PricingPolicyServiceTest.java @@ -98,6 +98,33 @@ class PricingPolicyServiceTest { .hasMessageContaining("No default pricing_policy row"); } + @Test + void nullTeamId_returnsDefaultPolicyDirectly_noCacheLookup() { + PricingPolicy result = service.getEffectivePolicy(null); + + assertThat(result).isEqualTo(defaultPolicy); + verify(extensionsRepo, never()).findById(any()); + // Team-less users go straight to the default — no cache pollution either. + assertThat(service.cacheSize()).isZero(); + } + + @Test + void nullTeamIdUncached_returnsDefaultPolicyDirectly() { + PricingPolicy result = service.getEffectivePolicyUncached(null); + + assertThat(result).isEqualTo(defaultPolicy); + verify(extensionsRepo, never()).findById(any()); + } + + @Test + void nullTeamId_noDefaultExists_throws() { + when(policyRepo.findFirstByIsDefaultTrue()).thenReturn(Optional.empty()); + + assertThatThrownBy(() -> service.getEffectivePolicy(null)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("No default pricing_policy row"); + } + @Test void secondCallHitsCache_noRepoLookup() { when(extensionsRepo.findById(42L)).thenReturn(Optional.empty());