feat(snap): download state leaves only and build state trie on client side

Author: Vovchyk

rskj-core/src/integrationTest/java/co/rsk/snap/SnapshotSyncIntegrationTest.java

@@ -18,13 +18,15 @@
  */
 package co.rsk.snap;
 
+import co.rsk.core.RskAddress;
 import co.rsk.util.*;
 import co.rsk.util.cli.NodeIntegrationTestCommandLine;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.squareup.okhttp.Response;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
+import org.ethereum.crypto.ECKey;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -170,14 +172,20 @@ private String readServerNodeId(Path serverDatabasePath) throws IOException {
     private void generateBlocks() throws IOException {
         List<String> accounts = OkHttpClientTestFixture.PRE_FUNDED_ACCOUNTS;
         Random rand = new Random(111);
+        int bulkSize = 10;
 
         for (int i = 0; i < 1001; i++) {
-            OkHttpClientTestFixture.FromToAddressPair[] pairs = IntStream.range(0, 10)
-                    .mapToObj(n -> of(accounts.get(rand.nextInt(accounts.size())), accounts.get(rand.nextInt(accounts.size()))))
+            int bulkPos = bulkSize * i;
+            OkHttpClientTestFixture.FromToAddressPair[] pairs = IntStream.range(0, bulkSize)
+                    .mapToObj(n -> of(accounts.get(rand.nextInt(accounts.size())), accountFromSeed("cow #" + bulkPos + n)))
                     .toArray(OkHttpClientTestFixture.FromToAddressPair[]::new);
             Response response = OkHttpClientTestFixture.sendBulkTransactions(portServerRpc, pairs);
             assertTrue(response.isSuccessful());
             response.body().close();
         }
     }
+
+    private static String accountFromSeed(String seed) {
+        return "0x" + new RskAddress(ECKey.fromPrivate(seed.getBytes()).getAddress());
+    }
 }

rskj-core/src/integrationTest/java/co/rsk/util/OkHttpClientTestFixture.java

