Add execution stepcount in header [X-Cairo-Steps]

jsonrpc/http.go

@@ -1,6 +1,7 @@
 package jsonrpc
 
 import (
+	"maps"
 	"net/http"
 
 	"github.com/NethermindEth/juno/utils"
@@ -46,8 +47,11 @@ func (h *HTTP) ServeHTTP(writer http.ResponseWriter, req *http.Request) {
 
 	req.Body = http.MaxBytesReader(writer, req.Body, MaxRequestBodySize)
 	h.listener.OnNewRequest("any")
-	resp, err := h.rpc.HandleReader(req.Context(), req.Body)
+	resp, header, err := h.rpc.HandleReader(req.Context(), req.Body)
+
 	writer.Header().Set("Content-Type", "application/json")
+	maps.Copy(writer.Header(), header) // overwrites duplicate headers
+
 	if err != nil {
 		h.log.Errorw("Handler failure", "err", err)
 		writer.WriteHeader(http.StatusInternalServerError)

jsonrpc/server.go

@@ -9,6 +9,7 @@
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"reflect"
 	"strings"
 	"sync"
@@ -183,11 +184,18 @@
 	if numArgs != len(method.Params) {
 		return errors.New("number of non-context function params and param names must match")
 	}
-	if handlerT.NumOut() != 2 {
-		return errors.New("handler must return 2 values")
+	outSize := handlerT.NumOut()
+	if outSize < 2 || outSize > 3 {
+		return errors.New("handler must return 2 or 3 values")
 	}
-	if handlerT.Out(1) != reflect.TypeOf(&Error{}) {
-		return errors.New("second return value must be a *jsonrpc.Error")
+	if outSize == 2 && handlerT.Out(1) != reflect.TypeOf(&Error{}) {
+		return errors.New("second return value must be a *jsonrpc.Error for 2 tuple handler")
+	} else if outSize == 3 && handlerT.Out(2) != reflect.TypeOf(&Error{}) {
+		return errors.New("third return value must be a *jsonrpc.Error for 3 tuple handler")
+	}
+
+	if outSize == 3 && handlerT.Out(1) != reflect.TypeOf(http.Header{}) {
+		return errors.New("second return value must be a http.Header for 3 tuple handler")
 	}
 
 	// The method is valid. Mutate the appropriate fields and register on the server.
@@ -255,7 +263,8 @@
 		activated: activated,
 	}
 	msgCtx := context.WithValue(ctx, ConnKey{}, conn)
-	resp, err := s.HandleReader(msgCtx, rw)
+	// header is unnecessary for read-writer(websocket)
+	resp, _, err := s.HandleReader(msgCtx, rw)
 	if err != nil {
 		conn.initialErr = err
 		return err
@@ -272,27 +281,31 @@
 // HandleReader processes a request to the server
 // It returns the response in a byte array, only returns an
 // error if it can not create the response byte array
-func (s *Server) HandleReader(ctx context.Context, reader io.Reader) ([]byte, error) {
+func (s *Server) HandleReader(ctx context.Context, reader io.Reader) ([]byte, http.Header, error) {
 	bufferedReader := bufio.NewReaderSize(reader, bufferSize)
 	requestIsBatch := isBatch(bufferedReader)
 	res := &response{
 		Version: "2.0",
 	}
 
+	header := http.Header{}
+
 	dec := json.NewDecoder(bufferedReader)
 	dec.UseNumber()
 
 	if !requestIsBatch {
 		req := new(Request)
 		if jsonErr := dec.Decode(req); jsonErr != nil {
 			res.Error = Err(InvalidJSON, jsonErr.Error())
-		} else if resObject, handleErr := s.handleRequest(ctx, req); handleErr != nil {
+		} else if resObject, httpHeader, handleErr := s.handleRequest(ctx, req); handleErr != nil {
 			if !errors.Is(handleErr, ErrInvalidID) {
 				res.ID = req.ID
 			}
 			res.Error = Err(InvalidRequest, handleErr.Error())
+			header = httpHeader
 		} else {
 			res = resObject
+			header = httpHeader
 		}
 	} else {
 		var batchReq []json.RawMessage
@@ -307,23 +320,27 @@
 	}
 
 	if res == nil {
-		return nil, nil
+		return nil, header, nil
 	}
-	return json.Marshal(res)
+
+	result, err := json.Marshal(res)
+	return result, header, err
 }
 
-func (s *Server) handleBatchRequest(ctx context.Context, batchReq []json.RawMessage) ([]byte, error) {
+func (s *Server) handleBatchRequest(ctx context.Context, batchReq []json.RawMessage) ([]byte, http.Header, error) {
 	var (
-		responses []json.RawMessage
 		mutex     sync.Mutex
+		responses []json.RawMessage
+		headers   []http.Header
 	)
 
-	addResponse := func(response any) {
+	addResponse := func(response any, header http.Header) {
 		if responseJSON, err := json.Marshal(response); err != nil {
 			s.log.Errorw("failed to marshal response", "err", err)
 		} else {
 			mutex.Lock()
 			responses = append(responses, responseJSON)
+			headers = append(headers, header)
 			mutex.Unlock()
 		}
 	}
@@ -341,15 +358,15 @@
 			addResponse(&response{
 				Version: "2.0",
 				Error:   Err(InvalidRequest, err.Error()),
-			})
+			}, http.Header{})
 			continue
 		}
 
 		wg.Add(1)
 		s.pool.Go(func() {
 			defer wg.Done()
 
-			resp, err := s.handleRequest(ctx, req)
+			resp, header, err := s.handleRequest(ctx, req)
 			if err != nil {
 				resp = &response{
 					Version: "2.0",
@@ -359,20 +376,33 @@
 					resp.ID = req.ID
 				}
 			}
-			// for notification request response is nil
+			// for notification request response is nil and header is irrelevant for now
 			if resp != nil {
-				addResponse(resp)
+				addResponse(resp, header)
 			}
 		})
 	}
 
 	wg.Wait()
+
+	// merge headers
+	finalHeaders := http.Header{}
+	for _, header := range headers {
+		for k, v := range header {
+			for _, e := range v {
+				finalHeaders.Add(k, e)
+			}
+		}
+	}
+
 	// according to the spec if there are no response objects server must not return empty array
 	if len(responses) == 0 {
-		return nil, nil
+		return nil, finalHeaders, nil
 	}
 
-	return json.Marshal(responses)
+	result, err := json.Marshal(responses)
+
+	return result, finalHeaders, err // todo: fix batch request aggregate header
 }
 
 func isBatch(reader *bufio.Reader) bool {
@@ -396,11 +426,13 @@
 	return i == nil || reflect.ValueOf(i).IsNil()
 }
 
-func (s *Server) handleRequest(ctx context.Context, req *Request) (*response, error) {
+func (s *Server) handleRequest(ctx context.Context, req *Request) (*response, http.Header, error) {
 	s.log.Tracew("Received request", "req", req)
+
+	header := http.Header{}
 	if err := req.isSane(); err != nil {
 		s.log.Tracew("Request sanity check failed", "err", err)
-		return nil, err
+		return nil, header, err
 	}
 
 	res := &response{
@@ -412,7 +444,7 @@
 	if !found {
 		res.Error = Err(MethodNotFound, nil)
 		s.log.Tracew("Method not found in request", "method", req.Method)
-		return res, nil
+		return res, header, nil
 	}
 
 	handlerTimer := time.Now()
@@ -421,7 +453,7 @@
 	if err != nil {
 		res.Error = Err(InvalidParams, err.Error())
 		s.log.Tracew("Error building arguments for RPC call", "err", err)
-		return res, nil
+		return res, header, nil
 	}
 	defer func() {
 		s.listener.OnRequestHandled(req.Method, time.Since(handlerTimer))
@@ -430,22 +462,28 @@
 	tuple := reflect.ValueOf(calledMethod.Handler).Call(args)
 	if res.ID == nil { // notification
 		s.log.Tracew("Notification received, no response expected")
-		return nil, nil
+		return nil, header, nil
+	}
+
+	errorIndex := 1
+	if len(tuple) == 3 {
+		errorIndex = 2
+		header = (tuple[1].Interface()).(http.Header)
 	}
 
-	if errAny := tuple[1].Interface(); !isNil(errAny) {
+	if errAny := tuple[errorIndex].Interface(); !isNil(errAny) {
 		res.Error = errAny.(*Error)
 		if res.Error.Code == InternalError {
 			s.listener.OnRequestFailed(req.Method, res.Error)
 			reqJSON, _ := json.Marshal(req)
 			errJSON, _ := json.Marshal(res.Error)
 			s.log.Debugw("Failed handing RPC request", "req", string(reqJSON), "res", string(errJSON))
 		}
-		return res, nil
+		return res, header, nil
 	}
 	res.Result = tuple[0].Interface()
 
-	return res, nil
+	return res, header, nil
 }
 
 func (s *Server) buildArguments(ctx context.Context, params any, method Method) ([]reflect.Value, error) {

jsonrpc/server_test.go

@@ -40,17 +40,27 @@ func TestServer_RegisterMethod(t *testing.T) {
 		"no return": {
 			handler:    func(param1, param2 int) {},
 			paramNames: []jsonrpc.Parameter{{Name: "param1"}, {Name: "param2"}},
-			want:       "handler must return 2 values",
+			want:       "handler must return 2 or 3 values",
 		},
 		"int return": {
 			handler:    func(param1, param2 int) (int, int) { return 0, 0 },
 			paramNames: []jsonrpc.Parameter{{Name: "param1"}, {Name: "param2"}},
-			want:       "second return value must be a *jsonrpc.Error",
+			want:       "second return value must be a *jsonrpc.Error for 2 tuple handler",
+		},
+		"no error return 3": {
+			handler:    func(param1, param2 int) (int, int, int) { return 0, 0, 0 },
+			paramNames: []jsonrpc.Parameter{{Name: "param1"}, {Name: "param2"}},
+			want:       "third return value must be a *jsonrpc.Error for 3 tuple handler",
+		},
+		"no header return 3": {
+			handler:    func(param1, param2 int) (int, int, *jsonrpc.Error) { return 0, 0, &jsonrpc.Error{} },
+			paramNames: []jsonrpc.Parameter{{Name: "param1"}, {Name: "param2"}},
+			want:       "second return value must be a http.Header for 3 tuple handler",
 		},
 		"no error return": {
 			handler:    func(param1, param2 int) (any, int) { return 0, 0 },
 			paramNames: []jsonrpc.Parameter{{Name: "param1"}, {Name: "param2"}},
-			want:       "second return value must be a *jsonrpc.Error",
+			want:       "second return value must be a *jsonrpc.Error for 2 tuple handler",
 		},
 	}
 
@@ -472,8 +482,9 @@ func TestHandle(t *testing.T) {
 			oldRequestFailedEventCount := len(listener.OnRequestFailedCalls)
 			oldRequestHandledCalls := len(listener.OnRequestHandledCalls)
 
-			res, err := server.HandleReader(context.Background(), strings.NewReader(test.req))
+			res, httpHeader, err := server.HandleReader(context.Background(), strings.NewReader(test.req))
 			require.NoError(t, err)
+			assert.NotNil(t, httpHeader)
 
 			if test.isBatch {
 				assertBatchResponse(t, test.res, string(res))
@@ -515,8 +526,9 @@ func BenchmarkHandle(b *testing.B) {
 
 	const request = `{"jsonrpc":"2.0","id":1,"method":"test"}`
 	for i := 0; i < b.N; i++ {
-		_, err := server.HandleReader(context.Background(), strings.NewReader(request))
+		_, header, err := server.HandleReader(context.Background(), strings.NewReader(request))
 		require.NoError(b, err)
+		require.NotNil(b, header)
 	}
 }
 
@@ -531,9 +543,10 @@ func TestCannotWriteToConnInHandler(t *testing.T) {
 			return 0, nil
 		},
 	}))
-	res, err := server.HandleReader(context.Background(), strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"test"}`))
+	res, header, err := server.HandleReader(context.Background(), strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"test"}`))
 	require.NoError(t, err)
 	require.Equal(t, `{"jsonrpc":"2.0","result":0,"id":1}`, string(res))
+	require.NotNil(t, header)
 }
 
 type fakeConn struct{}

node/throttled_vm.go

@@ -33,13 +33,14 @@
 func (tvm *ThrottledVM) Execute(txns []core.Transaction, declaredClasses []core.Class, paidFeesOnL1 []*felt.Felt,
 	blockInfo *vm.BlockInfo, state core.StateReader, network *utils.Network, skipChargeFee, skipValidate, errOnRevert,
 	useBlobData bool,
-) ([]*felt.Felt, []*felt.Felt, []vm.TransactionTrace, error) {
+) ([]*felt.Felt, []*felt.Felt, []vm.TransactionTrace, uint64, error) {
 	var ret []*felt.Felt
 	var traces []vm.TransactionTrace
 	var dataGasConsumed []*felt.Felt
-	return ret, dataGasConsumed, traces, tvm.Do(func(vm *vm.VM) error {
+	var numSteps uint64
+	return ret, dataGasConsumed, traces, numSteps, tvm.Do(func(vm *vm.VM) error {
 		var err error
-		ret, dataGasConsumed, traces, err = (*vm).Execute(txns, declaredClasses, paidFeesOnL1, blockInfo, state, network,
+		ret, dataGasConsumed, traces, numSteps, err = (*vm).Execute(txns, declaredClasses, paidFeesOnL1, blockInfo, state, network,
 			skipChargeFee, skipValidate, errOnRevert, useBlobData)
 		return err
 	})