This is an automated email from the ASF dual-hosted git repository. imbajin pushed a commit to branch cx/bump-server-api-version in repository https://gitbox.apache.org/repos/asf/hugegraph.git
commit 3d9d9544e3bcd7f2a29a0562f1a67b47a848d164 Author: dark <[email protected]> AuthorDate: Mon Aug 17 10:01:22 2026 +0800 fix(pd): validate raft peer addresses - normalize configured and runtime peer addresses - enforce DNS-aware IP authorization for raft traffic - refresh peer allowlists during membership changes - cover service updates and raft authorization integration --- .../org/apache/hugegraph/pd/raft/PeerUtil.java | 43 +- .../org/apache/hugegraph/pd/raft/RaftEngine.java | 148 ++++--- .../hugegraph/pd/raft/auth/IpAuthHandler.java | 429 +++++++++++++++++++- hugegraph-pd/hg-pd-service/pom.xml | 12 + .../org/apache/hugegraph/pd/service/PDService.java | 119 +++++- .../pd/service/PDServiceUpdateRaftTest.java | 195 +++++++++ .../apache/hugegraph/pd/core/PDCoreSuiteTest.java | 2 +- .../hugegraph/pd/raft/IpAuthHandlerTest.java | 133 ------- .../pd/raft/RaftEngineIpAuthIntegrationTest.java | 81 +++- .../hugegraph/pd/raft/auth/IpAuthHandlerTest.java | 439 +++++++++++++++++++++ 10 files changed, 1375 insertions(+), 226 deletions(-) diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/PeerUtil.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/PeerUtil.java index 265c7d4fc..bfffdf285 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/PeerUtil.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/PeerUtil.java @@ -17,15 +17,17 @@ package org.apache.hugegraph.pd.raft; -import com.alipay.sofa.jraft.JRaftUtils; -import com.alipay.sofa.jraft.entity.PeerId; -import org.apache.hugegraph.pd.common.KVPair; - import java.util.LinkedList; import java.util.List; import java.util.Objects; +import org.apache.hugegraph.pd.common.KVPair; + +import com.alipay.sofa.jraft.conf.Configuration; +import com.alipay.sofa.jraft.entity.PeerId; + public class PeerUtil { + public static boolean isPeerEquals(PeerId p1, PeerId p2) { if (p1 == null && p2 == null) { return true; @@ -40,19 +42,42 @@ public class PeerUtil { List<KVPair<String, PeerId>> result = new LinkedList<>(); if (conf != null && conf.length() > 0) { - for (var s : conf.split(",")) { + for (var s : conf.split(",", -1)) { + String role; + String peer; if (s.endsWith("/leader")) { - result.add(new KVPair<>("leader", JRaftUtils.getPeerId(s.substring(0, s.length() - 7)))); + role = "leader"; + peer = s.substring(0, s.length() - 7); } else if (s.endsWith("/learner")) { - result.add(new KVPair<>("learner", JRaftUtils.getPeerId(s.substring(0, s.length() - 8)))); + role = "learner"; + peer = s.substring(0, s.length() - 8); } else if (s.endsWith("/follower")) { - result.add(new KVPair<>("follower", JRaftUtils.getPeerId(s.substring(0, s.length() - 9)))); + role = "follower"; + peer = s.substring(0, s.length() - 9); } else { - result.add(new KVPair<>("follower", JRaftUtils.getPeerId(s))); + role = "follower"; + peer = s; } + result.add(new KVPair<>(role, parsePeer(peer))); } } return result; } + + public static Configuration parsePeerList(String peerList) { + Configuration configuration = new Configuration(); + for (String peer : peerList.split(",", -1)) { + configuration.addPeer(parsePeer(peer)); + } + return configuration; + } + + private static PeerId parsePeer(String value) { + PeerId peer = new PeerId(); + if (value.isEmpty() || !peer.parse(value)) { + throw new IllegalArgumentException("Invalid Raft peer: " + value); + } + return peer; + } } diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/RaftEngine.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/RaftEngine.java index 2b08de7d4..81543ee1e 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/RaftEngine.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/RaftEngine.java @@ -127,14 +127,28 @@ public class RaftEngine { final PeerId serverId = JRaftUtils.getPeerId(config.getAddress()); - rpcServer = createRaftRpcServer(config.getAddress(), initConf.getPeers()); - // construct raft group and start raft - this.raftGroupService = - new RaftGroupService(groupId, serverId, nodeOptions, rpcServer, true); - this.raftNode = raftGroupService.start(false); - log.info("RaftEngine start successfully: id = {}, peers list = {}", groupId, - nodeOptions.getInitialConf().getPeers()); - return this.raftNode != null; + try { + rpcServer = createRaftRpcServer(config.getAddress(), initConf.getPeers()); + // construct raft group and start raft + this.raftGroupService = + new RaftGroupService(groupId, serverId, nodeOptions, + rpcServer, true); + this.raftNode = raftGroupService.start(false); + if (this.raftNode == null) { + this.shutDown(); + return false; + } + log.info("RaftEngine start successfully: id = {}, peers list = {}", + groupId, nodeOptions.getInitialConf().getPeers()); + return true; + } catch (RuntimeException | Error e) { + try { + this.shutDown(); + } catch (RuntimeException | Error cleanupFailure) { + e.addSuppressed(cleanupFailure); + } + throw e; + } } /** @@ -143,13 +157,32 @@ public class RaftEngine { private RpcServer createRaftRpcServer(String raftAddr, List<PeerId> peers) { Endpoint endpoint = JRaftUtils.getEndPoint(raftAddr); RpcServer rpcServer = RaftRpcServerFactory.createRaftRpcServer(endpoint); - configureRaftServerIpWhitelist(peers, rpcServer); - RaftRpcProcessor.registerProcessor(rpcServer, this); - rpcServer.init(null); - return rpcServer; + try { + IpAuthHandler ipAuthHandler = IpAuthHandler.getInstance( + peers.stream() + .map(PeerId::getIp) + .collect(Collectors.toSet())); + configureRaftServerIpWhitelist(ipAuthHandler, rpcServer); + RaftRpcProcessor.registerProcessor(rpcServer, this); + if (!rpcServer.init(null)) { + throw new IllegalStateException( + "Failed to initialize Raft RPC server"); + } + return rpcServer; + } catch (RuntimeException | Error e) { + try { + rpcServer.shutdown(); + } catch (RuntimeException | Error cleanupFailure) { + e.addSuppressed(cleanupFailure); + } finally { + IpAuthHandler.shutdownInstance(); + } + throw e; + } } - private static void configureRaftServerIpWhitelist(List<PeerId> peers, RpcServer rpcServer) { + private static void configureRaftServerIpWhitelist( + IpAuthHandler ipAuthHandler, RpcServer rpcServer) { if (rpcServer instanceof BoltRpcServer) { ((BoltRpcServer) rpcServer).getServer().option( BoltServerOption.EXTENDED_NETTY_CHANNEL_HANDLER, @@ -157,11 +190,7 @@ public class RaftEngine { @Override public List<ChannelHandler> frontChannelHandlers() { return Collections.singletonList( - IpAuthHandler.getInstance( - peers.stream() - .map(PeerId::getIp) - .collect(Collectors.toSet()) - ) + ipAuthHandler ); } @@ -175,24 +204,38 @@ public class RaftEngine { } public void shutDown() { - if (this.raftGroupService != null) { - this.raftGroupService.shutdown(); - try { - this.raftGroupService.join(); - } catch (final InterruptedException e) { - this.raftNode = null; - ThrowUtil.throwException(e); + InterruptedException interrupted = null; + try { + if (this.raftGroupService != null) { + this.raftGroupService.shutdown(); + try { + this.raftGroupService.join(); + } catch (InterruptedException e) { + interrupted = e; + } } + } finally { this.raftGroupService = null; + try { + if (this.rpcServer != null) { + this.rpcServer.shutdown(); + } + } finally { + this.rpcServer = null; + try { + if (this.raftNode != null) { + this.raftNode.shutdown(); + } + } finally { + this.raftNode = null; + IpAuthHandler.shutdownInstance(); + } + } } - if (this.rpcServer != null) { - this.rpcServer.shutdown(); - this.rpcServer = null; - } - if (this.raftNode != null) { - this.raftNode.shutdown(); + if (interrupted != null) { + Thread.currentThread().interrupt(); + ThrowUtil.throwException(interrupted); } - this.raftNode = null; } public boolean isLeader() { @@ -352,32 +395,43 @@ public class RaftEngine { public Status changePeerList(String peerList) { AtomicReference<Status> result = new AtomicReference<>(); - Configuration newPeers = new Configuration(); try { + IpAuthHandler.validatePeerListShape(peerList); String[] peers = peerList.split(",", -1); if ((peers.length & 1) != 1) { throw new PDException(-1, "the number of peer list must be odd."); } - newPeers.parse(peerList); + Configuration newPeers = PeerUtil.parsePeerList(peerList); + Set<String> newIps = newPeers.getPeers() + .stream() + .map(PeerId::getIp) + .collect(Collectors.toSet()); + IpAuthHandler.validateAllowedEntries(newIps); + IpAuthHandler.requireActiveInstance(); CountDownLatch latch = new CountDownLatch(1); this.raftNode.changePeers(newPeers, status -> { - result.compareAndSet(null, status); - if (status != null && status.isOk()) { - IpAuthHandler handler = IpAuthHandler.getInstance(); - if (handler != null) { - Set<String> newIps = newPeers.getPeers() - .stream() - .map(PeerId::getIp) - .collect(Collectors.toSet()); - handler.refresh(newIps); + Status callbackStatus = status; + try { + if (status != null && status.isOk()) { + IpAuthHandler.refreshInstance(newIps); log.info("IpAuthHandler refreshed after peer list change to: {}", peerList); - } else { - log.warn("IpAuthHandler not initialized, skipping refresh for " - + "peer list: {}", peerList); + } else if (status == null) { + callbackStatus = new Status( + RaftError.EINTERNAL, + "changePeers returned no status"); } + } catch (RuntimeException e) { + callbackStatus = new Status( + RaftError.EINTERNAL, + "Raft peers changed but allowlist refresh failed: %s", + e.getMessage()); + log.error("Failed to refresh IpAuthHandler after peer list change to {}", + peerList, e); + } finally { + result.compareAndSet(null, callbackStatus); + latch.countDown(); } - latch.countDown(); }); boolean completed = latch.await(3L * config.getRpcTimeout(), TimeUnit.MILLISECONDS); if (!completed && result.get() == null) { diff --git a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandler.java b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandler.java index bdccb6dd7..e81c86ecd 100644 --- a/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandler.java +++ b/hugegraph-pd/hg-pd-core/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandler.java @@ -19,28 +19,119 @@ package org.apache.hugegraph.pd.raft.auth; import java.net.InetAddress; import java.net.InetSocketAddress; -import java.net.UnknownHostException; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; +import java.util.List; +import java.util.Map; import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import io.netty.channel.ChannelDuplexHandler; import io.netty.channel.ChannelHandler; import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.nio.NioDatagramChannel; +import io.netty.resolver.dns.DnsNameResolver; +import io.netty.resolver.dns.DnsNameResolverBuilder; import lombok.extern.slf4j.Slf4j; @Slf4j @ChannelHandler.Sharable public class IpAuthHandler extends ChannelDuplexHandler { + private static final long DNS_QUERY_TIMEOUT_MILLIS = 500L; + private static final long DNS_STALE_MILLIS = 30_000L; + private static final long DNS_REFRESH_MILLIS = 1_000L; + private static final int MAX_CONCURRENT_DNS_QUERIES = 8; + private static final int MAX_ALLOWED_ENTRIES = 127; + private static final int MAX_HOST_LENGTH = 253; + private static final int MAX_PEER_LIST_LENGTH = + MAX_ALLOWED_ENTRIES * (MAX_HOST_LENGTH + 16); + + private final HostResolver resolver; + private final long queryTimeoutMillis; + private final long staleMillis; + private final long refreshMillis; + private final Map<String, ResolvedEntry> resolvedByEntry; + private final Map<String, Query> inFlight; + private final Set<String> failedEntries; + private final ScheduledExecutorService refreshExecutor; + private boolean closed; + private int nextResolutionIndex; + private List<String> resolutionOrder; + private volatile Set<String> allowedEntries; private volatile Set<String> resolvedIps; private static volatile IpAuthHandler instance; private IpAuthHandler(Set<String> allowedIps) { - this.resolvedIps = resolveAll(allowedIps); + this(allowedIps, new NettyHostResolver(DNS_QUERY_TIMEOUT_MILLIS), true, + DNS_QUERY_TIMEOUT_MILLIS, DNS_STALE_MILLIS, + DNS_REFRESH_MILLIS); + } + + IpAuthHandler(Set<String> allowedIps, HostResolver resolver, + boolean scheduleRefresh, long queryTimeoutMillis, + long staleMillis, long refreshMillis) { + this.resolver = resolver; + this.queryTimeoutMillis = queryTimeoutMillis; + this.staleMillis = staleMillis; + this.refreshMillis = refreshMillis; + this.resolvedByEntry = new HashMap<>(); + this.inFlight = new HashMap<>(); + this.failedEntries = new HashSet<>(); + this.nextResolutionIndex = 0; + this.resolutionOrder = Collections.emptyList(); + try { + this.replaceAllowedEntries(allowedIps); + } catch (RuntimeException | Error e) { + try { + this.resolver.close(); + } catch (RuntimeException | Error cleanupFailure) { + e.addSuppressed(cleanupFailure); + } + throw e; + } + this.resolvedIps = this.allowedEntries; + this.closed = false; + if (scheduleRefresh) { + this.refreshExecutor = Executors.newSingleThreadScheduledExecutor(task -> { + Thread thread = new Thread(task, "pd-raft-dns-resolver"); + thread.setDaemon(true); + return thread; + }); + } else { + this.refreshExecutor = null; + } + try { + this.refreshResolvedIps(); + if (this.refreshExecutor != null) { + this.refreshExecutor.scheduleWithFixedDelay( + this::refreshSafely, this.refreshMillis, + this.refreshMillis, TimeUnit.MILLISECONDS); + } + } catch (RuntimeException | Error e) { + if (this.refreshExecutor != null) { + this.refreshExecutor.shutdownNow(); + } + try { + this.resolver.close(); + } catch (RuntimeException | Error cleanupFailure) { + e.addSuppressed(cleanupFailure); + } + throw e; + } } public static IpAuthHandler getInstance(Set<String> allowedIps) { + validateAllowedEntries(allowedIps); if (instance == null) { synchronized (IpAuthHandler.class) { if (instance == null) { @@ -59,17 +150,48 @@ public class IpAuthHandler extends ChannelDuplexHandler { return instance; } + public static IpAuthHandler requireActiveInstance() { + IpAuthHandler handler = instance; + if (handler == null || handler.isClosed()) { + throw new IllegalStateException( + "Raft peer IP allowlist is not active"); + } + return handler; + } + + public static void refreshInstance(Set<String> newAllowedIps) { + requireActiveInstance().refresh(newAllowedIps); + } + /** * Refreshes the resolved IP allowlist from a new set of hostnames or IPs. * Should be called when the Raft peer list changes via RaftEngine#changePeerList(). - * Note: DNS-only changes (e.g. container restart with new IP, same hostname) - * are not automatically detected and still require a process restart. + * DNS is also refreshed in the background so stable peer names can safely + * follow address changes without blocking a Netty event loop. */ - public void refresh(Set<String> newAllowedIps) { - this.resolvedIps = resolveAll(newAllowedIps); + public synchronized void refresh(Set<String> newAllowedIps) { + if (this.closed) { + throw new IllegalStateException( + "Raft peer IP allowlist is closed"); + } + this.replaceAllowedEntries(newAllowedIps); + this.resolvedByEntry.keySet().retainAll(this.allowedEntries); + this.failedEntries.retainAll(this.allowedEntries); + this.inFlight.entrySet().removeIf(entry -> { + if (!this.allowedEntries.contains(entry.getKey())) { + entry.getValue().cancel(); + return true; + } + return false; + }); + this.refreshResolvedIps(); log.info("IpAuthHandler allowlist refreshed, resolved {} entries", resolvedIps.size()); } + private synchronized boolean isClosed() { + return this.closed; + } + @Override public void channelActive(ChannelHandlerContext ctx) throws Exception { String clientIp = getClientIp(ctx); @@ -92,20 +214,301 @@ public class IpAuthHandler extends ChannelDuplexHandler { return resolved.isEmpty() || resolved.contains(ip); } - private static Set<String> resolveAll(Set<String> entries) { - Set<String> result = new HashSet<>(entries); + synchronized void refreshResolvedIps() { + this.refreshResolvedIps(true); + } + synchronized void refreshResolvedIps(boolean waitForResults) { + if (this.closed) { + return; + } + Set<String> entries = this.allowedEntries; + this.collectQueries(entries, false); + int attempted = 0; + while (this.inFlight.size() < MAX_CONCURRENT_DNS_QUERIES && + attempted < this.resolutionOrder.size()) { + String entry = this.resolutionOrder.get(this.nextResolutionIndex); + this.nextResolutionIndex = + (this.nextResolutionIndex + 1) % this.resolutionOrder.size(); + attempted++; + if (!this.inFlight.containsKey(entry)) { + this.inFlight.put( + entry, new Query(this.resolver.resolve(entry), + System.nanoTime())); + } + } + this.collectQueries(entries, waitForResults); + + long staleNanos = TimeUnit.MILLISECONDS.toNanos(this.staleMillis); + long now = System.nanoTime(); + this.resolvedByEntry.entrySet().removeIf( + entry -> now - entry.getValue().resolvedAtNanos > staleNanos); + Set<String> resolved = new HashSet<>(entries); + this.resolvedByEntry.values().forEach( + entry -> resolved.addAll(entry.addresses)); + this.resolvedIps = Collections.unmodifiableSet(resolved); + } + + private void collectQueries(Set<String> entries, + boolean waitForResults) { + long deadline = System.nanoTime() + + TimeUnit.MILLISECONDS.toNanos(this.queryTimeoutMillis); for (String entry : entries) { + Query query = this.inFlight.get(entry); + if (query == null) { + continue; + } + CompletableFuture<ResolvedQuery> future = query.future; try { - for (InetAddress addr : InetAddress.getAllByName(entry)) { - result.add(addr.getHostAddress()); + ResolvedQuery result; + if (future.isDone()) { + result = future.get(); + } else if (waitForResults) { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0L) { + expireQuery(entry, query); + continue; + } + result = future.get(remaining, TimeUnit.NANOSECONDS); + } else { + long elapsed = System.nanoTime() - query.startedAtNanos; + if (elapsed > TimeUnit.MILLISECONDS.toNanos( + this.queryTimeoutMillis)) { + expireQuery(entry, query); + } + continue; + } + if (result.completedAtNanos - query.startedAtNanos > + TimeUnit.MILLISECONDS.toNanos(this.queryTimeoutMillis)) { + expireQuery(entry, query); + continue; + } + this.resolvedByEntry.put( + entry, new ResolvedEntry(result.addresses, + System.nanoTime())); + this.inFlight.remove(entry); + if (this.failedEntries.remove(entry)) { + log.info("Raft peer address resolution recovered for '{}'", entry); } - } catch (UnknownHostException e) { - log.warn("Could not resolve allowlist entry '{}': {}", entry, e.getMessage()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + markResolutionFailure(entry, e); + throw new IllegalStateException( + "Raft peer address refresh interrupted", e); + } catch (ExecutionException e) { + this.inFlight.remove(entry); + markResolutionFailure(entry, e); + } catch (TimeoutException e) { + expireQuery(entry, query); + } catch (CancellationException e) { + this.inFlight.remove(entry); + markResolutionFailure(entry, e); + } + } + } + + private void expireQuery(String entry, Query query) { + query.cancel(); + this.inFlight.remove(entry); + markResolutionFailure( + entry, new TimeoutException("DNS refresh deadline")); + } + + private void markResolutionFailure(String entry, Exception failure) { + if (this.failedEntries.add(entry)) { + log.warn("Could not resolve Raft peer allowlist entry '{}': {}", + entry, failure.getMessage()); + } + } + + private void refreshSafely() { + try { + this.refreshResolvedIps(false); + } catch (RuntimeException e) { + log.error("Unexpected Raft peer allowlist refresh failure", e); + } + } + + private void replaceAllowedEntries(Set<String> entries) { + validateAllowedEntries(entries); + Set<String> copy = new HashSet<>(entries); + if (copy.equals(this.allowedEntries)) { + return; + } + this.allowedEntries = Collections.unmodifiableSet(copy); + this.resolutionOrder = new ArrayList<>(copy); + Collections.sort(this.resolutionOrder); + this.nextResolutionIndex = 0; + } + + public static void validateAllowedEntries(Set<String> entries) { + if (entries.size() > MAX_ALLOWED_ENTRIES) { + throw new IllegalArgumentException( + "Raft peer allowlist exceeds " + MAX_ALLOWED_ENTRIES + + " entries"); + } + for (String entry : entries) { + if (entry == null || entry.isEmpty() || + entry.length() > MAX_HOST_LENGTH) { + throw new IllegalArgumentException( + "Invalid Raft peer allowlist entry"); } } + } + + public static void validatePeerListShape(String peerList) { + if (peerList == null || peerList.isEmpty() || + peerList.length() > MAX_PEER_LIST_LENGTH) { + throw new IllegalArgumentException( + "Invalid Raft peer list length"); + } + int entries = 1; + for (int i = 0; i < peerList.length(); i++) { + if (peerList.charAt(i) == ',' && + ++entries > MAX_ALLOWED_ENTRIES) { + throw new IllegalArgumentException( + "Raft peer list exceeds " + MAX_ALLOWED_ENTRIES + + " entries"); + } + } + } + + synchronized void shutdown() { + if (this.closed) { + return; + } + this.closed = true; + if (this.refreshExecutor != null) { + this.refreshExecutor.shutdownNow(); + } + this.inFlight.values().forEach(Query::cancel); + this.inFlight.clear(); + this.resolver.close(); + } + + public static synchronized void shutdownInstance() { + if (instance != null) { + instance.shutdown(); + instance = null; + } + } + + @FunctionalInterface + interface HostResolver extends AutoCloseable { + + CompletableFuture<Set<String>> resolve(String host); + + @Override + default void close() { + // Most injected resolvers do not own resources. + } + } + + private static final class ResolvedEntry { + + private final Set<String> addresses; + private final long resolvedAtNanos; + + private ResolvedEntry(Set<String> addresses, + long resolvedAtNanos) { + this.addresses = addresses; + this.resolvedAtNanos = resolvedAtNanos; + } + } + + private static final class Query { + + private final CompletableFuture<Set<String>> source; + private final CompletableFuture<ResolvedQuery> future; + private final long startedAtNanos; - return Collections.unmodifiableSet(result); + private Query(CompletableFuture<Set<String>> source, + long startedAtNanos) { + this.source = source; + this.startedAtNanos = startedAtNanos; + this.future = source.thenApply( + addresses -> new ResolvedQuery(addresses, + System.nanoTime())); + } + + private void cancel() { + this.source.cancel(true); + this.future.cancel(true); + } + } + + private static final class ResolvedQuery { + + private final Set<String> addresses; + private final long completedAtNanos; + + private ResolvedQuery(Set<String> addresses, + long completedAtNanos) { + this.addresses = addresses; + this.completedAtNanos = completedAtNanos; + } + } + + private static final class NettyHostResolver implements HostResolver { + + private final NioEventLoopGroup eventLoopGroup; + private final DnsNameResolver resolver; + + private NettyHostResolver(long queryTimeoutMillis) { + this.eventLoopGroup = new NioEventLoopGroup(1, task -> { + Thread thread = new Thread(task, "pd-raft-dns-event-loop"); + thread.setDaemon(true); + return thread; + }); + try { + this.resolver = new DnsNameResolverBuilder( + this.eventLoopGroup.next()) + .channelType(NioDatagramChannel.class) + .ttl(0, 1) + .negativeTtl(0) + .queryTimeoutMillis(queryTimeoutMillis) + .build(); + } catch (RuntimeException | Error e) { + this.eventLoopGroup.shutdownGracefully( + 0L, DNS_QUERY_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS) + .awaitUninterruptibly( + DNS_QUERY_TIMEOUT_MILLIS); + throw e; + } + } + + @Override + public CompletableFuture<Set<String>> resolve(String host) { + io.netty.util.concurrent.Future<List<InetAddress>> query = + this.resolver.resolveAll(host); + CompletableFuture<Set<String>> result = new CompletableFuture<>(); + query.addListener(done -> { + if (!done.isSuccess()) { + result.completeExceptionally(done.cause()); + return; + } + Set<String> addresses = new HashSet<>(); + for (InetAddress address : query.getNow()) { + addresses.add(address.getHostAddress()); + } + result.complete(Collections.unmodifiableSet(addresses)); + }); + result.whenComplete((ignored, failure) -> { + if (result.isCancelled()) { + query.cancel(true); + } + }); + return result; + } + + @Override + public void close() { + this.resolver.close(); + this.eventLoopGroup.shutdownGracefully( + 0L, DNS_QUERY_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS) + .awaitUninterruptibly( + DNS_QUERY_TIMEOUT_MILLIS); + } } @Override diff --git a/hugegraph-pd/hg-pd-service/pom.xml b/hugegraph-pd/hg-pd-service/pom.xml index ee78863f3..7ffb9ccd6 100644 --- a/hugegraph-pd/hg-pd-service/pom.xml +++ b/hugegraph-pd/hg-pd-service/pom.xml @@ -162,6 +162,18 @@ <artifactId>log4j-jul</artifactId> <version>2.17.2</version> </dependency> + <dependency> + <groupId>junit</groupId> + <artifactId>junit</artifactId> + <version>${junit.version}</version> + <scope>test</scope> + </dependency> + <dependency> + <groupId>org.mockito</groupId> + <artifactId>mockito-core</artifactId> + <version>3.9.0</version> + <scope>test</scope> + </dependency> </dependencies> <build> <plugins> diff --git a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java index 94d136a84..b31be3bb1 100644 --- a/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java +++ b/hugegraph-pd/hg-pd-service/src/main/java/org/apache/hugegraph/pd/service/PDService.java @@ -27,8 +27,10 @@ import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import javax.annotation.PostConstruct; @@ -99,6 +101,7 @@ import com.alipay.sofa.jraft.JRaftUtils; import com.alipay.sofa.jraft.Status; import com.alipay.sofa.jraft.conf.Configuration; import com.alipay.sofa.jraft.entity.PeerId; +import com.alipay.sofa.jraft.error.RaftError; import io.grpc.ManagedChannel; import io.grpc.stub.StreamObserver; @@ -1683,7 +1686,20 @@ public class PDService extends PDGrpc.PDImplBase implements RaftStateListener { return; } - var list = PeerUtil.parseConfig(request.getConfig()); + List<KVPair<String, PeerId>> list; + try { + IpAuthHandler.validatePeerListShape(request.getConfig()); + list = PeerUtil.parseConfig(request.getConfig()); + } catch (IllegalArgumentException e) { + Pdpb.UpdatePdRaftResponse response = + Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6668, e.getMessage())) + .build(); + observer.onNext(response); + observer.onCompleted(); + return; + } log.info("update raft request: {}, list: {}", request.getConfig(), list); @@ -1732,28 +1748,93 @@ public class PDService extends PDGrpc.PDImplBase implements RaftStateListener { } } + Set<String> newIps = new HashSet<>(); + config.getPeers().forEach(peer -> newIps.add(peer.getIp())); + config.getLearners().forEach(peer -> newIps.add(peer.getIp())); + try { + IpAuthHandler.validateAllowedEntries(newIps); + IpAuthHandler.requireActiveInstance(); + } catch (IllegalArgumentException e) { + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6668, + e.getMessage())) + .build(); + break; + } catch (IllegalStateException e) { + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6670, + e.getMessage())) + .build(); + break; + } + log.info("pd raft update with new config: {}", config); - node.changePeers(config, status -> { - if (status.isOk()) { - log.info("updatePdRaft, change peers success"); - // Refresh IpAuthHandler so newly added peers are not blocked - IpAuthHandler handler = IpAuthHandler.getInstance(); - if (handler != null) { - Set<String> newIps = new HashSet<>(); - config.getPeers().forEach(p -> newIps.add(p.getIp())); - config.getLearners().forEach(p -> newIps.add(p.getIp())); - handler.refresh(newIps); - log.info("IpAuthHandler refreshed after updatePdRaft peer change"); - } else { - log.warn("IpAuthHandler not initialized, skipping refresh"); + CountDownLatch changeLatch = new CountDownLatch(1); + AtomicReference<Status> changeStatus = new AtomicReference<>(); + try { + node.changePeers(config, status -> { + Status callbackStatus = status; + try { + if (status != null && status.isOk()) { + log.info("updatePdRaft, change peers success"); + IpAuthHandler.refreshInstance(newIps); + log.info("IpAuthHandler refreshed after updatePdRaft peer change"); + } else if (status != null) { + log.error("changePeers status: {}, msg:{}, code: {}, raft error:{}", + status, status.getErrorMsg(), status.getCode(), + status.getRaftError()); + } else { + callbackStatus = new Status( + RaftError.EINTERNAL, + "changePeers returned no status"); + } + } catch (RuntimeException e) { + callbackStatus = new Status( + RaftError.EINTERNAL, + "Raft peers changed but allowlist refresh failed: %s", + e.getMessage()); + log.error("Raft peers changed but IpAuthHandler refresh failed", + e); + } finally { + changeStatus.set(callbackStatus); + changeLatch.countDown(); } - } else { - log.error("changePeers status: {}, msg:{}, code: {}, raft error:{}", - status, status.getErrorMsg(), status.getCode(), - status.getRaftError()); + }); + long timeout = 3L * pdConfig.getRaft().getRpcTimeout(); + if (!changeLatch.await(timeout, TimeUnit.MILLISECONDS)) { + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6669, + "changePeers timed out")) + .build(); + } else if (changeStatus.get() == null || + !changeStatus.get().isOk()) { + String message = changeStatus.get() == null ? + "changePeers returned no status" : + changeStatus.get().getErrorMsg(); + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6670, message)) + .build(); } - }); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6670, + "changePeers interrupted")) + .build(); + } catch (RuntimeException e) { + log.error("changePeers failed before callback", e); + response = Pdpb.UpdatePdRaftResponse.newBuilder() + .setHeader(newErrorHeader( + 6670, + e.getMessage())) + .build(); + } } while (false); observer.onNext(response); diff --git a/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/service/PDServiceUpdateRaftTest.java b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/service/PDServiceUpdateRaftTest.java new file mode 100644 index 000000000..d7ee1401c --- /dev/null +++ b/hugegraph-pd/hg-pd-service/src/test/java/org/apache/hugegraph/pd/service/PDServiceUpdateRaftTest.java @@ -0,0 +1,195 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.service; + +import java.util.Collections; + +import org.apache.hugegraph.pd.config.PDConfig; +import org.apache.hugegraph.pd.grpc.Pdpb; +import org.apache.hugegraph.pd.raft.RaftEngine; +import org.apache.hugegraph.pd.raft.auth.IpAuthHandler; +import org.apache.hugegraph.testutil.Whitebox; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.ArgumentCaptor; + +import com.alipay.sofa.jraft.Closure; +import com.alipay.sofa.jraft.Node; +import com.alipay.sofa.jraft.Status; +import com.alipay.sofa.jraft.conf.Configuration; +import com.alipay.sofa.jraft.entity.PeerId; +import com.alipay.sofa.jraft.error.RaftError; + +import io.grpc.stub.StreamObserver; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class PDServiceUpdateRaftTest { + + private Node originalRaftNode; + private Node mockNode; + private PDService service; + private PeerId leader; + + @Before + public void setUp() { + this.originalRaftNode = RaftEngine.getInstance().getRaftNode(); + IpAuthHandler.shutdownInstance(); + + this.leader = new PeerId(); + Assert.assertTrue(this.leader.parse("127.0.0.1:8610")); + this.mockNode = mock(Node.class); + when(this.mockNode.isLeader(true)).thenReturn(true); + when(this.mockNode.getLeaderId()).thenReturn(this.leader); + when(this.mockNode.listPeers()).thenReturn( + Collections.singletonList(this.leader)); + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", + this.mockNode); + IpAuthHandler.getInstance(Collections.singleton("127.0.0.1")); + + PDConfig pdConfig = new PDConfig(); + PDConfig.Raft raft = pdConfig.new Raft(); + raft.setRpcTimeout(1); + pdConfig.setRaft(raft); + this.service = new PDService(); + this.service.setInitConfig(pdConfig); + } + + @After + public void tearDown() { + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", + this.originalRaftNode); + IpAuthHandler.shutdownInstance(); + } + + @Test + public void testRejectsMalformedConfigBeforeRaft() { + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader,bad,127.0.0.2:8610/follower"); + + Assert.assertEquals(6668, response.getHeader().getError().getTypeValue()); + Assert.assertTrue(response.getHeader().getError().getMessage() + .contains("Invalid Raft peer")); + verify(this.mockNode, never()).changePeers( + any(Configuration.class), any(Closure.class)); + } + + @Test + public void testReturnsSuccessAfterRaftCallbackAndAllowlistRefresh() + throws Exception { + IpAuthHandler handler = IpAuthHandler.requireActiveInstance(); + handler.refresh(Collections.singleton("10.0.0.1")); + doAnswer(invocation -> { + Closure closure = invocation.getArgument(1); + closure.run(Status.OK()); + return null; + }).when(this.mockNode).changePeers(any(Configuration.class), + any(Closure.class)); + + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader"); + + Assert.assertEquals(Pdpb.ErrorType.OK, + response.getHeader().getError().getType()); + Assert.assertTrue(isIpAllowed(handler, "127.0.0.1")); + Assert.assertFalse(isIpAllowed(handler, "10.0.0.1")); + } + + @Test + public void testReturnsRaftFailureFromCallback() { + doAnswer(invocation -> { + Closure closure = invocation.getArgument(1); + closure.run(new Status(RaftError.EINTERNAL, "simulated failure")); + return null; + }).when(this.mockNode).changePeers(any(Configuration.class), + any(Closure.class)); + + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader"); + + Assert.assertEquals(6670, response.getHeader().getError().getTypeValue()); + Assert.assertTrue(response.getHeader().getError().getMessage() + .contains("simulated failure")); + } + + @Test + public void testReturnsTimeoutWhenRaftDoesNotCallback() { + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader"); + + Assert.assertEquals(6669, response.getHeader().getError().getTypeValue()); + Assert.assertTrue(response.getHeader().getError().getMessage() + .contains("timed out")); + } + + @Test + public void testRejectsMissingAllowlistBeforeRaft() { + IpAuthHandler.shutdownInstance(); + + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader"); + + Assert.assertEquals(6670, response.getHeader().getError().getTypeValue()); + Assert.assertTrue(response.getHeader().getError().getMessage() + .contains("not active")); + verify(this.mockNode, never()).changePeers( + any(Configuration.class), any(Closure.class)); + } + + @Test + public void testMapsSynchronousRaftFailure() { + doThrow(new IllegalStateException("node stopped")) + .when(this.mockNode) + .changePeers(any(Configuration.class), any(Closure.class)); + + Pdpb.UpdatePdRaftResponse response = update( + "127.0.0.1:8610/leader"); + + Assert.assertEquals(6670, response.getHeader().getError().getTypeValue()); + Assert.assertTrue(response.getHeader().getError().getMessage() + .contains("node stopped")); + } + + @SuppressWarnings("unchecked") + private Pdpb.UpdatePdRaftResponse update(String config) { + StreamObserver<Pdpb.UpdatePdRaftResponse> observer = + mock(StreamObserver.class); + this.service.updatePdRaft( + Pdpb.UpdatePdRaftRequest.newBuilder().setConfig(config).build(), + observer); + ArgumentCaptor<Pdpb.UpdatePdRaftResponse> response = + ArgumentCaptor.forClass(Pdpb.UpdatePdRaftResponse.class); + verify(observer).onNext(response.capture()); + verify(observer).onCompleted(); + return response.getValue(); + } + + private boolean isIpAllowed(IpAuthHandler handler, String ip) { + return Whitebox.invoke(IpAuthHandler.class, + new Class[]{String.class}, + "isIpAllowed", handler, ip); + } +} diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java index 95b044c76..613d08559 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/core/PDCoreSuiteTest.java @@ -19,7 +19,7 @@ package org.apache.hugegraph.pd.core; import org.apache.hugegraph.pd.core.meta.MetadataKeyHelperTest; import org.apache.hugegraph.pd.core.store.HgKVStoreImplTest; -import org.apache.hugegraph.pd.raft.IpAuthHandlerTest; +import org.apache.hugegraph.pd.raft.auth.IpAuthHandlerTest; import org.apache.hugegraph.pd.raft.RaftEngineIpAuthIntegrationTest; import org.apache.hugegraph.pd.raft.RaftEngineLeaderAddressTest; import org.junit.runner.RunWith; diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/IpAuthHandlerTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/IpAuthHandlerTest.java deleted file mode 100644 index 31647b6d3..000000000 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/IpAuthHandlerTest.java +++ /dev/null @@ -1,133 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.hugegraph.pd.raft; - -import java.net.InetAddress; -import java.util.Collections; -import java.util.HashSet; -import java.util.Set; - -import org.apache.hugegraph.pd.raft.auth.IpAuthHandler; -import org.apache.hugegraph.testutil.Whitebox; -import org.junit.After; -import org.junit.Assert; -import org.junit.Before; -import org.junit.Test; - -public class IpAuthHandlerTest { - - @Before - public void setUp() { - // Must reset BEFORE each test — earlier suite classes (e.g. ConfigServiceTest) - // initialize RaftEngine which creates the IpAuthHandler singleton with their - // own peer IPs. Without this reset, our getInstance() calls return the stale - // singleton and ignore the allowlist passed by the test. - Whitebox.setInternalState(IpAuthHandler.class, "instance", null); - } - - @After - public void tearDown() { - // Must reset AFTER each test — prevents our test singleton from leaking - // into later suite classes that also depend on IpAuthHandler state. - Whitebox.setInternalState(IpAuthHandler.class, "instance", null); - } - - private boolean isIpAllowed(IpAuthHandler handler, String ip) { - return Whitebox.invoke(IpAuthHandler.class, - new Class[]{String.class}, - "isIpAllowed", handler, ip); - } - - @Test - public void testHostnameResolvesToIp() throws Exception { - // "localhost" should resolve to one or more IPs via InetAddress.getAllByName() - // This verifies the core fix: hostname allowlists match numeric remote addresses - // Using dynamic resolution avoids hardcoding "127.0.0.1" which may not be - // returned on IPv6-only or custom resolver environments - IpAuthHandler handler = IpAuthHandler.getInstance( - Collections.singleton("localhost")); - InetAddress[] addresses = InetAddress.getAllByName("localhost"); - // All resolved addresses should be allowed — resolveAll() adds every address - // returned by getAllByName() so none should be blocked - Assert.assertTrue("Expected at least one resolved address", - addresses.length > 0); - for (InetAddress address : addresses) { - Assert.assertTrue( - "Expected " + address.getHostAddress() + " to be allowed", - isIpAllowed(handler, address.getHostAddress())); - } - } - - @Test - public void testUnresolvableHostnameDoesNotCrash() { - // Should log a warning and skip — no exception thrown during construction - // Uses .invalid TLD which is RFC-2606 reserved and guaranteed to never resolve - IpAuthHandler handler = IpAuthHandler.getInstance( - Collections.singleton("nonexistent.invalid")); - // Handler was still created successfully despite bad hostname - Assert.assertNotNull(handler); - // Unresolvable entry is skipped so no IPs should be allowed - Assert.assertFalse(isIpAllowed(handler, "127.0.0.1")); - Assert.assertFalse(isIpAllowed(handler, "192.168.0.1")); - } - - @Test - public void testRefreshUpdatesResolvedIps() { - // Start with 127.0.0.1 - IpAuthHandler handler = IpAuthHandler.getInstance( - Collections.singleton("127.0.0.1")); - Assert.assertTrue(isIpAllowed(handler, "127.0.0.1")); - - // Refresh with a different IP — verifies refresh() swaps the set correctly - Set<String> newIps = new HashSet<>(); - newIps.add("192.168.0.1"); - handler.refresh(newIps); - - // Old IP should no longer be allowed - Assert.assertFalse(isIpAllowed(handler, "127.0.0.1")); - // New IP should now be allowed - Assert.assertTrue(isIpAllowed(handler, "192.168.0.1")); - } - - @Test - public void testEmptyAllowlistAllowsAll() { - // Empty allowlist = no restriction configured = allow all connections - // This is intentional fallback behavior and must be explicitly tested - // because it is a security-relevant boundary - IpAuthHandler handler = IpAuthHandler.getInstance( - Collections.emptySet()); - Assert.assertTrue(isIpAllowed(handler, "1.2.3.4")); - Assert.assertTrue(isIpAllowed(handler, "192.168.99.99")); - } - - @Test - public void testGetInstanceReturnsSingletonIgnoresNewAllowlist() { - // First call creates the singleton with 127.0.0.1 - IpAuthHandler first = IpAuthHandler.getInstance( - Collections.singleton("127.0.0.1")); - // Second call with a different set must return the same instance - // and must NOT reinitialize or override the existing allowlist - IpAuthHandler second = IpAuthHandler.getInstance( - Collections.singleton("192.168.0.1")); - Assert.assertSame(first, second); - // Original allowlist still in effect - Assert.assertTrue(isIpAllowed(second, "127.0.0.1")); - // New set was ignored — 192.168.0.1 should not be allowed - Assert.assertFalse(isIpAllowed(second, "192.168.0.1")); - } -} diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/RaftEngineIpAuthIntegrationTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/RaftEngineIpAuthIntegrationTest.java index 1f9857df0..1aa292174 100644 --- a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/RaftEngineIpAuthIntegrationTest.java +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/RaftEngineIpAuthIntegrationTest.java @@ -19,6 +19,7 @@ package org.apache.hugegraph.pd.raft; import java.util.Collections; +import org.apache.hugegraph.pd.config.PDConfig; import org.apache.hugegraph.pd.raft.auth.IpAuthHandler; import org.apache.hugegraph.testutil.Whitebox; import org.junit.After; @@ -35,25 +36,35 @@ import com.alipay.sofa.jraft.error.RaftError; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; public class RaftEngineIpAuthIntegrationTest { private Node originalRaftNode; + private PDConfig.Raft originalConfig; @Before public void setUp() { // Save original raftNode so we can restore it after the test originalRaftNode = RaftEngine.getInstance().getRaftNode(); + originalConfig = Whitebox.getInternalState(RaftEngine.getInstance(), + "config"); + PDConfig pdConfig = new PDConfig(); + PDConfig.Raft config = pdConfig.new Raft(); + config.setRpcTimeout(100); + Whitebox.setInternalState(RaftEngine.getInstance(), "config", config); // Reset IpAuthHandler singleton for a clean state - Whitebox.setInternalState(IpAuthHandler.class, "instance", null); + IpAuthHandler.shutdownInstance(); } @After public void tearDown() { // Restore original raftNode Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", originalRaftNode); + Whitebox.setInternalState(RaftEngine.getInstance(), "config", originalConfig); // Reset IpAuthHandler singleton - Whitebox.setInternalState(IpAuthHandler.class, "instance", null); + IpAuthHandler.shutdownInstance(); } @Test @@ -80,9 +91,11 @@ public class RaftEngineIpAuthIntegrationTest { Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", mockNode); // Call changePeerList with new peer — must be odd count - RaftEngine.getInstance().changePeerList("127.0.0.1:8610"); + Status status = RaftEngine.getInstance().changePeerList( + "127.0.0.1:8610"); // Verify IpAuthHandler was refreshed with the new peer IP + Assert.assertTrue(status.isOk()); Assert.assertTrue(invokeIsIpAllowed(handler, "127.0.0.1")); // Old IP should no longer be allowed Assert.assertFalse(invokeIsIpAllowed(handler, "10.0.0.1")); @@ -109,13 +122,73 @@ public class RaftEngineIpAuthIntegrationTest { Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", mockNode); - RaftEngine.getInstance().changePeerList("127.0.0.1:8610"); + Status status = RaftEngine.getInstance().changePeerList( + "127.0.0.1:8610"); // Handler should NOT be refreshed — old IP still allowed + Assert.assertFalse(status.isOk()); Assert.assertTrue(invokeIsIpAllowed(handler, "10.0.0.1")); Assert.assertFalse(invokeIsIpAllowed(handler, "127.0.0.1")); } + @Test + public void testChangePeerListRejectsNullCallbackStatus() { + IpAuthHandler.getInstance(Collections.singleton("10.0.0.1")); + Node mockNode = mock(Node.class); + doAnswer(invocation -> { + Closure closure = invocation.getArgument(1); + closure.run(null); + return null; + }).when(mockNode).changePeers(any(Configuration.class), + any(Closure.class)); + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", + mockNode); + + Status status = RaftEngine.getInstance().changePeerList( + "127.0.0.1:8610"); + + Assert.assertNotNull(status); + Assert.assertFalse(status.isOk()); + Assert.assertTrue(status.getErrorMsg() + .contains("returned no status")); + } + + @Test + public void testChangePeerListRejectsOversizedAllowlistBeforeRaft() { + Node mockNode = mock(Node.class); + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", mockNode); + StringBuilder peers = new StringBuilder(); + for (int i = 0; i < 129; i++) { + if (i > 0) { + peers.append(','); + } + peers.append("pd-").append(i).append(":8610"); + } + + Status status = RaftEngine.getInstance().changePeerList( + peers.toString()); + + Assert.assertNotNull(status); + Assert.assertFalse(status.isOk()); + verify(mockNode, never()).changePeers( + any(Configuration.class), any(Closure.class)); + } + + @Test + public void testChangePeerListRejectsMalformedPeerBeforeRaft() { + Node mockNode = mock(Node.class); + Whitebox.setInternalState(RaftEngine.getInstance(), "raftNode", mockNode); + + Status status = RaftEngine.getInstance().changePeerList( + "127.0.0.1:8610,bad,127.0.0.2:8610"); + + Assert.assertNotNull(status); + Assert.assertFalse(status.isOk()); + Assert.assertTrue(status.getErrorMsg().contains("Invalid Raft peer")); + verify(mockNode, never()).changePeers( + any(Configuration.class), any(Closure.class)); + } + private boolean invokeIsIpAllowed(IpAuthHandler handler, String ip) { return Whitebox.invoke(IpAuthHandler.class, new Class[]{String.class}, diff --git a/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandlerTest.java b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandlerTest.java new file mode 100644 index 000000000..833d1eeaa --- /dev/null +++ b/hugegraph-pd/hg-pd-test/src/main/java/org/apache/hugegraph/pd/raft/auth/IpAuthHandlerTest.java @@ -0,0 +1,439 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hugegraph.pd.raft.auth; + +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.hugegraph.testutil.Whitebox; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class IpAuthHandlerTest { + + @Before + public void setUp() { + // Must reset BEFORE each test — earlier suite classes (e.g. ConfigServiceTest) + // initialize RaftEngine which creates the IpAuthHandler singleton with their + // own peer IPs. Without this reset, our getInstance() calls return the stale + // singleton and ignore the allowlist passed by the test. + IpAuthHandler.shutdownInstance(); + } + + @After + public void tearDown() { + // Must reset AFTER each test — prevents our test singleton from leaking + // into later suite classes that also depend on IpAuthHandler state. + IpAuthHandler handler = IpAuthHandler.getInstance(); + if (handler != null) { + IpAuthHandler.shutdownInstance(); + } + } + + private boolean isIpAllowed(IpAuthHandler handler, String ip) { + return Whitebox.invoke(IpAuthHandler.class, + new Class[]{String.class}, + "isIpAllowed", handler, ip); + } + + @Test + public void testHostnameResolvesToIp() throws Exception { + // "localhost" should resolve to one or more IPs via InetAddress.getAllByName() + // This verifies the core fix: hostname allowlists match numeric remote addresses + // Using dynamic resolution avoids hardcoding "127.0.0.1" which may not be + // returned on IPv6-only or custom resolver environments + IpAuthHandler handler = IpAuthHandler.getInstance( + Collections.singleton("localhost")); + InetAddress[] addresses = InetAddress.getAllByName("localhost"); + Assert.assertTrue("Expected at least one resolved address", + addresses.length > 0); + boolean matched = false; + for (InetAddress address : addresses) { + matched |= isIpAllowed(handler, address.getHostAddress()); + } + Assert.assertTrue("Expected a resolved address to be allowed", matched); + } + + @Test + public void testTransientDnsFailureRecoversOnRefresh() throws Exception { + AtomicInteger attempts = new AtomicInteger(); + InetAddress expected = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 1}); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("pd-1"), + host -> { + if (attempts.incrementAndGet() < 3) { + return failed(host); + } + return resolved(expected); + }, + false, 100L, 1_000L, 1_000L); + + Assert.assertFalse(isIpAllowed(handler, expected.getHostAddress())); + handler.refreshResolvedIps(); + handler.refreshResolvedIps(); + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + Assert.assertEquals(3, attempts.get()); + handler.shutdown(); + } + + @Test + public void testTransientDnsFailureKeepsLastKnownAddress() throws Exception { + AtomicInteger attempts = new AtomicInteger(); + InetAddress expected = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 1}); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("pd-1"), + host -> { + if (attempts.incrementAndGet() > 1) { + return failed(host); + } + return resolved(expected); + }, + false, 100L, 1_000L, 1_000L); + + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + handler.refreshResolvedIps(); + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + handler.shutdown(); + } + + @Test + public void testSlowPeerDoesNotBlockFollowingPeer() throws Exception { + InetAddress expected = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 2}); + Set<String> peers = new LinkedHashSet<>(); + peers.add("pd-slow"); + peers.add("pd-ready"); + IpAuthHandler handler = new IpAuthHandler( + peers, + host -> { + if ("pd-slow".equals(host)) { + return new CompletableFuture<>(); + } + return resolved(expected); + }, + false, 10L, 1_000L, 1_000L); + + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + handler.shutdown(); + } + + @Test + public void testExpiredAddressFailsClosed() throws Exception { + AtomicInteger attempts = new AtomicInteger(); + InetAddress expected = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 3}); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("pd-1"), + host -> { + if (attempts.incrementAndGet() > 1) { + return failed(host); + } + return resolved(expected); + }, + false, 100L, 1L, 1_000L); + + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + Thread.sleep(5L); + handler.refreshResolvedIps(); + Assert.assertFalse(isIpAllowed(handler, expected.getHostAddress())); + handler.shutdown(); + } + + @Test + public void testScheduledRefreshAddsLatePeerAndRotatesAddress() + throws Exception { + AtomicInteger attempts = new AtomicInteger(); + InetAddress first = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 4}); + InetAddress second = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 5}); + AtomicReference<InetAddress> current = new AtomicReference<>(first); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("pd-late"), + host -> { + if (attempts.incrementAndGet() == 1) { + return failed(host); + } + return resolved(current.get()); + }, + true, 20L, 1_000L, 10L); + try { + awaitAllowed(handler, first.getHostAddress()); + current.set(second); + awaitAllowed(handler, second.getHostAddress()); + Assert.assertFalse(isIpAllowed(handler, first.getHostAddress())); + } finally { + handler.shutdown(); + } + } + + @Test + public void testNeverCompletingPeersDoNotStarveReadyPeer() + throws Exception { + InetAddress expected = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 6}); + Set<String> peers = new LinkedHashSet<>(); + for (int i = 0; i < 8; i++) { + peers.add("00-pd-slow-" + i); + } + peers.add("99-pd-ready"); + IpAuthHandler handler = new IpAuthHandler( + peers, + host -> { + if (host.startsWith("00-pd-slow-")) { + return new CompletableFuture<>(); + } + return resolved(expected); + }, + false, 10L, 1_000L, 1_000L); + + handler.refresh(peers); + Assert.assertTrue(isIpAllowed(handler, expected.getHostAddress())); + handler.shutdown(); + } + + @Test + public void testRejectsOversizedAllowlist() { + Set<String> peers = new HashSet<>(); + for (int i = 0; i < 128; i++) { + peers.add("pd-" + i); + } + + try { + new IpAuthHandler(peers, host -> new CompletableFuture<>(), + false, 10L, 1_000L, 1_000L); + Assert.fail("Expected oversized allowlist rejection"); + } catch (IllegalArgumentException e) { + Assert.assertTrue(e.getMessage().contains("127")); + } + } + + @Test + public void testLateSuccessfulResultIsDiscarded() throws Exception { + AtomicInteger attempts = new AtomicInteger(); + InetAddress first = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 7}); + InetAddress late = InetAddress.getByAddress( + new byte[]{(byte) 192, (byte) 168, 0, 8}); + CompletableFuture<Set<String>> delayed = new CompletableFuture<>(); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("pd-1"), + host -> { + int attempt = attempts.incrementAndGet(); + if (attempt == 1) { + return resolved(first); + } + if (attempt == 2) { + return delayed; + } + return new CompletableFuture<>(); + }, + false, 10L, 1_000L, 1_000L); + + handler.refreshResolvedIps(false); + Thread.sleep(20L); + delayed.complete(resolved(late).get()); + handler.refreshResolvedIps(false); + + Assert.assertTrue(isIpAllowed(handler, first.getHostAddress())); + Assert.assertFalse(isIpAllowed(handler, late.getHostAddress())); + handler.shutdown(); + } + + @Test + public void testRefreshCollectsPreviousBatchBeforeStartingNext() + throws Exception { + Set<String> peers = new HashSet<>(); + Map<Integer, CompletableFuture<Set<String>>> delayed = + new HashMap<>(); + for (int i = 0; i < 17; i++) { + peers.add(String.format("pd-%02d", i)); + if (i >= 8 && i < 16) { + delayed.put(i, new CompletableFuture<>()); + } + } + IpAuthHandler handler = new IpAuthHandler( + peers, + host -> { + int index = Integer.parseInt(host.substring(3)); + CompletableFuture<Set<String>> future = delayed.get(index); + if (future != null) { + return future; + } + return resolved(address(index)); + }, + false, 100L, 1_000L, 1_000L); + + handler.refreshResolvedIps(false); + for (Map.Entry<Integer, CompletableFuture<Set<String>>> entry : + delayed.entrySet()) { + entry.getValue().complete(resolved(address(entry.getKey())).get()); + } + handler.refreshResolvedIps(false); + + Assert.assertTrue(isIpAllowed( + handler, address(16).getHostAddress())); + handler.shutdown(); + } + + @Test + public void testConstructorFailureClosesResolver() { + AtomicBoolean closed = new AtomicBoolean(); + IpAuthHandler.HostResolver resolver = new IpAuthHandler.HostResolver() { + + @Override + public CompletableFuture<Set<String>> resolve(String host) { + throw new IllegalStateException("simulated resolver failure"); + } + + @Override + public void close() { + closed.set(true); + } + }; + + try { + new IpAuthHandler(Collections.singleton("pd-1"), resolver, + false, 10L, 1_000L, 1_000L); + Assert.fail("Expected constructor failure"); + } catch (IllegalStateException e) { + Assert.assertEquals("simulated resolver failure", e.getMessage()); + } + Assert.assertTrue(closed.get()); + } + + @Test + public void testInterruptedRefreshFailsAndPreservesInterrupt() + throws Exception { + InetAddress initial = address(20); + IpAuthHandler handler = new IpAuthHandler( + Collections.singleton("ready"), + host -> { + if ("ready".equals(host)) { + return resolved(initial); + } + return new CompletableFuture<>(); + }, + false, 100L, 1_000L, 1_000L); + try { + Thread.currentThread().interrupt(); + handler.refresh(Collections.singleton("slow")); + Assert.fail("Expected interrupted refresh to fail"); + } catch (IllegalStateException e) { + Assert.assertTrue(e.getMessage().contains("interrupted")); + Assert.assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + handler.shutdown(); + } + } + + private void awaitAllowed(IpAuthHandler handler, String address) + throws InterruptedException { + long deadline = System.currentTimeMillis() + 1_000L; + while (!isIpAllowed(handler, address) && + System.currentTimeMillis() < deadline) { + Thread.sleep(10L); + } + Assert.assertTrue(isIpAllowed(handler, address)); + } + + @Test + public void testRefreshUpdatesResolvedIps() { + // Start with 127.0.0.1 + IpAuthHandler handler = IpAuthHandler.getInstance( + Collections.singleton("127.0.0.1")); + Assert.assertTrue(isIpAllowed(handler, "127.0.0.1")); + + // Refresh with a different IP — verifies refresh() swaps the set correctly + Set<String> newIps = new HashSet<>(); + newIps.add("192.168.0.1"); + handler.refresh(newIps); + + // Old IP should no longer be allowed + Assert.assertFalse(isIpAllowed(handler, "127.0.0.1")); + // New IP should now be allowed + Assert.assertTrue(isIpAllowed(handler, "192.168.0.1")); + } + + @Test + public void testEmptyAllowlistAllowsAll() { + // Empty allowlist = no restriction configured = allow all connections + // This is intentional fallback behavior and must be explicitly tested + // because it is a security-relevant boundary + IpAuthHandler handler = IpAuthHandler.getInstance( + Collections.emptySet()); + Assert.assertTrue(isIpAllowed(handler, "1.2.3.4")); + Assert.assertTrue(isIpAllowed(handler, "192.168.99.99")); + } + + @Test + public void testGetInstanceReturnsSingletonIgnoresNewAllowlist() { + // First call creates the singleton with 127.0.0.1 + IpAuthHandler first = IpAuthHandler.getInstance( + Collections.singleton("127.0.0.1")); + // Second call with a different set must return the same instance + // and must NOT reinitialize or override the existing allowlist + IpAuthHandler second = IpAuthHandler.getInstance( + Collections.singleton("192.168.0.1")); + Assert.assertSame(first, second); + // Original allowlist still in effect + Assert.assertTrue(isIpAllowed(second, "127.0.0.1")); + // New set was ignored — 192.168.0.1 should not be allowed + Assert.assertFalse(isIpAllowed(second, "192.168.0.1")); + } + + private static CompletableFuture<Set<String>> resolved( + InetAddress... addresses) { + Set<String> result = new HashSet<>(); + for (InetAddress address : addresses) { + result.add(address.getHostAddress()); + } + return CompletableFuture.completedFuture( + Collections.unmodifiableSet(result)); + } + + private static CompletableFuture<Set<String>> failed(String host) { + CompletableFuture<Set<String>> result = new CompletableFuture<>(); + result.completeExceptionally(new UnknownHostException(host)); + return result; + } + + private static InetAddress address(int suffix) { + try { + return InetAddress.getByAddress( + new byte[]{10, 0, 0, (byte) (suffix + 1)}); + } catch (UnknownHostException e) { + throw new AssertionError(e); + } + } +}
