compute workload id from quote

Author: avalonche

config/l2.go

@@ -33,7 +33,7 @@ type L2 struct {
 	MonitorTxReceipts                                bool              `yaml:"monitor_tx_receipts"`
 	MonitorWalletAddresses                           map[string]string `yaml:"monitor_wallet_addresses"`
 
-	ProbeTx               *ProbeTx               `yaml:"probe"`
+	ProbeTx *ProbeTx `yaml:"probe"`
 }
 
 const (

contracts/types.go

@@ -18,8 +18,8 @@ var (
 	FlashblocksLandedCount otelapi.Int64Gauge
 	FlashblocksMissedCount otelapi.Int64Gauge
 
-	FlashtestationsLandedCount otelapi.Int64Gauge
-	FlashtestationsMissedCount otelapi.Int64Gauge
+	FlashtestationsLandedCount     otelapi.Int64Gauge
+	FlashtestationsMissedCount     otelapi.Int64Gauge
 	RegisteredFlashtestationsCount otelapi.Int64Gauge
 
 	ReorgsCount otelapi.Int64Counter

metrics/exports.go

@@ -155,7 +155,7 @@ func setupFlashtestationsMissedCount(ctx context.Context, _ *config.ProbeTx) err
 	return nil
 }
 
-func setupRegisteredFlashtestationsCount(ctx context.Context, _*config.ProbeTx) error {
+func setupRegisteredFlashtestationsCount(ctx context.Context, _ *config.ProbeTx) error {
 	m, err := meter.Int64Gauge("registered_flashtestations_count",
 		otelapi.WithDescription("registered flashtestations count"),
 	)

metrics/metrics.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"encoding/binary"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -23,6 +24,7 @@ import (
 
 	"go.uber.org/zap"
 
+	"github.com/ethereum/go-ethereum/accounts/abi"
 	ethcommon "github.com/ethereum/go-ethereum/common"
 	ethtypes "github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -1109,7 +1111,7 @@ func (l2 *L2) checkAndResetProbeTxNonce(ctx context.Context) {
 func (l2 *L2) handleRegistrationTx(ctx context.Context, txHash ethcommon.Hash) {
 	l := logutils.LoggerFromContext(ctx)
 
-	teeAddress, err := l2.getTEEAddressFromTx(ctx, txHash)
+	teeAddress, rawQuote, err := l2.getTEEAddressAndQuoteFromTx(ctx, txHash)
 	if err != nil {
 		l.Warn("Failed to get register flashtestations transaction receipt",
 			zap.Error(err),
@@ -1118,32 +1120,74 @@ func (l2 *L2) handleRegistrationTx(ctx context.Context, txHash ethcommon.Hash) {
 		return
 	}
 
+	workloadId, err := ComputeWorkloadID(rawQuote)
+	if err != nil {
+		l.Warn("Failed to compute workload id",
+			zap.Error(err),
+			zap.String("tx", txHash.Hex()),
+		)
+		return
+	}
+
 	metrics.RegisteredFlashtestationsCount.Record(ctx, 1, otelapi.WithAttributes(
 		attribute.KeyValue{Key: "kind", Value: attribute.StringValue("l2")},
 		attribute.KeyValue{Key: "network_id", Value: attribute.Int64Value(l2.chainID.Int64())},
 		attribute.KeyValue{Key: "tee_address", Value: attribute.StringValue(teeAddress.Hex())},
+		attribute.KeyValue{Key: "workload_id", Value: attribute.StringValue(hex.EncodeToString(workloadId[:]))},
 	))
 
+	l.Info("TEE service registered",
+		zap.String("teeAddress", teeAddress.Hex()),
+		zap.String("workloadId", hex.EncodeToString(workloadId[:])),
+	)
+
 	return
 }
 
-// Extract TEE address from TEEServiceRegistered event
-func (l2 *L2) getTEEAddressFromTx(ctx context.Context, txHash ethcommon.Hash) (ethcommon.Address, error) {
+// Extract TEE address and raw quote from TEEServiceRegistered event
+func (l2 *L2) getTEEAddressAndQuoteFromTx(ctx context.Context, txHash ethcommon.Hash) (ethcommon.Address, []byte, error) {
 	receipt, err := l2.rpc.TransactionReceipt(ctx, txHash)
 	if err != nil {
-		return ethcommon.Address{}, err
+		return ethcommon.Address{}, nil, err
 	}
 
 	if receipt.Status == ethtypes.ReceiptStatusFailed {
-		return ethcommon.Address{}, fmt.Errorf("Register tee transaction did not succeeed %s", txHash.Hex())
+		return ethcommon.Address{}, nil, fmt.Errorf("Register tee transaction did not succeeed %s", txHash.Hex())
+	}
+
+	// Define the event arguments for decoding (non-indexed parameters only)
+	// event TEEServiceRegistered(address indexed teeAddress, bytes rawQuote, bool alreadyExists);
+	bytesType, _ := abi.NewType("bytes", "", nil)
+	boolType, _ := abi.NewType("bool", "", nil)
+
+	eventABI := abi.Arguments{
+		{Type: bytesType}, // rawQuote
+		{Type: boolType},  // alreadyExists
 	}
 
 	for _, log := range receipt.Logs {
 		if len(log.Topics) > 1 && log.Topics[0] == l2.flashtestationsRegistryEventSignature {
 			// TEE address is in Topics[1] (indexed parameter)
-			return ethcommon.BytesToAddress(log.Topics[1].Bytes()), nil
+			teeAddress := ethcommon.BytesToAddress(log.Topics[1].Bytes())
+
+			// Decode the data field to get rawQuote and alreadyExists
+			decoded, err := eventABI.Unpack(log.Data)
+			if err != nil {
+				return ethcommon.Address{}, nil, fmt.Errorf("failed to decode event data: %w", err)
+			}
+
+			if len(decoded) < 2 {
+				return ethcommon.Address{}, nil, fmt.Errorf("unexpected decoded data length: %d", len(decoded))
+			}
+
+			rawQuote, ok := decoded[0].([]byte)
+			if !ok {
+				return ethcommon.Address{}, nil, fmt.Errorf("failed to type assert rawQuote")
+			}
+
+			return teeAddress, rawQuote, nil
 		}
 	}
 
-	return ethcommon.Address{}, fmt.Errorf("TEEServiceRegistered event not found in tx %s", txHash.Hex())
+	return ethcommon.Address{}, nil, fmt.Errorf("TEEServiceRegistered event not found in tx %s", txHash.Hex())
 }

server/l2.go

@@ -0,0 +1,90 @@
+package server
+
+// Implementation based on https://github.com/flashbots/flashtestations/tree/7cc7f68492fe672a823dd2dead649793aac1f216
+import (
+	"encoding/binary"
+	"fmt"
+
+	"golang.org/x/crypto/sha3"
+)
+
+const (
+	// Raw TDX v4 quote structure constants
+	// Raw quote has a 48-byte header before the TD10ReportBody
+	HEADER_LENGTH      = 48
+	TD_REPORT10_LENGTH = 584
+
+	// TDX workload constants
+	TD_XFAM_FPU            = 0x0000000000000001
+	TD_XFAM_SSE            = 0x0000000000000002
+	TD_TDATTRS_VE_DISABLED = 0x0000000010000000
+	TD_TDATTRS_PKS         = 0x0000000040000000
+	TD_TDATTRS_KL          = 0x0000000080000000
+)
+
+// ComputeWorkloadID computes the workload ID from Automata's serialized verifier output
+// This corresponds to QuoteParser.parseV4Quote in Solidity
+// The workload ID uniquely identifies a TEE workload based on its measurement registers
+func ComputeWorkloadID(rawQuote []byte) ([32]byte, error) {
+	var workloadID [32]byte
+
+	// Validate quote length
+	if len(rawQuote) < HEADER_LENGTH+TD_REPORT10_LENGTH {
+		return workloadID, fmt.Errorf("invalid quote length: %d, expected at least %d",
+			len(rawQuote), HEADER_LENGTH+TD_REPORT10_LENGTH)
+	}
+
+	// Skip the 48-byte header to get to the TD10ReportBody
+	reportBody := rawQuote[HEADER_LENGTH:]
+
+	// Extract fields exactly as parseRawReportBody does in Solidity
+	// Using hardcoded offsets to match Solidity implementation exactly
+	mrTd := reportBody[136 : 136+48]
+	rtMr0 := reportBody[328 : 328+48]
+	rtMr1 := reportBody[376 : 376+48]
+	rtMr2 := reportBody[424 : 424+48]
+	rtMr3 := reportBody[472 : 472+48]
+	mrConfigId := reportBody[184 : 184+48]
+
+	// Extract xFAM and tdAttributes (8 bytes each)
+	// In Solidity, bytes8 is treated as big-endian for bitwise operations
+	xfam := binary.BigEndian.Uint64(reportBody[128 : 128+8])
+	tdAttributes := binary.BigEndian.Uint64(reportBody[120 : 120+8])
+
+	// Apply transformations as per the Solidity implementation
+	// expectedXfamBits = TD_XFAM_FPU | TD_XFAM_SSE
+	expectedXfamBits := uint64(TD_XFAM_FPU | TD_XFAM_SSE)
+
+	// ignoredTdAttributesBitmask = TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL
+	ignoredTdAttributesBitmask := uint64(TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL)
+
+	// Transform xFAM: xFAM ^ expectedXfamBits
+	transformedXfam := xfam ^ expectedXfamBits
+
+	// Transform tdAttributes: tdAttributes & ~ignoredTdAttributesBitmask
+	transformedTdAttributes := tdAttributes & ^ignoredTdAttributesBitmask
+
+	// Convert transformed values back to bytes (big-endian, to match Solidity bytes8)
+	xfamBytes := make([]byte, 8)
+	tdAttributesBytes := make([]byte, 8)
+	binary.BigEndian.PutUint64(xfamBytes, transformedXfam)
+	binary.BigEndian.PutUint64(tdAttributesBytes, transformedTdAttributes)
+
+	// Concatenate all fields
+	var concatenated []byte
+	concatenated = append(concatenated, mrTd...)
+	concatenated = append(concatenated, rtMr0...)
+	concatenated = append(concatenated, rtMr1...)
+	concatenated = append(concatenated, rtMr2...)
+	concatenated = append(concatenated, rtMr3...)
+	concatenated = append(concatenated, mrConfigId...)
+	concatenated = append(concatenated, xfamBytes...)
+	concatenated = append(concatenated, tdAttributesBytes...)
+
+	// Compute keccak256 hash
+	hash := sha3.NewLegacyKeccak256()
+	hash.Write(concatenated)
+	copy(workloadID[:], hash.Sum(nil))
+
+	return workloadID, nil
+}