refactor(operator): use pointers for cons key

imua-xyz · Apr 3, 2024 · 48f12a5 · 48f12a5
1 parent b310aa8
commit 48f12a5
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 94 deletions.
diff --git a/x/operator/keeper/consensus_keys.go b/x/operator/keeper/consensus_keys.go
@@ -24,12 +24,13 @@ func (k *Keeper) Logger(ctx sdk.Context) log.Logger {
 // and chain id. By doing this, an operator is consenting to be an operator on the given chain.
 // If a key already exists, it will be overwritten and the change in voting power will flow
 // through to the validator set.
+// TODO: Rationalize with k.OptIn
 func (k *Keeper) SetOperatorConsKeyForChainID(
 	ctx sdk.Context,
 	opAccAddr sdk.AccAddress,
 	chainID string,
 	// should be tm-ed25519
-	consKey tmprotocrypto.PublicKey,
+	consKey *tmprotocrypto.PublicKey,
 ) error {
 	// check if we are an operator
 	if !k.IsOperator(ctx, opAccAddr) {
@@ -44,23 +45,22 @@ func (k *Keeper) SetOperatorConsKeyForChainID(
 		return assetstypes.ErrUnknownAppChainID
 	}
 	// if opting out, do not allow key replacement
+	// TODO: merge with the functionality in state_update.go
 	if k.IsOperatorOptingOutFromChainID(ctx, opAccAddr, chainID) {
 		return types.ErrAlreadyOptingOut
 	}
 	// convert to bytes
 	bz, err := consKey.Marshal()
 	if err != nil {
 		return errorsmod.Wrap(
-			err,
-			"SetOperatorConsKeyForChainID: error occurred when marshal public key",
+			err, "SetOperatorConsKeyForChainID: cannot marshal public key",
 		)
 	}
 	// convert to address for reverse lookup
 	consAddr, err := types.TMCryptoPublicKeyToConsAddr(consKey)
 	if err != nil {
 		return errorsmod.Wrap(
-			err,
-			"SetOperatorConsKeyForChainID: error occurred when convert public key to consensus address",
+			err, "SetOperatorConsKeyForChainID: cannot convert pub key to consensus address",
 		)
 	}
 	// check if the key is already in use by another operator
@@ -94,24 +94,32 @@ func (k *Keeper) SetOperatorConsKeyForChainID(
 			panic(err)
 		}
 		if !alreadyRecorded {
-			if err := k.setOperatorPrevConsKeyForChainID(ctx, opAccAddr, chainID, prevKey); err != nil {
+			if err := k.setOperatorPrevConsKeyForChainID(
+				ctx, opAccAddr, chainID, prevKey,
+			); err != nil {
 				// this should not happen
 				panic(err)
 			}
 		}
 	}
-	// 	k.setOperatorConsKeyForChainID(ctx, opAccAddr, chainID, bz)
-	// 	return nil
-	// }
+	k.setOperatorConsKeyForChainID(ctx, opAccAddr, consAddr, chainID, bz)
+	if found {
+		if !alreadyRecorded {
+			k.Hooks().AfterOperatorKeyReplacement(ctx, opAccAddr, prevKey, consKey, chainID)
+		}
+	} else {
+		k.Hooks().AfterOperatorOptIn(ctx, opAccAddr, chainID, consKey)
+	}
+	return nil
+}
 
-	// // setOperatorConsKeyForChainID is the internal private version. It performs
-	// // no error checking of the input.
-	// func (k Keeper) setOperatorConsKeyForChainID(
-	// 	ctx sdk.Context,
-	// 	opAccAddr sdk.AccAddress,
-	// 	chainID string,
-	// 	bz []byte,
-	// ) {
+// setOperatorConsKeyForChainID is the internal private version. It performs
+// no error checking of the input. The caller must do the error checking
+// and then call this function.
+func (k Keeper) setOperatorConsKeyForChainID(
+	ctx sdk.Context, opAccAddr sdk.AccAddress, consAddr sdk.ConsAddress,
+	chainID string, bz []byte,
+) {
 	store := ctx.KVStore(k.storeKey)
 	// forward lookup
 	// given operator address and chain id, find the consensus key,
@@ -132,14 +140,6 @@ func (k *Keeper) SetOperatorConsKeyForChainID(
 	// this pruning will be triggered by the app chain module and will not be
 	// recorded here.
 	store.Set(types.KeyForChainIDAndConsKeyToOperator(chainID, consAddr), opAccAddr.Bytes())
-	if found {
-		if !alreadyRecorded {
-			k.Hooks().AfterOperatorKeyReplacement(ctx, opAccAddr, prevKey, consKey, chainID)
-		}
-	} else {
-		k.Hooks().AfterOperatorOptIn(ctx, opAccAddr, chainID, consKey)
-	}
-	return nil
 }
 
 // setOperatorPrevConsKeyForChainID sets the previous (consensus) public key for the given
@@ -150,7 +150,7 @@ func (k *Keeper) setOperatorPrevConsKeyForChainID(
 	ctx sdk.Context,
 	opAccAddr sdk.AccAddress,
 	chainID string,
-	prevKey tmprotocrypto.PublicKey,
+	prevKey *tmprotocrypto.PublicKey,
 ) error {
 	bz, err := prevKey.Marshal()
 	if err != nil {
@@ -169,7 +169,7 @@ func (k *Keeper) setOperatorPrevConsKeyForChainID(
 // to 0 in the validator update.
 func (k *Keeper) GetOperatorPrevConsKeyForChainID(
 	ctx sdk.Context, opAccAddr sdk.AccAddress, chainID string,
-) (found bool, key tmprotocrypto.PublicKey, err error) {
+) (found bool, key *tmprotocrypto.PublicKey, err error) {
 	// check if we are an operator
 	if !k.IsOperator(ctx, opAccAddr) {
 		err = delegationtypes.ErrOperatorNotExist
@@ -184,19 +184,19 @@ func (k *Keeper) GetOperatorPrevConsKeyForChainID(
 	return
 }
 
-// getOperatorPrevConsKeyForChainID is the internal version of
-// GetOperatorPrevConsKeyForChainID.
+// getOperatorPrevConsKeyForChainID is the internal version of GetOperatorPrevConsKeyForChainID.
 // It performs no error checking of the input.
 func (k *Keeper) getOperatorPrevConsKeyForChainID(
 	ctx sdk.Context,
 	opAccAddr sdk.AccAddress,
 	chainID string,
-) (found bool, key tmprotocrypto.PublicKey, err error) {
+) (found bool, key *tmprotocrypto.PublicKey, err error) {
 	store := ctx.KVStore(k.storeKey)
 	res := store.Get(types.KeyForOperatorAndChainIDToPrevConsKey(opAccAddr, chainID))
 	if res == nil {
 		return
 	}
+	key = &tmprotocrypto.PublicKey{}
 	if err = key.Unmarshal(res); err != nil {
 		return
 	}
@@ -210,7 +210,7 @@ func (k *Keeper) GetOperatorConsKeyForChainID(
 	ctx sdk.Context,
 	opAccAddr sdk.AccAddress,
 	chainID string,
-) (found bool, key tmprotocrypto.PublicKey, err error) {
+) (found bool, key *tmprotocrypto.PublicKey, err error) {
 	// check if we are an operator
 	if !k.IsOperator(ctx, opAccAddr) {
 		err = delegationtypes.ErrOperatorNotExist
@@ -232,25 +232,27 @@ func (k *Keeper) getOperatorConsKeyForChainID(
 	ctx sdk.Context,
 	opAccAddr sdk.AccAddress,
 	chainID string,
-) (found bool, key tmprotocrypto.PublicKey, err error) {
+) (found bool, key *tmprotocrypto.PublicKey, err error) {
 	store := ctx.KVStore(k.storeKey)
 	res := store.Get(types.KeyForOperatorAndChainIDToConsKey(opAccAddr, chainID))
 	if res == nil {
 		return
 	}
+	key = &tmprotocrypto.PublicKey{}
 	if err = key.Unmarshal(res); err != nil {
 		return
 	}
 	return true, key, nil
 }
 
-// GetChainIDsAndKeysForOperator gets the chain ids for which the given operator address has set a
+// GetChainIDsAndKeysForOperator gets the chain ids for which the given operator address has set
+// a
 // (consensus) public key. TODO: would it be better to make this a key per operator?
 // This is intentionally an array of strings because I don't see the utility for the vote power
 // or the public key here. If we need it, we can add it later.
 func (k *Keeper) GetChainIDsAndKeysForOperator(
 	ctx sdk.Context, opAccAddr sdk.AccAddress,
-) (chainIDs []string, consKeys []tmprotocrypto.PublicKey) {
+) (chainIDs []string, consKeys []*tmprotocrypto.PublicKey) {
 	// check if we are an operator
 	if !k.IsOperator(ctx, opAccAddr) {
 		return
@@ -268,7 +270,7 @@ func (k *Keeper) GetChainIDsAndKeysForOperator(
 	for ; iterator.Valid(); iterator.Next() {
 		// the key returned is the full key, with the prefix. drop the prefix and the length.
 		chainID := string(iterator.Key()[len(prefix)+8:])
-		var key tmprotocrypto.PublicKey
+		var key *tmprotocrypto.PublicKey
 		if err := key.Unmarshal(iterator.Value()); err != nil {
 			// grave error because we are the ones who stored this information in the first
 			// place
@@ -285,7 +287,7 @@ func (k *Keeper) GetChainIDsAndKeysForOperator(
 // jailed or frozen operators.
 func (k *Keeper) GetOperatorsForChainID(
 	ctx sdk.Context, chainID string,
-) (addrs []sdk.AccAddress, pubKeys []tmprotocrypto.PublicKey) {
+) (addrs []sdk.AccAddress, pubKeys []*tmprotocrypto.PublicKey) {
 	if !k.assetsKeeper.AppChainInfoIsExist(ctx, chainID) {
 		return nil, nil
 	}
@@ -306,7 +308,7 @@ func (k *Keeper) GetOperatorsForChainID(
 		// so just drop it and convert to sdk.AccAddress
 		addr := iterator.Key()[len(prefix):]
 		res := iterator.Value()
-		var ret tmprotocrypto.PublicKey
+		var ret *tmprotocrypto.PublicKey
 		if err := ret.Unmarshal(res); err != nil {
 			// grave error
 			panic(err)
@@ -426,6 +428,6 @@ func (k *Keeper) Jail(sdk.Context, sdk.ConsAddress, string) {}
 
 func (k *Keeper) GetActiveOperatorsForChainID(
 	sdk.Context, string,
-) ([]sdk.AccAddress, []tmprotocrypto.PublicKey) {
+) ([]sdk.AccAddress, []*tmprotocrypto.PublicKey) {
 	return nil, nil
 }
diff --git a/x/operator/keeper/grpc_query.go b/x/operator/keeper/grpc_query.go
@@ -35,6 +35,6 @@ func (k *Keeper) QueryOperatorConsKeyForChainID(
 		return nil, errors.New("no key assigned")
 	}
 	return &operatortypes.QueryOperatorConsKeyResponse{
-		PublicKey: key,
+		PublicKey: *key,
 	}, nil
 }
diff --git a/x/operator/keeper/msg_server.go b/x/operator/keeper/msg_server.go
@@ -2,9 +2,6 @@ package keeper
 
 import (
 	context "context"
-	"encoding/base64"
-
-	tmprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 
 	"github.com/ExocoreNetwork/exocore/x/operator/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -34,7 +31,7 @@ func (k *Keeper) OptInToCosmosChain(
 	if err != nil {
 		return nil, err
 	}
-	key, err := stringToPubKey(req.PublicKey)
+	key, err := types.StringToPubKey(req.PublicKey)
 	if err != nil {
 		return nil, err
 	}
@@ -65,16 +62,3 @@ func (k *Keeper) InitOptOutFromCosmosChain(
 	}
 	return &types.InitOptOutFromCosmosChainResponse{}, nil
 }
-
-func stringToPubKey(pubKey string) (key tmprotocrypto.PublicKey, err error) {
-	pubKeyBytes, err := base64.StdEncoding.DecodeString(pubKey)
-	if err != nil {
-		return
-	}
-	subscriberTMConsKey := tmprotocrypto.PublicKey{
-		Sum: &tmprotocrypto.PublicKey_Ed25519{
-			Ed25519: pubKeyBytes,
-		},
-	}
-	return subscriberTMConsKey, nil
-}
diff --git a/x/operator/types/consensus_keys.go b/x/operator/types/consensus_keys.go
@@ -41,13 +41,13 @@ func HexStringToPubKey(key string) (*tmprotocrypto.PublicKey, error) {
 }
 
 // StringToPubKey converts a base64-encoded public key to a tendermint public key.
-// Typically, this function is fed an input from ParseConsumerKeyFromJson.
-func StringToPubKey(pubKey string) (key tmprotocrypto.PublicKey, err error) {
+// Typically, this function is fed an input from ParseConsumerKeyFromJSON.
+func StringToPubKey(pubKey string) (key *tmprotocrypto.PublicKey, err error) {
 	pubKeyBytes, err := base64.StdEncoding.DecodeString(pubKey)
 	if err != nil {
 		return
 	}
-	subscriberTMConsKey := tmprotocrypto.PublicKey{
+	subscriberTMConsKey := &tmprotocrypto.PublicKey{
 		Sum: &tmprotocrypto.PublicKey_Ed25519{
 			Ed25519: pubKeyBytes,
 		},