diff --git a/wallet/recovery_test.go b/wallet/recovery_test.go new file mode 100644 index 0000000..8a66934 --- /dev/null +++ b/wallet/recovery_test.go @@ -0,0 +1,245 @@ +package wallet_test + +import ( + "runtime" + "testing" + + "github.com/roasbeef/btcwallet/wallet" +) + +// Harness holds the BranchRecoveryState being tested, the recovery window being +// used, provides access to the test object, and tracks the expected horizon +// and next unfound values. +type Harness struct { + t *testing.T + brs *wallet.BranchRecoveryState + recoveryWindow uint32 + expHorizon uint32 + expNextUnfound uint32 +} + +type ( + // Stepper is a generic interface that performs an action or assertion + // against a test Harness. + Stepper interface { + // Apply performs an action or assertion against branch recovery + // state held by the Harness. The step index is provided so + // that any failures can report which Step failed. + Apply(step int, harness *Harness) + } + + // InitialiDelta is a Step that verifies our first attempt to expand the + // branch recovery state's horizons tells us to derive a number of + // adddresses equal to the recovery window. + InitialDelta struct{} + + // CheckDelta is a Step that expands the branch recovery state's + // horizon, and checks that the returned delta meets our expected + // `delta`. + CheckDelta struct { + delta uint32 + } + + // CheckNumInvalid is a Step that asserts that the branch recovery + // state reports `total` invalid children with the current horizon. + CheckNumInvalid struct { + total uint32 + } + + // MarkInvalid is a Step that marks the `child` as invalid in the branch + // recovery state. + MarkInvalid struct { + child uint32 + } + + // ReportFound is a Step that reports `child` as being found to the + // branch recovery state. + ReportFound struct { + child uint32 + } +) + +// Apply extends the current horizon of the branch recovery state, and checks +// that the returned delta is equal to the test's recovery window. If the +// assertions pass, the harness's expected horizon is increased by the returned +// delta. +// +// NOTE: This should be used before applying any CheckDelta steps. +func (_ InitialDelta) Apply(i int, h *Harness) { + curHorizon, delta := h.brs.ExtendHorizon() + assertHorizon(h.t, i, curHorizon, h.expHorizon) + assertDelta(h.t, i, delta, h.recoveryWindow) + h.expHorizon += delta +} + +// Apply extends the current horizon of the branch recovery state, and checks +// that the returned delta is equal to the CheckDelta's child value. +func (d CheckDelta) Apply(i int, h *Harness) { + curHorizon, delta := h.brs.ExtendHorizon() + assertHorizon(h.t, i, curHorizon, h.expHorizon) + assertDelta(h.t, i, delta, d.delta) + h.expHorizon += delta +} + +// Apply queries the branch recovery state for the number of invalid children +// that lie between the last found address and the current horizon, and compares +// that to the CheckNumInvalid's total. +func (m CheckNumInvalid) Apply(i int, h *Harness) { + assertNumInvalid(h.t, i, h.brs.NumInvalidInHorizon(), m.total) +} + +// Apply marks the MarkInvalid's child index as invalid in the branch recovery +// state, and increments the harness's expected horizon. +func (m MarkInvalid) Apply(i int, h *Harness) { + h.brs.MarkInvalidChild(m.child) + h.expHorizon++ +} + +// Apply reports the ReportFound's child index as found in the branch recovery +// state. If the child index meets or exceeds our expected next unfound value, +// the expected value will be modified to be the child index + 1. Afterwards, +// this step asserts that the branch recovery state's next reported unfound +// value matches our potentially-updated value. +func (r ReportFound) Apply(i int, h *Harness) { + h.brs.ReportFound(r.child) + if r.child >= h.expNextUnfound { + h.expNextUnfound = r.child + 1 + } + assertNextUnfound(h.t, i, h.brs.NextUnfound(), h.expNextUnfound) +} + +// Compile-time checks to ensure our steps implement the Step interface. +var _ Stepper = InitialDelta{} +var _ Stepper = CheckDelta{} +var _ Stepper = CheckNumInvalid{} +var _ Stepper = MarkInvalid{} +var _ Stepper = ReportFound{} + +// TestBranchRecoveryState walks the BranchRecoveryState through a sequence of +// steps, verifying that: +// - the horizon is properly expanded in response to found addrs +// - report found children below or equal to previously found causes no change +// - marking invalid children expands the horizon +func TestBranchRecoveryState(t *testing.T) { + + const recoveryWindow = 10 + + recoverySteps := []Stepper{ + // First, check that expanding our horizon returns exactly the + // recovery window (10). + InitialDelta{}, + + // Expected horizon: 10. + + // Report finding the 2nd addr, this should cause our horizon + // to expand by 2. + ReportFound{1}, + CheckDelta{2}, + + // Expected horizon: 12. + + // Sanity check that expanding again reports zero delta, as + // nothing has changed. + CheckDelta{0}, + + // Now, report finding the 6th addr, which should expand our + // horizon to 16 with a detla of 4. + ReportFound{5}, + CheckDelta{4}, + + // Expected horizon: 16. + + // Sanity check that expanding again reports zero delta, as + // nothing has changed. + CheckDelta{0}, + + // Report finding child index 5 again, nothing should change. + ReportFound{5}, + CheckDelta{0}, + + // Report finding a lower index that what was last found, + // nothing should change. + ReportFound{4}, + CheckDelta{0}, + + // Moving on, report finding the 11th addr, which should extend + // our horizon to 21. + ReportFound{10}, + CheckDelta{5}, + + // Expected horizon: 21. + + // Before testing the lookahead expansion when encountering + // invalid child keys, check that we are correctly starting with + // no invalid keys. + CheckNumInvalid{0}, + + // Now that the window has been expanded, simulate deriving + // invalid keys in range of addrs that are being derived for the + // first time. The horizon will be incremented by one, as the + // recovery manager is expected to try and derive at least the + // next address. + MarkInvalid{17}, + CheckNumInvalid{1}, + CheckDelta{0}, + + // Expected horizon: 22. + + // Check that deriving a second invalid key shows both invalid + // indexes currently within the horizon. + MarkInvalid{18}, + CheckNumInvalid{2}, + CheckDelta{0}, + + // Expected horizon: 23. + + // Lastly, report finding the addr immediately after our two + // invalid keys. This should return our number of invalid keys + // within the horizon back to 0. + ReportFound{19}, + CheckNumInvalid{0}, + + // As the 20-th key was just marked found, our horizon will need + // to expand to 30. With the horizon at 23, the delta returned + // should be 7. + CheckDelta{7}, + CheckDelta{0}, + + // Expected horizon: 30. + } + + brs := wallet.NewBranchRecoveryState(recoveryWindow) + harness := &Harness{ + t: t, + brs: brs, + recoveryWindow: recoveryWindow, + } + + for i, step := range recoverySteps { + step.Apply(i, harness) + } +} + +func assertHorizon(t *testing.T, i int, have, want uint32) { + assertHaveWant(t, i, "incorrect horizon", have, want) +} + +func assertDelta(t *testing.T, i int, have, want uint32) { + assertHaveWant(t, i, "incorrect delta", have, want) +} + +func assertNextUnfound(t *testing.T, i int, have, want uint32) { + assertHaveWant(t, i, "incorrect next unfound", have, want) +} + +func assertNumInvalid(t *testing.T, i int, have, want uint32) { + assertHaveWant(t, i, "incorrect num invalid children", have, want) +} + +func assertHaveWant(t *testing.T, i int, msg string, have, want uint32) { + _, _, line, _ := runtime.Caller(2) + if want != have { + t.Fatalf("[line: %d, step: %d] %s: got %d, want %d", + line, i, msg, have, want) + } +}