Janady13/missionmatch / internal/platform/federated/federated.go

1package federated 2 3// federated.go defines the cross tenant federated aggregator boundary (Clip, Noise, Emit) and the HARD 4// server side spend gate that backs Charge with the dp_budget_ledger. Everything here is DEAD CODE unless 5// MM_FEDERATED_DP_ENABLED is set: the constructor for an active aggregator refuses to build while the flag is 6// false, so there is no live emission path today, and the noise sampler is an activation blocker that returns 7// an error rather than sampling. The package has no non test caller; cmd/teach-aggregate in Phase 5 will be 8// the gated caller. 9// 10// THE GATE IS NEVER CLIENT TRUSTED. Charge reads the PERSISTED accumulated RDP for a (model_version, 11// tenant_id) budget key from the ledger, composes the candidate release into it, reconverts to epsilon at the 12// configured delta, and REFUSES the release (writing nothing) if it would exceed MM_FEDERATED_DP_EPSILON. 13// Only an in budget release persists a new ledger row, and it does so through the MIGRATOR pool, because 14// dp_budget_ledger is a global migrator owned table the app role can only SELECT. No client supplied epsilon, 15// no client supplied accumulation, ever enters the decision. 16// 17// AUDIT RESPONSIBILITY (documented, load bearing). A successful, in budget release is audited as 18// federated.dp.released through audit.Writer.System (NULL org, a cross tenant global event) BY THE GATE 19// ITSELF, inside Charge, after the ledger write succeeds. The gate owns the audit so a release can never be 20// persisted without its audit record. The event carries the model version, the contributing tenant count, 21// the converted epsilon and delta, and the worst case accumulated RDP, never any PII, candidate id, or raw 22// per tenant vector. If the audit cannot be written the release fails closed. 23 24import ( 25 "context" 26 "encoding/json" 27 "errors" 28 "fmt" 29 "math" 30 "strconv" 31 32 "github.com/google/uuid" 33 "github.com/jackc/pgx/v5" 34 "github.com/shopspring/decimal" 35 36 "github.com/missionworkforce/missionmatch/internal/platform/audit" 37 "github.com/missionworkforce/missionmatch/internal/platform/config" 38 "github.com/missionworkforce/missionmatch/internal/platform/database/dbgen" 39) 40 41// ErrDisabled is returned by the active aggregator constructor when MM_FEDERATED_DP_ENABLED is false. It is 42// the guard that keeps the whole emission path dead in tenant LIVE: an active aggregator cannot even be 43// constructed while the flag is off. 44var ErrDisabled = errors.New("federated: MM_FEDERATED_DP_ENABLED is false, the active aggregator is disabled") 45 46// Aggregator is the cross tenant federated aggregator boundary. It is the seam for the central DP v1 47// implementation and its TEE or secure aggregation upgrade. The three steps are the DP-FedEmb pattern: Clip 48// each per tenant aggregate to a fixed L2 sensitivity, Noise the clipped aggregate with calibrated Gaussian 49// noise, then Emit the noised aggregate through the spend gate (which charges the budget and persists the 50// release). Today no implementation can be constructed while the flag is off, so none of these run. 51type Aggregator interface { 52 // Clip bounds a per tenant aggregate vector to L2 norm clipNorm (the sensitivity C), returning the 53 // clipped vector. Clipping is the deterministic step that establishes the per tenant sensitivity the 54 // accountant calibrates noise against. It does not require the secure sampler and is safe to compute. 55 Clip(vec []float64, clipNorm float64) []float64 56 // Noise adds calibrated Gaussian noise to a clipped vector. THIS IS THE ACTIVATION BLOCKER: it must use a 57 // Mironov safe discrete or snapped Gaussian sampler behind the FIPS boundary, which is not implemented, so 58 // the only implementation returns ErrNotImplemented. The naive float Gaussian is intentionally absent so 59 // it can never run on an activated path. 60 Noise(vec []float64, noiseSigma float64) ([]float64, error) 61 // Emit charges the release against the budget through the gate and, only if in budget, persists the 62 // release to the ledger and audits it. It returns the epsilon at the configured delta after the release, 63 // or ErrBudgetExceeded if the budget is spent. 64 Emit(ctx context.Context, release Release) (epsilonNow float64, err error) 65} 66 67// Release is one cross tenant differential privacy release to emit. ModelVersion is the shared model version 68// the release contributes to (pinned from engine.ModelVersion by the caller, never a literal). TenantID is 69// nil for a cross tenant cohort release, which has no single owning tenant. ContributingTenantCount is the 70// number of live non test distinct organizations in the cohort, sourced from the production flag, and is 71// recorded so the (epsilon, delta) claim is always reported with its cohort size. Gaussian carries the noise 72// multiplier inputs the accountant charges. The noised aggregate vector itself is NOT carried here because 73// the gate accounts and audits counts only; the vector lives in the shared model artifact, never the ledger 74// or the audit log. 75type Release struct { 76 ModelVersion string 77 TenantID *uuid.UUID 78 ContributingTenantCount int 79 Gaussian GaussianRelease 80} 81 82// LedgerStore is the migrator pool seam the gate reads the running budget through and writes a release 83// through. The live implementation is backed by the MIGRATOR pool (the app role has SELECT only on 84// dp_budget_ledger); a test substitutes a fake in memory store. The package depends on this interface, not 85// the concrete dbgen.DBTX, so the gate is testable without a database and the migrator pool is supplied by 86// the caller (the cmd/teach-aggregate two pool wiring in Phase 5), never the app role. 87type LedgerStore interface { 88 // LatestForKey returns the most recent ledger row for the (modelVersion, tenantID) budget key, carrying 89 // the running accumulated RDP composed so far. It returns ok false when the budget is fresh (no prior 90 // release), in which case the accumulation starts at zero. 91 LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (row dbgen.DpBudgetLedger, ok bool, err error) 92 // Insert persists one in budget release row and returns it. The gate calls this ONLY after it has 93 // confirmed the release stays within budget, so a persisted row is by construction an in budget release. 94 Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error) 95} 96 97// systemAuditor is the off pool audit seam the gate writes federated.dp.released through. *audit.Writer 98// satisfies it via System, so production wiring passes a real writer unchanged; a test substitutes a fake to 99// assert the release is audited without a database. The gate depends on this interface, not the concrete 100// writer, for the same testability reason the ledger uses LedgerStore. 101type systemAuditor interface { 102 System(ctx context.Context, e audit.Event) error 103} 104 105// SpendGate is the hard server side differential privacy spend gate. It is the ledger backed Accountant: each 106// Charge reads the persisted accumulated RDP for the release's budget key, composes the candidate release, 107// reconverts to epsilon at the configured delta, refuses (writing nothing) if over the budget, and otherwise 108// persists the new release row through the migrator pool and audits it as federated.dp.released. The gate is 109// the only path that writes the ledger, and it is never client trusted: the accumulation is read from the 110// persisted ledger, never supplied by a caller. 111type SpendGate struct { 112 store LedgerStore // the migrator pool seam (dp_budget_ledger reads and writes) 113 auditW systemAuditor // audit.Writer.System for the federated.dp.released event (NULL org, global) 114 epsilon float64 // the epsilon budget target, from config (MM_FEDERATED_DP_EPSILON), never a constant 115 delta float64 // the delta target, from config (MM_FEDERATED_DP_DELTA), never a constant 116 enabled bool // MM_FEDERATED_DP_ENABLED; false makes every Charge a no op refusal (dead code guard) 117} 118 119// NewSpendGate builds the hard spend gate from configuration. The flag, epsilon, and delta come from the 120// Federated config block, never from code constants, so a privacy officer signs off and tunes them without a 121// rebuild. store is the MIGRATOR pool seam (never the app pool, which cannot write the global ledger) and 122// auditW writes the federated.dp.released system event (pass a real *audit.Writer in production). The gate is 123// built regardless of the flag, but when the flag is false every Charge refuses with ErrDisabled and writes 124// nothing, so the emission path stays dead. A caller that wants a hard construction guard should use 125// NewActiveAggregator, which refuses to build at all while the flag is off. 126func NewSpendGate(cfg config.Federated, store LedgerStore, auditW *audit.Writer) *SpendGate { 127 return &SpendGate{ 128 store: store, 129 auditW: auditW, 130 epsilon: cfg.DPEpsilon, 131 delta: cfg.DPDelta, 132 enabled: cfg.DPEnabled, 133 } 134} 135 136// chargeRelease is the shared budget decision for one release. It loads the persisted accumulation, composes 137// the candidate release, and returns the post release accountant plus the per order accumulation, or 138// ErrBudgetExceeded. It performs NO write and NO audit, so it is the pure decision the Emit path and any 139// dry run share. The flag guard lives in Charge and Emit, not here. 140func (g *SpendGate) chargeRelease(ctx context.Context, release Release) (acct *rdpAccountant, epsilonNow float64, err error) { 141 acct, err = g.loadAccountant(ctx, release.ModelVersion, release.TenantID) 142 if err != nil { 143 return nil, 0, err 144 } 145 epsilonNow, err = acct.Charge(release.Gaussian) 146 if err != nil { 147 return nil, 0, err 148 } 149 return acct, epsilonNow, nil 150} 151 152// loadAccountant builds an RDP accountant seeded from the persisted accumulation for the (modelVersion, 153// tenantID) budget key. A fresh budget (no prior release) starts at zero; a resumed budget decodes the per 154// order accumulation out of the stored rdp_orders jsonb so composition continues across process restarts. 155func (g *SpendGate) loadAccountant(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (*rdpAccountant, error) { 156 prior, ok, err := g.store.LatestForKey(ctx, modelVersion, tenantID) 157 if err != nil { 158 return nil, fmt.Errorf("federated: read budget ledger: %w", err) 159 } 160 acct, err := newRDPAccountant(g.epsilon, g.delta, nil) 161 if err != nil { 162 return nil, err 163 } 164 if ok { 165 seed, err := decodeRDPOrders(prior.RdpOrders, acct.orders) 166 if err != nil { 167 return nil, err 168 } 169 acct, err = newRDPAccountant(g.epsilon, g.delta, seed) 170 if err != nil { 171 return nil, err 172 } 173 } 174 return acct, nil 175} 176 177// Charge is the hard spend gate. While the flag is off it refuses with ErrDisabled and writes nothing (the 178// dead code guard). While the flag is on it reads the persisted accumulated RDP for the budget key, composes 179// the candidate release, refuses with ErrBudgetExceeded (writing nothing) if over budget, and otherwise 180// persists the new release row through the migrator pool and audits federated.dp.released through 181// audit.Writer.System. The persisted accumulated_rdp and epsilon_at_delta are shopspring decimals so the 182// audit record is bit stable; the grid math ran in float64. A persisted row is by construction an in budget, 183// already audited release. 184func (g *SpendGate) Charge(ctx context.Context, release Release) (epsilonNow float64, err error) { 185 if !g.enabled { 186 return 0, ErrDisabled 187 } 188 acct, epsilonNow, err := g.chargeRelease(ctx, release) 189 if err != nil { 190 return 0, err 191 } 192 193 ordersJSON, err := encodeRDPOrders(acct.accumulatedRDPDecimal()) 194 if err != nil { 195 return 0, err 196 } 197 row, err := g.store.Insert(ctx, dbgen.InsertDpBudgetLedgerParams{ 198 ModelVersion: release.ModelVersion, 199 TenantID: release.TenantID, 200 RdpOrders: ordersJSON, 201 AccumulatedRdp: acct.worstRDPDecimal(), 202 EpsilonAtDelta: acct.epsilonDecimal(), 203 Delta: decimal.NewFromFloat(g.delta), 204 ContributingTenantCount: int32(release.ContributingTenantCount), 205 }) 206 if err != nil { 207 return 0, fmt.Errorf("federated: persist budget release: %w", err) 208 } 209 210 if err := g.auditRelease(ctx, row); err != nil { 211 // Fail closed: a release the audit cannot record is treated as failed even though the row persisted, 212 // because the dp_budget_ledger is append only and the release is already counted against the budget. 213 return 0, fmt.Errorf("federated: audit dp release: %w", err) 214 } 215 return epsilonNow, nil 216} 217 218// Remaining returns the epsilon headroom left on a budget key. It is a read only budget check (no write, no 219// audit) so it is safe to call regardless of the flag, but it returns zero headroom while the flag is off so 220// nothing reads a live budget on a disabled deploy. tenantID is nil for the cross tenant cohort key. 221func (g *SpendGate) Remaining(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (float64, error) { 222 if !g.enabled { 223 return 0, nil 224 } 225 acct, err := g.loadAccountant(ctx, modelVersion, tenantID) 226 if err != nil { 227 return 0, err 228 } 229 return acct.Remaining(), nil 230} 231 232// auditRelease writes the federated.dp.released system event (NULL org, a cross tenant global event). It 233// carries the model version, the contributing tenant count, the converted epsilon and delta, and the worst 234// case accumulated RDP, never any PII, candidate id, or raw vector. 235func (g *SpendGate) auditRelease(ctx context.Context, row dbgen.DpBudgetLedger) error { 236 payload := map[string]any{ 237 "model_version": row.ModelVersion, 238 "contributing_tenant_count": row.ContributingTenantCount, 239 "epsilon_at_delta": row.EpsilonAtDelta.String(), 240 "delta": row.Delta.String(), 241 "accumulated_rdp": row.AccumulatedRdp.String(), 242 } 243 body, err := json.Marshal(payload) 244 if err != nil { 245 return err 246 } 247 return g.auditW.System(ctx, audit.Event{ 248 Action: "federated.dp.released", 249 EntityType: "dp_budget_ledger", 250 EntityID: row.ID.String(), 251 NewValues: body, 252 }) 253} 254 255// activeAggregator is the only Aggregator implementation. It is constructed ONLY when the flag is on (see 256// NewActiveAggregator), so its emission path is unreachable in tenant LIVE. Clip is deterministic and safe; 257// Noise is the activation blocker; Emit routes through the spend gate. 258type activeAggregator struct { 259 gate *SpendGate 260} 261 262// NewActiveAggregator is the HARD construction guard for the emission path. It refuses to build an active 263// aggregator (returning ErrDisabled) while MM_FEDERATED_DP_ENABLED is false, so no live emission path can 264// even be assembled in tenant LIVE. When the flag is on it returns an aggregator whose Emit routes through 265// the spend gate; its Noise still returns ErrNotImplemented until the Mironov safe sampler is written, so an 266// activated path cannot actually sample noise either, only assemble the pieces. 267func NewActiveAggregator(cfg config.Federated, gate *SpendGate) (Aggregator, error) { 268 if !cfg.DPEnabled { 269 return nil, ErrDisabled 270 } 271 return &activeAggregator{gate: gate}, nil 272} 273 274// Clip bounds vec to L2 norm clipNorm. It is the deterministic sensitivity step and does not sample noise, so 275// it is safe to compute even on a flag off path (no constructor returns this type while the flag is off, but 276// the method itself touches no randomness). A non positive clipNorm or a vector already inside the ball is 277// returned unchanged. 278func (a *activeAggregator) Clip(vec []float64, clipNorm float64) []float64 { 279 if clipNorm <= 0 { 280 return vec 281 } 282 var sumSq float64 283 for _, v := range vec { 284 sumSq += v * v 285 } 286 norm := math.Sqrt(sumSq) 287 if norm <= clipNorm || norm == 0 { 288 return vec 289 } 290 scale := clipNorm / norm 291 out := make([]float64, len(vec)) 292 for i, v := range vec { 293 out[i] = v * scale 294 } 295 return out 296} 297 298// Noise is the HARD activation blocker. It returns ErrNotImplemented and samples nothing. The Mironov safe 299// discrete or snapped Gaussian sampler (the only acceptable activated sampler, behind the FIPS boundary) is a 300// documented TODO; the naive float Gaussian is intentionally absent from the whole package so it can never 301// run. See docs/AI_RESEARCH.md gotcha 7 (floating point DP attacks, Mironov 2012). 302// 303// TODO activation blocker: implement a Mironov safe discrete or snapped Gaussian sampler that draws only from 304// crypto/rand under GOFIPS140 (no new crypto dependency, in module primitives only) before the flag is ever 305// flipped on. Until then this returns an error and no noise is added. 306func (a *activeAggregator) Noise(vec []float64, noiseSigma float64) ([]float64, error) { 307 return nil, ErrNotImplemented 308} 309 310// Emit routes the release through the hard spend gate, which charges the budget, persists an in budget 311// release, and audits it. It returns the epsilon at the configured delta after the release or 312// ErrBudgetExceeded. 313func (a *activeAggregator) Emit(ctx context.Context, release Release) (float64, error) { 314 return a.gate.Charge(ctx, release) 315} 316 317// migratorLedgerStore backs LedgerStore with a migrator scoped dbgen.DBTX (the migrator pool). It is a thin 318// adapter so the gate depends on the interface, not the concrete DBTX, and a test substitutes a fake. The 319// caller passes the MIGRATOR pool here, never the app pool, because the app role cannot write the global 320// dp_budget_ledger. 321type migratorLedgerStore struct{ db dbgen.DBTX } 322 323// NewMigratorLedgerStore wraps a migrator scoped dbgen.DBTX as the ledger store. The caller passes the 324// migrator pool, never the app pool. 325func NewMigratorLedgerStore(migrator dbgen.DBTX) LedgerStore { 326 return migratorLedgerStore{db: migrator} 327} 328 329func (s migratorLedgerStore) LatestForKey(ctx context.Context, modelVersion string, tenantID *uuid.UUID) (dbgen.DpBudgetLedger, bool, error) { 330 row, err := dbgen.New(s.db).GetLatestDpBudgetLedger(ctx, dbgen.GetLatestDpBudgetLedgerParams{ 331 ModelVersion: modelVersion, 332 TenantID: tenantID, 333 }) 334 if err != nil { 335 if errors.Is(err, pgx.ErrNoRows) { 336 return dbgen.DpBudgetLedger{}, false, nil 337 } 338 return dbgen.DpBudgetLedger{}, false, err 339 } 340 return row, true, nil 341} 342 343func (s migratorLedgerStore) Insert(ctx context.Context, arg dbgen.InsertDpBudgetLedgerParams) (dbgen.DpBudgetLedger, error) { 344 return dbgen.New(s.db).InsertDpBudgetLedger(ctx, arg) 345} 346 347// encodeRDPOrders serializes the per order accumulated RDP decimal map to the rdp_orders jsonb, keyed by the 348// Renyi order as a string so the json object is stable and the audit record reproduces exactly. 349func encodeRDPOrders(m map[float64]decimal.Decimal) ([]byte, error) { 350 obj := make(map[string]string, len(m)) 351 for order, v := range m { 352 obj[strconv.FormatFloat(order, 'g', -1, 64)] = v.String() 353 } 354 return json.Marshal(obj) 355} 356 357// decodeRDPOrders parses the rdp_orders jsonb back into a per order accumulated RDP slice index aligned with 358// the grid, so composition resumes from a persisted ledger row. An order in the grid that is absent from the 359// stored map seeds to zero (a fresh order); the float64 grid math then resumes from the persisted values. 360func decodeRDPOrders(raw []byte, orders []float64) ([]float64, error) { 361 out := make([]float64, len(orders)) 362 if len(raw) == 0 { 363 return out, nil 364 } 365 var obj map[string]string 366 if err := json.Unmarshal(raw, &obj); err != nil { 367 return nil, fmt.Errorf("federated: decode rdp_orders: %w", err) 368 } 369 for i, order := range orders { 370 key := strconv.FormatFloat(order, 'g', -1, 64) 371 s, ok := obj[key] 372 if !ok { 373 continue 374 } 375 v, err := strconv.ParseFloat(s, 64) 376 if err != nil { 377 return nil, fmt.Errorf("federated: parse rdp_orders value: %w", err) 378 } 379 out[i] = v 380 } 381 return out, nil 382} 383