add unjail functionality

proto/babylon/btcstaking/v1/events.proto

@@ -46,13 +46,21 @@ message EventPowerDistUpdate {
     bytes pk = 1 [ (gogoproto.customtype) = "github.com/babylonlabs-io/babylon/types.BIP340PubKey" ];
   }
 
+  // EventUnjailedFinalityProvider defines an event that a jailed finality provider
+  // is unjailed after the jailing period is passed
+  message EventUnjailedFinalityProvider {
+    bytes pk = 1 [ (gogoproto.customtype) = "github.com/babylonlabs-io/babylon/types.BIP340PubKey" ];
+  }
+
   // ev is the event that affects voting power distribution
   oneof ev {
     // slashed_fp means a finality provider is slashed
     EventSlashedFinalityProvider slashed_fp = 1;
     // jailed_fp means a finality provider is jailed
     EventJailedFinalityProvider jailed_fp = 2;
+    // unjailed_fp means a jailed finality provider is unjailed
+    EventUnjailedFinalityProvider unjailed_fp = 3;
     // btc_del_state_update means a BTC delegation's state is updated
-    EventBTCDelegationStateUpdate btc_del_state_update = 3;
+    EventBTCDelegationStateUpdate btc_del_state_update = 4;
   }
 }

x/btcstaking/keeper/finality_providers.go

@@ -139,6 +139,35 @@ func (k Keeper) JailFinalityProvider(ctx context.Context, fpBTCPK []byte) error
 	return nil
 }
 
+// UnjailFinalityProvider reverts the Jailed flag of a finality provider
+func (k Keeper) UnjailFinalityProvider(ctx context.Context, fpBTCPK []byte) error {
+	// ensure finality provider exists
+	fp, err := k.GetFinalityProvider(ctx, fpBTCPK)
+	if err != nil {
+		return err
+	}
+
+	// ensure finality provider is not jailed yet
+	if fp.IsJailed() {
+		return types.ErrFpAlreadyJailed
+	}
+
+	fp.Jailed = false
+	k.setFinalityProvider(ctx, fp)
+
+	btcTip := k.btclcKeeper.GetTipInfo(ctx)
+	if btcTip == nil {
+		return fmt.Errorf("failed to get current BTC tip")
+	}
+
+	// record unjailed event. The next `BeginBlock` will consume this
+	// event for updating the finality provider set
+	powerUpdateEvent := types.NewEventPowerDistUpdateWithUnjailedFP(fp.BtcPk)
+	k.addPowerDistUpdateEvent(ctx, btcTip.Height, powerUpdateEvent)
+
+	return nil
+}
+
 // finalityProviderStore returns the KVStore of the finality provider set
 // prefix: FinalityProviderKey
 // key: Bitcoin secp256k1 PK

x/btcstaking/keeper/power_dist_change.go

@@ -160,6 +160,8 @@ func (k Keeper) ProcessAllPowerDistUpdateEvents(
 	slashedFPs := map[string]struct{}{}
 	// a map where key is jailed finality providers' BTC PK
 	jailedFPs := map[string]struct{}{}
+	// a map where key is unjailed finality providers' BTC PK
+	unjailedFPs := map[string]struct{}{}
 
 	/*
 		filter and classify all events into new/expired BTC delegations and jailed/slashed FPs
@@ -189,6 +191,9 @@ func (k Keeper) ProcessAllPowerDistUpdateEvents(
 		case *types.EventPowerDistUpdate_JailedFp:
 			// record jailed fps
 			jailedFPs[typedEvent.JailedFp.Pk.MarshalHex()] = struct{}{}
+		case *types.EventPowerDistUpdate_UnjailedFp:
+			// record unjailed fps
+			unjailedFPs[typedEvent.UnjailedFp.Pk.MarshalHex()] = struct{}{}
 		}
 	}
 
@@ -225,6 +230,11 @@ func (k Keeper) ProcessAllPowerDistUpdateEvents(
 			fp.IsJailed = true
 		}
 
+		// set IsJailed to be false if the fp is unjailed
+		if _, ok := unjailedFPs[fpBTCPKHex]; ok {
+			fp.IsJailed = false
+		}
+
 		// add all BTC delegations that are not unbonded to the new finality provider
 		for j := range dc.FinalityProviders[i].BtcDels {
 			btcDel := *dc.FinalityProviders[i].BtcDels[j]

x/btcstaking/types/events.go

@@ -31,3 +31,13 @@ func NewEventPowerDistUpdateWithJailedFP(fpBTCPK *bbn.BIP340PubKey) *EventPowerD
 		},
 	}
 }
+
+func NewEventPowerDistUpdateWithUnjailedFP(fpBTCPK *bbn.BIP340PubKey) *EventPowerDistUpdate {
+	return &EventPowerDistUpdate{
+		Ev: &EventPowerDistUpdate_UnjailedFp{
+			UnjailedFp: &EventPowerDistUpdate_EventUnjailedFinalityProvider{
+				Pk: fpBTCPK,
+			},
+		},
+	}
+}