diff --git a/instantout/reservation/interfaces.go b/instantout/reservation/interfaces.go index 04bf830d..23658d5a 100644 --- a/instantout/reservation/interfaces.go +++ b/instantout/reservation/interfaces.go @@ -8,14 +8,18 @@ import ( ) var ( - ErrReservationAlreadyExists = fmt.Errorf("reservation already exists") - ErrReservationNotFound = fmt.Errorf("reservation not found") + ErrReservationAlreadyExists = fmt.Errorf("reservation already exists") + ErrReservationNotFound = fmt.Errorf("reservation not found") + ErrTooManyActiveReservations = fmt.Errorf( + "too many active reservations", + ) ) const ( - KeyFamily = int32(42068) - DefaultConfTarget = int32(3) - IdLength = 32 + KeyFamily = int32(42068) + DefaultConfTarget = int32(3) + IdLength = 32 + maxActiveReservations = 1000 ) // Store is the interface that stores the reservations. diff --git a/instantout/reservation/manager.go b/instantout/reservation/manager.go index 930dff7b..2f4c80c4 100644 --- a/instantout/reservation/manager.go +++ b/instantout/reservation/manager.go @@ -26,6 +26,28 @@ type Manager struct { activeReservations map[ID]*FSM } +// finalStateObserver removes a reservation FSM from the active set once it +// reaches a terminal state. +type finalStateObserver struct { + manager *Manager + id ID + fsm *FSM +} + +// Notify implements the fsm.Observer interface. +func (o *finalStateObserver) Notify(notification fsm.Notification) { + if !isFinalState(notification.NextState) { + return + } + + o.manager.Lock() + defer o.manager.Unlock() + + if o.manager.activeReservations[o.id] == o.fsm { + delete(o.manager.activeReservations, o.id) + } +} + // NewManager creates a new reservation manager. func NewManager(cfg *Config) *Manager { return &Manager{ @@ -134,9 +156,19 @@ func (m *Manager) newReservation(ctx context.Context, currentHeight uint32, m.Unlock() return nil, ErrReservationAlreadyExists } + if len(m.activeReservations) >= maxActiveReservations { + m.Unlock() + return nil, ErrTooManyActiveReservations + } m.activeReservations[reservationID] = reservationFSM m.Unlock() + reservationFSM.RegisterObserver(&finalStateObserver{ + manager: m, + id: reservationID, + fsm: reservationFSM, + }) + initContext := &InitReservationContext{ reservationID: reservationID, serverPubkey: serverKey, @@ -187,6 +219,16 @@ func (m *Manager) RecoverReservations(ctx context.Context) error { return err } + activeCount := 0 + for _, reservation := range reservations { + if !isFinalState(reservation.State) { + activeCount++ + } + } + if activeCount > maxActiveReservations { + return ErrTooManyActiveReservations + } + for _, reservation := range reservations { if isFinalState(reservation.State) { continue @@ -199,6 +241,11 @@ func (m *Manager) RecoverReservations(ctx context.Context) error { reservationFSM := NewFSMFromReservation(m.cfg, reservation) m.activeReservations[reservation.ID] = reservationFSM + reservationFSM.RegisterObserver(&finalStateObserver{ + manager: m, + id: reservation.ID, + fsm: reservationFSM, + }) // As SendEvent can block, we'll start a goroutine to process // the event. diff --git a/instantout/reservation/manager_test.go b/instantout/reservation/manager_test.go index 0437955e..bf1d1a46 100644 --- a/instantout/reservation/manager_test.go +++ b/instantout/reservation/manager_test.go @@ -97,6 +97,11 @@ func TestManager(t *testing.T) { // We'll now expect the reservation to be expired. err = reservationFSM.DefaultObserver.WaitForState(ctxb, 5*time.Second, Spent) require.NoError(t, err) + + testContext.manager.Lock() + _, ok := testContext.manager.activeReservations[defaultReservationId] + testContext.manager.Unlock() + require.False(t, ok) } // TestManagerContinuesAfterInvalidNotification verifies that a malformed @@ -173,6 +178,38 @@ func TestManagerRejectsDuplicateReservation(t *testing.T) { ) } +// TestManagerLimitsActiveReservations verifies that server notifications +// cannot grow the active FSM set without bound. +func TestManagerLimitsActiveReservations(t *testing.T) { + testContext := newManagerTestContext(t) + + for i := range maxActiveReservations { + var id ID + id[0] = byte(i) + id[1] = byte(i >> 8) + testContext.manager.activeReservations[id] = NewFSM( + testContext.manager.cfg, + ) + } + + reservationFSM, err := testContext.manager.newReservation( + t.Context(), uint32(testContext.mockLnd.Height), + &swapserverrpc.ServerReservationNotification{ + ReservationId: defaultReservationId[:], + Value: uint64(defaultValue), + ServerKey: defaultPubkeyBytes, + Expiry: uint32(testContext.mockLnd.Height) + + defaultExpiry, + }, + ) + require.ErrorIs(t, err, ErrTooManyActiveReservations) + require.Nil(t, reservationFSM) + require.Len( + t, testContext.manager.activeReservations, + maxActiveReservations, + ) +} + // ManagerTestContext is a helper struct that contains all the necessary // components to test the reservation manager. type ManagerTestContext struct {