Add new endpoint for generating messages and refactor message generation

go/common/viewingkey/viewing_key.go

@@ -3,6 +3,7 @@ package viewingkey
 import (
 	"crypto/ecdsa"
 	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"math/big"
@@ -100,8 +101,10 @@ func (psc PersonalSignChecker) CheckSignature(encryptionToken string, signature
 
 	// create all possible hashes (for all the supported versions) of the message (needed for signature verification)
 	for _, version := range PersonalSignMessageSupportedVersions {
-		message := GeneratePersonalSignMessage(encryptionToken, chainID, version)
-		messageHash := accounts.TextHash([]byte(message))
+		messageHash, err := GenerateMessage(encryptionToken, chainID, version, PersonalSign, true)
+		if err != nil {
+			return nil, fmt.Errorf("cannot generate message. Cause %w", err)
+		}
 
 		// current signature is valid - return account address
 		address, err := CheckSignatureAndReturnAccountAddress(messageHash, signature)
@@ -118,9 +121,9 @@ func (e EIP712Checker) CheckSignature(encryptionToken string, signature []byte,
 		return nil, fmt.Errorf("invalid signaure length: %d", len(signature))
 	}
 
-	rawDataOptions, err := GenerateAuthenticationEIP712RawDataOptions(encryptionToken, chainID)
+	messageHash, err := GenerateMessage(encryptionToken, chainID, 1, EIP712Signature, true)
 	if err != nil {
-		return nil, fmt.Errorf("cannot generate eip712 message. Cause %w", err)
+		return nil, fmt.Errorf("cannot generate message. Cause %w", err)
 	}
 
 	// We transform the V from 27/28 to 0/1. This same change is made in Geth internals, for legacy reasons to be able
@@ -129,16 +132,12 @@ func (e EIP712Checker) CheckSignature(encryptionToken string, signature []byte,
 		signature[64] -= 27
 	}
 
-	for _, rawData := range rawDataOptions {
-		// create a hash of structured message (needed for signature verification)
-		hashBytes := crypto.Keccak256(rawData)
-
-		// current signature is valid - return account address
-		address, err := CheckSignatureAndReturnAccountAddress(hashBytes, signature)
-		if err == nil {
-			return address, nil
-		}
+	// current signature is valid - return account address
+	address, err := CheckSignatureAndReturnAccountAddress(messageHash, signature)
+	if err == nil {
+		return address, nil
 	}
+
 	return nil, errors.New("EIP 712 signature verification failed")
 }
 
@@ -249,10 +248,6 @@ func GenerateSignMessage(vkPubKey []byte) string {
 	return SignedMsgPrefix + hex.EncodeToString(vkPubKey)
 }
 
-func GeneratePersonalSignMessage(encryptionToken string, chainID int64, version int) string {
-	return fmt.Sprintf(PersonalSignMessageFormat, encryptionToken, chainID, version)
-}
-
 // getBytesFromTypedData creates EIP-712 compliant hash from typedData.
 // It involves hashing the message with its structure, hashing domain separator,
 // and then encoding both hashes with specific EIP-712 bytes to construct the final message format.
@@ -271,14 +266,71 @@ func getBytesFromTypedData(typedData apitypes.TypedData) ([]byte, error) {
 	return rawData, nil
 }
 
-// GenerateAuthenticationEIP712RawDataOptions generates all the options or raw data messages (bytes)
-// for an EIP-712 message used to authenticate an address with user
-// (currently only one option is supported, but function leaves room for future expansion of options)
-func GenerateAuthenticationEIP712RawDataOptions(userID string, chainID int64) ([][]byte, error) {
-	if len(userID) != UserIDHexLength {
-		return nil, fmt.Errorf("userID hex length must be %d, received %d", UserIDHexLength, len(userID))
+// CalculateUserIDHex CalculateUserID calculates userID from a public key
+// (we truncate it, because we want it to have length 20) and encode to hex strings
+func CalculateUserIDHex(publicKeyBytes []byte) string {
+	return hex.EncodeToString(CalculateUserID(publicKeyBytes))
+}
+
+// CalculateUserID calculates userID from a public key (we truncate it, because we want it to have length 20)
+func CalculateUserID(publicKeyBytes []byte) []byte {
+	return crypto.Keccak256Hash(publicKeyBytes).Bytes()[:20]
+}
+
+// CheckSignatureAndReturnAccountAddress checks if the signature is valid for hash of the message and checks if
+// signer is an address provided to the function.
+// It returns an address if the signature is valid and nil otherwise
+func CheckSignatureAndReturnAccountAddress(hashBytes []byte, signature []byte) (*gethcommon.Address, error) {
+	pubKeyBytes, err := crypto.Ecrecover(hashBytes, signature)
+	if err != nil {
+		return nil, err
+	}
+
+	pubKey, err := crypto.UnmarshalPubkey(pubKeyBytes)
+	if err != nil {
+		return nil, err
+	}
+
+	r := new(big.Int).SetBytes(signature[:32])
+	s := new(big.Int).SetBytes(signature[32:64])
+
+	// Verify the signature and return the result (all the checks above passed)
+	isSigValid := ecdsa.Verify(pubKey, hashBytes, r, s)
+	if isSigValid {
+		address := crypto.PubkeyToAddress(*pubKey)
+		return &address, nil
 	}
-	encryptionToken := "0x" + userID
+	return nil, fmt.Errorf("invalid signature")
+}
+
+type MessageGenerator interface {
+	generateMessage(encryptionToken string, chainID int64, version int, hash bool) ([]byte, error)
+}
+
+type (
+	PersonalMessageGenerator struct{}
+	EIP712MessageGenerator   struct{}
+)
+
+var messageGenerators = map[SignatureType]MessageGenerator{
+	PersonalSign:    PersonalMessageGenerator{},
+	EIP712Signature: EIP712MessageGenerator{},
+}
+
+// GenerateMessage generates a message for the given encryptionToken, chainID, version and signatureType
+func (p PersonalMessageGenerator) generateMessage(encryptionToken string, chainID int64, version int, hash bool) ([]byte, error) {
+	textMessage := fmt.Sprintf(PersonalSignMessageFormat, encryptionToken, chainID, version)
+	if hash {
+		return accounts.TextHash([]byte(textMessage)), nil
+	}
+	return []byte(textMessage), nil
+}
+
+func (e EIP712MessageGenerator) generateMessage(encryptionToken string, chainID int64, _ int, hash bool) ([]byte, error) {
+	if len(encryptionToken) != UserIDHexLength {
+		return nil, fmt.Errorf("userID hex length must be %d, received %d", UserIDHexLength, len(encryptionToken))
+	}
+	encryptionToken = "0x" + encryptionToken
 
 	domain := apitypes.TypedDataDomain{
 		Name:    EIP712DomainNameValue,
@@ -308,49 +360,27 @@ func GenerateAuthenticationEIP712RawDataOptions(userID string, chainID int64) ([
 		Message:     message,
 	}
 
-	rawDataOptions := make([][]byte, 0)
 	rawData, err := getBytesFromTypedData(newTypeElement)
 	if err != nil {
 		return nil, err
 	}
-	rawDataOptions = append(rawDataOptions, rawData)
-
-	return rawDataOptions, nil
-}
-
-// CalculateUserIDHex CalculateUserID calculates userID from a public key
-// (we truncate it, because we want it to have length 20) and encode to hex strings
-func CalculateUserIDHex(publicKeyBytes []byte) string {
-	return hex.EncodeToString(CalculateUserID(publicKeyBytes))
-}
-
-// CalculateUserID calculates userID from a public key (we truncate it, because we want it to have length 20)
-func CalculateUserID(publicKeyBytes []byte) []byte {
-	return crypto.Keccak256Hash(publicKeyBytes).Bytes()[:20]
-}
 
-// CheckSignatureAndReturnAccountAddress checks if the signature is valid for hash of the message and checks if
-// signer is an address provided to the function.
-// It returns an address if the signature is valid and nil otherwise
-func CheckSignatureAndReturnAccountAddress(hashBytes []byte, signature []byte) (*gethcommon.Address, error) {
-	pubKeyBytes, err := crypto.Ecrecover(hashBytes, signature)
+	// add the JSON message to the list of messages
+	jsonData, err := json.Marshal(newTypeElement)
 	if err != nil {
 		return nil, err
 	}
 
-	pubKey, err := crypto.UnmarshalPubkey(pubKeyBytes)
-	if err != nil {
-		return nil, err
+	if hash {
+		return crypto.Keccak256(rawData), nil
 	}
+	return jsonData, nil
+}
 
-	r := new(big.Int).SetBytes(signature[:32])
-	s := new(big.Int).SetBytes(signature[32:64])
-
-	// Verify the signature and return the result (all the checks above passed)
-	isSigValid := ecdsa.Verify(pubKey, hashBytes, r, s)
-	if isSigValid {
-		address := crypto.PubkeyToAddress(*pubKey)
-		return &address, nil
+func GenerateMessage(encryptionToken string, chainID int64, version int, signatureType SignatureType, hash bool) ([]byte, error) {
+	generator, exists := messageGenerators[signatureType]
+	if !exists {
+		return nil, fmt.Errorf("unsupported signature type")
 	}
-	return nil, fmt.Errorf("invalid signature")
+	return generator.generateMessage(encryptionToken, chainID, version, hash)
 }

go/enclave/vkhandler/vk_handler_test.go

@@ -7,7 +7,6 @@ import (
 
 	gethcommon "github.com/ethereum/go-ethereum/common"
 
-	"github.com/ethereum/go-ethereum/accounts"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/crypto/ecies"
 	"github.com/stretchr/testify/assert"
@@ -49,13 +48,14 @@ func TestCheckSignature(t *testing.T) {
 	userPrivKey, _, userID, userAddress := generateRandomUserKeys()
 
 	// Generate all message types and create map with the corresponding signature type
-	// Test EIP712 message format
-	EIP712MessageDataOptions, err := viewingkey.GenerateAuthenticationEIP712RawDataOptions(userID, chainID)
+	EIP712MessageHash, err := viewingkey.GenerateMessage(userID, chainID, 0, viewingkey.EIP712Signature, true)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	PersonalSignMessageHash, err := viewingkey.GenerateMessage(userID, chainID, viewingkey.PersonalSignMessageSupportedVersions[0], viewingkey.PersonalSign, true)
 	if err != nil {
 		t.Fatalf(err.Error())
 	}
-	EIP712MessageHash := crypto.Keccak256(EIP712MessageDataOptions[0])
-	PersonalSignMessageHash := accounts.TextHash([]byte(viewingkey.GeneratePersonalSignMessage(userID, chainID, viewingkey.PersonalSignMessageSupportedVersions[0])))
 
 	messages := map[string]MessageWithSignatureType{
 		"EIP712MessageHash": {
@@ -86,12 +86,15 @@ func TestVerifyViewingKey(t *testing.T) {
 	userPrivKey, vkPrivKey, userID, userAddress := generateRandomUserKeys()
 	// Generate all message types and create map with the corresponding signature type
 	// Test EIP712 message format
-	EIP712MessageDataOptions, err := viewingkey.GenerateAuthenticationEIP712RawDataOptions(userID, chainID)
+
+	EIP712MessageHash, err := viewingkey.GenerateMessage(userID, chainID, viewingkey.PersonalSignMessageSupportedVersions[0], viewingkey.EIP712Signature, true)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	PersonalSignMessageHash, err := viewingkey.GenerateMessage(userID, chainID, viewingkey.PersonalSignMessageSupportedVersions[0], viewingkey.PersonalSign, true)
 	if err != nil {
 		t.Fatalf(err.Error())
 	}
-	EIP712MessageHash := crypto.Keccak256(EIP712MessageDataOptions[0])
-	PersonalSignMessageHash := accounts.TextHash([]byte(viewingkey.GeneratePersonalSignMessage(userID, chainID, viewingkey.PersonalSignMessageSupportedVersions[0])))
 
 	messages := map[string]MessageWithSignatureType{
 		"EIP712MessageHash": {

tools/walletextension/api/routes.go

@@ -46,6 +46,10 @@ func NewHTTPRoutes(walletExt *walletextension.WalletExtension) []Route {
 			Name: common.APIVersion1 + common.PathJoin,
 			Func: httpHandler(walletExt, joinRequestHandler),
 		},
+		{
+			Name: common.APIVersion1 + common.PathGetMessage,
+			Func: httpHandler(walletExt, getMessageRequestHandler),
+		},
 		{
 			Name: common.APIVersion1 + common.PathAuthenticate,
 			Func: httpHandler(walletExt, authenticateRequestHandler),
@@ -490,3 +494,84 @@ func versionRequestHandler(walletExt *walletextension.WalletExtension, userConn
 		walletExt.Logger().Error("error writing success response", log.ErrKey, err)
 	}
 }
+
+// ethRequestHandler parses the user eth request, passes it on to the WE to proxy it and processes the response
+func getMessageRequestHandler(walletExt *walletextension.WalletExtension, conn userconn.UserConn) {
+	// read the request
+	body, err := conn.ReadRequest()
+	if err != nil {
+		handleError(conn, walletExt.Logger(), fmt.Errorf("error reading request: %w", err))
+		return
+	}
+
+	// read body of the request
+	var reqJSONMap map[string]interface{}
+	err = json.Unmarshal(body, &reqJSONMap)
+	if err != nil {
+		handleError(conn, walletExt.Logger(), fmt.Errorf("could not unmarshal address request - %w", err))
+		return
+	}
+
+	// get address from the request
+	encryptionToken, ok := reqJSONMap[common.JSONKeyEncryptionToken]
+	if !ok || len(encryptionToken.(string)) != common.MessageUserIDLen {
+		handleError(conn, walletExt.Logger(), fmt.Errorf("unable to read encryptionToken field from the request or it is not of correct length"))
+		return
+	}
+
+	// get formats from the request, if present
+	var formatsSlice []string
+	if formatsInterface, ok := reqJSONMap[common.JSONKeyFormats]; ok {
+		formats, ok := formatsInterface.([]interface{})
+		if !ok {
+			handleError(conn, walletExt.Logger(), fmt.Errorf("formats field is not an array"))
+			return
+		}
+
+		for _, f := range formats {
+			formatStr, ok := f.(string)
+			if !ok {
+				handleError(conn, walletExt.Logger(), fmt.Errorf("format value is not a string"))
+				return
+			}
+			formatsSlice = append(formatsSlice, formatStr)
+		}
+	}
+
+	message, err := walletExt.GetMessage(encryptionToken.(string), formatsSlice)
+	if err != nil {
+		handleError(conn, walletExt.Logger(), fmt.Errorf("internal error"))
+		walletExt.Logger().Error("error getting message", log.ErrKey, err)
+		return
+	}
+
+	// create the response structure
+	type JSONResponse struct {
+		Message string `json:"message"`
+		Type    string `json:"type"`
+	}
+
+	// get string representation of the message format
+	messageFormat := common.GetBestFormat(formatsSlice)
+	messageFormatString, err := common.GetSignatureTypeString(messageFormat)
+	if err != nil {
+		handleError(conn, walletExt.Logger(), fmt.Errorf("internal error"))
+		return
+	}
+
+	response := JSONResponse{
+		Message: message,
+		Type:    messageFormatString,
+	}
+
+	responseBytes, err := json.Marshal(response)
+	if err != nil {
+		walletExt.Logger().Error("error marshaling JSON response", log.ErrKey, err)
+		return
+	}
+
+	err = conn.WriteResponse(responseBytes)
+	if err != nil {
+		walletExt.Logger().Error("error writing success response", log.ErrKey, err)
+	}
+}

tools/walletextension/common/common.go

@@ -99,3 +99,23 @@ func NewFileLogger() gethlog.Logger {
 
 	return logger
 }
+
+// GetBestFormat returns the best format for a message based on available formats that are supported by the user
+func GetBestFormat(formatsSlice []string) viewingkey.SignatureType {
+	// If "Personal" is the only format available, choose it
+	if len(formatsSlice) == 1 && formatsSlice[0] == "Personal" {
+		return viewingkey.PersonalSign
+	}
+
+	// otherwise, choose EIP712
+	return viewingkey.EIP712Signature
+}
+
+func GetSignatureTypeString(expectedSignatureType viewingkey.SignatureType) (string, error) {
+	for key, value := range SignatureTypeMap {
+		if value == expectedSignatureType {
+			return key, nil
+		}
+	}
+	return "", fmt.Errorf("unable to find signature type")
+}