From a774dc9bb0aa7bd23ff6f3c0db21960d9b69e1fd Mon Sep 17 00:00:00 2001
From: agnusmor <100322135+agnusmor@users.noreply.github.com>
Date: Wed, 7 Aug 2024 11:45:00 +0200
Subject: [PATCH] Fix timestamp when querying zkevm_getBatchByNumer and the
 batch is only in trusted state (#3750)

* fix timestamp when querying zkevm_getBatchByNumer and the batch is only in trusted state

* control timestamp to be returned for a batch in GetBatchTimestamp

* fix ErrNotFound

* fix state.ErrNotFound check

* fix comments

* fix GetRawBatchTimestamps

* fix comments

* fix GetRawBatchTimestamps return err
---
 state/batch.go                | 23 ++++++++++++++++++++---
 state/pgstatestorage/batch.go |  2 +-
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/state/batch.go b/state/batch.go
index f5325a08c7..c52d2eb278 100644
--- a/state/batch.go
+++ b/state/batch.go
@@ -568,10 +568,13 @@ func (s *State) GetLastBatch(ctx context.Context, dbTx pgx.Tx) (*Batch, error) {
 	return batches[0], nil
 }
 
-// GetBatchTimestamp returns the batch timestamp.
+// GetBatchTimestamp returns the batch timestamp
+// If batch >= etrog
 //
-//	   for >= etrog is stored on virtual_batch.batch_timestamp
-//		  previous batches is stored on batch.timestamp
+//	if the batch it's virtualized it will return virtual_batch.timestamp_batch_etrog field value
+//	if the batch if's only trusted and it has L2 blocks it will return the timestamp of the last L2 block, otherwise it will return batchTimestamp
+//
+// If batch < etrog it will return batchTimestamp value
 func (s *State) GetBatchTimestamp(ctx context.Context, batchNumber uint64, forcedForkId *uint64, dbTx pgx.Tx) (*time.Time, error) {
 	var forkid uint64
 	if forcedForkId != nil {
@@ -584,6 +587,20 @@ func (s *State) GetBatchTimestamp(ctx context.Context, batchNumber uint64, force
 		return nil, err
 	}
 	if forkid >= FORKID_ETROG {
+		if virtualTimestamp == nil {
+			lastL2Block, err := s.GetLastL2BlockByBatchNumber(ctx, batchNumber, dbTx)
+			if err != nil && !errors.Is(err, ErrNotFound) {
+				return nil, err
+			}
+
+			// If the batch has L2 blocks we will return the timestamp of the last L2 block as the timestamp of the batch
+			// else we will return the batchTimestamp value (timestamp of batch creation)
+			if lastL2Block != nil {
+				return &lastL2Block.ReceivedAt, nil
+			}
+
+			return batchTimestamp, nil
+		}
 		return virtualTimestamp, nil
 	}
 	return batchTimestamp, nil
diff --git a/state/pgstatestorage/batch.go b/state/pgstatestorage/batch.go
index b0b1aa6389..bbe12a355a 100644
--- a/state/pgstatestorage/batch.go
+++ b/state/pgstatestorage/batch.go
@@ -995,7 +995,7 @@ func (p *PostgresStorage) BuildChangeL2Block(deltaTimestamp uint32, l1InfoTreeIn
 }
 
 // GetRawBatchTimestamps returns the timestamp of the batch with the given number.
-// it returns batch_num.tstamp and virtual_batch.batch_timestamp
+// it returns batch.timestamp and virtual_batch.timestamp_batch_etrog
 func (p *PostgresStorage) GetRawBatchTimestamps(ctx context.Context, batchNumber uint64, dbTx pgx.Tx) (*time.Time, *time.Time, error) {
 	const sql = `
 	SELECT b.timestamp AS batch_timestamp, v.timestamp_batch_etrog AS virtual_batch_timestamp