chore: backport register signer integration v3.5.6 (#1002)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: lukeiannucci

arbnode/batch_poster.go

@@ -80,9 +80,8 @@ var (
 
 	batchPosterFailureCounter = metrics.NewRegisteredCounter("arb/batchPoster/action/failure", nil)
 
-	usableBytesInBlob              = big.NewInt(int64(len(kzg4844.Blob{}) * 31 / 32))
-	blobTxBlobGasPerBlob           = big.NewInt(params.BlobTxBlobGasPerBlob)
-	FatalErrUnableToRegisterSigner = errors.New("unable to register signer")
+	usableBytesInBlob    = big.NewInt(int64(len(kzg4844.Blob{}) * 31 / 32))
+	blobTxBlobGasPerBlob = big.NewInt(params.BlobTxBlobGasPerBlob)
 )
 
 const (
@@ -1633,13 +1632,9 @@ func (b *BatchPoster) maybePostSequencerBatch(ctx context.Context) (bool, error)
 		return false, fmt.Errorf("batch was reverted, not posting any more batches")
 	}
 	if espressoSubmitter := b.streamer.espressoSubmitter; espressoSubmitter != nil {
-		registered := espressoSubmitter.GetKeyManager().HasRegistered()
-		if !registered {
-			log.Warn("ephemeral keys are not yet registered in Espresso TEE Contract")
-			err := espressoSubmitter.RegisterService()
-			if err != nil {
-				return false, fmt.Errorf("%w: %w", FatalErrUnableToRegisterSigner, err)
-			}
+		isRegistered, err := espressoSubmitter.GetKeyManager().CheckRegistration()
+		if !isRegistered {
+			return false, err
 		}
 	}
 
@@ -2314,7 +2309,7 @@ func (b *BatchPoster) Start(ctxIn context.Context) {
 				// Shutting down. No need to print the context canceled error.
 				return 0
 			}
-			if errors.Is(err, FatalErrUnableToRegisterSigner) {
+			if errors.Is(err, espresso_key_manager.FatalErrUnableToRegisterSigner) {
 				log.Warn(
 					"Espresso signer registration failed consecutively. Stopping.",
 					"retries", espressotee.EspressoMaxRetries,

arbnode/espresso_caff_node.go

@@ -507,8 +507,13 @@ func (n *EspressoCaffNode) Start(ctx context.Context) error {
 	}
 
 	if n.keyManager != nil {
-		registered := n.keyManager.HasRegistered()
-		if !registered {
+		if err := n.keyManager.Init(); err != nil {
+			return err
+		}
+		if state := n.keyManager.GetKeyManagerState(); state == espresso_key_manager.Registered {
+			log.Info("Caff node address is already registered on chain!")
+		} else {
+			log.Info("caff node completed init, trying to register signer")
 			if err := n.keyManager.RegisterService(); err != nil {
 				return err
 			}

espresso/key-manager/key_manager.go

@@ -3,6 +3,7 @@ package keymanager
 import (
 	"context"
 	"crypto/ecdsa"
+	"crypto/rand"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -29,20 +30,38 @@ const (
 	EMPTY = espressotee.EMPTY
 )
 
+var FatalErrUnableToRegisterSigner = errors.New("unable to register signer")
+
+type KeyManagerState int
+
+const (
+	Init KeyManagerState = iota
+	PendingRegistration
+	Registered
+)
+
+type state struct {
+	currentState KeyManagerState
+	attestation  []byte
+	data         []byte
+}
+
 // This is a private key derived from the test test test ... test junk BIP-39 mnemonic. It is a well known private key, so it should be fine to hardcode for tests.
 // I found it here: https://ethereum.stackexchange.com/questions/147078/hardhat-which-file-is-initial-state-in-such-as-the-mnemonic-and-20-accounts
 const TEST_PERSISTENT_KEY = "ac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
 const quoteFile = "/dev/attestation/quote"
 const userDataAttestationFile = "/dev/attestation/user_report_data"
 
 type EspressoKeyManagerInterface interface {
-	HasRegistered() bool
-	Register(getAttestationFunc func([]byte) ([]byte, error)) error
 	RegisterService() error
+	Init() error
 	GetCurrentKey() *ecdsa.PublicKey
 	SignPayload(message []byte) ([]byte, error)
 	SignMessage(message []byte) ([]byte, error)
 	TeeType() espressotee.TEE
+	GetKeyManagerState() KeyManagerState
+	CheckRegistration() (bool, error)
+	InitRegistration(getAttestationFunc func([]byte) ([]byte, error)) error
 }
 
 var _ EspressoKeyManagerInterface = &EspressoKeyManager{}
@@ -56,8 +75,8 @@ type EspressoKeyManager struct {
 	teeType     espressotee.TEE
 	serviceType espressotee.ServiceType
 
-	hasRegistered                          bool
 	espressoNitroAttestationVerifierClient *attestationverifierclient.EspressoAttestationVerifierClient
+	state                                  state
 }
 
 func NewEspressoKeyManager(
@@ -78,7 +97,12 @@ func NewEspressoKeyManager(
 	// is provided, we use that one. Otherwise, we read the enclave private key from the attestation path.
 	// Note: The current implementation only supports reading the key during key manager construction
 	// for the batch poster. Support for reading the caff node key will be added in a later PR.
-	if keyPairAttestationsPath != "" && chainID != 0 {
+	if teeType == espressotee.SGX {
+		privKey, err = ecdsa.GenerateKey(crypto.S256(), rand.Reader)
+		if err != nil {
+			panic(err)
+		}
+	} else if keyPairAttestationsPath != "" && chainID != 0 {
 		// Read enclave private key
 		privKey, err = espresso_tee_utils.ReadEnclavePrivateKey(keyPairAttestationsPath, chainID)
 		if err != nil {
@@ -120,15 +144,24 @@ func NewEspressoKeyManager(
 		teeType:                                teeType,
 		serviceType:                            serviceType,
 		espressoNitroAttestationVerifierClient: espressoNitroAttestationVerifierClient,
+		state: state{
+			currentState: Init,
+			attestation:  []byte{},
+			data:         []byte{},
+		},
 	}
 }
 
-func (k *EspressoKeyManager) HasRegistered() bool {
-	return k.hasRegistered
+func (k *EspressoKeyManager) hasRegistered() bool {
+	return k.state.currentState == Registered
+}
+
+func (k *EspressoKeyManager) GetKeyManagerState() KeyManagerState {
+	return k.state.currentState
 }
 
 func (k *EspressoKeyManager) VerifyRegistered() (bool, error) {
-	if k.hasRegistered {
+	if k.hasRegistered() {
 		return true, nil
 	}
 	pubKey, ok := k.privKey.Public().(*ecdsa.PublicKey)
@@ -143,6 +176,30 @@ func (k *EspressoKeyManager) VerifyRegistered() (bool, error) {
 	return ok, nil
 }
 
+func (k *EspressoKeyManager) CheckRegistration() (bool, error) {
+	state := k.state.currentState
+	switch state {
+	case Init:
+		log.Warn("ephemeral keys are not yet registered in Espresso TEE Contract, KeyManager in Init phase. Waiting for Zk proof to be generated")
+		err := k.Init()
+		if err != nil {
+			return false, fmt.Errorf("unable to init keymanager: %w", err)
+		}
+		return false, nil
+	case PendingRegistration:
+		log.Warn("ephemeral keys are not yet registered in Espresso TEE Contract, KeyManager in Registration phase")
+		err := k.RegisterService()
+		if err != nil {
+			return false, fmt.Errorf("%w: %w", FatalErrUnableToRegisterSigner, err)
+		}
+		return false, nil
+	case Registered:
+		return true, nil
+	default:
+		return false, fmt.Errorf("key manager in an unknown state: %v", state)
+	}
+}
+
 /*
  * This function will get the attestation in order to properly register the signing address on chain for a given TEE type
  */
@@ -201,20 +258,28 @@ func (k *EspressoKeyManager) PrepareRegisterService(getAttestationFunc func([]by
 	}
 }
 
-func (k *EspressoKeyManager) Register(getAttestationFunc func([]byte) ([]byte, error)) error {
-	if k.hasRegistered {
+func (k *EspressoKeyManager) InitRegistration(getAttestationFunc func([]byte) ([]byte, error)) error {
+	if k.hasRegistered() {
 		log.Info("EspressoKeyManager already registered")
 		return nil
 	}
 
+	// In tests we use TESTS tee type but the contract only accepts SGX tee type
+	if k.teeType == TESTS && k.serviceType != espressotee.Test {
+		k.teeType = SGX
+	}
+
 	// Check on-chain if we have already registered
 	hasRegistered, err := k.VerifyRegistered()
 	if err != nil {
 		return err
 	}
 	if hasRegistered {
-		k.hasRegistered = true
-		log.Info("Signer already registered on-chain")
+		k.state.currentState = Registered
+		k.state.attestation = []byte{}
+		k.state.data = []byte{}
+		signerAddr := crypto.PubkeyToAddress(k.privKey.PublicKey)
+		log.Info("Signer already registered on-chain", "signer address", signerAddr.Hex())
 		return nil
 	}
 
@@ -223,7 +288,18 @@ func (k *EspressoKeyManager) Register(getAttestationFunc func([]byte) ([]byte, e
 	if err != nil {
 		return err
 	}
-	err = k.espressoTEEVerifierCaller.RegisterService(k.dataPoster, attestation, data, uint8(k.teeType), k.serviceType)
+	k.state.currentState = PendingRegistration
+	k.state.attestation = attestation
+	k.state.data = data
+	return nil
+}
+
+func (k *EspressoKeyManager) RegisterService() error {
+	currentState := k.GetKeyManagerState()
+	if currentState != PendingRegistration {
+		return fmt.Errorf("invalid state to register signer: got %v, want PendingRegistration", currentState)
+	}
+	err := k.espressoTEEVerifierCaller.RegisterService(k.dataPoster, k.state.attestation, k.state.data, uint8(k.teeType), k.serviceType)
 	if err != nil {
 		return err
 	}
@@ -232,18 +308,18 @@ func (k *EspressoKeyManager) Register(getAttestationFunc func([]byte) ([]byte, e
 	log.Info("Register signer transaction sent", "signer address", signerAddr.Hex())
 
 	// Verify our address is actually registered in contract
-	hasRegistered, err = k.VerifyRegistered()
+	hasRegistered, err := k.VerifyRegistered()
 	if err != nil {
 		return err
 	}
 	if !hasRegistered {
 		return errors.New("address is not registered in contract even after successful transaction and retries")
 	}
 
-	k.hasRegistered = true
-	if k.teeType == TESTS {
-		k.teeType = SGX
-	}
+	// We are registered free up the memory
+	k.state.currentState = Registered
+	k.state.attestation = []byte{}
+	k.state.data = []byte{}
 	log.Info("Signer registration confirmed on-chain")
 	return nil
 }
@@ -265,15 +341,15 @@ func (k *EspressoKeyManager) SignMessage(message []byte) ([]byte, error) {
 	return arbutil.SignMessage(message, k.privKey)
 }
 
-func (k *EspressoKeyManager) RegisterService() error {
+func (k *EspressoKeyManager) Init() error {
 	teeType := k.TeeType()
 	switch teeType {
 	case SGX:
-		return k.Register(k.getAttestationQuote)
+		return k.InitRegistration(k.getAttestationQuote)
 	case NITRO:
-		return k.Register(k.getNitroAttestation)
+		return k.InitRegistration(k.getNitroAttestation)
 	case TESTS:
-		return k.Register(k.noOpSignerFunc)
+		return k.InitRegistration(k.noOpSignerFunc)
 	default:
 		return fmt.Errorf("unsupported tee Type: %d", teeType)
 	}

espresso/key-manager/key_manager_test.go

@@ -67,8 +67,8 @@ func TestEspressoKeyManager(t *testing.T) {
 		require.NotNil(t, km, "Key manager should not be nil")
 		assert.NotEmpty(t, km.GetCurrentKey(), "Public key should be set")
 		// assert.NotNil(t, km.privKey, "Private key should be set")
-		registered := km.HasRegistered()
-		assert.False(t, registered, "Should not be registered initially")
+		state := km.GetKeyManagerState()
+		assert.Equal(t, espresso_key_manager.Init, state, "Should not be registered initially")
 	})
 
 	// Test HasRegistered and Registry
@@ -78,8 +78,8 @@ func TestEspressoKeyManager(t *testing.T) {
 		mockEspressoTEEVerifierClient.On("RegisteredServices", mock.Anything, mock.Anything, mock.Anything).Return(false, nil).Once()
 		mockEspressoTEEVerifierClient.On("RegisteredServices", mock.Anything, mock.Anything, mock.Anything).Return(true, nil).Once()
 		km := espresso_key_manager.NewEspressoKeyManager(mockEspressoTEEVerifierClient, dataposter, dataSigner, espresso_key_manager.SGX, espressotee.Test, persistentPrivKey, "", "", 0)
-		registered := km.HasRegistered()
-		assert.False(t, registered, "Should start unregistered")
+		state := km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.Init, "Should start unregistered")
 
 		// Mock sign function
 		called := false
@@ -91,18 +91,18 @@ func TestEspressoKeyManager(t *testing.T) {
 			return []byte("mock-signature"), nil
 		}
 
-		// First registration
-		err := km.Register(getAttestationFunc)
+		// Register: should call sign function
+		err := km.InitRegistration(getAttestationFunc)
 		require.NoError(t, err, "Registry should succeed")
 		assert.True(t, called, "Sign function should be called")
-		registered = km.HasRegistered()
-		assert.True(t, registered, "Should be registered after call")
+		state = km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.PendingRegistration, "Should be pending registration after call")
 
 		// Second call (already registered)
-		called = false
-		err = km.Register(getAttestationFunc)
-		require.NoError(t, err, "Registry should succeed when already registered")
-		assert.False(t, called, "Sign function should not be called again")
+		err = km.RegisterService()
+		require.NoError(t, err, "Registry should succeed")
+		state = km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.Registered, "Should be registered after call")
 	})
 
 	// Test GetCurrentKey
@@ -159,8 +159,8 @@ func TestEspressoKeyManager(t *testing.T) {
 		mockEspressoTEEVerifierClient.On("RegisteredServices", mock.Anything, mock.Anything, mock.Anything).Return(false, nil).Once()
 		mockEspressoTEEVerifierClient.On("RegisteredServices", mock.Anything, mock.Anything, mock.Anything).Return(true, nil).Once()
 		km := espresso_key_manager.NewEspressoKeyManager(mockEspressoTEEVerifierClient, dataposter, dataSigner, espresso_key_manager.TESTS, espressotee.Test, persistentPrivKey, "", "", 0)
-		registered := km.HasRegistered()
-		assert.False(t, registered, "Should start unregistered")
+		state := km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.Init, "Should start unregistered")
 
 		// Mock sign function
 		called := false
@@ -171,20 +171,18 @@ func TestEspressoKeyManager(t *testing.T) {
 			return []byte{}, nil
 		}
 
-		// First registration
-		err := km.Register(getAttestationFunc)
+		// Register: should call sign function
+		err := km.InitRegistration(getAttestationFunc)
 		require.NoError(t, err, "Registry should succeed")
 		assert.True(t, called, "Sign function should be called")
-		registered = km.HasRegistered()
-		assert.True(t, registered, "Should be registered after call")
+		state = km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.PendingRegistration, "Should be pending registration after call")
 
 		// Second call (already registered)
-		called = false
-		err = km.Register(getAttestationFunc)
-		require.NoError(t, err, "Registry should succeed when already registered")
-		assert.False(t, called, "Sign function should not be called again")
-		registered = km.HasRegistered()
-		assert.True(t, registered, "Register function should still return true")
+		err = km.RegisterService()
+		require.NoError(t, err, "Registry should succeed")
+		state = km.GetKeyManagerState()
+		assert.Equal(t, state, espresso_key_manager.Registered, "Should be registered after call")
 	})
 
 	// Test Sign