feat: update security parameters when updating the addressbook

Signed-off-by: Matt Riben <matt.riben@swirldslabs.com>

sdk/src/main/java/com/hedera/hashgraph/sdk/Client.java

@@ -416,6 +416,7 @@ public synchronized Client setNetworkFromAddressBook(NodeAddressBook addressBook
             addressBook.nodeAddresses,
             isTransportSecurity() ? PORT_NODE_TLS : PORT_NODE_PLAIN
         ));
+        network.setAddressBook(addressBook);
         return this;
     }
 

sdk/src/main/java/com/hedera/hashgraph/sdk/Network.java

@@ -23,6 +23,8 @@
 import com.google.errorprone.annotations.Var;
 import com.google.protobuf.ByteString;
 
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -158,6 +160,44 @@ private Network setLedgerIdInternal(@Nullable LedgerId ledgerId, @Nullable Map<A
         return this;
     }
 
+    void setAddressBook(NodeAddressBook addressBook) {
+        Map<AccountId, NodeAddress> newAddressBook = addressBook.getNodeAddresses().stream()
+            .filter(nodeAddress -> Objects.nonNull(nodeAddress.getAccountId()))
+            .collect(Collectors.toMap(NodeAddress::getAccountId, Function.identity(),
+                /*
+                * Here we index by AccountId ignoring any subsequent entries with the same AccountId.
+                *
+                * Currently, this seems to be needed when reloading predefined address book for testnet which contains
+                * multiple entries with the same AccountId.
+                *
+                * If it becomes necessary to better handle such cases, either the one-to-one mapping from AccountId to
+                * single NodeAddress should be abandoned or NodeAddresses with the same AccountId may need to be merged.
+                * */
+                (a, b) -> a));
+        /*
+        * Here we preserve the certificate hash in the case where one is previously defined and no new one is provided.
+        *
+        * Currently, this seems to be needed since the downloaded address book lacks the certificate hash. However,
+        * it is expected the certificate hash will be provided in the future in which case this workaround will no
+        * longer be necessary.
+        * */
+        if (null != this.addressBook) {
+            for (Map.Entry<AccountId, NodeAddress> entry : newAddressBook.entrySet()) {
+                NodeAddress previous = this.addressBook.get(entry.getKey());
+                if (null != previous) {
+                    ByteString certHash = entry.getValue().getCertHash();
+                    if (null == certHash || certHash.isEmpty()) {
+                        entry.getValue().setCertHash(previous.certHash);
+                    }
+                }
+            }
+        }
+        this.addressBook = newAddressBook;
+        for (var node : nodes) {
+            node.setAddressBookEntry(this.addressBook.get(node.getAccountId()));
+        }
+    }
+
     @Nullable
     private static Map<AccountId, NodeAddress> getAddressBookForLedger(@Nullable LedgerId ledgerId) {
         return (ledgerId == null || !ledgerId.isKnownNetwork()) ?

sdk/src/test/java/com/hedera/hashgraph/sdk/ClientTest.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.params.provider.NullSource;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import com.google.protobuf.ByteString;
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.InputStream;
@@ -38,7 +39,9 @@
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
 
+import static com.hedera.hashgraph.sdk.BaseNodeAddress.PORT_NODE_PLAIN;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -360,4 +363,64 @@ void testExecuteSyncTimeout(String timeoutSite) throws Exception {
 
         client.close();
     }
+
+    com.hedera.hashgraph.sdk.proto.NodeAddress nodeAddress(long accountNum, String rsaPubKeyHex, byte[] certHash, byte[] ipv4) {
+        com.hedera.hashgraph.sdk.proto.NodeAddress.Builder builder = com.hedera.hashgraph.sdk.proto.NodeAddress.newBuilder()
+            .setNodeAccountId(com.hedera.hashgraph.sdk.proto.AccountID.newBuilder()
+                .setAccountNum(accountNum)
+                .build())
+            .addServiceEndpoint(com.hedera.hashgraph.sdk.proto.ServiceEndpoint.newBuilder()
+                .setIpAddressV4(ByteString.copyFrom(ipv4))
+                .setPort(PORT_NODE_PLAIN)
+                .build())
+            .setRSAPubKey(rsaPubKeyHex);
+        if (certHash != null) {
+            builder.setNodeCertHash(ByteString.copyFrom(certHash));
+        }
+        return builder.build();
+    }
+
+    @Test
+    @DisplayName("setNetworkFromAddressBook() updates security parameters in the client")
+    void setNetworkFromAddressBook() throws Exception {
+        try (Client client = Client.forNetwork(Map.of())) {
+            Function<Integer, NodeAddress> nodeAddress = accountNum -> client.network.network.get(new AccountId(accountNum)).get(0).getAddressBookEntry();
+
+            // reconfigure client network from addressbook (add new nodes)
+            client.setNetworkFromAddressBook(NodeAddressBook.fromBytes(com.hedera.hashgraph.sdk.proto.NodeAddressBook.newBuilder()
+                .addNodeAddress(nodeAddress(10001, "10001", new byte[] {1, 0, 1}, new byte[] {10, 0, 0, 1}))
+                .addNodeAddress(nodeAddress(10002, "10002", new byte[] {1, 0, 2}, new byte[] {10, 0, 0, 2}))
+                .build().toByteString()));
+
+            // verify security parameters in client
+            assertThat(nodeAddress.apply(10001).certHash).isEqualTo(ByteString.copyFrom(new byte[]{1, 0, 1}));
+            assertThat(nodeAddress.apply(10001).publicKey).isEqualTo("10001");
+            assertThat(nodeAddress.apply(10002).certHash).isEqualTo(ByteString.copyFrom(new byte[]{1, 0, 2}));
+            assertThat(nodeAddress.apply(10002).publicKey).isEqualTo("10002");
+
+            // reconfigure client network from addressbook without `certHash`
+            client.setNetworkFromAddressBook(NodeAddressBook.fromBytes(com.hedera.hashgraph.sdk.proto.NodeAddressBook.newBuilder()
+                .addNodeAddress(nodeAddress(10001, "10001", null, new byte[] {10, 0, 0, 1}))
+                .addNodeAddress(nodeAddress(10002, "10002", null, new byte[] {10, 0, 0, 2}))
+                .build().toByteString()));
+
+            // verify security parameters in client (unchanged)
+            assertThat(nodeAddress.apply(10001).certHash).isEqualTo(ByteString.copyFrom(new byte[]{1, 0, 1}));
+            assertThat(nodeAddress.apply(10001).publicKey).isEqualTo("10001");
+            assertThat(nodeAddress.apply(10002).certHash).isEqualTo(ByteString.copyFrom(new byte[]{1, 0, 2}));
+            assertThat(nodeAddress.apply(10002).publicKey).isEqualTo("10002");
+
+            // reconfigure client network from addressbook (update existing nodes)
+            client.setNetworkFromAddressBook(NodeAddressBook.fromBytes(com.hedera.hashgraph.sdk.proto.NodeAddressBook.newBuilder()
+                .addNodeAddress(nodeAddress(10001, "810001", new byte[] {8, 1, 0, 1}, new byte[] {10, 0, 0, 1}))
+                .addNodeAddress(nodeAddress(10002, "810002", new byte[] {8, 1, 0, 2}, new byte[] {10, 0, 0, 2}))
+                .build().toByteString()));
+
+            // verify security parameters in client
+            assertThat(nodeAddress.apply(10001).certHash).isEqualTo(ByteString.copyFrom(new byte[]{8, 1, 0, 1}));
+            assertThat(nodeAddress.apply(10001).publicKey).isEqualTo("810001");
+            assertThat(nodeAddress.apply(10002).certHash).isEqualTo(ByteString.copyFrom(new byte[]{8, 1, 0, 2}));
+            assertThat(nodeAddress.apply(10002).publicKey).isEqualTo("810002");
+        }
+    }
 }