add tests

Author: gfukushima

src/test/java/org/ethereum/beacon/discovery/SimpleIdentitySchemaInterpreter.java

@@ -4,6 +4,7 @@
 
 package org.ethereum.beacon.discovery;
 
+import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
@@ -35,12 +36,30 @@ public static NodeRecord createNodeRecord(final int nodeId) {
 
   public static NodeRecord createNodeRecord(
       final Bytes nodeId, final InetSocketAddress udpAddress) {
+    if (udpAddress.getAddress() instanceof Inet6Address) {
+      return createNodeRecord(
+          nodeId,
+          new EnrField(EnrField.IP_V6, Bytes.wrap(udpAddress.getAddress().getAddress())),
+          new EnrField(EnrField.UDP_V6, udpAddress.getPort()));
+    }
     return createNodeRecord(
         nodeId,
         new EnrField(EnrField.IP_V4, Bytes.wrap(udpAddress.getAddress().getAddress())),
         new EnrField(EnrField.UDP, udpAddress.getPort()));
   }
 
+  public static NodeRecord createDualStackNodeRecord(
+      final Bytes nodeId,
+      final InetSocketAddress ipv4Address,
+      final InetSocketAddress ipv6Address) {
+    return createNodeRecord(
+        nodeId,
+        new EnrField(EnrField.IP_V4, Bytes.wrap(ipv4Address.getAddress().getAddress())),
+        new EnrField(EnrField.UDP, ipv4Address.getPort()),
+        new EnrField(EnrField.IP_V6, Bytes.wrap(ipv6Address.getAddress().getAddress())),
+        new EnrField(EnrField.UDP_V6, ipv6Address.getPort()));
+  }
+
   public static NodeRecord createNodeRecord(final Bytes nodeId, final EnrField... extraFields) {
     final List<EnrField> fields = new ArrayList<>(List.of(extraFields));
     fields.add(new EnrField(EnrField.ID, IdentitySchema.V4));
@@ -86,8 +105,18 @@ public Optional<InetSocketAddress> getUdpAddress(final NodeRecord nodeRecord) {
   }
 
   @Override
-  public Optional<InetSocketAddress> getUdp6Address(NodeRecord nodeRecord) {
-    return Optional.empty();
+  public Optional<InetSocketAddress> getUdp6Address(final NodeRecord nodeRecord) {
+    try {
+      final Bytes ipBytes = (Bytes) nodeRecord.get(EnrField.IP_V6);
+      if (ipBytes == null) {
+        return Optional.empty();
+      }
+      final InetAddress ipAddress = InetAddress.getByAddress(ipBytes.toArrayUnsafe());
+      final int port = (int) nodeRecord.get(EnrField.UDP_V6);
+      return Optional.of(new InetSocketAddress(ipAddress, port));
+    } catch (UnknownHostException e) {
+      return Optional.empty();
+    }
   }
 
   @Override
@@ -117,7 +146,31 @@ public NodeRecord createWithNewAddress(
       final Optional<Integer> newTcpPort,
       final Optional<Integer> newQuicPort,
       final Signer signer) {
-    final NodeRecord newRecord = createNodeRecord(getNodeId(nodeRecord), newAddress);
+    final List<EnrField> fields = new ArrayList<>();
+    // Preserve fields from the other IP family
+    if (newAddress.getAddress() instanceof Inet6Address) {
+      // Updating IPv6 — preserve IPv4 fields if present
+      if (nodeRecord.get(EnrField.IP_V4) != null) {
+        fields.add(new EnrField(EnrField.IP_V4, nodeRecord.get(EnrField.IP_V4)));
+        fields.add(new EnrField(EnrField.UDP, nodeRecord.get(EnrField.UDP)));
+      }
+      fields.add(
+          new EnrField(EnrField.IP_V6, Bytes.wrap(newAddress.getAddress().getAddress())));
+      fields.add(new EnrField(EnrField.UDP_V6, newAddress.getPort()));
+    } else {
+      // Updating IPv4 — preserve IPv6 fields if present
+      fields.add(
+          new EnrField(EnrField.IP_V4, Bytes.wrap(newAddress.getAddress().getAddress())));
+      fields.add(new EnrField(EnrField.UDP, newAddress.getPort()));
+      if (nodeRecord.get(EnrField.IP_V6) != null) {
+        fields.add(new EnrField(EnrField.IP_V6, nodeRecord.get(EnrField.IP_V6)));
+        fields.add(new EnrField(EnrField.UDP_V6, nodeRecord.get(EnrField.UDP_V6)));
+      }
+    }
+    fields.add(new EnrField(EnrField.ID, IdentitySchema.V4));
+    fields.add(new EnrField(EnrField.PKEY_SECP256K1, getNodeId(nodeRecord)));
+    final NodeRecord newRecord =
+        NodeRecord.fromValues(this, nodeRecord.getSeq().add(1), fields);
     sign(newRecord, signer);
     return newRecord;
   }

src/test/java/org/ethereum/beacon/discovery/message/handler/DefaultExternalAddressSelectorTest.java

@@ -27,6 +27,8 @@ class DefaultExternalAddressSelectorTest {
   private static final InetSocketAddress ADDRESS1 = new InetSocketAddress("127.0.0.1", 2000);
   private static final InetSocketAddress ADDRESS2 = new InetSocketAddress("127.0.0.2", 2002);
   private static final InetSocketAddress ADDRESS3 = new InetSocketAddress("127.0.0.3", 2003);
+  private static final InetSocketAddress IPV6_ADDRESS1 = new InetSocketAddress("::1", 3000);
+  private static final InetSocketAddress IPV6_ADDRESS2 = new InetSocketAddress("::2", 3002);
   private static final Instant START_TIME = Instant.ofEpochSecond(1_000_000);
   private final Bytes nodeId = Bytes.fromHexString("0x1234567890");
   private final NodeRecord originalNodeRecord =
@@ -141,7 +143,58 @@ void shouldRemoveStaleAddresses() {
     assertSelectedAddress(ADDRESS3);
   }
 
-  private void assertSelectedAddress(final InetSocketAddress address2) {
-    assertThat(localNodeRecordStore.getLocalNodeRecord().getUdpAddress()).contains(address2);
+  @Test
+  void shouldAutoDiscoverBothIpv4AndIpv6Addresses() {
+    // IPv4 gets enough confirmations
+    for (int i = 0; i < MIN_CONFIRMATIONS; i++) {
+      selector.onExternalAddressReport(Optional.empty(), ADDRESS2, START_TIME);
+    }
+    assertSelectedAddress(ADDRESS2);
+
+    // IPv6 also gets enough confirmations
+    for (int i = 0; i < MIN_CONFIRMATIONS; i++) {
+      selector.onExternalAddressReport(Optional.empty(), IPV6_ADDRESS1, START_TIME);
+    }
+    // Both families should be present in the ENR
+    assertSelectedAddress(ADDRESS2);
+    assertSelectedIpv6Address(IPV6_ADDRESS1);
+  }
+
+  @Test
+  void shouldAutoDiscoverIpv6EvenWhenIpv4HasMoreVotes() {
+    // IPv4 accumulates many more votes than IPv6
+    for (int i = 0; i < MIN_CONFIRMATIONS * 10; i++) {
+      selector.onExternalAddressReport(Optional.empty(), ADDRESS2, START_TIME);
+    }
+    assertSelectedAddress(ADDRESS2);
+
+    // IPv6 gets just enough confirmations
+    for (int i = 0; i < MIN_CONFIRMATIONS; i++) {
+      selector.onExternalAddressReport(Optional.empty(), IPV6_ADDRESS1, START_TIME);
+    }
+    // IPv6 should still be discovered despite IPv4 having far more votes
+    assertSelectedAddress(ADDRESS2);
+    assertSelectedIpv6Address(IPV6_ADDRESS1);
+  }
+
+  @Test
+  void shouldSelectMostVotedIpv6AddressIndependentlyFromIpv4() {
+    // Two competing IPv6 addresses
+    for (int i = 0; i < MIN_CONFIRMATIONS; i++) {
+      selector.onExternalAddressReport(Optional.empty(), IPV6_ADDRESS1, START_TIME);
+    }
+    for (int i = 0; i < MIN_CONFIRMATIONS + 1; i++) {
+      selector.onExternalAddressReport(Optional.empty(), IPV6_ADDRESS2, START_TIME);
+    }
+    // IPV6_ADDRESS2 should win (more votes)
+    assertSelectedIpv6Address(IPV6_ADDRESS2);
+  }
+
+  private void assertSelectedAddress(final InetSocketAddress address) {
+    assertThat(localNodeRecordStore.getLocalNodeRecord().getUdpAddress()).contains(address);
+  }
+
+  private void assertSelectedIpv6Address(final InetSocketAddress address) {
+    assertThat(localNodeRecordStore.getLocalNodeRecord().getUdp6Address()).contains(address);
   }
 }