minimal diff to be able to return the task response bytes instead of just the task response digest

services/bls_aggregation/blsagg.go

@@ -45,6 +45,7 @@ var (
 type BlsAggregationServiceResponse struct {
 	Err                error                    // if Err is not nil, the other fields are not valid
 	TaskIndex          types.TaskIndex          // unique identifier of the task
+	TaskResponse       types.TaskResponse       // the task response that was signed
 	TaskResponseDigest types.TaskResponseDigest // digest of the task response that was signed
 	// The below 8 fields are the data needed to build the IBLSSignatureChecker.NonSignerStakesAndSignature struct
 	// users of this service will need to build the struct themselves by converting the bls points
@@ -97,7 +98,7 @@ type BlsAggregationService interface {
 	ProcessNewSignature(
 		ctx context.Context,
 		taskIndex types.TaskIndex,
-		taskResponseDigest types.TaskResponseDigest,
+		taskResponse types.TaskResponse,
 		blsSignature *bls.Signature,
 		operatorId types.OperatorId,
 	) error
@@ -134,17 +135,24 @@ type BlsAggregatorService struct {
 	taskChansMutex     sync.RWMutex
 	avsRegistryService avsregistry.AvsRegistryService
 	logger             logging.Logger
+
+	// taskResponseMap is a map of taskResponseDigest to taskResponse
+	taskResponseMap map[types.TaskResponseDigest]types.TaskResponse
+
+	hashFunction types.TaskResponseHashFunction
 }
 
 var _ BlsAggregationService = (*BlsAggregatorService)(nil)
 
-func NewBlsAggregatorService(avsRegistryService avsregistry.AvsRegistryService, logger logging.Logger) *BlsAggregatorService {
+func NewBlsAggregatorService(avsRegistryService avsregistry.AvsRegistryService, hashFunction types.TaskResponseHashFunction, logger logging.Logger) *BlsAggregatorService {
 	return &BlsAggregatorService{
 		aggregatedResponsesC: make(chan BlsAggregationServiceResponse),
 		signedTaskRespsCs:    make(map[types.TaskIndex]chan types.SignedTaskResponseDigest),
 		taskChansMutex:       sync.RWMutex{},
 		avsRegistryService:   avsRegistryService,
 		logger:               logger,
+		taskResponseMap:      make(map[types.TaskResponseDigest]types.TaskResponse),
+		hashFunction:         hashFunction,
 	}
 }
 
@@ -179,7 +187,7 @@ func (a *BlsAggregatorService) InitializeNewTask(
 func (a *BlsAggregatorService) ProcessNewSignature(
 	ctx context.Context,
 	taskIndex types.TaskIndex,
-	taskResponseDigest types.TaskResponseDigest,
+	taskResponse types.TaskResponse,
 	blsSignature *bls.Signature,
 	operatorId types.OperatorId,
 ) error {
@@ -189,9 +197,20 @@ func (a *BlsAggregatorService) ProcessNewSignature(
 	if !taskInitialized {
 		return TaskNotFoundErrorFn(taskIndex)
 	}
+	// compute the taskResponseDigest using the hash function
+	taskResponseDigest := a.hashFunction(taskResponse)
+
+	// check if the taskResponseDigest is already in the map
+	_, taskResponseExists := a.taskResponseMap[taskResponseDigest]
+	if !taskResponseExists {
+		// Store the TaskResponse in our mapping
+		a.taskResponseMap[taskResponseDigest] = taskResponse
+	}
+
 	signatureVerificationErrorC := make(chan error)
 	// send the task to the goroutine processing this task
 	// and return the error (if any) returned by the signature verification routine
+
 	select {
 	// we need to send this as part of select because if the goroutine is processing another SignedTaskResponseDigest
 	// and cannot receive this one, we want the context to be able to cancel the request
@@ -255,6 +274,9 @@ func (a *BlsAggregatorService) singleTaskAggregatorGoroutineFunc(
 		select {
 		case signedTaskResponseDigest := <-signedTaskRespsC:
 			a.logger.Debug("Task goroutine received new signed task response digest", "taskIndex", taskIndex, "signedTaskResponseDigest", signedTaskResponseDigest)
+			// Retrieve the TaskResponse from the map
+			taskResponse := a.taskResponseMap[signedTaskResponseDigest.TaskResponseDigest]
+
 			err := a.verifySignature(taskIndex, signedTaskResponseDigest, operatorsAvsStateDict)
 			signedTaskResponseDigest.SignatureVerificationErrorC <- err
 			if err != nil {
@@ -316,9 +338,11 @@ func (a *BlsAggregatorService) singleTaskAggregatorGoroutineFunc(
 					}
 					return
 				}
+
 				blsAggregationServiceResponse := BlsAggregationServiceResponse{
 					Err:                          nil,
 					TaskIndex:                    taskIndex,
+					TaskResponse:                 taskResponse,
 					TaskResponseDigest:           signedTaskResponseDigest.TaskResponseDigest,
 					NonSignersPubkeysG1:          nonSignersG1Pubkeys,
 					QuorumApksG1:                 quorumApksG1,
@@ -371,7 +395,7 @@ func (a *BlsAggregatorService) verifySignature(
 		return OperatorNotPartOfTaskQuorumErrorFn(signedTaskResponseDigest.OperatorId, taskIndex)
 	}
 
-	// 0. verify that the msg actually came from the correct operator
+	// verify that the msg actually came from the correct operator
 	operatorG2Pubkey := operatorsAvsStateDict[signedTaskResponseDigest.OperatorId].OperatorInfo.Pubkeys.G2Pubkey
 	if operatorG2Pubkey == nil {
 		a.logger.Error("Operator G2 pubkey not found", "operatorId", signedTaskResponseDigest.OperatorId, "taskId", taskIndex)
@@ -382,6 +406,9 @@ func (a *BlsAggregatorService) verifySignature(
 		"taskResponseDigest", signedTaskResponseDigest.TaskResponseDigest,
 		"blsSignature", signedTaskResponseDigest.BlsSignature,
 	)
+
+	// if the operator signs a digest that is not the digest of the TaskResponse submitted in ProcessNewTask
+	// then the signature will not be verified
 	signatureVerified, err := signedTaskResponseDigest.BlsSignature.Verify(operatorG2Pubkey, signedTaskResponseDigest.TaskResponseDigest)
 	if err != nil {
 		return SignatureVerificationError(err)