bug fix marshal unmarshal batch data

sh-cha · sh-cha · commit 57e379ead3aa · 2024-09-09T19:41:04.000+09:00
diff --git a/executor/batch/handler.go b/executor/batch/handler.go
@@ -1,8 +1,8 @@
 package batch
 
 import (
+	"bytes"
 	"context"
-	"crypto/sha256"
 	"fmt"
 	"io"
 	"time"
@@ -138,6 +138,8 @@ func (bs *BatchSubmitter) prepareBatch(blockHeight uint64) error {
 		bs.localBatchInfo.BatchFileSize = 0
 		bs.localBatchInfo.Start = blockHeight
 		bs.localBatchInfo.End = 0
+
+		bs.batchWriter.Reset(bs.batchFile)
 	}
 	return nil
 }
@@ -158,7 +160,11 @@ func (bs *BatchSubmitter) finalizeBatch(ctx context.Context, blockHeight uint64)
 	if err != nil {
 		return errors.Wrap(err, "failed to write raw commit")
 	}
-	fileSize, err := bs.batchFileSize()
+	err = bs.batchWriter.Close()
+	if err != nil {
+		return errors.Wrap(err, "failed to close batch writer")
+	}
+	fileSize, err := bs.batchFileSize(false)
 	if err != nil {
 		return err
 	}
@@ -169,7 +175,7 @@ func (bs *BatchSubmitter) finalizeBatch(ctx context.Context, blockHeight uint64)
 
 	// TODO: improve this logic to avoid hold all the batch data in memory
 	chunks := make([][]byte, 0)
