Skip to content

Commit

Permalink
chore: optimizations for the encryptor
Browse files Browse the repository at this point in the history
  • Loading branch information
@Mikelle
Mikelle committed May 21, 2024
1 parent 16e93c1 commit 11726b4
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 41 deletions.
4 changes: 3 additions & 1 deletion p2p/pkg/preconfirmation/preconfirmation.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,12 +297,14 @@ func (p *Preconfirmation) handleBid(
case providerapiv1.BidResponse_STATUS_REJECTED:
return status.Errorf(codes.Internal, "bid rejected")
case providerapiv1.BidResponse_STATUS_ACCEPTED:
constructStartTime := time.Now()
preConfirmation, encryptedPreConfirmation, err := p.encryptor.ConstructEncryptedPreConfirmation(bid)
if err != nil {
return status.Errorf(codes.Internal, "failed to constuct encrypted preconfirmation: %v", err)
}
constructDuration := time.Since(constructStartTime)
p.logger.Info("constructed encrypted preconfirmation", "duration", constructDuration)

p.logger.Info("sending preconfirmation", "preConfirmation", encryptedPreConfirmation)
err = stream.WriteMsg(ctx, encryptedPreConfirmation)
if err != nil {
return status.Errorf(codes.Internal, "failed to send preconfirmation: %v", err)
Expand Down
75 changes: 43 additions & 32 deletions p2p/pkg/signer/preconfencryptor/encryptor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"crypto/ecdh"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"math/big"

Expand All @@ -16,6 +15,7 @@ import (
preconfpb "github.com/primev/mev-commit/p2p/gen/go/preconfirmation/v1"
p2pcrypto "github.com/primev/mev-commit/p2p/pkg/crypto"
"github.com/primev/mev-commit/x/keysigner"
"google.golang.org/protobuf/proto"
)

var (
Expand All @@ -25,6 +25,10 @@ var (
ErrInvalidHash = errors.New("bidhash doesn't match bid payload")
ErrAlreadySignedPreConfirmation = errors.New("preConfirmation is already hashed or signed")
ErrInvalidCommitment = errors.New("commitment is incorrect")
ErrMissingRequiredFields = errors.New("missing required fields")
ErrNoAesKeyFound = errors.New("no AES key found for bidder")
ErrInvalidBidAmt = errors.New("invalid bid amount")
ErrBidNotFound = errors.New("bid not found")
)

type Encryptor interface {
Expand All @@ -42,6 +46,9 @@ type Store interface {

type encryptor struct {
keySigner keysigner.KeySigner
address []byte // set for the provider
nikePrvKey *ecdh.PrivateKey // set for the provider
aesKey []byte // set for the bidder
store Store
bidHashesToBid *lru.Cache[string, *preconfpb.Bid]
}
Expand All @@ -51,8 +58,22 @@ func NewEncryptor(ks keysigner.KeySigner, store Store) (*encryptor, error) {
if err != nil {
return nil, err
}
address := ks.GetAddress()
// those keys are set up during the libp2p.New initialization.
aesKey, err := store.GetAESKey(address)
if err != nil {
return nil, err
}
nikePrvKey, err := store.GetNikePrivateKey()
if err != nil {
return nil, err
}

return &encryptor{
keySigner: ks,
address: address.Bytes(), // set for the provider
nikePrvKey: nikePrvKey,
aesKey: aesKey, // set for the bidder
store: store,
bidHashesToBid: bidHashesToBidCache,
}, nil
Expand All @@ -66,7 +87,7 @@ func (e *encryptor) ConstructEncryptedBid(
decayEndTimeStamp int64,
) (*preconfpb.Bid, *preconfpb.EncryptedBid, *ecdh.PrivateKey, error) {
if txHash == "" || bidAmt == "" || blockNumber == 0 {
return nil, nil, nil, errors.New("missing required fields")
return nil, nil, nil, ErrMissingRequiredFields
}

bid := &preconfpb.Bid{
Expand All @@ -87,36 +108,25 @@ func (e *encryptor) ConstructEncryptedBid(
return nil, nil, nil, err
}

if sig[64] == 0 || sig[64] == 1 {
sig[64] += 27 // Transform V from 0/1 to 27/28
}
transformSignatureVValue(sig)

nikePrivateKey, err := ecdh.P256().GenerateKey(rand.Reader)
if err != nil {
return nil, nil, nil, err
}

nikePublicKey := nikePrivateKey.PublicKey()

bid.NikePublicKey = nikePublicKey.Bytes()
bid.NikePublicKey = nikePrivateKey.PublicKey().Bytes()
bid.Digest = bidHash
bid.Signature = sig

bidDataBytes, err := json.Marshal(bid)
bidDataBytes, err := proto.Marshal(bid)
if err != nil {
return nil, nil, nil, err
}

e.bidHashesToBid.Add(hex.EncodeToString(bidHash), bid)

aesKey, err := e.store.GetAESKey(e.keySigner.GetAddress())
if err != nil {
return nil, nil, nil, err
}
if aesKey == nil {
return nil, nil, nil, errors.New("no AES key found for bidder")
}
encryptedBidData, err := p2pcrypto.EncryptWithAESGCM(aesKey, bidDataBytes)
encryptedBidData, err := p2pcrypto.EncryptWithAESGCM(e.aesKey, bidDataBytes)
if err != nil {
return nil, nil, nil, err
}
Expand All @@ -135,19 +145,15 @@ func (e *encryptor) ConstructEncryptedPreConfirmation(bid *preconfpb.Bid) (*prec
return nil, nil, err
}

nikePrvKey, err := e.store.GetNikePrivateKey()
if err != nil {
return nil, nil, err
}
sharedSecretProviderSk, err := nikePrvKey.ECDH(bidDataPublicKey)
sharedSecretProviderSk, err := e.nikePrvKey.ECDH(bidDataPublicKey)
if err != nil {
return nil, nil, err
}

preConfirmation := &preconfpb.PreConfirmation{
Bid: bid,
SharedSecret: sharedSecretProviderSk,
ProviderAddress: e.keySigner.GetAddress().Bytes(),
ProviderAddress: e.address,
}

preConfirmationHash, err := GetPreConfirmationHash(preConfirmation)
Expand All @@ -160,9 +166,7 @@ func (e *encryptor) ConstructEncryptedPreConfirmation(bid *preconfpb.Bid) (*prec
return nil, nil, err
}

if sig[64] == 0 || sig[64] == 1 {
sig[64] += 27 // Transform V from 0/1 to 27/28
}
transformSignatureVValue(sig)

preConfirmation.Digest = preConfirmationHash
preConfirmation.Signature = sig
Expand Down Expand Up @@ -196,15 +200,15 @@ func (e *encryptor) DecryptBidData(bidderAddress common.Address, bid *preconfpb.
return nil, err
}
if aesKey == nil {
return nil, errors.New("no AES key found for bidder")
return nil, ErrNoAesKeyFound
}
decryptedBytes, err := p2pcrypto.DecryptWithAESGCM(aesKey, bid.Ciphertext)
if err != nil {
return nil, err
}

var bidData preconfpb.Bid
if err := json.Unmarshal(decryptedBytes, &bidData); err != nil {
if err := proto.Unmarshal(decryptedBytes, &bidData); err != nil {
return nil, err
}

Expand All @@ -216,7 +220,8 @@ func (e *encryptor) DecryptBidData(bidderAddress common.Address, bid *preconfpb.
func (e *encryptor) VerifyEncryptedPreConfirmation(
providerNikePK *ecdh.PublicKey,
bidderNikeSC *ecdh.PrivateKey,
bidHash []byte, c *preconfpb.EncryptedPreConfirmation,
bidHash []byte,
c *preconfpb.EncryptedPreConfirmation,
) ([]byte, *common.Address, error) {
if c.Signature == nil {
return nil, nil, ErrMissingHashSignature
Expand All @@ -225,7 +230,7 @@ func (e *encryptor) VerifyEncryptedPreConfirmation(
bidHashStr := hex.EncodeToString(bidHash)
bid, ok := e.bidHashesToBid.Get(bidHashStr)
if !ok {
return nil, nil, errors.New("bid not found")
return nil, nil, ErrBidNotFound
}
sharedSecredBidderSk, err := bidderNikeSC.ECDH(providerNikePK)
if err != nil {
Expand Down Expand Up @@ -302,7 +307,7 @@ func GetBidHash(bid *preconfpb.Bid) ([]byte, error) {

bidAmt, ok := big.NewInt(0).SetString(bid.BidAmount, 10)
if !ok {
return nil, errors.New("invalid bid amount")
return nil, ErrInvalidBidAmt
}

// EIP712_MESSAGE_TYPEHASH
Expand Down Expand Up @@ -345,7 +350,7 @@ func GetPreConfirmationHash(c *preconfpb.PreConfirmation) ([]byte, error) {

bidAmt, ok := big.NewInt(0).SetString(c.Bid.BidAmount, 10)
if !ok {
return nil, errors.New("invalid bid amount")
return nil, ErrInvalidBidAmt
}

// EIP712_MESSAGE_TYPEHASH
Expand Down Expand Up @@ -374,3 +379,9 @@ func GetPreConfirmationHash(c *preconfpb.PreConfirmation) ([]byte, error) {
// Create the final hash
return crypto.Keccak256Hash(rawData).Bytes(), nil
}

func transformSignatureVValue(sig []byte) {
if sig[64] == 0 || sig[64] == 1 {
sig[64] += 27 // Transform V from 0/1 to 27/28
}
}
17 changes: 9 additions & 8 deletions p2p/pkg/signer/preconfencryptor/encryptor_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,27 +106,28 @@ func TestBids(t *testing.T) {
if err != nil {
t.Fatal(err)
}
providerEncryptor, err := preconfencryptor.NewEncryptor(keySigner, providerStore)
providerNikePrivateKey, err := ecdh.P256().GenerateKey(rand.Reader)
if err != nil {
t.Fatal(err)
}
start := time.Now().UnixMilli()
end := start + 100000

bid, encryptedBid, nikePrivateKey, err := bidderEncryptor.ConstructEncryptedBid("0xkartik", "10", 2, start, end)
err = providerStore.SetNikePrivateKey(providerNikePrivateKey)
if err != nil {
t.Fatal(err)
}

decryptedBid, err := providerEncryptor.DecryptBidData(bidderAddress, encryptedBid)
providerEncryptor, err := preconfencryptor.NewEncryptor(keySigner, providerStore)
if err != nil {
t.Fatal(err)
}
providerNikePrivateKey, err := ecdh.P256().GenerateKey(rand.Reader)
start := time.Now().UnixMilli()
end := start + 100000

bid, encryptedBid, nikePrivateKey, err := bidderEncryptor.ConstructEncryptedBid("0xkartik", "10", 2, start, end)
if err != nil {
t.Fatal(err)
}
err = providerStore.SetNikePrivateKey(providerNikePrivateKey)

decryptedBid, err := providerEncryptor.DecryptBidData(bidderAddress, encryptedBid)
if err != nil {
t.Fatal(err)
}
Expand Down

0 comments on commit 11726b4

Please sign in to comment.