chore: address the review comments

store/v2/multi/snapshot.go

@@ -107,9 +107,12 @@ func (rs *Store) Restore(
 	var subStore *substore
 	var storeSchemaReceived = false
 
+	var snapshotItem snapshottypes.SnapshotItem
+
+loop:
 	for {
-		snapshotItem := &snapshottypes.SnapshotItem{}
-		err := protoReader.ReadMsg(snapshotItem)
+		snapshotItem = snapshottypes.SnapshotItem{}
+		err := protoReader.ReadMsg(&snapshotItem)
 		if err == io.EOF {
 			break
 		} else if err != nil {
@@ -153,15 +156,15 @@ func (rs *Store) Restore(
 			subStore.Set(item.KV.Key, item.KV.Value)
 
 		default:
-			return *snapshotItem, nil
+			break loop
 		}
 	}
 
 	// commit the all key/values to store
 	_, err := rs.commit(height)
 	if err != nil {
-		return snapshottypes.SnapshotItem{}, sdkerrors.Wrap(err, fmt.Sprintf("error during commit the store at height %d", height))
+		return snapshotItem, sdkerrors.Wrap(err, fmt.Sprintf("error during commit the store at height %d", height))
 	}
 
-	return snapshottypes.SnapshotItem{}, nil
+	return snapshotItem, nil
 }

store/v2/multi/snapshot_test.go

@@ -184,19 +184,32 @@ func TestMultistoreSnapshotRestore(t *testing.T) {
 	// check for target store restore
 	require.Equal(t, target.LastCommitID().Version, int64(0))
 
+	dummyExtensionItem := snapshottypes.SnapshotItem{
+		Item: &snapshottypes.SnapshotItem_Extension{
+			Extension: &snapshottypes.SnapshotExtensionMeta{
+				Name:   "test",
+				Format: 1,
+			},
+		},
+	}
+
 	chunks := make(chan io.ReadCloser, 100)
 	go func() {
 		streamWriter := snapshots.NewStreamWriter(chunks)
 		require.NotNil(t, streamWriter)
 		defer streamWriter.Close()
 		err := source.Snapshot(version, streamWriter)
 		require.NoError(t, err)
+		// write an extension metadata
+		err = streamWriter.WriteMsg(&dummyExtensionItem)
+		require.NoError(t, err)
 	}()
 
 	streamReader, err := snapshots.NewStreamReader(chunks)
 	require.NoError(t, err)
-	_, err = target.Restore(version, snapshottypes.CurrentFormat, streamReader)
+	nextItem, err := target.Restore(version, snapshottypes.CurrentFormat, streamReader)
 	require.NoError(t, err)
+	require.Equal(t, *dummyExtensionItem.GetExtension(), *nextItem.GetExtension())
 
 	assert.Equal(t, source.LastCommitID(), target.LastCommitID())