feat: Sync shasta fork with raiko (vibe-kanban 03860811)

taikoxyz/raiko#642
请审查这个 pull request，我们当前分支也是要实现一模一样的功能，你需要评估 raiko2 的 rust 实现和我们 go 版本的实现的差异（目前可能有部分差异），功能上请以 raiko2 当前的 pr 实现为准，为我们 gaiko 制定重构和补齐功能的计划，开始实现之前请让我 review 方案

Author: johntaiko

internal/prover/proof.go

@@ -20,11 +20,12 @@ import (
 )
 
 type ProofResponse struct {
-	Proof           hexutil.Bytes  `json:"proof"`
-	Quote           hexutil.Bytes  `json:"quote"`
-	PublicKey       hexutil.Bytes  `json:"public_key"`
-	InstanceAddress common.Address `json:"instance_address"`
-	Input           common.Hash    `json:"input"`
+	Proof           hexutil.Bytes           `json:"proof"`
+	Quote           hexutil.Bytes           `json:"quote"`
+	PublicKey       hexutil.Bytes           `json:"public_key"`
+	InstanceAddress common.Address          `json:"instance_address"`
+	Input           common.Hash             `json:"input"`
+	ExtraData       *witness.ProofCarryData `json:"extra_data,omitempty"`
 }
 
 func (p *ProofResponse) Output(w io.Writer) error {
@@ -150,6 +151,23 @@ func genOneshotProof(
 		return err
 	}
 
+	var extraData *witness.ProofCarryData
+	if guestInput.BlockProposed().IsShasta() {
+		transitionInput, ok := guestInput.Transition().(witness.TransitionInputData)
+		if !ok {
+			if ptr, ok := guestInput.Transition().(*witness.TransitionInputData); ok && ptr != nil {
+				transitionInput = *ptr
+			} else {
+				return fmt.Errorf("invalid shasta transition type: %T", guestInput.Transition())
+			}
+		}
+		extraData = &witness.ProofCarryData{
+			ChainID:         guestInput.ChainID(),
+			Verifier:        guestInput.ForkVerifierAddress(args.ProofType),
+			TransitionInput: transitionInput,
+		}
+	}
+
 	sign, err := Sign(piHash.Bytes(), prevPrivKey)
 	if err != nil {
 		return err
@@ -167,5 +185,6 @@ func genOneshotProof(
 		PublicKey:       crypto.FromECDSAPub(&prevPrivKey.PublicKey),
 		InstanceAddress: newInstance,
 		Input:           piHash,
+		ExtraData:       extraData,
 	}).Output(args.ProofWriter)
 }

internal/prover/shasta.go

@@ -35,26 +35,41 @@ func genShastaAggregateProof(
 		return err
 	}
 	log.Info("receive input: ", "input", input)
