fix(snapshot): harden restore and final flush
- commit
- 55b4f0299a40e0ff1246416d00e341472bac4d56
- author
- brice · · 2026-07-12 16:56 UTC
- parent
- 6b789ff
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
6 files · +342 -53
M internal/usecase/daemon/daemon.go +2 -0
| @@ -109,6 +109,8 @@ type Daemon struct { | |||
| 109 | 109 | snapshotWorkerCtx context.Context | |
| 110 | 110 | snapshotWorkerCancel context.CancelFunc | |
| 111 | 111 | snapshotWorkerDone chan struct{} | |
| 112 | + | snapshotWorkerFlush chan struct{} | |
| 113 | + | snapshotWorkerClosing bool | |
| 112 | 114 | snapshotWorkerInFlight *snapshotCapture | |
| 113 | 115 | restoreDone chan struct{} | |
| 114 | 116 | restoreOnce sync.Once | |
| @@ -109,6 +109,8 @@ type Daemon struct { | |||||
| 109 | snapshotWorkerCtx context.Context | 109 | snapshotWorkerCtx context.Context | ||
| 110 | snapshotWorkerCancel context.CancelFunc | 110 | snapshotWorkerCancel context.CancelFunc | ||
| 111 | snapshotWorkerDone chan struct{} | 111 | snapshotWorkerDone chan struct{} | ||
| 112 | + | snapshotWorkerFlush chan struct{} | |||
| 113 | + | snapshotWorkerClosing bool | |||
| 112 | snapshotWorkerInFlight *snapshotCapture | 114 | snapshotWorkerInFlight *snapshotCapture | ||
| 113 | restoreDone chan struct{} | 115 | restoreDone chan struct{} | ||
| 114 | restoreOnce sync.Once | 116 | restoreOnce sync.Once | ||
M internal/usecase/daemon/snapshot.go +74 -21
| @@ -47,7 +47,10 @@ type snapshotCapturePane struct { | |||
| 47 | 47 | process *snapcodec.Process | |
| 48 | 48 | } | |
| 49 | 49 | ||
| 50 | - | const snapshotInterval = 30 * time.Second | |
| 50 | + | const ( | |
| 51 | + | snapshotInterval = 30 * time.Second | |
| 52 | + | snapshotFinalFlushTimeout = time.Second | |
| 53 | + | ) | |
| 51 | 54 | ||
| 52 | 55 | func (d *Daemon) closeRestoreDone() { | |
| 53 | 56 | d.restoreOnce.Do(func() { close(d.restoreDone) }) | |
| @@ -183,6 +186,9 @@ func (d *Daemon) restoreSession(ctx context.Context, snap snapcodec.Session) err | |||
| 183 | 186 | p.ctx, p.cancel = context.WithCancel(tb.ctx) | |
| 184 | 187 | p.rect = contentRect | |
| 185 | 188 | if err := restorePaneTerminal(p, paneSnap); err != nil { | |
| 189 | + | // p is not in tb.panes yet, so closeOpened cannot reach it. | |
| 190 | + | p.cancel() | |
| 191 | + | _ = pty.Close() | |
| 186 | 192 | closeOpened() | |
| 187 | 193 | return err | |
| 188 | 194 | } | |
| @@ -320,7 +326,7 @@ func (d *Daemon) scheduleSnapshot(sess *session) bool { | |||
| 320 | 326 | func (d *Daemon) enqueueSnapshotCapture(capture *snapshotCapture) bool { | |
| 321 | 327 | d.snapshotWorkerMu.Lock() | |
| 322 | 328 | defer d.snapshotWorkerMu.Unlock() | |
| 323 | - | if d.snapshotWorkerCancel == nil || d.snapshotWorkerCtx == nil || d.snapshotWorkerCtx.Err() != nil { | |
| 329 | + | if d.snapshotWorkerClosing || d.snapshotWorkerCancel == nil || d.snapshotWorkerCtx == nil || d.snapshotWorkerCtx.Err() != nil { | |
| 324 | 330 | return false | |
| 325 | 331 | } | |
| 326 | 332 | select { | |
| @@ -457,36 +463,61 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||
| 457 | 463 | d.snapshotWorkerMu.Unlock() | |
| 458 | 464 | return | |
| 459 | 465 | } | |
| 460 | - | workerCtx, cancel := context.WithCancel(ctx) | |
| 466 | + | // Shutdown owns the worker lifetime so it can flush captures after Serve | |
| 467 | + | // cancels its parent context. The worker is always stopped explicitly. | |
| 468 | + | workerCtx, cancel := context.WithCancel(context.Background()) | |
| 461 | 469 | d.snapshotWorkerID++ | |
| 462 | 470 | workerID := d.snapshotWorkerID | |
| 463 | 471 | d.snapshotWorkerCtx = workerCtx | |
| 464 | 472 | d.snapshotWorkerCancel = cancel | |
| 465 | 473 | d.snapshotWorkerDone = make(chan struct{}) | |
| 474 | + | d.snapshotWorkerFlush = make(chan struct{}) | |
| 475 | + | d.snapshotWorkerClosing = false | |
| 466 | 476 | d.snapshotWorkerInFlight = nil | |
| 467 | 477 | done := d.snapshotWorkerDone | |
| 478 | + | flush := d.snapshotWorkerFlush | |
| 468 | 479 | d.snapshotWorkerMu.Unlock() | |
| 469 | 480 | go func() { | |
| 470 | 481 | defer close(done) | |
| 482 | + | write := func(capture *snapshotCapture) bool { | |
| 483 | + | if workerCtx.Err() != nil || !d.setSnapshotWorkerInFlight(workerID, capture) { | |
| 484 | + | d.finishSnapshotCapture(capture, false) | |
| 485 | + | return false | |
| 486 | + | } | |
| 487 | + | data, err := d.encodeSnapshotCapture(capture) | |
| 488 | + | if err == nil && workerCtx.Err() == nil { | |
| 489 | + | err = d.snaps.Write(capture.name, data) | |
| 490 | + | } | |
| 491 | + | d.clearSnapshotWorkerInFlight(workerID, capture) | |
| 492 | + | if err != nil && workerCtx.Err() == nil { | |
| 493 | + | d.log.Warn("writing session snapshot failed", "err", err, "session", capture.name) | |
| 494 | + | } | |
| 495 | + | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | |
| 496 | + | return workerCtx.Err() == nil | |
| 497 | + | } | |
| 498 | + | flushQueued := func() { | |
| 499 | + | for { | |
| 500 | + | select { | |
| 501 | + | case capture := <-d.snapshotJobs: | |
| 502 | + | if !write(capture) { | |
| 503 | + | return | |
| 504 | + | } | |
| 505 | + | default: | |
| 506 | + | return | |
| 507 | + | } | |
| 508 | + | } | |
| 509 | + | } | |
| 471 | 510 | for { | |
| 472 | 511 | select { | |
| 473 | 512 | case <-workerCtx.Done(): | |
| 474 | 513 | return | |
| 514 | + | case <-flush: | |
| 515 | + | flushQueued() | |
| 516 | + | return | |
| 475 | 517 | case capture := <-d.snapshotJobs: | |
| 476 | - | // A queued capture must not begin persistence after shutdown. | |
| 477 | - | if workerCtx.Err() != nil || !d.setSnapshotWorkerInFlight(workerID, capture) { | |
| 478 | - | d.finishSnapshotCapture(capture, false) | |
| 518 | + | if !write(capture) { | |
| 479 | 519 | return | |
| 480 | 520 | } | |
| 481 | - | data, err := d.encodeSnapshotCapture(capture) | |
| 482 | - | if err == nil && workerCtx.Err() == nil { | |
| 483 | - | err = d.snaps.Write(capture.name, data) | |
| 484 | - | } | |
| 485 | - | d.clearSnapshotWorkerInFlight(workerID, capture) | |
| 486 | - | if err != nil && workerCtx.Err() == nil { | |
| 487 | - | d.log.Warn("writing session snapshot failed", "err", err, "session", capture.name) | |
| 488 | - | } | |
| 489 | - | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | |
| 490 | 521 | } | |
| 491 | 522 | } | |
| 492 | 523 | }() | |
| @@ -516,19 +547,39 @@ func (d *Daemon) stopSnapshotEncodeWorker() { | |||
| 516 | 547 | } | |
| 517 | 548 | d.snapshotWorkerMu.Lock() | |
| 518 | 549 | cancel := d.snapshotWorkerCancel | |
| 519 | - | if cancel == nil { | |
| 550 | + | if cancel == nil || d.snapshotWorkerClosing { | |
| 520 | 551 | d.snapshotWorkerMu.Unlock() | |
| 521 | 552 | return | |
| 522 | 553 | } | |
| 523 | - | // The queue stays open, so producers can safely race teardown. Detach an | |
| 524 | - | // in-flight uncancellable Write instead of waiting for it; the process owns | |
| 525 | - | // that goroutine until the store returns. | |
| 554 | + | // Stop accepting producers, then let the single worker drain its bounded | |
| 555 | + | // queue. A non-context-aware store may still block indefinitely, so the | |
| 556 | + | // injected clock bounds how long teardown waits before detaching it. | |
| 557 | + | d.snapshotWorkerClosing = true | |
| 558 | + | flush := d.snapshotWorkerFlush | |
| 559 | + | done := d.snapshotWorkerDone | |
| 560 | + | close(flush) | |
| 561 | + | timer := d.clock.NewTimer(snapshotFinalFlushTimeout) | |
| 562 | + | d.snapshotWorkerMu.Unlock() | |
| 563 | + | ||
| 564 | + | select { | |
| 565 | + | case <-done: | |
| 566 | + | timer.Stop() | |
| 567 | + | d.finishStoppedSnapshotWorker(false) | |
| 568 | + | case <-timer.C(): | |
| 569 | + | cancel() | |
| 570 | + | d.finishStoppedSnapshotWorker(true) | |
| 571 | + | } | |
| 572 | + | } | |
| 573 | + | ||
| 574 | + | func (d *Daemon) finishStoppedSnapshotWorker(abandoned bool) { | |
| 575 | + | d.snapshotWorkerMu.Lock() | |
| 526 | 576 | inFlight := d.snapshotWorkerInFlight | |
| 527 | 577 | d.snapshotWorkerCtx = nil | |
| 528 | 578 | d.snapshotWorkerCancel = nil | |
| 529 | 579 | d.snapshotWorkerDone = nil | |
| 580 | + | d.snapshotWorkerFlush = nil | |
| 581 | + | d.snapshotWorkerClosing = false | |
| 530 | 582 | d.snapshotWorkerInFlight = nil | |
| 531 | - | cancel() | |
| 532 | 583 | queued := make([]*snapshotCapture, 0, len(d.snapshotJobs)) | |
| 533 | 584 | for { | |
| 534 | 585 | select { | |
| @@ -536,7 +587,9 @@ func (d *Daemon) stopSnapshotEncodeWorker() { | |||
| 536 | 587 | queued = append(queued, capture) | |
| 537 | 588 | default: | |
| 538 | 589 | d.snapshotWorkerMu.Unlock() | |
| 539 | - | d.finishSnapshotCapture(inFlight, false) | |
| 590 | + | if abandoned { | |
| 591 | + | d.finishSnapshotCapture(inFlight, false) | |
| 592 | + | } | |
| 540 | 593 | for _, capture := range queued { | |
| 541 | 594 | d.finishSnapshotCapture(capture, false) | |
| 542 | 595 | } | |
| @@ -47,7 +47,10 @@ type snapshotCapturePane struct { | |||||
| 47 | process *snapcodec.Process | 47 | process *snapcodec.Process | ||
| 48 | } | 48 | } | ||
| 49 | 49 | ||||
| 50 | - | const snapshotInterval = 30 * time.Second | 50 | + | const ( |
| 51 | + | snapshotInterval = 30 * time.Second | |||
| 52 | + | snapshotFinalFlushTimeout = time.Second | |||
| 53 | + | ) | |||
| 51 | 54 | ||||
| 52 | func (d *Daemon) closeRestoreDone() { | 55 | func (d *Daemon) closeRestoreDone() { | ||
| 53 | d.restoreOnce.Do(func() { close(d.restoreDone) }) | 56 | d.restoreOnce.Do(func() { close(d.restoreDone) }) | ||
| @@ -183,6 +186,9 @@ func (d *Daemon) restoreSession(ctx context.Context, snap snapcodec.Session) err | |||||
| 183 | p.ctx, p.cancel = context.WithCancel(tb.ctx) | 186 | p.ctx, p.cancel = context.WithCancel(tb.ctx) | ||
| 184 | p.rect = contentRect | 187 | p.rect = contentRect | ||
| 185 | if err := restorePaneTerminal(p, paneSnap); err != nil { | 188 | if err := restorePaneTerminal(p, paneSnap); err != nil { | ||
| 189 | + | // p is not in tb.panes yet, so closeOpened cannot reach it. | |||
| 190 | + | p.cancel() | |||
| 191 | + | _ = pty.Close() | |||
| 186 | closeOpened() | 192 | closeOpened() | ||
| 187 | return err | 193 | return err | ||
| 188 | } | 194 | } | ||
| @@ -320,7 +326,7 @@ func (d *Daemon) scheduleSnapshot(sess *session) bool { | |||||
| 320 | func (d *Daemon) enqueueSnapshotCapture(capture *snapshotCapture) bool { | 326 | func (d *Daemon) enqueueSnapshotCapture(capture *snapshotCapture) bool { | ||
| 321 | d.snapshotWorkerMu.Lock() | 327 | d.snapshotWorkerMu.Lock() | ||
| 322 | defer d.snapshotWorkerMu.Unlock() | 328 | defer d.snapshotWorkerMu.Unlock() | ||
| 323 | - | if d.snapshotWorkerCancel == nil || d.snapshotWorkerCtx == nil || d.snapshotWorkerCtx.Err() != nil { | 329 | + | if d.snapshotWorkerClosing || d.snapshotWorkerCancel == nil || d.snapshotWorkerCtx == nil || d.snapshotWorkerCtx.Err() != nil { |
| 324 | return false | 330 | return false | ||
| 325 | } | 331 | } | ||
| 326 | select { | 332 | select { | ||
| @@ -457,36 +463,61 @@ func (d *Daemon) startSnapshotEncodeWorker(ctx context.Context) { | |||||
| 457 | d.snapshotWorkerMu.Unlock() | 463 | d.snapshotWorkerMu.Unlock() | ||
| 458 | return | 464 | return | ||
| 459 | } | 465 | } | ||
| 460 | - | workerCtx, cancel := context.WithCancel(ctx) | 466 | + | // Shutdown owns the worker lifetime so it can flush captures after Serve |
| 467 | + | // cancels its parent context. The worker is always stopped explicitly. | |||
| 468 | + | workerCtx, cancel := context.WithCancel(context.Background()) | |||
| 461 | d.snapshotWorkerID++ | 469 | d.snapshotWorkerID++ | ||
| 462 | workerID := d.snapshotWorkerID | 470 | workerID := d.snapshotWorkerID | ||
| 463 | d.snapshotWorkerCtx = workerCtx | 471 | d.snapshotWorkerCtx = workerCtx | ||
| 464 | d.snapshotWorkerCancel = cancel | 472 | d.snapshotWorkerCancel = cancel | ||
| 465 | d.snapshotWorkerDone = make(chan struct{}) | 473 | d.snapshotWorkerDone = make(chan struct{}) | ||
| 474 | + | d.snapshotWorkerFlush = make(chan struct{}) | |||
| 475 | + | d.snapshotWorkerClosing = false | |||
| 466 | d.snapshotWorkerInFlight = nil | 476 | d.snapshotWorkerInFlight = nil | ||
| 467 | done := d.snapshotWorkerDone | 477 | done := d.snapshotWorkerDone | ||
| 478 | + | flush := d.snapshotWorkerFlush | |||
| 468 | d.snapshotWorkerMu.Unlock() | 479 | d.snapshotWorkerMu.Unlock() | ||
| 469 | go func() { | 480 | go func() { | ||
| 470 | defer close(done) | 481 | defer close(done) | ||
| 482 | + | write := func(capture *snapshotCapture) bool { | |||
| 483 | + | if workerCtx.Err() != nil || !d.setSnapshotWorkerInFlight(workerID, capture) { | |||
| 484 | + | d.finishSnapshotCapture(capture, false) | |||
| 485 | + | return false | |||
| 486 | + | } | |||
| 487 | + | data, err := d.encodeSnapshotCapture(capture) | |||
| 488 | + | if err == nil && workerCtx.Err() == nil { | |||
| 489 | + | err = d.snaps.Write(capture.name, data) | |||
| 490 | + | } | |||
| 491 | + | d.clearSnapshotWorkerInFlight(workerID, capture) | |||
| 492 | + | if err != nil && workerCtx.Err() == nil { | |||
| 493 | + | d.log.Warn("writing session snapshot failed", "err", err, "session", capture.name) | |||
| 494 | + | } | |||
| 495 | + | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | |||
| 496 | + | return workerCtx.Err() == nil | |||
| 497 | + | } | |||
| 498 | + | flushQueued := func() { | |||
| 499 | + | for { | |||
| 500 | + | select { | |||
| 501 | + | case capture := <-d.snapshotJobs: | |||
| 502 | + | if !write(capture) { | |||
| 503 | + | return | |||
| 504 | + | } | |||
| 505 | + | default: | |||
| 506 | + | return | |||
| 507 | + | } | |||
| 508 | + | } | |||
| 509 | + | } | |||
| 471 | for { | 510 | for { | ||
| 472 | select { | 511 | select { | ||
| 473 | case <-workerCtx.Done(): | 512 | case <-workerCtx.Done(): | ||
| 474 | return | 513 | return | ||
| 514 | + | case <-flush: | |||
| 515 | + | flushQueued() | |||
| 516 | + | return | |||
| 475 | case capture := <-d.snapshotJobs: | 517 | case capture := <-d.snapshotJobs: | ||
| 476 | - | // A queued capture must not begin persistence after shutdown. | 518 | + | if !write(capture) { |
| 477 | - | if workerCtx.Err() != nil || !d.setSnapshotWorkerInFlight(workerID, capture) { | |||
| 478 | - | d.finishSnapshotCapture(capture, false) | |||
| 479 | return | 519 | return | ||
| 480 | } | 520 | } | ||
| 481 | - | data, err := d.encodeSnapshotCapture(capture) | |||
| 482 | - | if err == nil && workerCtx.Err() == nil { | |||
| 483 | - | err = d.snaps.Write(capture.name, data) | |||
| 484 | - | } | |||
| 485 | - | d.clearSnapshotWorkerInFlight(workerID, capture) | |||
| 486 | - | if err != nil && workerCtx.Err() == nil { | |||
| 487 | - | d.log.Warn("writing session snapshot failed", "err", err, "session", capture.name) | |||
| 488 | - | } | |||
| 489 | - | d.finishSnapshotCapture(capture, err == nil && workerCtx.Err() == nil) | |||
| 490 | } | 521 | } | ||
| 491 | } | 522 | } | ||
| 492 | }() | 523 | }() | ||
| @@ -516,19 +547,39 @@ func (d *Daemon) stopSnapshotEncodeWorker() { | |||||
| 516 | } | 547 | } | ||
| 517 | d.snapshotWorkerMu.Lock() | 548 | d.snapshotWorkerMu.Lock() | ||
| 518 | cancel := d.snapshotWorkerCancel | 549 | cancel := d.snapshotWorkerCancel | ||
| 519 | - | if cancel == nil { | 550 | + | if cancel == nil || d.snapshotWorkerClosing { |
| 520 | d.snapshotWorkerMu.Unlock() | 551 | d.snapshotWorkerMu.Unlock() | ||
| 521 | return | 552 | return | ||
| 522 | } | 553 | } | ||
| 523 | - | // The queue stays open, so producers can safely race teardown. Detach an | 554 | + | // Stop accepting producers, then let the single worker drain its bounded |
| 524 | - | // in-flight uncancellable Write instead of waiting for it; the process owns | 555 | + | // queue. A non-context-aware store may still block indefinitely, so the |
| 525 | - | // that goroutine until the store returns. | 556 | + | // injected clock bounds how long teardown waits before detaching it. |
| 557 | + | d.snapshotWorkerClosing = true | |||
| 558 | + | flush := d.snapshotWorkerFlush | |||
| 559 | + | done := d.snapshotWorkerDone | |||
| 560 | + | close(flush) | |||
| 561 | + | timer := d.clock.NewTimer(snapshotFinalFlushTimeout) | |||
| 562 | + | d.snapshotWorkerMu.Unlock() | |||
| 563 | + | ||||
| 564 | + | select { | |||
| 565 | + | case <-done: | |||
| 566 | + | timer.Stop() | |||
| 567 | + | d.finishStoppedSnapshotWorker(false) | |||
| 568 | + | case <-timer.C(): | |||
| 569 | + | cancel() | |||
| 570 | + | d.finishStoppedSnapshotWorker(true) | |||
| 571 | + | } | |||
| 572 | + | } | |||
| 573 | + | ||||
| 574 | + | func (d *Daemon) finishStoppedSnapshotWorker(abandoned bool) { | |||
| 575 | + | d.snapshotWorkerMu.Lock() | |||
| 526 | inFlight := d.snapshotWorkerInFlight | 576 | inFlight := d.snapshotWorkerInFlight | ||
| 527 | d.snapshotWorkerCtx = nil | 577 | d.snapshotWorkerCtx = nil | ||
| 528 | d.snapshotWorkerCancel = nil | 578 | d.snapshotWorkerCancel = nil | ||
| 529 | d.snapshotWorkerDone = nil | 579 | d.snapshotWorkerDone = nil | ||
| 580 | + | d.snapshotWorkerFlush = nil | |||
| 581 | + | d.snapshotWorkerClosing = false | |||
| 530 | d.snapshotWorkerInFlight = nil | 582 | d.snapshotWorkerInFlight = nil | ||
| 531 | - | cancel() | |||
| 532 | queued := make([]*snapshotCapture, 0, len(d.snapshotJobs)) | 583 | queued := make([]*snapshotCapture, 0, len(d.snapshotJobs)) | ||
| 533 | for { | 584 | for { | ||
| 534 | select { | 585 | select { | ||
| @@ -536,7 +587,9 @@ func (d *Daemon) stopSnapshotEncodeWorker() { | |||||
| 536 | queued = append(queued, capture) | 587 | queued = append(queued, capture) | ||
| 537 | default: | 588 | default: | ||
| 538 | d.snapshotWorkerMu.Unlock() | 589 | d.snapshotWorkerMu.Unlock() | ||
| 539 | - | d.finishSnapshotCapture(inFlight, false) | 590 | + | if abandoned { |
| 591 | + | d.finishSnapshotCapture(inFlight, false) | |||
| 592 | + | } | |||
| 540 | for _, capture := range queued { | 593 | for _, capture := range queued { | ||
| 541 | d.finishSnapshotCapture(capture, false) | 594 | d.finishSnapshotCapture(capture, false) | ||
| 542 | } | 595 | } | ||
M internal/usecase/daemon/snapshot_restore_test.go +28 -1
| @@ -59,6 +59,21 @@ func TestRestorePaneTerminalRejectsMissingOrMalformedCanonicalBlobs(t *testing.T | |||
| 59 | 59 | } | |
| 60 | 60 | } | |
| 61 | 61 | ||
| 62 | + | func TestRestoreSessionClosesOpenedPTYWhenTerminalRestoreFails(t *testing.T) { | |
| 63 | + | factory := &restorePTYFactory{} | |
| 64 | + | d := newTestDaemon(t, factory, stubClock{}) | |
| 65 | + | pane := emptyTerminalPane(t, "pane-1", "/tmp") | |
| 66 | + | pane.Tail = []byte("malformed") | |
| 67 | + | ||
| 68 | + | err := d.restoreSession(context.Background(), snapcodec.Session{Name: "bad", Tabs: []snapcodec.Tab{{ | |
| 69 | + | Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{pane}, | |
| 70 | + | }}}) | |
| 71 | + | ||
| 72 | + | require.Error(t, err) | |
| 73 | + | require.Len(t, factory.opens, 1) | |
| 74 | + | require.Equal(t, 1, factory.opens[0].pty.closeCount()) | |
| 75 | + | } | |
| 76 | + | ||
| 62 | 77 | func TestRestoreSnapshotsRestoresLayoutCwdAndRows(t *testing.T) { | |
| 63 | 78 | store := &restoreSnapshotStore{} | |
| 64 | 79 | store.blobs = []ports.SnapshotBlob{{Name: "work", Data: mustSnapshotBytes(t, snapcodec.Session{ | |
| @@ -532,6 +547,7 @@ func (f *restorePTYFactory) Open(ctx context.Context, _ string, _ []string, env | |||
| 532 | 547 | type restorePTY struct { | |
| 533 | 548 | mu sync.Mutex | |
| 534 | 549 | writes []string | |
| 550 | + | closes int | |
| 535 | 551 | writeErr error | |
| 536 | 552 | done chan struct{} | |
| 537 | 553 | once sync.Once | |
| @@ -551,7 +567,18 @@ func (p *restorePTY) Write(b []byte) (int, error) { | |||
| 551 | 567 | } | |
| 552 | 568 | return len(b), nil | |
| 553 | 569 | } | |
| 554 | - | func (p *restorePTY) Close() error { p.once.Do(func() { close(p.done) }); return nil } | |
| 570 | + | func (p *restorePTY) Close() error { | |
| 571 | + | p.mu.Lock() | |
| 572 | + | p.closes++ | |
| 573 | + | p.mu.Unlock() | |
| 574 | + | p.once.Do(func() { close(p.done) }) | |
| 575 | + | return nil | |
| 576 | + | } | |
| 577 | + | func (p *restorePTY) closeCount() int { | |
| 578 | + | p.mu.Lock() | |
| 579 | + | defer p.mu.Unlock() | |
| 580 | + | return p.closes | |
| 581 | + | } | |
| 555 | 582 | func (p *restorePTY) Resize(domain.Size) error { return nil } | |
| 556 | 583 | func (p *restorePTY) Pid() int { return 0 } | |
| 557 | 584 | func (p *restorePTY) ForegroundPgid() (int, error) { return 0, nil } | |
| @@ -59,6 +59,21 @@ func TestRestorePaneTerminalRejectsMissingOrMalformedCanonicalBlobs(t *testing.T | |||||
| 59 | } | 59 | } | ||
| 60 | } | 60 | } | ||
| 61 | 61 | ||||
| 62 | + | func TestRestoreSessionClosesOpenedPTYWhenTerminalRestoreFails(t *testing.T) { | |||
| 63 | + | factory := &restorePTYFactory{} | |||
| 64 | + | d := newTestDaemon(t, factory, stubClock{}) | |||
| 65 | + | pane := emptyTerminalPane(t, "pane-1", "/tmp") | |||
| 66 | + | pane.Tail = []byte("malformed") | |||
| 67 | + | ||||
| 68 | + | err := d.restoreSession(context.Background(), snapcodec.Session{Name: "bad", Tabs: []snapcodec.Tab{{ | |||
| 69 | + | Cols: 80, Rows: 24, Tree: layout.NewTree("pane-1"), Panes: []snapcodec.Pane{pane}, | |||
| 70 | + | }}}) | |||
| 71 | + | ||||
| 72 | + | require.Error(t, err) | |||
| 73 | + | require.Len(t, factory.opens, 1) | |||
| 74 | + | require.Equal(t, 1, factory.opens[0].pty.closeCount()) | |||
| 75 | + | } | |||
| 76 | + | ||||
| 62 | func TestRestoreSnapshotsRestoresLayoutCwdAndRows(t *testing.T) { | 77 | func TestRestoreSnapshotsRestoresLayoutCwdAndRows(t *testing.T) { | ||
| 63 | store := &restoreSnapshotStore{} | 78 | store := &restoreSnapshotStore{} | ||
| 64 | store.blobs = []ports.SnapshotBlob{{Name: "work", Data: mustSnapshotBytes(t, snapcodec.Session{ | 79 | store.blobs = []ports.SnapshotBlob{{Name: "work", Data: mustSnapshotBytes(t, snapcodec.Session{ | ||
| @@ -532,6 +547,7 @@ func (f *restorePTYFactory) Open(ctx context.Context, _ string, _ []string, env | |||||
| 532 | type restorePTY struct { | 547 | type restorePTY struct { | ||
| 533 | mu sync.Mutex | 548 | mu sync.Mutex | ||
| 534 | writes []string | 549 | writes []string | ||
| 550 | + | closes int | |||
| 535 | writeErr error | 551 | writeErr error | ||
| 536 | done chan struct{} | 552 | done chan struct{} | ||
| 537 | once sync.Once | 553 | once sync.Once | ||
| @@ -551,7 +567,18 @@ func (p *restorePTY) Write(b []byte) (int, error) { | |||||
| 551 | } | 567 | } | ||
| 552 | return len(b), nil | 568 | return len(b), nil | ||
| 553 | } | 569 | } | ||
| 554 | - | func (p *restorePTY) Close() error { p.once.Do(func() { close(p.done) }); return nil } | 570 | + | func (p *restorePTY) Close() error { |
| 571 | + | p.mu.Lock() | |||
| 572 | + | p.closes++ | |||
| 573 | + | p.mu.Unlock() | |||
| 574 | + | p.once.Do(func() { close(p.done) }) | |||
| 575 | + | return nil | |||
| 576 | + | } | |||
| 577 | + | func (p *restorePTY) closeCount() int { | |||
| 578 | + | p.mu.Lock() | |||
| 579 | + | defer p.mu.Unlock() | |||
| 580 | + | return p.closes | |||
| 581 | + | } | |||
| 555 | func (p *restorePTY) Resize(domain.Size) error { return nil } | 582 | func (p *restorePTY) Resize(domain.Size) error { return nil } | ||
| 556 | func (p *restorePTY) Pid() int { return 0 } | 583 | func (p *restorePTY) Pid() int { return 0 } | ||
| 557 | func (p *restorePTY) ForegroundPgid() (int, error) { return 0, nil } | 584 | func (p *restorePTY) ForegroundPgid() (int, error) { return 0, nil } | ||
M internal/usecase/daemon/snapshot_test.go +64 -6
| @@ -19,6 +19,7 @@ import ( | |||
| 19 | 19 | "github.com/bnema/vev/internal/usecase/layout" | |
| 20 | 20 | snapcodec "github.com/bnema/vev/internal/usecase/snapshot" | |
| 21 | 21 | "github.com/bnema/vev/pkg/renderer" | |
| 22 | + | "github.com/bnema/vev/pkg/vt" | |
| 22 | 23 | ) | |
| 23 | 24 | ||
| 24 | 25 | func TestSnapshotEncodeWorkerDefersMarshalUntilAfterCapture(t *testing.T) { | |
| @@ -152,9 +153,55 @@ func exerciseBlockedSnapshotWorker(t *testing.T, blockEncode bool) { | |||
| 152 | 153 | require.True(t, sess.snapDirty.Load(), "PTY output published while blocked must remain retryable") | |
| 153 | 154 | } | |
| 154 | 155 | ||
| 156 | + | func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *testing.T) { | |
| 157 | + | clock := newFinalFlushClock() | |
| 158 | + | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |
| 159 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |
| 160 | + | WithSnapshotStore(store)(d) | |
| 161 | + | ctx, cancel := context.WithCancel(t.Context()) | |
| 162 | + | defer cancel() | |
| 163 | + | d.startSnapshotEncodeWorker(ctx) | |
| 164 | + | ||
| 165 | + | entered := make(chan struct{}) | |
| 166 | + | release := make(chan struct{}) | |
| 167 | + | d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) { | |
| 168 | + | close(entered) | |
| 169 | + | <-release | |
| 170 | + | return snapcodec.Marshal(s) | |
| 171 | + | } | |
| 172 | + | sess := newSnapshotTestSession(t, "final", false, "/work") | |
| 173 | + | pty := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |
| 174 | + | pty.EXPECT().Close().Return(nil).Once() | |
| 175 | + | sess.tabs[0].panes["pane-1"].screen.Write([]byte("\rnewest")) | |
| 176 | + | d.sessions[sess.id] = sess | |
| 177 | + | ||
| 178 | + | require.NoError(t, d.killSession(sess, 0, false)) | |
| 179 | + | <-entered | |
| 180 | + | stopped := make(chan struct{}) | |
| 181 | + | go func() { | |
| 182 | + | d.stopSnapshotEncodeWorker() | |
| 183 | + | close(stopped) | |
| 184 | + | }() | |
| 185 | + | close(release) | |
| 186 | + | <-stopped | |
| 187 | + | awaitSnapshotIdle(t, sess) | |
| 188 | + | ||
| 189 | + | select { | |
| 190 | + | case data := <-store.writes: | |
| 191 | + | snap, err := snapcodec.Unmarshal(data) | |
| 192 | + | require.NoError(t, err) | |
| 193 | + | frame, err := vt.UnmarshalVisible(snap.Tabs[0].Panes[0].Visible) | |
| 194 | + | require.NoError(t, err) | |
| 195 | + | require.Contains(t, rowText(frame.Row(0)), "newest") | |
| 196 | + | default: | |
| 197 | + | t.Fatal("final snapshot was dropped during worker teardown") | |
| 198 | + | } | |
| 199 | + | } | |
| 200 | + | ||
| 155 | 201 | func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) { | |
| 156 | 202 | store := newNeverReturningSnapshotStore() | |
| 157 | - | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | |
| 203 | + | clock := newFinalFlushClock() | |
| 204 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |
| 158 | 205 | WithSnapshotStore(store)(d) | |
| 159 | 206 | ctx, cancel := context.WithCancel(t.Context()) | |
| 160 | 207 | defer cancel() | |
| @@ -169,11 +216,9 @@ func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) | |||
| 169 | 216 | d.stopSnapshotEncodeWorker() | |
| 170 | 217 | close(stopped) | |
| 171 | 218 | }() | |
| 172 | - | select { | |
| 173 | - | case <-stopped: | |
| 174 | - | case <-time.After(time.Second): | |
| 175 | - | t.Fatal("snapshot worker shutdown waited for uncancellable Write") | |
| 176 | - | } | |
| 219 | + | timer := <-clock.timers | |
| 220 | + | timer.ch <- time.Time{} | |
| 221 | + | <-stopped | |
| 177 | 222 | awaitSnapshotIdle(t, sess) | |
| 178 | 223 | require.True(t, sess.snapDirty.Load(), "an abandoned write must remain retryable") | |
| 179 | 224 | require.NotPanics(t, func() { require.False(t, d.scheduleSnapshot(sess)) }) | |
| @@ -870,6 +915,19 @@ func (c *manualSnapshotClock) fireNext(t *testing.T) { | |||
| 870 | 915 | } | |
| 871 | 916 | } | |
| 872 | 917 | ||
| 918 | + | type finalFlushClock struct{ timers chan *manualSnapshotTimer } | |
| 919 | + | ||
| 920 | + | func newFinalFlushClock() *finalFlushClock { | |
| 921 | + | return &finalFlushClock{timers: make(chan *manualSnapshotTimer, 1)} | |
| 922 | + | } | |
| 923 | + | ||
| 924 | + | func (c *finalFlushClock) Now() time.Time { return time.Unix(0, 0) } | |
| 925 | + | func (c *finalFlushClock) NewTimer(time.Duration) ports.Timer { | |
| 926 | + | t := &manualSnapshotTimer{ch: make(chan time.Time, 1)} | |
| 927 | + | c.timers <- t | |
| 928 | + | return t | |
| 929 | + | } | |
| 930 | + | ||
| 873 | 931 | type manualSnapshotTimer struct{ ch chan time.Time } | |
| 874 | 932 | ||
| 875 | 933 | func (t *manualSnapshotTimer) C() <-chan time.Time { return t.ch } | |
| @@ -19,6 +19,7 @@ import ( | |||||
| 19 | "github.com/bnema/vev/internal/usecase/layout" | 19 | "github.com/bnema/vev/internal/usecase/layout" | ||
| 20 | snapcodec "github.com/bnema/vev/internal/usecase/snapshot" | 20 | snapcodec "github.com/bnema/vev/internal/usecase/snapshot" | ||
| 21 | "github.com/bnema/vev/pkg/renderer" | 21 | "github.com/bnema/vev/pkg/renderer" | ||
| 22 | + | "github.com/bnema/vev/pkg/vt" | |||
| 22 | ) | 23 | ) | ||
| 23 | 24 | ||||
| 24 | func TestSnapshotEncodeWorkerDefersMarshalUntilAfterCapture(t *testing.T) { | 25 | func TestSnapshotEncodeWorkerDefersMarshalUntilAfterCapture(t *testing.T) { | ||
| @@ -152,9 +153,55 @@ func exerciseBlockedSnapshotWorker(t *testing.T, blockEncode bool) { | |||||
| 152 | require.True(t, sess.snapDirty.Load(), "PTY output published while blocked must remain retryable") | 153 | require.True(t, sess.snapDirty.Load(), "PTY output published while blocked must remain retryable") | ||
| 153 | } | 154 | } | ||
| 154 | 155 | ||||
| 156 | + | func TestFinalSnapshotFlushPersistsNewestNamedSessionBeforeWorkerTeardown(t *testing.T) { | |||
| 157 | + | clock := newFinalFlushClock() | |||
| 158 | + | store := &channelSnapshotStore{writes: make(chan []byte, 1)} | |||
| 159 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |||
| 160 | + | WithSnapshotStore(store)(d) | |||
| 161 | + | ctx, cancel := context.WithCancel(t.Context()) | |||
| 162 | + | defer cancel() | |||
| 163 | + | d.startSnapshotEncodeWorker(ctx) | |||
| 164 | + | ||||
| 165 | + | entered := make(chan struct{}) | |||
| 166 | + | release := make(chan struct{}) | |||
| 167 | + | d.snapshotMarshal = func(s snapcodec.Session) ([]byte, error) { | |||
| 168 | + | close(entered) | |||
| 169 | + | <-release | |||
| 170 | + | return snapcodec.Marshal(s) | |||
| 171 | + | } | |||
| 172 | + | sess := newSnapshotTestSession(t, "final", false, "/work") | |||
| 173 | + | pty := sess.tabs[0].panes["pane-1"].pty.(*portsmocks.MockPTY) | |||
| 174 | + | pty.EXPECT().Close().Return(nil).Once() | |||
| 175 | + | sess.tabs[0].panes["pane-1"].screen.Write([]byte("\rnewest")) | |||
| 176 | + | d.sessions[sess.id] = sess | |||
| 177 | + | ||||
| 178 | + | require.NoError(t, d.killSession(sess, 0, false)) | |||
| 179 | + | <-entered | |||
| 180 | + | stopped := make(chan struct{}) | |||
| 181 | + | go func() { | |||
| 182 | + | d.stopSnapshotEncodeWorker() | |||
| 183 | + | close(stopped) | |||
| 184 | + | }() | |||
| 185 | + | close(release) | |||
| 186 | + | <-stopped | |||
| 187 | + | awaitSnapshotIdle(t, sess) | |||
| 188 | + | ||||
| 189 | + | select { | |||
| 190 | + | case data := <-store.writes: | |||
| 191 | + | snap, err := snapcodec.Unmarshal(data) | |||
| 192 | + | require.NoError(t, err) | |||
| 193 | + | frame, err := vt.UnmarshalVisible(snap.Tabs[0].Panes[0].Visible) | |||
| 194 | + | require.NoError(t, err) | |||
| 195 | + | require.Contains(t, rowText(frame.Row(0)), "newest") | |||
| 196 | + | default: | |||
| 197 | + | t.Fatal("final snapshot was dropped during worker teardown") | |||
| 198 | + | } | |||
| 199 | + | } | |||
| 200 | + | ||||
| 155 | func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) { | 201 | func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) { | ||
| 156 | store := newNeverReturningSnapshotStore() | 202 | store := newNeverReturningSnapshotStore() | ||
| 157 | - | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), stubClock{}) | 203 | + | clock := newFinalFlushClock() |
| 204 | + | d := newTestDaemon(t, portsmocks.NewMockPTYFactory(t), clock) | |||
| 158 | WithSnapshotStore(store)(d) | 205 | WithSnapshotStore(store)(d) | ||
| 159 | ctx, cancel := context.WithCancel(t.Context()) | 206 | ctx, cancel := context.WithCancel(t.Context()) | ||
| 160 | defer cancel() | 207 | defer cancel() | ||
| @@ -169,11 +216,9 @@ func TestStopSnapshotEncodeWorkerDoesNotWaitForUncancellableWrite(t *testing.T) | |||||
| 169 | d.stopSnapshotEncodeWorker() | 216 | d.stopSnapshotEncodeWorker() | ||
| 170 | close(stopped) | 217 | close(stopped) | ||
| 171 | }() | 218 | }() | ||
| 172 | - | select { | 219 | + | timer := <-clock.timers |
| 173 | - | case <-stopped: | 220 | + | timer.ch <- time.Time{} |
| 174 | - | case <-time.After(time.Second): | 221 | + | <-stopped |
| 175 | - | t.Fatal("snapshot worker shutdown waited for uncancellable Write") | |||
| 176 | - | } | |||
| 177 | awaitSnapshotIdle(t, sess) | 222 | awaitSnapshotIdle(t, sess) | ||
| 178 | require.True(t, sess.snapDirty.Load(), "an abandoned write must remain retryable") | 223 | require.True(t, sess.snapDirty.Load(), "an abandoned write must remain retryable") | ||
| 179 | require.NotPanics(t, func() { require.False(t, d.scheduleSnapshot(sess)) }) | 224 | require.NotPanics(t, func() { require.False(t, d.scheduleSnapshot(sess)) }) | ||
| @@ -870,6 +915,19 @@ func (c *manualSnapshotClock) fireNext(t *testing.T) { | |||||
| 870 | } | 915 | } | ||
| 871 | } | 916 | } | ||
| 872 | 917 | ||||
| 918 | + | type finalFlushClock struct{ timers chan *manualSnapshotTimer } | |||
| 919 | + | ||||
| 920 | + | func newFinalFlushClock() *finalFlushClock { | |||
| 921 | + | return &finalFlushClock{timers: make(chan *manualSnapshotTimer, 1)} | |||
| 922 | + | } | |||
| 923 | + | ||||
| 924 | + | func (c *finalFlushClock) Now() time.Time { return time.Unix(0, 0) } | |||
| 925 | + | func (c *finalFlushClock) NewTimer(time.Duration) ports.Timer { | |||
| 926 | + | t := &manualSnapshotTimer{ch: make(chan time.Time, 1)} | |||
| 927 | + | c.timers <- t | |||
| 928 | + | return t | |||
| 929 | + | } | |||
| 930 | + | ||||
| 873 | type manualSnapshotTimer struct{ ch chan time.Time } | 931 | type manualSnapshotTimer struct{ ch chan time.Time } | ||
| 874 | 932 | ||||
| 875 | func (t *manualSnapshotTimer) C() <-chan time.Time { return t.ch } | 933 | func (t *manualSnapshotTimer) C() <-chan time.Time { return t.ch } | ||
M internal/usecase/snapshot/codec.go +94 -25
| @@ -22,13 +22,17 @@ var ( | |||
| 22 | 22 | ) | |
| 23 | 23 | ||
| 24 | 24 | const ( | |
| 25 | - | magic = "VEVS" | |
| 26 | - | version = uint16(3) | |
| 27 | - | maxDecodedBodySize = 256 << 20 | |
| 28 | - | maxSnapshotBlobs = 1 << 16 | |
| 29 | - | maxSnapshotBytes = 256 << 20 | |
| 30 | - | maxSnapshotRows = 1 << 20 | |
| 31 | - | maxSnapshotCells = 16 << 20 | |
| 25 | + | magic = "VEVS" | |
| 26 | + | version = uint16(3) | |
| 27 | + | maxDecodedBodySize = 256 << 20 | |
| 28 | + | maxSnapshotBlobs = 1 << 16 | |
| 29 | + | maxSnapshotBytes = 256 << 20 | |
| 30 | + | maxSnapshotRows = 1 << 20 | |
| 31 | + | maxSnapshotCells = 16 << 20 | |
| 32 | + | maxSnapshotTreeNodes = 1 << 15 | |
| 33 | + | // 256 MiB accommodates a full 10k-row history at practical terminal widths | |
| 34 | + | // while bounding hostile manifests before decoded slices are allocated. | |
| 35 | + | maxSnapshotDecodedAllocation = 256 << 20 | |
| 32 | 36 | ) | |
| 33 | 37 | ||
| 34 | 38 | // Marshal encodes the strict v3, flags-zero durable envelope. | |
| @@ -574,8 +578,9 @@ func preflightSession(body []byte) error { | |||
| 574 | 578 | } | |
| 575 | 579 | var totals vt.DecodeStats | |
| 576 | 580 | blobs := uint64(0) | |
| 581 | + | budget := preflightBudget{} | |
| 577 | 582 | for range tabs { | |
| 578 | - | if err := preflightTab(&r, &totals, &blobs); err != nil { | |
| 583 | + | if err := preflightTab(&r, &totals, &blobs, &budget); err != nil { | |
| 579 | 584 | return fmt.Errorf("preflight tab: %w", err) | |
| 580 | 585 | } | |
| 581 | 586 | } | |
| @@ -600,7 +605,38 @@ func skipString(r *payloadReader) error { | |||
| 600 | 605 | return nil | |
| 601 | 606 | } | |
| 602 | 607 | ||
| 603 | - | func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error { | |
| 608 | + | type preflightBudget struct { | |
| 609 | + | nodes uint64 | |
| 610 | + | alloc uint64 | |
| 611 | + | } | |
| 612 | + | ||
| 613 | + | func (b *preflightBudget) addAllocation(amount uint64) bool { | |
| 614 | + | if amount > maxSnapshotDecodedAllocation-b.alloc { | |
| 615 | + | return false | |
| 616 | + | } | |
| 617 | + | b.alloc += amount | |
| 618 | + | return true | |
| 619 | + | } | |
| 620 | + | ||
| 621 | + | func (b *preflightBudget) addProduct(count, size uint64) bool { | |
| 622 | + | if count != 0 && size > math.MaxUint64/count { | |
| 623 | + | return false | |
| 624 | + | } | |
| 625 | + | return b.addAllocation(count * size) | |
| 626 | + | } | |
| 627 | + | ||
| 628 | + | func (b *preflightBudget) addNode() bool { | |
| 629 | + | if b.nodes >= maxSnapshotTreeNodes || !b.addAllocation(128) { | |
| 630 | + | return false | |
| 631 | + | } | |
| 632 | + | b.nodes++ | |
| 633 | + | return true | |
| 634 | + | } | |
| 635 | + | ||
| 636 | + | func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, budget *preflightBudget) error { | |
| 637 | + | if !budget.addAllocation(256) { | |
| 638 | + | return fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |
| 639 | + | } | |
| 604 | 640 | if err := skipString(r); err != nil { | |
| 605 | 641 | return err | |
| 606 | 642 | } | |
| @@ -616,7 +652,7 @@ func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error | |||
| 616 | 652 | if err := skipString(r); err != nil { | |
| 617 | 653 | return err | |
| 618 | 654 | } | |
| 619 | - | if err := preflightNode(r, 0); err != nil { | |
| 655 | + | if err := preflightNode(r, 0, budget); err != nil { | |
| 620 | 656 | return err | |
| 621 | 657 | } | |
| 622 | 658 | panes, err := r.getUint16() | |
| @@ -626,18 +662,22 @@ func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error | |||
| 626 | 662 | if uint64(panes) > uint64(len(r.b))/2 { | |
| 627 | 663 | return ErrShortPayload | |
| 628 | 664 | } | |
| 665 | + | ids := make(map[layout.PaneID]struct{}) | |
| 629 | 666 | for range panes { | |
| 630 | - | if err := preflightPane(r, totals, blobs); err != nil { | |
| 667 | + | if err := preflightPane(r, totals, blobs, ids, budget); err != nil { | |
| 631 | 668 | return err | |
| 632 | 669 | } | |
| 633 | 670 | } | |
| 634 | 671 | return nil | |
| 635 | 672 | } | |
| 636 | 673 | ||
| 637 | - | func preflightNode(r *payloadReader, depth int) error { | |
| 674 | + | func preflightNode(r *payloadReader, depth int, budget *preflightBudget) error { | |
| 638 | 675 | if depth > 64 { | |
| 639 | 676 | return fmt.Errorf("%w: tree depth", ErrInvalidData) | |
| 640 | 677 | } | |
| 678 | + | if !budget.addNode() { | |
| 679 | + | return fmt.Errorf("%w: tree node or decoded allocation budget", ErrInvalidData) | |
| 680 | + | } | |
| 641 | 681 | kind, err := r.getUint8() | |
| 642 | 682 | if err != nil { | |
| 643 | 683 | return err | |
| @@ -666,17 +706,29 @@ func preflightNode(r *payloadReader, depth int) error { | |||
| 666 | 706 | return ErrShortPayload | |
| 667 | 707 | } | |
| 668 | 708 | for range children { | |
| 669 | - | if err := preflightNode(r, depth+1); err != nil { | |
| 709 | + | if err := preflightNode(r, depth+1, budget); err != nil { | |
| 670 | 710 | return err | |
| 671 | 711 | } | |
| 672 | 712 | } | |
| 673 | 713 | return nil | |
| 674 | 714 | } | |
| 675 | 715 | ||
| 676 | - | func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error { | |
| 677 | - | if err := skipString(r); err != nil { | |
| 716 | + | func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, ids map[layout.PaneID]struct{}, budget *preflightBudget) error { | |
| 717 | + | if !budget.addAllocation(256) { | |
| 718 | + | return fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |
| 719 | + | } | |
| 720 | + | id, err := r.getString() | |
| 721 | + | if err != nil { | |
| 678 | 722 | return err | |
| 679 | 723 | } | |
| 724 | + | if id == "" { | |
| 725 | + | return fmt.Errorf("%w: empty pane ID", ErrInvalidData) | |
| 726 | + | } | |
| 727 | + | paneID := layout.PaneID(id) | |
| 728 | + | if _, exists := ids[paneID]; exists { | |
| 729 | + | return fmt.Errorf("%w: duplicate pane ID", ErrInvalidData) | |
| 730 | + | } | |
| 731 | + | ids[paneID] = struct{}{} | |
| 680 | 732 | if err := skipString(r); err != nil { | |
| 681 | 733 | return err | |
| 682 | 734 | } | |
| @@ -691,29 +743,43 @@ func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) erro | |||
| 691 | 743 | return fmt.Errorf("%w: sealed chunks", ErrInvalidData) | |
| 692 | 744 | } | |
| 693 | 745 | for range n { | |
| 694 | - | if err := preflightBlob(r, totals, blobs, true); err != nil { | |
| 746 | + | stats, err := preflightBlob(r, totals, blobs, true, budget) | |
| 747 | + | if err != nil { | |
| 695 | 748 | return err | |
| 696 | 749 | } | |
| 750 | + | if stats.Chunks != 1 || stats.Rows == 0 { | |
| 751 | + | return fmt.Errorf("%w: sealed blob role", ErrInvalidData) | |
| 752 | + | } | |
| 697 | 753 | } | |
| 698 | - | if err := preflightBlob(r, totals, blobs, true); err != nil { | |
| 754 | + | tail, err := preflightBlob(r, totals, blobs, true, budget) | |
| 755 | + | if err != nil { | |
| 699 | 756 | return err | |
| 700 | 757 | } | |
| 701 | - | if err := preflightBlob(r, totals, blobs, false); err != nil { | |
| 758 | + | if tail.Chunks != 0 || tail.Rows != 0 || tail.Cells != 0 { | |
| 759 | + | return fmt.Errorf("%w: tail blob role", ErrInvalidData) | |
| 760 | + | } | |
| 761 | + | visible, err := preflightBlob(r, totals, blobs, false, budget) | |
| 762 | + | if err != nil { | |
| 702 | 763 | return err | |
| 703 | 764 | } | |
| 765 | + | // PreflightVisibleBlob above validates width×height and exact payload | |
| 766 | + | // cardinality. A visible terminal frame itself must also be nonempty. | |
| 767 | + | if visible.Rows == 0 || visible.Cells == 0 { | |
| 768 | + | return fmt.Errorf("%w: visible geometry", ErrInvalidData) | |
| 769 | + | } | |
| 704 | 770 | return preflightProcess(r) | |
| 705 | 771 | } | |
| 706 | 772 | ||
| 707 | - | func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, history bool) error { | |
| 773 | + | func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, history bool, budget *preflightBudget) (vt.DecodeStats, error) { | |
| 708 | 774 | n, err := r.getUint32() | |
| 709 | 775 | if err != nil { | |
| 710 | - | return err | |
| 776 | + | return vt.DecodeStats{}, err | |
| 711 | 777 | } | |
| 712 | 778 | if n == 0 || uint64(n) > uint64(len(r.b)) || uint64(n) > maxSnapshotBytes { | |
| 713 | - | return fmt.Errorf("%w: VT blob length", ErrInvalidData) | |
| 779 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT blob length", ErrInvalidData) | |
| 714 | 780 | } | |
| 715 | 781 | if *blobs >= maxSnapshotBlobs { | |
| 716 | - | return fmt.Errorf("%w: too many VT blobs", ErrInvalidData) | |
| 782 | + | return vt.DecodeStats{}, fmt.Errorf("%w: too many VT blobs", ErrInvalidData) | |
| 717 | 783 | } | |
| 718 | 784 | blob := r.b[:n] | |
| 719 | 785 | r.b = r.b[n:] | |
| @@ -724,13 +790,16 @@ func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, hist | |||
| 724 | 790 | stats, err = vt.PreflightVisibleBlob(blob) | |
| 725 | 791 | } | |
| 726 | 792 | if err != nil { | |
| 727 | - | return fmt.Errorf("%w: VT blob", ErrInvalidData) | |
| 793 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT blob", ErrInvalidData) | |
| 728 | 794 | } | |
| 729 | 795 | if !totals.Add(stats) { | |
| 730 | - | return fmt.Errorf("%w: VT budget overflow", ErrInvalidData) | |
| 796 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT budget overflow", ErrInvalidData) | |
| 797 | + | } | |
| 798 | + | if !budget.addAllocation(uint64(n)) || !budget.addProduct(stats.Chunks, 64) || !budget.addProduct(stats.Rows, 32) || !budget.addProduct(stats.Cells, 64) { | |
| 799 | + | return vt.DecodeStats{}, fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |
| 731 | 800 | } | |
| 732 | 801 | *blobs++ | |
| 733 | - | return nil | |
| 802 | + | return stats, nil | |
| 734 | 803 | } | |
| 735 | 804 | ||
| 736 | 805 | func preflightProcess(r *payloadReader) error { | |
| @@ -22,13 +22,17 @@ var ( | |||||
| 22 | ) | 22 | ) | ||
| 23 | 23 | ||||
| 24 | const ( | 24 | const ( | ||
| 25 | - | magic = "VEVS" | 25 | + | magic = "VEVS" |
| 26 | - | version = uint16(3) | 26 | + | version = uint16(3) |
| 27 | - | maxDecodedBodySize = 256 << 20 | 27 | + | maxDecodedBodySize = 256 << 20 |
| 28 | - | maxSnapshotBlobs = 1 << 16 | 28 | + | maxSnapshotBlobs = 1 << 16 |
| 29 | - | maxSnapshotBytes = 256 << 20 | 29 | + | maxSnapshotBytes = 256 << 20 |
| 30 | - | maxSnapshotRows = 1 << 20 | 30 | + | maxSnapshotRows = 1 << 20 |
| 31 | - | maxSnapshotCells = 16 << 20 | 31 | + | maxSnapshotCells = 16 << 20 |
| 32 | + | maxSnapshotTreeNodes = 1 << 15 | |||
| 33 | + | // 256 MiB accommodates a full 10k-row history at practical terminal widths | |||
| 34 | + | // while bounding hostile manifests before decoded slices are allocated. | |||
| 35 | + | maxSnapshotDecodedAllocation = 256 << 20 | |||
| 32 | ) | 36 | ) | ||
| 33 | 37 | ||||
| 34 | // Marshal encodes the strict v3, flags-zero durable envelope. | 38 | // Marshal encodes the strict v3, flags-zero durable envelope. | ||
| @@ -574,8 +578,9 @@ func preflightSession(body []byte) error { | |||||
| 574 | } | 578 | } | ||
| 575 | var totals vt.DecodeStats | 579 | var totals vt.DecodeStats | ||
| 576 | blobs := uint64(0) | 580 | blobs := uint64(0) | ||
| 581 | + | budget := preflightBudget{} | |||
| 577 | for range tabs { | 582 | for range tabs { | ||
| 578 | - | if err := preflightTab(&r, &totals, &blobs); err != nil { | 583 | + | if err := preflightTab(&r, &totals, &blobs, &budget); err != nil { |
| 579 | return fmt.Errorf("preflight tab: %w", err) | 584 | return fmt.Errorf("preflight tab: %w", err) | ||
| 580 | } | 585 | } | ||
| 581 | } | 586 | } | ||
| @@ -600,7 +605,38 @@ func skipString(r *payloadReader) error { | |||||
| 600 | return nil | 605 | return nil | ||
| 601 | } | 606 | } | ||
| 602 | 607 | ||||
| 603 | - | func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error { | 608 | + | type preflightBudget struct { |
| 609 | + | nodes uint64 | |||
| 610 | + | alloc uint64 | |||
| 611 | + | } | |||
| 612 | + | ||||
| 613 | + | func (b *preflightBudget) addAllocation(amount uint64) bool { | |||
| 614 | + | if amount > maxSnapshotDecodedAllocation-b.alloc { | |||
| 615 | + | return false | |||
| 616 | + | } | |||
| 617 | + | b.alloc += amount | |||
| 618 | + | return true | |||
| 619 | + | } | |||
| 620 | + | ||||
| 621 | + | func (b *preflightBudget) addProduct(count, size uint64) bool { | |||
| 622 | + | if count != 0 && size > math.MaxUint64/count { | |||
| 623 | + | return false | |||
| 624 | + | } | |||
| 625 | + | return b.addAllocation(count * size) | |||
| 626 | + | } | |||
| 627 | + | ||||
| 628 | + | func (b *preflightBudget) addNode() bool { | |||
| 629 | + | if b.nodes >= maxSnapshotTreeNodes || !b.addAllocation(128) { | |||
| 630 | + | return false | |||
| 631 | + | } | |||
| 632 | + | b.nodes++ | |||
| 633 | + | return true | |||
| 634 | + | } | |||
| 635 | + | ||||
| 636 | + | func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, budget *preflightBudget) error { | |||
| 637 | + | if !budget.addAllocation(256) { | |||
| 638 | + | return fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |||
| 639 | + | } | |||
| 604 | if err := skipString(r); err != nil { | 640 | if err := skipString(r); err != nil { | ||
| 605 | return err | 641 | return err | ||
| 606 | } | 642 | } | ||
| @@ -616,7 +652,7 @@ func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error | |||||
| 616 | if err := skipString(r); err != nil { | 652 | if err := skipString(r); err != nil { | ||
| 617 | return err | 653 | return err | ||
| 618 | } | 654 | } | ||
| 619 | - | if err := preflightNode(r, 0); err != nil { | 655 | + | if err := preflightNode(r, 0, budget); err != nil { |
| 620 | return err | 656 | return err | ||
| 621 | } | 657 | } | ||
| 622 | panes, err := r.getUint16() | 658 | panes, err := r.getUint16() | ||
| @@ -626,18 +662,22 @@ func preflightTab(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error | |||||
| 626 | if uint64(panes) > uint64(len(r.b))/2 { | 662 | if uint64(panes) > uint64(len(r.b))/2 { | ||
| 627 | return ErrShortPayload | 663 | return ErrShortPayload | ||
| 628 | } | 664 | } | ||
| 665 | + | ids := make(map[layout.PaneID]struct{}) | |||
| 629 | for range panes { | 666 | for range panes { | ||
| 630 | - | if err := preflightPane(r, totals, blobs); err != nil { | 667 | + | if err := preflightPane(r, totals, blobs, ids, budget); err != nil { |
| 631 | return err | 668 | return err | ||
| 632 | } | 669 | } | ||
| 633 | } | 670 | } | ||
| 634 | return nil | 671 | return nil | ||
| 635 | } | 672 | } | ||
| 636 | 673 | ||||
| 637 | - | func preflightNode(r *payloadReader, depth int) error { | 674 | + | func preflightNode(r *payloadReader, depth int, budget *preflightBudget) error { |
| 638 | if depth > 64 { | 675 | if depth > 64 { | ||
| 639 | return fmt.Errorf("%w: tree depth", ErrInvalidData) | 676 | return fmt.Errorf("%w: tree depth", ErrInvalidData) | ||
| 640 | } | 677 | } | ||
| 678 | + | if !budget.addNode() { | |||
| 679 | + | return fmt.Errorf("%w: tree node or decoded allocation budget", ErrInvalidData) | |||
| 680 | + | } | |||
| 641 | kind, err := r.getUint8() | 681 | kind, err := r.getUint8() | ||
| 642 | if err != nil { | 682 | if err != nil { | ||
| 643 | return err | 683 | return err | ||
| @@ -666,17 +706,29 @@ func preflightNode(r *payloadReader, depth int) error { | |||||
| 666 | return ErrShortPayload | 706 | return ErrShortPayload | ||
| 667 | } | 707 | } | ||
| 668 | for range children { | 708 | for range children { | ||
| 669 | - | if err := preflightNode(r, depth+1); err != nil { | 709 | + | if err := preflightNode(r, depth+1, budget); err != nil { |
| 670 | return err | 710 | return err | ||
| 671 | } | 711 | } | ||
| 672 | } | 712 | } | ||
| 673 | return nil | 713 | return nil | ||
| 674 | } | 714 | } | ||
| 675 | 715 | ||||
| 676 | - | func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) error { | 716 | + | func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, ids map[layout.PaneID]struct{}, budget *preflightBudget) error { |
| 677 | - | if err := skipString(r); err != nil { | 717 | + | if !budget.addAllocation(256) { |
| 718 | + | return fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |||
| 719 | + | } | |||
| 720 | + | id, err := r.getString() | |||
| 721 | + | if err != nil { | |||
| 678 | return err | 722 | return err | ||
| 679 | } | 723 | } | ||
| 724 | + | if id == "" { | |||
| 725 | + | return fmt.Errorf("%w: empty pane ID", ErrInvalidData) | |||
| 726 | + | } | |||
| 727 | + | paneID := layout.PaneID(id) | |||
| 728 | + | if _, exists := ids[paneID]; exists { | |||
| 729 | + | return fmt.Errorf("%w: duplicate pane ID", ErrInvalidData) | |||
| 730 | + | } | |||
| 731 | + | ids[paneID] = struct{}{} | |||
| 680 | if err := skipString(r); err != nil { | 732 | if err := skipString(r); err != nil { | ||
| 681 | return err | 733 | return err | ||
| 682 | } | 734 | } | ||
| @@ -691,29 +743,43 @@ func preflightPane(r *payloadReader, totals *vt.DecodeStats, blobs *uint64) erro | |||||
| 691 | return fmt.Errorf("%w: sealed chunks", ErrInvalidData) | 743 | return fmt.Errorf("%w: sealed chunks", ErrInvalidData) | ||
| 692 | } | 744 | } | ||
| 693 | for range n { | 745 | for range n { | ||
| 694 | - | if err := preflightBlob(r, totals, blobs, true); err != nil { | 746 | + | stats, err := preflightBlob(r, totals, blobs, true, budget) |
| 747 | + | if err != nil { | |||
| 695 | return err | 748 | return err | ||
| 696 | } | 749 | } | ||
| 750 | + | if stats.Chunks != 1 || stats.Rows == 0 { | |||
| 751 | + | return fmt.Errorf("%w: sealed blob role", ErrInvalidData) | |||
| 752 | + | } | |||
| 697 | } | 753 | } | ||
| 698 | - | if err := preflightBlob(r, totals, blobs, true); err != nil { | 754 | + | tail, err := preflightBlob(r, totals, blobs, true, budget) |
| 755 | + | if err != nil { | |||
| 699 | return err | 756 | return err | ||
| 700 | } | 757 | } | ||
| 701 | - | if err := preflightBlob(r, totals, blobs, false); err != nil { | 758 | + | if tail.Chunks != 0 || tail.Rows != 0 || tail.Cells != 0 { |
| 759 | + | return fmt.Errorf("%w: tail blob role", ErrInvalidData) | |||
| 760 | + | } | |||
| 761 | + | visible, err := preflightBlob(r, totals, blobs, false, budget) | |||
| 762 | + | if err != nil { | |||
| 702 | return err | 763 | return err | ||
| 703 | } | 764 | } | ||
| 765 | + | // PreflightVisibleBlob above validates width×height and exact payload | |||
| 766 | + | // cardinality. A visible terminal frame itself must also be nonempty. | |||
| 767 | + | if visible.Rows == 0 || visible.Cells == 0 { | |||
| 768 | + | return fmt.Errorf("%w: visible geometry", ErrInvalidData) | |||
| 769 | + | } | |||
| 704 | return preflightProcess(r) | 770 | return preflightProcess(r) | ||
| 705 | } | 771 | } | ||
| 706 | 772 | ||||
| 707 | - | func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, history bool) error { | 773 | + | func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, history bool, budget *preflightBudget) (vt.DecodeStats, error) { |
| 708 | n, err := r.getUint32() | 774 | n, err := r.getUint32() | ||
| 709 | if err != nil { | 775 | if err != nil { | ||
| 710 | - | return err | 776 | + | return vt.DecodeStats{}, err |
| 711 | } | 777 | } | ||
| 712 | if n == 0 || uint64(n) > uint64(len(r.b)) || uint64(n) > maxSnapshotBytes { | 778 | if n == 0 || uint64(n) > uint64(len(r.b)) || uint64(n) > maxSnapshotBytes { | ||
| 713 | - | return fmt.Errorf("%w: VT blob length", ErrInvalidData) | 779 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT blob length", ErrInvalidData) |
| 714 | } | 780 | } | ||
| 715 | if *blobs >= maxSnapshotBlobs { | 781 | if *blobs >= maxSnapshotBlobs { | ||
| 716 | - | return fmt.Errorf("%w: too many VT blobs", ErrInvalidData) | 782 | + | return vt.DecodeStats{}, fmt.Errorf("%w: too many VT blobs", ErrInvalidData) |
| 717 | } | 783 | } | ||
| 718 | blob := r.b[:n] | 784 | blob := r.b[:n] | ||
| 719 | r.b = r.b[n:] | 785 | r.b = r.b[n:] | ||
| @@ -724,13 +790,16 @@ func preflightBlob(r *payloadReader, totals *vt.DecodeStats, blobs *uint64, hist | |||||
| 724 | stats, err = vt.PreflightVisibleBlob(blob) | 790 | stats, err = vt.PreflightVisibleBlob(blob) | ||
| 725 | } | 791 | } | ||
| 726 | if err != nil { | 792 | if err != nil { | ||
| 727 | - | return fmt.Errorf("%w: VT blob", ErrInvalidData) | 793 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT blob", ErrInvalidData) |
| 728 | } | 794 | } | ||
| 729 | if !totals.Add(stats) { | 795 | if !totals.Add(stats) { | ||
| 730 | - | return fmt.Errorf("%w: VT budget overflow", ErrInvalidData) | 796 | + | return vt.DecodeStats{}, fmt.Errorf("%w: VT budget overflow", ErrInvalidData) |
| 797 | + | } | |||
| 798 | + | if !budget.addAllocation(uint64(n)) || !budget.addProduct(stats.Chunks, 64) || !budget.addProduct(stats.Rows, 32) || !budget.addProduct(stats.Cells, 64) { | |||
| 799 | + | return vt.DecodeStats{}, fmt.Errorf("%w: decoded allocation budget", ErrInvalidData) | |||
| 731 | } | 800 | } | ||
| 732 | *blobs++ | 801 | *blobs++ | ||
| 733 | - | return nil | 802 | + | return stats, nil |
| 734 | } | 803 | } | ||
| 735 | 804 | ||||
| 736 | func preflightProcess(r *payloadReader) error { | 805 | func preflightProcess(r *payloadReader) error { | ||
M internal/usecase/snapshot/codec_test.go +80 -0
| @@ -125,6 +125,24 @@ func TestUnmarshalRejectsEveryV3PrefixAndTrailingGarbage(t *testing.T) { | |||
| 125 | 125 | reject(t, append(encoded, 0), ErrTrailingBytes) | |
| 126 | 126 | } | |
| 127 | 127 | ||
| 128 | + | func TestUnmarshalRejectsEmptyAndDuplicatePaneIDsDuringPreflight(t *testing.T) { | |
| 129 | + | sealed, tail := historyBlobs(t, nil) | |
| 130 | + | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |
| 131 | + | for _, tc := range []struct { | |
| 132 | + | name string | |
| 133 | + | panes []Pane | |
| 134 | + | }{ | |
| 135 | + | {name: "empty", panes: []Pane{{ID: "", Tail: tail, Visible: visible}}}, | |
| 136 | + | {name: "duplicate", panes: []Pane{{ID: "p", SealedChunks: sealed, Tail: tail, Visible: visible}, {ID: "p", Tail: tail, Visible: visible}}}, | |
| 137 | + | } { | |
| 138 | + | t.Run(tc.name, func(t *testing.T) { | |
| 139 | + | data, err := Marshal(Session{Name: "s", Tabs: []Tab{{Panes: tc.panes}}}) | |
| 140 | + | requireNoError(t, err) | |
| 141 | + | reject(t, data, ErrInvalidData) | |
| 142 | + | }) | |
| 143 | + | } | |
| 144 | + | } | |
| 145 | + | ||
| 128 | 146 | func TestUnmarshalRejectsInvalidTreeReference(t *testing.T) { | |
| 129 | 147 | sealed, tail := historyBlobs(t, nil) | |
| 130 | 148 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |
| @@ -137,6 +155,33 @@ func TestUnmarshalRejectsInvalidTreeReference(t *testing.T) { | |||
| 137 | 155 | } | |
| 138 | 156 | } | |
| 139 | 157 | ||
| 158 | + | func TestUnmarshalRejectsNonCanonicalBlobRolesDuringPreflight(t *testing.T) { | |
| 159 | + | sealed, tail := historyBlobs(t, [][]renderer.Cell{{{Rune: 'a'}}, {{Rune: 'b'}}, {{Rune: 'c'}}}) | |
| 160 | + | if len(sealed) == 0 { | |
| 161 | + | t.Fatal("expected a sealed chunk") | |
| 162 | + | } | |
| 163 | + | multiChunk := append([]byte(nil), sealed[0][:5]...) | |
| 164 | + | multiChunk = binary.BigEndian.AppendUint32(multiChunk, 2) | |
| 165 | + | multiChunk = append(multiChunk, sealed[0][9:]...) | |
| 166 | + | multiChunk = append(multiChunk, sealed[0][9:]...) | |
| 167 | + | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |
| 168 | + | emptyVisible := visibleBlob(t, nil) | |
| 169 | + | for _, tc := range []struct { | |
| 170 | + | name string | |
| 171 | + | tab Tab | |
| 172 | + | }{ | |
| 173 | + | {name: "sealed multi-chunk", tab: Tab{Cols: 1, Rows: 1, Panes: []Pane{{ID: "p", SealedChunks: [][]byte{multiChunk}, Tail: tail, Visible: visible}}}}, | |
| 174 | + | {name: "nonempty tail", tab: Tab{Cols: 1, Rows: 1, Panes: []Pane{{ID: "p", Tail: sealed[0], Visible: visible}}}}, | |
| 175 | + | {name: "empty visible geometry", tab: Tab{Panes: []Pane{{ID: "p", Tail: tail, Visible: emptyVisible}}}}, | |
| 176 | + | } { | |
| 177 | + | t.Run(tc.name, func(t *testing.T) { | |
| 178 | + | data, err := Marshal(Session{Name: "s", Tabs: []Tab{tc.tab}}) | |
| 179 | + | requireNoError(t, err) | |
| 180 | + | reject(t, data, ErrInvalidData) | |
| 181 | + | }) | |
| 182 | + | } | |
| 183 | + | } | |
| 184 | + | ||
| 140 | 185 | func TestUnmarshalRejectsMalformedTerminalBlobWithoutPanic(t *testing.T) { | |
| 141 | 186 | tail := mustHistoryTail(t) | |
| 142 | 187 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |
| @@ -236,6 +281,20 @@ func TestUnmarshalRejectsProcessWithoutArgv(t *testing.T) { | |||
| 236 | 281 | } | |
| 237 | 282 | } | |
| 238 | 283 | ||
| 284 | + | func TestPreflightBudgetRejectsAggregateDecodedAllocation(t *testing.T) { | |
| 285 | + | budget := preflightBudget{alloc: maxSnapshotDecodedAllocation - 1} | |
| 286 | + | if budget.addProduct(1, 2) { | |
| 287 | + | t.Fatal("allocation budget accepted an over-limit aggregate") | |
| 288 | + | } | |
| 289 | + | } | |
| 290 | + | ||
| 291 | + | func TestPreflightRejectsBroadTreeBeforeSecondPassAllocation(t *testing.T) { | |
| 292 | + | body := v3WideTreeManifest(32_769) | |
| 293 | + | if err := preflightSession(body); !errors.Is(err, ErrInvalidData) { | |
| 294 | + | t.Fatalf("preflightSession() error = %v, want %v", err, ErrInvalidData) | |
| 295 | + | } | |
| 296 | + | } | |
| 297 | + | ||
| 239 | 298 | func TestUnmarshalRejectsAggregateOpaqueBlobDeclarations(t *testing.T) { | |
| 240 | 299 | body := v3Manifest(math.MaxUint32, nil, nil, math.MaxUint32, nil, math.MaxUint32, nil) | |
| 241 | 300 | reject(t, v3Envelope(version, 0, body), ErrInvalidData) | |
| @@ -318,6 +377,27 @@ func v3EnvelopeWithLength(v, flags uint16, body []byte, length uint32) []byte { | |||
| 318 | 377 | } | |
| 319 | 378 | ||
| 320 | 379 | // v3Manifest creates just enough session/tab/pane framing for hostile declarations. | |
| 380 | + | // v3WideTreeManifest declares a broad but compact tree without allocating a | |
| 381 | + | // matching layout.Node graph. It exercises preflight's manifest budget. | |
| 382 | + | func v3WideTreeManifest(children uint16) []byte { | |
| 383 | + | var b []byte | |
| 384 | + | putS := func(s string) { b = binary.BigEndian.AppendUint16(b, uint16(len(s))); b = append(b, s...) } | |
| 385 | + | putS("s") | |
| 386 | + | b = append(b, make([]byte, 8+2)...) | |
| 387 | + | b = binary.BigEndian.AppendUint16(b, 1) | |
| 388 | + | putS("") | |
| 389 | + | b = append(b, make([]byte, 2+2+8)...) | |
| 390 | + | putS("") | |
| 391 | + | b = append(b, 1, byte(layout.Horizontal)) // split node | |
| 392 | + | b = binary.BigEndian.AppendUint16(b, children) | |
| 393 | + | for range children { | |
| 394 | + | b = append(b, 0) // leaf | |
| 395 | + | putS("") | |
| 396 | + | } | |
| 397 | + | b = binary.BigEndian.AppendUint16(b, 0) // panes | |
| 398 | + | return b | |
| 399 | + | } | |
| 400 | + | ||
| 321 | 401 | func v3Manifest(sealed uint32, sealedLens []uint32, sealedData []byte, tailLen uint32, tailData []byte, visibleLen uint32, visibleData []byte) []byte { | |
| 322 | 402 | var b []byte | |
| 323 | 403 | putS := func(s string) { b = binary.BigEndian.AppendUint16(b, uint16(len(s))); b = append(b, s...) } | |
| @@ -125,6 +125,24 @@ func TestUnmarshalRejectsEveryV3PrefixAndTrailingGarbage(t *testing.T) { | |||||
| 125 | reject(t, append(encoded, 0), ErrTrailingBytes) | 125 | reject(t, append(encoded, 0), ErrTrailingBytes) | ||
| 126 | } | 126 | } | ||
| 127 | 127 | ||||
| 128 | + | func TestUnmarshalRejectsEmptyAndDuplicatePaneIDsDuringPreflight(t *testing.T) { | |||
| 129 | + | sealed, tail := historyBlobs(t, nil) | |||
| 130 | + | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |||
| 131 | + | for _, tc := range []struct { | |||
| 132 | + | name string | |||
| 133 | + | panes []Pane | |||
| 134 | + | }{ | |||
| 135 | + | {name: "empty", panes: []Pane{{ID: "", Tail: tail, Visible: visible}}}, | |||
| 136 | + | {name: "duplicate", panes: []Pane{{ID: "p", SealedChunks: sealed, Tail: tail, Visible: visible}, {ID: "p", Tail: tail, Visible: visible}}}, | |||
| 137 | + | } { | |||
| 138 | + | t.Run(tc.name, func(t *testing.T) { | |||
| 139 | + | data, err := Marshal(Session{Name: "s", Tabs: []Tab{{Panes: tc.panes}}}) | |||
| 140 | + | requireNoError(t, err) | |||
| 141 | + | reject(t, data, ErrInvalidData) | |||
| 142 | + | }) | |||
| 143 | + | } | |||
| 144 | + | } | |||
| 145 | + | ||||
| 128 | func TestUnmarshalRejectsInvalidTreeReference(t *testing.T) { | 146 | func TestUnmarshalRejectsInvalidTreeReference(t *testing.T) { | ||
| 129 | sealed, tail := historyBlobs(t, nil) | 147 | sealed, tail := historyBlobs(t, nil) | ||
| 130 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | 148 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | ||
| @@ -137,6 +155,33 @@ func TestUnmarshalRejectsInvalidTreeReference(t *testing.T) { | |||||
| 137 | } | 155 | } | ||
| 138 | } | 156 | } | ||
| 139 | 157 | ||||
| 158 | + | func TestUnmarshalRejectsNonCanonicalBlobRolesDuringPreflight(t *testing.T) { | |||
| 159 | + | sealed, tail := historyBlobs(t, [][]renderer.Cell{{{Rune: 'a'}}, {{Rune: 'b'}}, {{Rune: 'c'}}}) | |||
| 160 | + | if len(sealed) == 0 { | |||
| 161 | + | t.Fatal("expected a sealed chunk") | |||
| 162 | + | } | |||
| 163 | + | multiChunk := append([]byte(nil), sealed[0][:5]...) | |||
| 164 | + | multiChunk = binary.BigEndian.AppendUint32(multiChunk, 2) | |||
| 165 | + | multiChunk = append(multiChunk, sealed[0][9:]...) | |||
| 166 | + | multiChunk = append(multiChunk, sealed[0][9:]...) | |||
| 167 | + | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | |||
| 168 | + | emptyVisible := visibleBlob(t, nil) | |||
| 169 | + | for _, tc := range []struct { | |||
| 170 | + | name string | |||
| 171 | + | tab Tab | |||
| 172 | + | }{ | |||
| 173 | + | {name: "sealed multi-chunk", tab: Tab{Cols: 1, Rows: 1, Panes: []Pane{{ID: "p", SealedChunks: [][]byte{multiChunk}, Tail: tail, Visible: visible}}}}, | |||
| 174 | + | {name: "nonempty tail", tab: Tab{Cols: 1, Rows: 1, Panes: []Pane{{ID: "p", Tail: sealed[0], Visible: visible}}}}, | |||
| 175 | + | {name: "empty visible geometry", tab: Tab{Panes: []Pane{{ID: "p", Tail: tail, Visible: emptyVisible}}}}, | |||
| 176 | + | } { | |||
| 177 | + | t.Run(tc.name, func(t *testing.T) { | |||
| 178 | + | data, err := Marshal(Session{Name: "s", Tabs: []Tab{tc.tab}}) | |||
| 179 | + | requireNoError(t, err) | |||
| 180 | + | reject(t, data, ErrInvalidData) | |||
| 181 | + | }) | |||
| 182 | + | } | |||
| 183 | + | } | |||
| 184 | + | ||||
| 140 | func TestUnmarshalRejectsMalformedTerminalBlobWithoutPanic(t *testing.T) { | 185 | func TestUnmarshalRejectsMalformedTerminalBlobWithoutPanic(t *testing.T) { | ||
| 141 | tail := mustHistoryTail(t) | 186 | tail := mustHistoryTail(t) | ||
| 142 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | 187 | visible := visibleBlob(t, [][]renderer.Cell{{{Rune: 'v'}}}) | ||
| @@ -236,6 +281,20 @@ func TestUnmarshalRejectsProcessWithoutArgv(t *testing.T) { | |||||
| 236 | } | 281 | } | ||
| 237 | } | 282 | } | ||
| 238 | 283 | ||||
| 284 | + | func TestPreflightBudgetRejectsAggregateDecodedAllocation(t *testing.T) { | |||
| 285 | + | budget := preflightBudget{alloc: maxSnapshotDecodedAllocation - 1} | |||
| 286 | + | if budget.addProduct(1, 2) { | |||
| 287 | + | t.Fatal("allocation budget accepted an over-limit aggregate") | |||
| 288 | + | } | |||
| 289 | + | } | |||
| 290 | + | ||||
| 291 | + | func TestPreflightRejectsBroadTreeBeforeSecondPassAllocation(t *testing.T) { | |||
| 292 | + | body := v3WideTreeManifest(32_769) | |||
| 293 | + | if err := preflightSession(body); !errors.Is(err, ErrInvalidData) { | |||
| 294 | + | t.Fatalf("preflightSession() error = %v, want %v", err, ErrInvalidData) | |||
| 295 | + | } | |||
| 296 | + | } | |||
| 297 | + | ||||
| 239 | func TestUnmarshalRejectsAggregateOpaqueBlobDeclarations(t *testing.T) { | 298 | func TestUnmarshalRejectsAggregateOpaqueBlobDeclarations(t *testing.T) { | ||
| 240 | body := v3Manifest(math.MaxUint32, nil, nil, math.MaxUint32, nil, math.MaxUint32, nil) | 299 | body := v3Manifest(math.MaxUint32, nil, nil, math.MaxUint32, nil, math.MaxUint32, nil) | ||
| 241 | reject(t, v3Envelope(version, 0, body), ErrInvalidData) | 300 | reject(t, v3Envelope(version, 0, body), ErrInvalidData) | ||
| @@ -318,6 +377,27 @@ func v3EnvelopeWithLength(v, flags uint16, body []byte, length uint32) []byte { | |||||
| 318 | } | 377 | } | ||
| 319 | 378 | ||||
| 320 | // v3Manifest creates just enough session/tab/pane framing for hostile declarations. | 379 | // v3Manifest creates just enough session/tab/pane framing for hostile declarations. | ||
| 380 | + | // v3WideTreeManifest declares a broad but compact tree without allocating a | |||
| 381 | + | // matching layout.Node graph. It exercises preflight's manifest budget. | |||
| 382 | + | func v3WideTreeManifest(children uint16) []byte { | |||
| 383 | + | var b []byte | |||
| 384 | + | putS := func(s string) { b = binary.BigEndian.AppendUint16(b, uint16(len(s))); b = append(b, s...) } | |||
| 385 | + | putS("s") | |||
| 386 | + | b = append(b, make([]byte, 8+2)...) | |||
| 387 | + | b = binary.BigEndian.AppendUint16(b, 1) | |||
| 388 | + | putS("") | |||
| 389 | + | b = append(b, make([]byte, 2+2+8)...) | |||
| 390 | + | putS("") | |||
| 391 | + | b = append(b, 1, byte(layout.Horizontal)) // split node | |||
| 392 | + | b = binary.BigEndian.AppendUint16(b, children) | |||
| 393 | + | for range children { | |||
| 394 | + | b = append(b, 0) // leaf | |||
| 395 | + | putS("") | |||
| 396 | + | } | |||
| 397 | + | b = binary.BigEndian.AppendUint16(b, 0) // panes | |||
| 398 | + | return b | |||
| 399 | + | } | |||
| 400 | + | ||||
| 321 | func v3Manifest(sealed uint32, sealedLens []uint32, sealedData []byte, tailLen uint32, tailData []byte, visibleLen uint32, visibleData []byte) []byte { | 401 | func v3Manifest(sealed uint32, sealedLens []uint32, sealedData []byte, tailLen uint32, tailData []byte, visibleLen uint32, visibleData []byte) []byte { | ||
| 322 | var b []byte | 402 | var b []byte | ||
| 323 | putS := func(s string) { b = binary.BigEndian.AppendUint16(b, uint16(len(s))); b = append(b, s...) } | 403 | putS := func(s string) { b = binary.BigEndian.AppendUint16(b, uint16(len(s))); b = append(b, s...) } | ||