@@ -163,8 +163,8 @@ public static Response sendBulkTransactions(int rpcPort, FromToAddressPair... fr
         Objects.requireNonNull(fromToAddresses);
 
         String gas = "0x9C40";
-        String gasPrice = "0x10";
-        String value = "0x500";
+        String gasPrice = "0x1";
+        String value = "0x1";
 
         String[] placeholders = new String[]{
                 "<ADDRESS_FROM>", "<ADDRESS_TO>", "<GAS>",

rskj-core/src/main/java/co/rsk/RskContext.java

@@ -2090,6 +2090,7 @@ private SnapshotProcessor getSnapshotProcessor() {
                     getPeersInformation(),
                     getBlockStore(),
                     getTransactionPool(),
+                    getStateRootHandler(),
                     getSnapBlockParentDependantValidationRule(),
                     getSnapBlockValidationRule(),
                     new BlockHeaderParentCompositeRule(

rskj-core/src/main/java/co/rsk/metrics/profilers/MetricKind.java

@@ -60,4 +60,6 @@ public enum MetricKind {
     SNAP_BLOCKS_RESPONSE,
     SNAP_STATE_CHUNK_REQUEST,
     SNAP_STATE_CHUNK_RESPONSE,
+    SNAP_STATE_V2_CHUNK_REQUEST,
+    SNAP_STATE_V2_CHUNK_RESPONSE,
 }

rskj-core/src/main/java/co/rsk/net/SnapshotProcessor.java

@@ -264,6 +264,19 @@ public void processStateChunkResponse(Peer peer, SnapStateChunkResponseMessage r
         }
     }
 
+    public void processStateChunkResponse(Peer peer, SnapStateChunkV2ResponseMessage responseMessage) {
+        peersInformation.getOrRegisterPeer(peer);
+
+        long messageId = responseMessage.getId();
+        MessageType messageType = responseMessage.getMessageType();
+        if (isPending(messageId, messageType)) {
+            removePendingMessage(messageId, messageType);
+            syncState.onSnapStateChunk(peer, responseMessage);
+        } else {
+            notifyUnexpectedMessageToPeerScoring(peer, "snap state chunk");
+        }
+    }
+
     @Override
     public void sendSkeletonRequest(Peer peer, long height) {
         logger.debug("Send skeleton request to node {} height {}", peer.getPeerNodeID(), height);
@@ -318,7 +331,7 @@ public void startBlockForwardSyncing(Peer peer) {
 
     @Override
     public void startSnapCapablePeerSelection() {
-        logger.info("Start peer selection");
+        logger.info("Start snap peer selection");
         setSyncState(new SnapCapablePeerSelectionSyncState(this, syncConfiguration, peersInformation, blockHeaderValidationRule, difficultyRule));
     }
 

rskj-core/src/main/java/co/rsk/net/SyncProcessor.java

@@ -294,6 +294,18 @@ public Message createMessage(BlockFactory blockFactory, RLPList list) {
             return SnapBlocksResponseMessage.decodeMessage(blockFactory, list);
         }
     },
+    SNAP_STATE_CHUNK_V2_REQUEST_MESSAGE(26) {
+        @Override
+        public Message createMessage(BlockFactory blockFactory, RLPList list) {
+            return SnapStateChunkV2RequestMessage.decodeMessage(list);
+        }
+    },
+    SNAP_STATE_CHUNK_V2_RESPONSE_MESSAGE(27) {
+        @Override
+        public Message createMessage(BlockFactory blockFactory, RLPList list) {
+            return SnapStateChunkV2ResponseMessage.decodeMessage(list);
+        }
+    },
     ;
 
     private final int type;

rskj-core/src/main/java/co/rsk/net/messages/MessageType.java

@@ -219,6 +219,16 @@ public void apply(SnapStateChunkResponseMessage message) {
         this.syncProcessor.processStateChunkResponse(sender, message);
     }
 
+    public void apply(SnapStateChunkV2RequestMessage message) {
+        logger.debug("snapshot chunk v2 request : {}", message.getId());
+        this.snapshotProcessor.processStateChunkRequest(sender, message);
+    }
+
+    public void apply(SnapStateChunkV2ResponseMessage message) {
+        logger.debug("snapshot chunk v2 response : {}", message.getId());
+        this.syncProcessor.processStateChunkResponse(sender, message);
+    }
+
     public void apply(TransactionsMessage message) {
         if (blockProcessor.hasBetterBlockToSync()) {
             loggerMessageProcess.debug("Message[{}] not processed.", message.getMessageType());

rskj-core/src/main/java/co/rsk/net/messages/MessageVisitor.java

@@ -0,0 +1,85 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2025 RSK Labs Ltd.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package co.rsk.net.messages;
+
+import org.bouncycastle.util.BigIntegers;
+import org.ethereum.util.RLP;
+import org.ethereum.util.RLPList;
+
+public class SnapStateChunkV2RequestMessage extends MessageWithId {
+
+    private final long id;
+
+    private final byte[] blockHash;
+
+    private final byte[] fromKey;
+
+    public SnapStateChunkV2RequestMessage(long id, byte[] blockHash, byte[] fromKey) {
+        this.id = id;
+        this.blockHash = blockHash;
+        this.fromKey = fromKey;
+    }
+
+    @Override
+    public long getId() {
+        return this.id;
+    }
+
+    public byte[] getBlockHash() {
+        return blockHash;
+    }
+
+    public byte[] getFromKey() {
+        return fromKey;
+    }
+
+    @Override
+    protected byte[] getEncodedMessageWithoutId() {
+        byte[] rlpBlockHash = RLP.encodeElement(this.blockHash);
+        byte[] rlpFromKey = RLP.encodeElement(this.fromKey);
+        return RLP.encodeList(rlpBlockHash, rlpFromKey);
+    }
+
+    @Override
+    public MessageType getMessageType() {
+        return MessageType.SNAP_STATE_CHUNK_V2_REQUEST_MESSAGE;
+    }
+
+    @Override
+    public MessageType getResponseMessageType() {
+        return MessageType.SNAP_STATE_CHUNK_V2_RESPONSE_MESSAGE;
+    }
+
+    @Override
+    public void accept(MessageVisitor v) {
+        v.apply(this);
+    }
+
+    public static Message decodeMessage(RLPList list) {
+        byte[] rlpId = list.get(0).getRLPData();
+
+        long id = rlpId == null ? 0 : BigIntegers.fromUnsignedByteArray(rlpId).longValue();
+
+        RLPList message = (RLPList)RLP.decode2(list.get(1).getRLPData()).get(0);
+        byte[] rlpBlockHash = message.get(0).getRLPData();
+        byte[] rlpFromKey = message.get(1).getRLPData();
+
+        return new SnapStateChunkV2RequestMessage(id, rlpBlockHash, rlpFromKey);
+    }
+}

rskj-core/src/main/java/co/rsk/net/messages/SnapStateChunkV2RequestMessage.java

@@ -0,0 +1,72 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2025 RSK Labs Ltd.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package co.rsk.net.messages;
+
+import co.rsk.trie.TrieChunk;
+import org.bouncycastle.util.BigIntegers;
+import org.ethereum.util.RLP;
+import org.ethereum.util.RLPList;
+
+public class SnapStateChunkV2ResponseMessage extends MessageWithId {
+
+    private final long id;
+
+    private final TrieChunk chunk;
+
+    public SnapStateChunkV2ResponseMessage(long id, TrieChunk chunk) {
+        this.id = id;
+        this.chunk = chunk;
+    }
+
+    @Override
+    public long getId() {
+        return this.id;
+    }
+
+    public TrieChunk getChunk() {
+        return chunk;
+    }
+
+    @Override
+    protected byte[] getEncodedMessageWithoutId() {
+        return chunk.encode();
+    }
+
+    @Override
+    public MessageType getMessageType() {
+        return MessageType.SNAP_STATE_CHUNK_V2_RESPONSE_MESSAGE;
+    }
+
+    @Override
+    public void accept(MessageVisitor v) {
+        v.apply(this);
+    }
+
+    public static Message decodeMessage(RLPList list) {
+        byte[] rlpId = list.get(0).getRLPData();
+
+        long id = rlpId == null ? 0 : BigIntegers.fromUnsignedByteArray(rlpId).longValue();
+
+        RLPList message = (RLPList)RLP.decode2(list.get(1).getRLPData()).get(0);
+
+        TrieChunk chunk = TrieChunk.decode(message);
+
+        return new SnapStateChunkV2ResponseMessage(id, chunk);
+    }
+}