Handroll validation input preimages JSON marshal/unmarshal

validator/server_api/json.go

@@ -1,13 +1,173 @@
 package server_api
 
 import (
+	"bytes"
 	"encoding/base64"
+	"fmt"
+	"io"
 
 	"github.com/ethereum/go-ethereum/common"
 
 	"github.com/offchainlabs/nitro/validator"
 )
 
+type PreimagesMapJson struct {
+	Map map[common.Hash][]byte
+}
+
+func (m *PreimagesMapJson) MarshalJSON() ([]byte, error) {
+	size := 2 // {}
+	encoding := base64.StdEncoding
+	for key, value := range m.Map {
+		size += 5 // "":""
+		size += encoding.EncodedLen(len(key))
+		size += encoding.EncodedLen(len(value))
+	}
+	if len(m.Map) > 0 {
+		// commas
+		size += len(m.Map) - 1
+	}
+	out := make([]byte, size)
+	i := 0
+	out[i] = '{'
+	i++
+	for key, value := range m.Map {
+		if i > 1 {
+			out[i] = ','
+			i++
+		}
+		out[i] = '"'
+		i++
+		encoding.Encode(out[i:], key[:])
+		i += encoding.EncodedLen(len(key))
+		out[i] = '"'
+		i++
+		out[i] = ':'
+		i++
+		out[i] = '"'
+		i++
+		encoding.Encode(out[i:], value)
+		i += encoding.EncodedLen(len(value))
+		out[i] = '"'
+		i++
+	}
+	out[i] = '}'
+	i++
+	if i != len(out) {
+		return nil, fmt.Errorf("preimage map wrote %v bytes but expected to write %v", i, len(out))
+	}
+	return out, nil
+}
+
+func readNonWhitespace(data *[]byte) (byte, error) {
+	c := byte('\t')
+	for c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r' || c == ' ' {
+		if len(*data) == 0 {
+			return 0, io.ErrUnexpectedEOF
+		}
+		c = (*data)[0]
+		*data = (*data)[1:]
+	}
+	return c, nil
+}
+
+func expectCharacter(data *[]byte, expected rune) error {
+	got, err := readNonWhitespace(data)
+	if err != nil {
+		return fmt.Errorf("while looking for '%v' got %w", expected, err)
+	}
+	if rune(got) != expected {
+		return fmt.Errorf("while looking for '%v' got '%v'", expected, rune(got))
+	}
+	return nil
+}
+
+func getStrLen(data []byte) (int, error) {
+	// We don't allow strings to contain an escape sequence.
+	// Searching for a backslash here would be duplicated work.
+	// If the returned string length includes a backslash, base64 decoding will fail and error there.
+	strLen := bytes.IndexByte(data, '"')
+	if strLen == -1 {
+		return 0, fmt.Errorf("%w: hit end of preimages map looking for end quote", io.ErrUnexpectedEOF)
+	}
+	return strLen, nil
+}
+
+func (m *PreimagesMapJson) UnmarshalJSON(data []byte) error {
+	err := expectCharacter(&data, '{')
+	if err != nil {
+		return err
+	}
+	m.Map = make(map[common.Hash][]byte)
+	encoding := base64.StdEncoding
+	// Used to store base64 decoded data
+	// Returned unmarshalled preimage slices will just be parts of this one
+	buf := make([]byte, encoding.DecodedLen(len(data)))
+	for {
+		c, err := readNonWhitespace(&data)
+		if err != nil {
+			return fmt.Errorf("while looking for key in preimages map got %w", err)
+		}
+		if len(m.Map) == 0 && c == '}' {
+			break
+		} else if c != '"' {
+			return fmt.Errorf("expected '\"' to begin key in preimages map but got '%v'", c)
+		}
+		strLen, err := getStrLen(data)
+		if err != nil {
+			return err
+		}
+		maxKeyLen := encoding.DecodedLen(strLen)
+		if maxKeyLen > len(buf) {
+			return fmt.Errorf("preimage key base64 possible length %v is greater than buffer size of %v", maxKeyLen, len(buf))
+		}
+		keyLen, err := encoding.Decode(buf, data[:strLen])
+		if err != nil {
+			return fmt.Errorf("error base64 decoding preimage key: %w", err)
+		}
+		var key common.Hash
+		if keyLen != len(key) {
+			return fmt.Errorf("expected preimage to be %v bytes long, but got %v bytes", len(key), keyLen)
+		}
+		copy(key[:], buf[:len(key)])
+		// We don't need to advance buf here because we already copied the data we needed out of it
+		data = data[strLen+1:]
+		err = expectCharacter(&data, ':')
+		if err != nil {
+			return err
+		}
+		err = expectCharacter(&data, '"')
+		if err != nil {
+			return err
+		}
+		strLen, err = getStrLen(data)
+		if err != nil {
+			return err
+		}
+		maxValueLen := encoding.DecodedLen(strLen)
+		if maxValueLen > len(buf) {
+			return fmt.Errorf("preimage value base64 possible length %v is greater than buffer size of %v", maxValueLen, len(buf))
+		}
+		valueLen, err := encoding.Decode(buf, data[:strLen])
+		if err != nil {
+			return fmt.Errorf("error base64 decoding preimage value: %w", err)
+		}
+		m.Map[key] = buf[:valueLen]
+		buf = buf[valueLen:]
+		data = data[strLen+1:]
+		c, err = readNonWhitespace(&data)
+		if err != nil {
+			return fmt.Errorf("after value in preimages map got %w", err)
+		}
+		if c == '}' {
+			break
+		} else if c != ',' {
+			return fmt.Errorf("expected ',' or '}' after value in preimages map but got '%v'", c)
+		}
+	}
+	return nil
+}
+
 type BatchInfoJson struct {
 	Number  uint64
 	DataB64 string
@@ -17,7 +177,7 @@ type ValidationInputJson struct {
 	Id            uint64
 	HasDelayedMsg bool
 	DelayedMsgNr  uint64
-	PreimagesB64  map[string]string
+	PreimagesB64  PreimagesMapJson
 	BatchInfo     []BatchInfoJson
 	DelayedMsgB64 string
 	StartState    validator.GoGlobalState
@@ -30,12 +190,7 @@ func ValidationInputToJson(entry *validator.ValidationInput) *ValidationInputJso
 		DelayedMsgNr:  entry.DelayedMsgNr,
 		DelayedMsgB64: base64.StdEncoding.EncodeToString(entry.DelayedMsg),
 		StartState:    entry.StartState,
-		PreimagesB64:  make(map[string]string),
-	}
-	for hash, data := range entry.Preimages {
-		encHash := base64.StdEncoding.EncodeToString(hash.Bytes())
-		encData := base64.StdEncoding.EncodeToString(data)
-		res.PreimagesB64[encHash] = encData
+		PreimagesB64:  PreimagesMapJson{entry.Preimages},
 	}
 	for _, binfo := range entry.BatchInfo {
 		encData := base64.StdEncoding.EncodeToString(binfo.Data)
@@ -50,24 +205,13 @@ func ValidationInputFromJson(entry *ValidationInputJson) (*validator.ValidationI
 		HasDelayedMsg: entry.HasDelayedMsg,
 		DelayedMsgNr:  entry.DelayedMsgNr,
 		StartState:    entry.StartState,
-		Preimages:     make(map[common.Hash][]byte),
+		Preimages:     entry.PreimagesB64.Map,
 	}
 	delayed, err := base64.StdEncoding.DecodeString(entry.DelayedMsgB64)
 	if err != nil {
 		return nil, err
 	}
 	valInput.DelayedMsg = delayed
-	for encHash, encData := range entry.PreimagesB64 {
-		hash, err := base64.StdEncoding.DecodeString(encHash)
-		if err != nil {
-			return nil, err
-		}
-		data, err := base64.StdEncoding.DecodeString(encData)
-		if err != nil {
-			return nil, err
-		}
-		valInput.Preimages[common.BytesToHash(hash)] = data
-	}
 	for _, binfo := range entry.BatchInfo {
 		data, err := base64.StdEncoding.DecodeString(binfo.DataB64)
 		if err != nil {

validator/server_api/json_test.go

@@ -0,0 +1,54 @@
+package server_api
+
+import (
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/offchainlabs/nitro/util/testhelpers"
+)
+
+func Require(t *testing.T, err error, printables ...interface{}) {
+	t.Helper()
+	testhelpers.RequireImpl(t, err, printables...)
+}
+
+func TestPreimagesMapJson(t *testing.T) {
+	t.Parallel()
+	for _, preimages := range []PreimagesMapJson{
+		{},
+		{make(map[common.Hash][]byte)},
+		{map[common.Hash][]byte{
+			{}: {},
+		}},
+		{map[common.Hash][]byte{
+			{1}: {1},
+			{2}: {1, 2},
+			{3}: {1, 2, 3},
+		}},
+	} {
+		t.Run(fmt.Sprintf("%v preimages", len(preimages.Map)), func(t *testing.T) {
+			// These test cases are fast enough that t.Parallel() probably isn't worth it
+			serialized, err := preimages.MarshalJSON()
+			Require(t, err, "Failed to marshal preimagesj")
+
+			// Make sure that `serialized` is a valid JSON map
+			stringMap := make(map[string]string)
+			err = json.Unmarshal(serialized, &stringMap)
+			Require(t, err, "Failed to unmarshal preimages as string map")
+			if len(stringMap) != len(preimages.Map) {
+				t.Errorf("Got %v entries in string map but only had %v preimages", len(stringMap), len(preimages.Map))
+			}
+
+			var deserialized PreimagesMapJson
+			err = deserialized.UnmarshalJSON(serialized)
+			Require(t, err)
+
+			if (len(preimages.Map) > 0 || len(deserialized.Map) > 0) && !reflect.DeepEqual(preimages, deserialized) {
+				t.Errorf("Preimages map %v serialized to %v but then deserialized to different map %v", preimages, string(serialized), deserialized)
+			}
+		})
+	}
+}