-	oldInstance := common.BytesToAddress(input.Proofs[0].Proof[4:24])
-	curInstance := oldInstance
-	transitionHashes := make([]common.Hash, 0, len(input.Proofs))
+	if len(input.Proofs) == 0 {
+		return fmt.Errorf("no proofs provided")
+	}
+	if len(input.Proofs) != len(input.ProofCarryDataVec) {
+		return fmt.Errorf("proofs and proof_carry_data_vec length mismatch")
+	}
+	if !witness.ValidateShastaProofCarryDataVec(input.ProofCarryDataVec) {
+		return fmt.Errorf("invalid shasta proof carry data")
+	}
+
+	expectedInstance := common.BytesToAddress(input.Proofs[0].Proof[4:24])
+	if expectedInstance != newInstance {
+		return fmt.Errorf("invalid instance: %#x", expectedInstance)
+	}
 	var pubKey *ecdsa.PublicKey
 	for i, proof := range input.Proofs {
+		if proof.Input != witness.HashShastaSubproofInput(&input.ProofCarryDataVec[i]) {
+			return fmt.Errorf("invalid shasta proof input at index %d", i)
+		}
 		pubKey, err = SigToPub(proof.Input.Bytes(), proof.Proof[24:])
 		if err != nil {
 			return err
 		}
-		if crypto.PubkeyToAddress(*pubKey) != curInstance {
+		if crypto.PubkeyToAddress(*pubKey) != expectedInstance {
 			return fmt.Errorf("invalid proof[%d]", i)
 		}
-		curInstance = common.BytesToAddress(proof.Proof[4:24])
-		transitionHashes = append(transitionHashes, proof.Input)
-	}
-	if newInstance != curInstance {
-		return fmt.Errorf("invalid instance: %#x", curInstance)
+		if common.BytesToAddress(proof.Proof[4:24]) != expectedInstance {
+			return fmt.Errorf("invalid proof[%d] instance", i)
+		}
 	}
 
-	aggHash := witness.HashShastaAggregation(transitionHashes, input.ChainID, input.VerifierAddress, newInstance)
+	aggHash, err := witness.ShastaPCDAggregationHash(input.ProofCarryDataVec, newInstance)
+	if err != nil {
+		return err
+	}
 	sign, err := Sign(aggHash.Bytes(), prevPrivKey)
 	if err != nil {
 		return err

internal/witness/abi.go

@@ -2,9 +2,14 @@ package witness
 
 import (
 	"errors"
+	"math/big"
+	"reflect"
+	"strings"
 
 	"github.com/ethereum/go-ethereum/accounts/abi"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/taikoxyz/taiko-mono/packages/taiko-client/bindings/encoding"
+	"github.com/taikoxyz/taiko-mono/packages/taiko-client/bindings/shasta"
 )
 
 var (
@@ -66,6 +71,7 @@ var (
 	blockProposedEvent            = encoding.TaikoL1ABI.Events["BlockProposed"]
 	blockProposedV2Event          = encoding.TaikoL1ABI.Events["BlockProposedV2"]
 	anchorV3Method                = encoding.TaikoAnchorABI.Methods["anchorV3"]
+	shastaAnchorV4Method          abi.Method
 )
 
 func init() {
@@ -89,6 +95,12 @@ func init() {
 		panic(err)
 	}
 	blockMetadataV2ComponentsArgs = abi.Arguments{arg}
+
+	shastaAnchorABI, err := abi.JSON(strings.NewReader(shasta.ShastaAnchorABI))
+	if err != nil {
+		panic(err)
+	}
+	shastaAnchorV4Method = shastaAnchorABI.Methods["anchorV4"]
 }
 
 // ABIEncoder is an interface for solidity structs encoding
@@ -131,3 +143,106 @@ func decodeAnchorV3ArgsSignalSlots(input []byte) ([][32]byte, error) {
 	}
 	return signalSlots, nil
 }
+
+func decodeShastaAnchorCheckpoint(input []byte) (*ShastaCheckpoint, error) {
+	if shastaAnchorV4Method.Name == "" {
+		return nil, errors.New("shasta anchor ABI not initialized")
+	}
+	args := map[string]any{}
+	if err := shastaAnchorV4Method.Inputs.UnpackIntoMap(args, input); err != nil {
+		return nil, err
+	}
+
+	var param any
+	if value, ok := args["_blockParams"]; ok {
+		param = value
+	} else if value, ok := args["_checkpoint"]; ok {
+		param = value
+	} else {
+		return nil, errors.New("shasta anchor params not found")
+	}
+
+	return shastaCheckpointFromTuple(param)
+}
+
+func shastaCheckpointFromTuple(value any) (*ShastaCheckpoint, error) {
+	rv := reflect.ValueOf(value)
+	if rv.Kind() == reflect.Pointer {
+		if rv.IsNil() {
+			return nil, errors.New("nil shasta anchor param")
+		}
+		rv = rv.Elem()
+	}
+	if rv.Kind() != reflect.Struct {
+		return nil, errors.New("unexpected shasta anchor param type")
+	}
+
+	numberField := rv.FieldByName("AnchorBlockNumber")
+	if !numberField.IsValid() {
+		numberField = rv.FieldByName("BlockNumber")
+	}
+	hashField := rv.FieldByName("AnchorBlockHash")
+	if !hashField.IsValid() {
+		hashField = rv.FieldByName("BlockHash")
+	}
+	stateField := rv.FieldByName("AnchorStateRoot")
+	if !stateField.IsValid() {
+		stateField = rv.FieldByName("StateRoot")
+	}
+
+	if !numberField.IsValid() || !hashField.IsValid() || !stateField.IsValid() {
+		return nil, errors.New("invalid shasta anchor param fields")
+	}
+
+	var blockNumber uint64
+	switch v := numberField.Interface().(type) {
+	case *big.Int:
+		blockNumber = v.Uint64()
+	case uint64:
+		blockNumber = v
+	case uint32:
+		blockNumber = uint64(v)
+	default:
+		if numberField.Kind() == reflect.Uint64 {
+			blockNumber = numberField.Uint()
+		} else {
+			return nil, errors.New("unsupported block number type")
+		}
+	}
+
+	blockHash, err := toCommonHash(hashField)
+	if err != nil {
+		return nil, err
+	}
+	stateRoot, err := toCommonHash(stateField)
+	if err != nil {
+		return nil, err
+	}
+
+	return &ShastaCheckpoint{
+		BlockNumber: blockNumber,
+		BlockHash:   blockHash,
+		StateRoot:   stateRoot,
+	}, nil
+}
+
+func toCommonHash(value reflect.Value) (common.Hash, error) {
+	if !value.IsValid() {
+		return common.Hash{}, errors.New("invalid hash value")
+	}
+	switch v := value.Interface().(type) {
+	case [32]byte:
+		return common.BytesToHash(v[:]), nil
+	case common.Hash:
+		return v, nil
+	default:
+		if value.Kind() == reflect.Array && value.Len() == 32 && value.Type().Elem().Kind() == reflect.Uint8 {
+			var out [32]byte
+			for i := 0; i < 32; i++ {
+				out[i] = byte(value.Index(i).Uint())
+			}
+			return common.BytesToHash(out[:]), nil
+		}
+		return common.Hash{}, errors.New("unsupported hash type")
+	}
+}

internal/witness/aggregate.go

@@ -12,7 +12,6 @@ type RawAggregationGuestInput struct {
 }
 
 type ShastaRawAggregationGuestInput struct {
-	Proofs          []*RawProof    `json:"proofs"`
-	ChainID         uint64         `json:"chain_id"`
-	VerifierAddress common.Address `json:"verifier_address"`
+	Proofs            []*RawProof      `json:"proofs"`
+	ProofCarryDataVec []ProofCarryData `json:"proof_carry_data_vec"`
 }