-	for offset := int64(0); ; offset += int64(bs.batchCfg.MaxChunkSize) {
+	for offset := int64(0); ; {
 		readLength, err := bs.batchFile.ReadAt(batchBuffer, offset)
 		if err != nil && err != io.EOF {
 			return err
@@ -178,15 +184,15 @@ func (bs *BatchSubmitter) finalizeBatch(ctx context.Context, blockHeight uint64)
 		}
 
 		// trim the buffer to the actual read length
-		chunk := make([]byte, readLength)
-		copy(chunk, batchBuffer[:readLength])
+		chunk := bytes.Clone(batchBuffer[:readLength])
 		chunks = append(chunks, chunk)
 
-		checksum := sha256.Sum256(batchBuffer)
+		checksum := executortypes.GetChecksumFromChunk(chunk)
 		checksums = append(checksums, checksum[:])
 		if uint64(readLength) < bs.batchCfg.MaxChunkSize {
 			break
 		}
+		offset += int64(readLength)
 	}
 
 	headerData := executortypes.MarshalBatchDataHeader(
@@ -236,7 +242,7 @@ func (bs *BatchSubmitter) finalizeBatch(ctx context.Context, blockHeight uint64)
 }
 
 func (bs *BatchSubmitter) checkBatch(ctx context.Context, blockHeight uint64, latestHeight uint64, blockTime time.Time) error {
-	fileSize, err := bs.batchFileSize()
+	fileSize, err := bs.batchFileSize(true)
 	if err != nil {
 		return err
 	}
@@ -263,13 +269,15 @@ func (bs *BatchSubmitter) checkBatch(ctx context.Context, blockHeight uint64, la
 	return nil
 }
 
-func (bs *BatchSubmitter) batchFileSize() (int64, error) {
+func (bs *BatchSubmitter) batchFileSize(flush bool) (int64, error) {
 	if bs.batchFile == nil {
 		return 0, errors.New("batch file is not initialized")
 	}
-	err := bs.batchWriter.Flush()
-	if err != nil {
-		return 0, errors.Wrap(err, "failed to flush batch writer")
+	if flush {
+		err := bs.batchWriter.Flush()
+		if err != nil {
+			return 0, errors.Wrap(err, "failed to flush batch writer")
+		}
 	}
 
 	info, err := bs.batchFile.Stat()
diff --git a/executor/types/batch.go b/executor/types/batch.go
@@ -2,8 +2,9 @@ package types
 
 import (
 	"context"
+	"crypto/sha256"
 	"encoding/binary"
-	"errors"
+	"fmt"
 	"time"
 
 	btypes "github.com/initia-labs/opinit-bots/node/broadcaster/types"
@@ -39,41 +40,67 @@ const (
 	BatchDataTypeChunk
 )
 
+type BatchDataHeader struct {
+	Start     uint64
+	End       uint64
+	Checksums [][]byte
+}
+
+type BatchDataChunk struct {
+	Start     uint64
+	End       uint64
+	Index     uint64
+	Length    uint64
+	ChunkData []byte
+}
+
+func GetChecksumFromChunk(chunk []byte) [32]byte {
+	return sha256.Sum256(chunk)
+}
+
 func MarshalBatchDataHeader(
 	start uint64,
 	end uint64,
 	checksums [][]byte,
 ) []byte {
 	data := make([]byte, 1)
 	data[0] = byte(BatchDataTypeHeader)
-	data = binary.AppendUvarint(data, start)
-	data = binary.AppendUvarint(data, end)
-	data = binary.AppendUvarint(data, uint64(len(checksums)))
+	data = binary.BigEndian.AppendUint64(data, start)
+	data = binary.BigEndian.AppendUint64(data, end)
+	data = binary.BigEndian.AppendUint64(data, uint64(len(checksums)))
 	for _, checksum := range checksums {
 		data = append(data, checksum...)
 	}
 	return data
 }
 
-func UnmarshalBatchDataHeader(data []byte) (start uint64, end uint64, length uint64, checksums [][]byte, err error) {
+func UnmarshalBatchDataHeader(data []byte) (BatchDataHeader, error) {
 	if len(data) < 25 {
-		err = errors.New("invalid data length")
-		return
+		err := fmt.Errorf("invalid data length: %d, expected > 25", len(data))
+		return BatchDataHeader{}, err
 	}
-	start, _ = binary.Uvarint(data[1:9])
-	end, _ = binary.Uvarint(data[9:17])
-	length, _ = binary.Uvarint(data[17:25])
-	checksums = make([][]byte, 0, length)
-
-	if len(data)-25%32 != 0 || (uint64(len(data)-25)/32) != length {
-		err = errors.New("invalid checksum data")
-		return
+	start := binary.BigEndian.Uint64(data[1:9])
+	end := binary.BigEndian.Uint64(data[9:17])
+	if start > end {
+		return BatchDataHeader{}, fmt.Errorf("invalid start: %d, end: %d", start, end)
 	}
 
+	length := binary.BigEndian.Uint64(data[17:25])
+	if (len(data)-25)%32 != 0 || (uint64(len(data)-25)/32) != length {
+		err := fmt.Errorf("invalid checksum length: %d, data length: %d", length, len(data)-25)
+		return BatchDataHeader{}, err
+	}
+
+	checksums := make([][]byte, 0, length)
 	for i := 25; i < len(data); i += 32 {
 		checksums = append(checksums, data[i:i+32])
 	}
-	return
+
+	return BatchDataHeader{
+		Start:     start,
+		End:       end,
+		Checksums: checksums,
+	}, nil
 }
 
 func MarshalBatchDataChunk(
@@ -85,23 +112,33 @@ func MarshalBatchDataChunk(
 ) []byte {
 	data := make([]byte, 1)
 	data[0] = byte(BatchDataTypeChunk)
-	data = binary.AppendUvarint(data, start)
-	data = binary.AppendUvarint(data, end)
-	data = binary.AppendUvarint(data, index)
-	data = binary.AppendUvarint(data, length)
+	data = binary.BigEndian.AppendUint64(data, start)
+	data = binary.BigEndian.AppendUint64(data, end)
+	data = binary.BigEndian.AppendUint64(data, index)
+	data = binary.BigEndian.AppendUint64(data, length)
 	data = append(data, chunkData...)
 	return data
 }
 
-func UnmarshalBatchDataChunk(data []byte) (start uint64, end uint64, index uint64, length uint64, chunkData []byte, err error) {
+func UnmarshalBatchDataChunk(data []byte) (BatchDataChunk, error) {
 	if len(data) < 33 {
-		err = errors.New("invalid data length")
-		return
+		err := fmt.Errorf("invalid data length: %d, expected > 33", len(data))
+		return BatchDataChunk{}, err
 	}
-	start, _ = binary.Uvarint(data[1:9])
-	end, _ = binary.Uvarint(data[9:17])
-	index, _ = binary.Uvarint(data[17:25])
-	length, _ = binary.Uvarint(data[25:33])
-	chunkData = data[33:]
-	return
+	start := binary.BigEndian.Uint64(data[1:9])
+	end := binary.BigEndian.Uint64(data[9:17])
+	if start > end {
+		return BatchDataChunk{}, fmt.Errorf("invalid start: %d, end: %d", start, end)
+	}
+	index := binary.BigEndian.Uint64(data[17:25])
+	length := binary.BigEndian.Uint64(data[25:33])
+	chunkData := data[33:]
+
+	return BatchDataChunk{
+		Start:     start,
+		End:       end,
+		Index:     index,
+		Length:    length,
+		ChunkData: chunkData,
+	}, nil
 }
diff --git a/executor/types/batch_test.go b/executor/types/batch_test.go
@@ -0,0 +1,38 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBatchDataHeader(t *testing.T) {
+	start := uint64(1)
+	end := uint64(100)
+
+	chunks := [][]byte{
+		[]byte("chunk1"),
+		[]byte("chunk2"),
+		[]byte("chunk3"),
+	}
+
+	checksums := make([][]byte, 0, len(chunks))
+	for _, chunk := range chunks {
+		checksum := GetChecksumFromChunk(chunk)
+		checksums = append(checksums, checksum[:])
+	}
+
+	headerData := MarshalBatchDataHeader(
+		start,
+		end,
+		checksums)
+	require.Equal(t, 1+8+8+8+3*32, len(headerData))
+
+	header, err := UnmarshalBatchDataHeader(headerData)
+	require.NoError(t, err)
+
+	require.Equal(t, start, header.Start)
+	require.Equal(t, end, header.End)
+	require.Equal(t, checksums, header.Checksums)
+	require.Equal(t, len(chunks), len(header.Checksums))
+}
