This is an automated email from the ASF dual-hosted git repository.
yuluo-yx pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
The following commit(s) were added to refs/heads/master by this push:
new e57c728364 feat: fix ssh connection bug (#4198)
e57c728364 is described below
commit e57c728364545f069c5cb8085ce82f8ac5b7aa02
Author: shown <[email protected]>
AuthorDate: Sun Jul 12 16:56:05 2026 +0800
feat: fix ssh connection bug (#4198)
---
.../collector/collect/ssh/SshCollectImpl.java | 62 ++-
.../collector/collect/ssh/SshCollectImplTest.java | 167 +++++++
.../collect/common/ssh/CommonSshClient.java | 15 +-
.../collect/common/ssh/SshCircuitBreaker.java | 102 ++++
.../collector/collect/common/ssh/SshHelper.java | 483 +++++++++++++-----
.../collect/common/ssh/CommonSshClientTest.java | 22 +-
.../collect/common/ssh/SshCircuitBreakerTest.java | 111 +++++
.../collect/common/ssh/SshHelperTest.java | 551 +++++++++++++++++++++
8 files changed, 1344 insertions(+), 169 deletions(-)
diff --git
a/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImpl.java
b/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImpl.java
index 870d912558..41b2987446 100644
---
a/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImpl.java
+++
b/hertzbeat-collector/hertzbeat-collector-basic/src/main/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImpl.java
@@ -19,6 +19,7 @@ package org.apache.hertzbeat.collector.collect.ssh;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
+import java.io.InterruptedIOException;
import java.net.ConnectException;
import java.net.SocketTimeoutException;
import java.security.GeneralSecurityException;
@@ -32,8 +33,6 @@ import java.util.Objects;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.hertzbeat.collector.collect.AbstractCollect;
-import org.apache.hertzbeat.collector.collect.common.cache.CacheIdentifier;
-import
org.apache.hertzbeat.collector.collect.common.cache.GlobalConnectionCache;
import org.apache.hertzbeat.collector.collect.common.ssh.CommonSshBlacklist;
import org.apache.hertzbeat.collector.collect.common.ssh.SshHelper;
import org.apache.hertzbeat.collector.constants.CollectorConstants;
@@ -49,6 +48,7 @@ import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.SshException;
import org.apache.sshd.common.channel.exception.SshChannelOpenException;
+import org.apache.sshd.common.future.CloseFuture;
import org.apache.sshd.common.util.io.output.NoCloseOutputStream;
import org.springframework.util.StringUtils;
@@ -64,7 +64,6 @@ public class SshCollectImpl extends AbstractCollect {
private static final String PARSE_TYPE_LOG = "log";
private static final int DEFAULT_TIMEOUT = 10_000;
- private final GlobalConnectionCache connectionCommonCache =
GlobalConnectionCache.getInstance();
@Override
public void preCheck(Metrics metrics) throws IllegalArgumentException {
@@ -79,7 +78,8 @@ public class SshCollectImpl extends AbstractCollect {
long startTime = System.currentTimeMillis();
SshProtocol sshProtocol = metrics.getSsh();
boolean reuseConnection =
Boolean.parseBoolean(sshProtocol.getReuseConnection());
- boolean useProxy = Boolean.parseBoolean(sshProtocol.getUseProxy());
+ boolean useProxy = Boolean.parseBoolean(sshProtocol.getUseProxy())
+ && StringUtils.hasText(sshProtocol.getProxyHost());
int timeout = CollectUtil.getTimeout(sshProtocol.getTimeout(),
DEFAULT_TIMEOUT);
ClientChannel channel = null;
ClientSession clientSession = null;
@@ -147,17 +147,14 @@ public class SshCollectImpl extends AbstractCollect {
builder.setCode(CollectRep.Code.FAIL);
builder.setMsg(errorMsg);
} finally {
- if (channel != null && channel.isOpen()) {
- try {
- // Close the SSH channel with the 'false' parameter to
ensure the session is not kept alive.
- long st = System.currentTimeMillis();
- channel.close(false).addListener(future ->
- log.debug("channel is closed in {} ms",
System.currentTimeMillis() - st));
- } catch (Exception e) {
- log.error(e.getMessage(), e);
- }
+ boolean channelClosed = true;
+ try {
+ channelClosed = closeChannel(channel, timeout);
+ } catch (Exception e) {
+ channelClosed = false;
+ log.error("Failed to close SSH channel", e);
}
- if (clientSession != null && !reuseConnection && !useProxy) {
+ if (clientSession != null && (!channelClosed || (!reuseConnection
&& !useProxy))) {
try {
clientSession.close();
} catch (Exception e) {
@@ -167,6 +164,35 @@ public class SshCollectImpl extends AbstractCollect {
}
}
+ static boolean closeChannel(ClientChannel channel, int timeout) throws
IOException {
+ if (channel == null || channel.isClosed()) {
+ return true;
+ }
+ long startTime = System.currentTimeMillis();
+ try {
+ CloseFuture closeFuture = channel.close(false);
+ if (!closeFuture.await(timeout)) {
+ log.warn("SSH channel graceful close timed out after {} ms,
forcing local cleanup", timeout);
+ closeFuture = channel.close(true);
+ if (!closeFuture.await(timeout)) {
+ log.warn("SSH channel immediate close timed out after {}
ms", timeout);
+ return false;
+ }
+ }
+ } catch (InterruptedIOException e) {
+ try {
+ channel.close(true);
+ } catch (RuntimeException closeException) {
+ e.addSuppressed(closeException);
+ } finally {
+ Thread.currentThread().interrupt();
+ }
+ throw e;
+ }
+ log.debug("SSH channel closed in {} ms", System.currentTimeMillis() -
startTime);
+ return true;
+ }
+
@Override
public String supportProtocol() {
return DispatchConstants.PROTOCOL_SSH;
@@ -273,14 +299,6 @@ public class SshCollectImpl extends AbstractCollect {
}
}
- private void removeConnectSessionCache(SshProtocol sshProtocol) {
- CacheIdentifier identifier = CacheIdentifier.builder()
- .ip(sshProtocol.getHost()).port(sshProtocol.getPort())
-
.username(sshProtocol.getUsername()).password(sshProtocol.getPassword())
- .build();
- connectionCommonCache.removeCache(identifier);
- }
-
private ClientSession getConnectSession(SshProtocol sshProtocol, int
timeout, boolean reuseConnection, boolean useProxy)
throws IOException, GeneralSecurityException {
return SshHelper.getConnectSession(sshProtocol, timeout,
reuseConnection, useProxy);
diff --git
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImplTest.java
b/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImplTest.java
index 466d1b9c6f..2810bb354d 100644
---
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImplTest.java
+++
b/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/ssh/SshCollectImplTest.java
@@ -19,14 +19,32 @@ package org.apache.hertzbeat.collector.collect.ssh;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import java.io.IOException;
+import java.io.InterruptedIOException;
+import java.util.List;
+import org.apache.hertzbeat.collector.collect.common.ssh.SshHelper;
import org.apache.hertzbeat.collector.dispatch.DispatchConstants;
import org.apache.hertzbeat.common.entity.job.Metrics;
import org.apache.hertzbeat.common.entity.job.protocol.SshProtocol;
import org.apache.hertzbeat.common.entity.message.CollectRep;
+import org.apache.sshd.client.channel.ChannelExec;
+import org.apache.sshd.client.channel.ClientChannel;
+import org.apache.sshd.client.future.OpenFuture;
+import org.apache.sshd.client.session.ClientSession;
+import org.apache.sshd.common.SshException;
+import org.apache.sshd.common.future.CloseFuture;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
/**
* Test case for {@link SshCollectImpl}
@@ -74,4 +92,153 @@ class SshCollectImplTest {
void supportProtocol() {
assertEquals(DispatchConstants.PROTOCOL_SSH,
sshCollect.supportProtocol());
}
+
+ @Test
+ void closeChannelWaitsForGracefulClose() throws IOException {
+ ClientChannel channel = mock(ClientChannel.class);
+ CloseFuture closeFuture = mock(CloseFuture.class);
+ when(channel.close(false)).thenReturn(closeFuture);
+ when(closeFuture.await(1_000)).thenReturn(true);
+
+ assertTrue(SshCollectImpl.closeChannel(channel, 1_000));
+ verify(channel).close(false);
+ verify(channel, never()).close(true);
+ }
+
+ @Test
+ void closeChannelForcesCleanupAfterGracefulTimeout() throws IOException {
+ ClientChannel channel = mock(ClientChannel.class);
+ CloseFuture gracefulClose = mock(CloseFuture.class);
+ CloseFuture immediateClose = mock(CloseFuture.class);
+ when(channel.close(false)).thenReturn(gracefulClose);
+ when(gracefulClose.await(1_000)).thenReturn(false);
+ when(channel.close(true)).thenReturn(immediateClose);
+ when(immediateClose.await(1_000)).thenReturn(true);
+
+ assertTrue(SshCollectImpl.closeChannel(channel, 1_000));
+ verify(channel).close(false);
+ verify(channel).close(true);
+ }
+
+ @Test
+ void closeChannelForcesCleanupAndPreservesInterrupt() throws IOException {
+ ClientChannel channel = mock(ClientChannel.class);
+ CloseFuture gracefulClose = mock(CloseFuture.class);
+ when(channel.close(false)).thenReturn(gracefulClose);
+ when(gracefulClose.await(1_000)).thenThrow(new
InterruptedIOException("interrupted"));
+
+ try {
+ assertThrows(InterruptedIOException.class,
+ () -> SshCollectImpl.closeChannel(channel, 1_000));
+ assertTrue(Thread.currentThread().isInterrupted());
+ verify(channel).close(true);
+ } finally {
+ Thread.interrupted();
+ }
+ }
+
+ @Test
+ void closeChannelSkipsAlreadyClosedChannel() throws IOException {
+ ClientChannel channel = mock(ClientChannel.class);
+ when(channel.isClosed()).thenReturn(true);
+
+ assertTrue(SshCollectImpl.closeChannel(channel, 1_000));
+ verify(channel, never()).close(false);
+ verify(channel, never()).close(true);
+ }
+
+ @Test
+ void closeChannelSuppressesForcedCloseFailureAndPreservesInterrupt()
throws IOException {
+ ClientChannel channel = mock(ClientChannel.class);
+ CloseFuture gracefulClose = mock(CloseFuture.class);
+ InterruptedIOException interrupted = new
InterruptedIOException("interrupted");
+ IllegalStateException closeFailure = new IllegalStateException("force
close failed");
+ when(channel.close(false)).thenReturn(gracefulClose);
+ when(gracefulClose.await(1_000)).thenThrow(interrupted);
+ when(channel.close(true)).thenThrow(closeFailure);
+
+ try {
+ InterruptedIOException thrown =
assertThrows(InterruptedIOException.class,
+ () -> SshCollectImpl.closeChannel(channel, 1_000));
+ assertSame(interrupted, thrown);
+ assertEquals(List.of(closeFailure),
List.of(thrown.getSuppressed()));
+ assertTrue(Thread.currentThread().isInterrupted());
+ } finally {
+ Thread.interrupted();
+ }
+ }
+
+ @Test
+ void collectFinallyAwaitsGracefulChannelClose() throws Exception {
+ int timeout = 1_000;
+ SshProtocol protocol = protocol(timeout);
+ Metrics metrics = Metrics.builder().ssh(protocol).build();
+ ClientSession clientSession = mock(ClientSession.class);
+ ChannelExec channel = mock(ChannelExec.class);
+ OpenFuture openFuture = mock(OpenFuture.class);
+ CloseFuture closeFuture = mock(CloseFuture.class);
+ when(clientSession.createExecChannel("echo ok")).thenReturn(channel);
+ when(channel.open()).thenReturn(openFuture);
+ when(openFuture.verify(timeout)).thenThrow(new SshException("channel
open failed"));
+ when(channel.close(false)).thenReturn(closeFuture);
+ when(closeFuture.await(timeout)).thenReturn(true);
+
+ try (MockedStatic<SshHelper> sshHelper = mockStatic(SshHelper.class)) {
+ sshHelper.when(() -> SshHelper.getConnectSession(protocol,
timeout, true, false))
+ .thenReturn(clientSession);
+
+ sshCollect.collect(builder, metrics);
+ }
+
+ assertEquals(CollectRep.Code.UN_CONNECTABLE, builder.getCode());
+ verify(channel).close(false);
+ verify(closeFuture).await(timeout);
+ verify(channel, never()).close(true);
+ verify(clientSession, never()).close();
+ }
+
+ @Test
+ void collectClosesExactSessionWhenChannelCannotClose() throws Exception {
+ int timeout = 1_000;
+ SshProtocol protocol = protocol(timeout);
+ Metrics metrics = Metrics.builder().ssh(protocol).build();
+ ClientSession clientSession = mock(ClientSession.class);
+ ChannelExec channel = mock(ChannelExec.class);
+ OpenFuture openFuture = mock(OpenFuture.class);
+ CloseFuture gracefulClose = mock(CloseFuture.class);
+ CloseFuture immediateClose = mock(CloseFuture.class);
+ when(clientSession.createExecChannel("echo ok")).thenReturn(channel);
+ when(channel.open()).thenReturn(openFuture);
+ when(openFuture.verify(timeout)).thenThrow(new SshException("channel
open failed"));
+ when(channel.close(false)).thenReturn(gracefulClose);
+ when(gracefulClose.await(timeout)).thenReturn(false);
+ when(channel.close(true)).thenReturn(immediateClose);
+ when(immediateClose.await(timeout)).thenReturn(false);
+
+ try (MockedStatic<SshHelper> sshHelper = mockStatic(SshHelper.class)) {
+ sshHelper.when(() -> SshHelper.getConnectSession(protocol,
timeout, true, false))
+ .thenReturn(clientSession);
+
+ sshCollect.collect(builder, metrics);
+ }
+
+ verify(channel).close(false);
+ verify(gracefulClose).await(timeout);
+ verify(channel).close(true);
+ verify(immediateClose).await(timeout);
+ verify(clientSession).close();
+ }
+
+ private SshProtocol protocol(int timeout) {
+ return SshProtocol.builder()
+ .host("target.example.com")
+ .port("22")
+ .username("root")
+ .password("password")
+ .timeout(String.valueOf(timeout))
+ .reuseConnection("true")
+ .useProxy("false")
+ .script("echo ok")
+ .build();
+ }
}
diff --git
a/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClient.java
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClient.java
index 4387415dd3..03ca9b1ec8 100644
---
a/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClient.java
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClient.java
@@ -17,12 +17,12 @@
package org.apache.hertzbeat.collector.collect.common.ssh;
+import java.time.Duration;
import lombok.extern.slf4j.Slf4j;
import org.apache.sshd.client.ClientBuilder;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.keyverifier.AcceptAllServerKeyVerifier;
import org.apache.sshd.common.NamedFactory;
-import org.apache.sshd.common.PropertyResolverUtils;
import org.apache.sshd.common.kex.BuiltinDHFactories;
import org.apache.sshd.core.CoreModuleProperties;
import org.apache.sshd.server.forward.AcceptAllForwardingFilter;
@@ -34,19 +34,16 @@ import
org.apache.sshd.server.forward.AcceptAllForwardingFilter;
public class CommonSshClient {
private static final SshClient SSH_CLIENT;
-
+
static {
SSH_CLIENT = SshClient.setUpDefaultClient();
// accept all server key verifier, will print warn log : Server at {}
presented unverified {} key: {}
AcceptAllServerKeyVerifier verifier =
AcceptAllServerKeyVerifier.INSTANCE;
SSH_CLIENT.setServerKeyVerifier(verifier);
- // set connection heartbeat interval time 2000ms, wait for heartbeat
response timeout 300_000ms
- PropertyResolverUtils.updateProperty(
- SSH_CLIENT, CoreModuleProperties.HEARTBEAT_INTERVAL.getName(),
2000);
- PropertyResolverUtils.updateProperty(
- SSH_CLIENT,
CoreModuleProperties.HEARTBEAT_NO_REPLY_MAX.getName(), 30);
- PropertyResolverUtils.updateProperty(
- SSH_CLIENT, CoreModuleProperties.SOCKET_KEEPALIVE.getName(),
true);
+ // Reduce heartbeat overhead while keeping unresponsive sessions
bounded to about 90 seconds.
+ CoreModuleProperties.HEARTBEAT_INTERVAL.set(SSH_CLIENT,
Duration.ofSeconds(30));
+ CoreModuleProperties.HEARTBEAT_NO_REPLY_MAX.set(SSH_CLIENT, 2);
+ CoreModuleProperties.SOCKET_KEEPALIVE.set(SSH_CLIENT, true);
// set support all KeyExchange
SSH_CLIENT.setKeyExchangeFactories(NamedFactory.setUpTransformedFactories(
false,
diff --git
a/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreaker.java
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreaker.java
new file mode 100644
index 0000000000..f69c1fa5ee
--- /dev/null
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreaker.java
@@ -0,0 +1,102 @@
+/*
+ * 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.hertzbeat.collector.collect.common.ssh;
+
+import com.google.common.base.Ticker;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.Striped;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Stops repeated SSH connection attempts after consecutive authentication
failures.
+ */
+@Slf4j
+final class SshCircuitBreaker {
+
+ static final int MAX_FAILURES = 3;
+ static final Duration COOLDOWN = Duration.ofMinutes(5);
+
+ private final int maxFailures;
+ private final Cache<Target, Integer> failureCounts;
+ private final Striped<Lock> connectionLocks;
+
+ SshCircuitBreaker() {
+ this(MAX_FAILURES, COOLDOWN, Ticker.systemTicker());
+ }
+
+ SshCircuitBreaker(int maxFailures, Duration cooldown, Ticker ticker) {
+ if (maxFailures <= 0) {
+ throw new IllegalArgumentException("maxFailures must be positive");
+ }
+ if (cooldown.isZero() || cooldown.isNegative()) {
+ throw new IllegalArgumentException("cooldown must be positive");
+ }
+ this.maxFailures = maxFailures;
+ this.failureCounts = CacheBuilder.newBuilder()
+ .expireAfterWrite(cooldown.toNanos(), TimeUnit.NANOSECONDS)
+ .ticker(ticker)
+ .build();
+ this.connectionLocks = Striped.lock(1_024);
+ }
+
+ Target buildTarget(String host, String port, String username, List<String>
credentials) {
+ return new Target(host, port, username, List.copyOf(credentials));
+ }
+
+ Lock connectionLock(Target target) {
+ return connectionLocks.get(target);
+ }
+
+ void checkOpen(Target target) throws IOException {
+ Integer failureCount = failureCounts.getIfPresent(target);
+ if (failureCount != null && failureCount >= maxFailures) {
+ throw new IOException("SSH authentication is temporarily blocked
for " + target.displayName()
+ + " after " + failureCount + " consecutive failures");
+ }
+ }
+
+ void recordFailure(Target target) {
+ int failureCount = failureCounts.asMap().merge(target, 1,
Integer::sum);
+ if (failureCount == maxFailures) {
+ log.warn("SSH authentication is temporarily blocked for {} after
{} consecutive failures",
+ target.displayName(), failureCount);
+ }
+ }
+
+ void recordSuccess(Target target) {
+ failureCounts.invalidate(target);
+ }
+
+ record Target(String host, String port, String username, List<String>
credentials) {
+
+ private String displayName() {
+ return username + "@" + host + ":" + port;
+ }
+
+ @Override
+ public String toString() {
+ return displayName();
+ }
+ }
+}
diff --git
a/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelper.java
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelper.java
index 70094feea7..7db6e3d79c 100644
---
a/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelper.java
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/main/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelper.java
@@ -17,29 +17,39 @@
package org.apache.hertzbeat.collector.collect.common.ssh;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
import java.io.InputStream;
+import java.io.InterruptedIOException;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.security.GeneralSecurityException;
import java.security.KeyPair;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
+import java.util.HexFormat;
import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import lombok.extern.slf4j.Slf4j;
import org.apache.hertzbeat.collector.collect.common.cache.AbstractConnection;
import org.apache.hertzbeat.collector.collect.common.cache.CacheIdentifier;
import
org.apache.hertzbeat.collector.collect.common.cache.GlobalConnectionCache;
import org.apache.hertzbeat.collector.collect.common.cache.SshConnect;
-import org.apache.hertzbeat.collector.util.PrivateKeyUtils;
import org.apache.hertzbeat.common.entity.job.protocol.SshProtocol;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.config.hosts.HostConfigEntry;
import org.apache.sshd.client.session.ClientSession;
+import org.apache.sshd.common.RuntimeSshException;
+import org.apache.sshd.common.SshConstants;
+import org.apache.sshd.common.SshException;
import org.apache.sshd.common.config.keys.FilePasswordProvider;
import org.apache.sshd.common.util.security.SecurityUtils;
import org.springframework.util.StringUtils;
-import java.io.FileInputStream;
-import java.io.IOException;
-import java.security.GeneralSecurityException;
-import java.util.Optional;
-import java.util.concurrent.TimeUnit;
-
/**
* ssh helper
*/
@@ -47,160 +57,377 @@ import java.util.concurrent.TimeUnit;
public class SshHelper {
private static final GlobalConnectionCache CONNECTION_COMMON_CACHE =
GlobalConnectionCache.getInstance();
+ private static final SshCircuitBreaker SSH_CIRCUIT_BREAKER = new
SshCircuitBreaker();
+ private static final ReentrantReadWriteLock CLIENT_IDENTITY_LOCK = new
ReentrantReadWriteLock(true);
public static ClientSession getConnectSession(String host, String port,
String username, String password, String privateKey,
String privateKeyPassphrase,
int timeout, boolean reuseConnection)
throws IOException, GeneralSecurityException {
- CacheIdentifier identifier = CacheIdentifier.builder()
- .ip(host).port(port)
- .username(username).password(password)
+ SshProtocol sshProtocol = SshProtocol.builder()
+ .host(host)
+ .port(port)
+ .username(username)
+ .password(password)
+ .privateKey(privateKey)
+ .privateKeyPassphrase(privateKeyPassphrase)
.build();
- ClientSession clientSession = null;
- if (reuseConnection) {
- Optional<AbstractConnection<?>> cacheOption =
CONNECTION_COMMON_CACHE.getCache(identifier, true);
- if (cacheOption.isPresent()) {
- SshConnect sshConnect = (SshConnect) cacheOption.get();
- clientSession = sshConnect.getConnection();
- try {
- if (clientSession == null || clientSession.isClosed() ||
clientSession.isClosing()) {
- clientSession = null;
- CONNECTION_COMMON_CACHE.removeCache(identifier);
- }
- } catch (Exception e) {
- log.warn(e.getMessage());
- clientSession = null;
- CONNECTION_COMMON_CACHE.removeCache(identifier);
- }
- }
+ return getConnectSession(sshProtocol, timeout, reuseConnection, false);
+ }
+
+ public static ClientSession getConnectSession(SshProtocol sshProtocol, int
timeout, boolean reuseConnection, boolean useProxy)
+ throws IOException, GeneralSecurityException {
+ boolean proxyEnabled = useProxy &&
StringUtils.hasText(sshProtocol.getProxyHost());
+ boolean cacheConnection = reuseConnection || proxyEnabled;
+ List<String> targetIdentity = connectionIdentity(
+ sshProtocol.getPassword(), sshProtocol.getPrivateKey(),
sshProtocol.getPrivateKeyPassphrase());
+ List<String> connectionIdentity = connectionIdentity(targetIdentity,
sshProtocol, proxyEnabled);
+ CacheIdentifier identifier = buildCacheIdentifier(
+ sshProtocol.getHost(), sshProtocol.getPort(),
sshProtocol.getUsername(), connectionIdentity);
+ SshCircuitBreaker.Target targetCircuit =
SSH_CIRCUIT_BREAKER.buildTarget(
+ sshProtocol.getHost(), sshProtocol.getPort(),
sshProtocol.getUsername(), targetIdentity);
+ SshCircuitBreaker.Target proxyCircuit = proxyEnabled
+ ? SSH_CIRCUIT_BREAKER.buildTarget(
+ sshProtocol.getProxyHost(),
sshProtocol.getProxyPort(), sshProtocol.getProxyUsername(),
+ proxyIdentity(sshProtocol))
+ : null;
+ // Keep the final ProxyJump session cached with its intermediate
session chain.
+ if (cacheConnection) {
+ ClientSession clientSession = getCachedSession(identifier);
if (clientSession != null) {
return clientSession;
}
}
- SshClient sshClient = CommonSshClient.getSshClient();
- clientSession = sshClient.connect(username, host,
Integer.parseInt(port))
- .verify(timeout, TimeUnit.MILLISECONDS).getSession();
- if (StringUtils.hasText(password)) {
- clientSession.addPasswordIdentity(password);
- } else if (StringUtils.hasText(privateKey)) {
- var resourceKey = PrivateKeyUtils.writePrivateKey(host,
privateKey);
- FilePasswordProvider passwordProvider = (session, resource, index)
-> {
- if (StringUtils.hasText(privateKeyPassphrase)) {
- return privateKeyPassphrase;
+ Lock connectionLock =
SSH_CIRCUIT_BREAKER.connectionLock(targetCircuit);
+ lockInterruptibly(connectionLock, "SSH connection");
+ try {
+ if (cacheConnection) {
+ ClientSession clientSession = getCachedSession(identifier);
+ if (clientSession != null) {
+ return clientSession;
}
- return null;
- };
- SecurityUtils.loadKeyPairIdentities(null, () -> resourceKey, new
FileInputStream(resourceKey), passwordProvider)
- .forEach(clientSession::addPublicKeyIdentity);
- } // else auth with localhost private public key certificates
+ }
+ SSH_CIRCUIT_BREAKER.checkOpen(targetCircuit);
+ if (proxyCircuit != null) {
+ SSH_CIRCUIT_BREAKER.checkOpen(proxyCircuit);
+ }
+ SshClient sshClient = CommonSshClient.getSshClient();
+ ClientSession clientSession = proxyEnabled
+ ? connectThroughProxy(sshClient, sshProtocol, timeout,
proxyCircuit)
+ : connectDirect(sshClient, sshProtocol.getUsername(),
sshProtocol.getHost(),
+ Integer.parseInt(sshProtocol.getPort()), timeout);
+ return authenticate(clientSession, sshProtocol.getHost(),
sshProtocol.getPassword(),
+ sshProtocol.getPrivateKey(),
sshProtocol.getPrivateKeyPassphrase(),
+ timeout, cacheConnection, identifier, targetCircuit);
+ } finally {
+ connectionLock.unlock();
+ }
+ }
- // auth
- if (!clientSession.auth().verify(timeout,
TimeUnit.MILLISECONDS).isSuccess()) {
- clientSession.close();
- throw new IllegalArgumentException("ssh auth failed.");
+ private static ClientSession authenticate(
+ ClientSession clientSession, String host, String password, String
privateKey,
+ String privateKeyPassphrase, int timeout, boolean cacheConnection,
+ CacheIdentifier identifier, SshCircuitBreaker.Target circuitTarget)
+ throws IOException, GeneralSecurityException {
+ boolean authenticated = false;
+ boolean identityReady = false;
+ try {
+ addIdentity(clientSession, host, password, privateKey,
privateKeyPassphrase);
+ identityReady = true;
+ clientSession.auth().verify(timeout, TimeUnit.MILLISECONDS);
+ authenticated = true;
+ SSH_CIRCUIT_BREAKER.recordSuccess(circuitTarget);
+ if (cacheConnection) {
+ CONNECTION_COMMON_CACHE.addCache(identifier, new
SshConnect(clientSession));
+ }
+ return clientSession;
+ } catch (IOException | GeneralSecurityException | RuntimeException e) {
+ if (!authenticated && ((!identityReady &&
isCredentialSetupFailure(e))
+ || isAuthenticationRejection(e))) {
+ SSH_CIRCUIT_BREAKER.recordFailure(circuitTarget);
+ }
+ closeSessionOnFailure(clientSession, e);
+ throw e;
}
- if (reuseConnection) {
- SshConnect sshConnect = new SshConnect(clientSession);
- CONNECTION_COMMON_CACHE.addCache(identifier, sshConnect);
+ }
+
+ private static ClientSession connectDirect(
+ SshClient sshClient, String username, String host, int port, int
timeout) throws IOException {
+ Lock identityLock = CLIENT_IDENTITY_LOCK.readLock();
+ lockInterruptibly(identityLock, "SSH client identity");
+ try {
+ return sshClient.connect(username, host, port)
+ .verify(timeout, TimeUnit.MILLISECONDS).getSession();
+ } finally {
+ identityLock.unlock();
}
- return clientSession;
}
- public static ClientSession getConnectSession(SshProtocol sshProtocol, int
timeout, boolean reuseConnection, boolean useProxy)
+ private static ClientSession connectThroughProxy(
+ SshClient sshClient, SshProtocol sshProtocol, int timeout,
SshCircuitBreaker.Target proxyCircuit)
+ throws IOException, GeneralSecurityException {
+ Lock identityLock = CLIENT_IDENTITY_LOCK.writeLock();
+ lockInterruptibly(identityLock, "SSH client identity");
+ try {
+ SSH_CIRCUIT_BREAKER.checkOpen(proxyCircuit);
+ try {
+ ClientSession clientSession =
connectThroughProxyWithIdentity(sshClient, sshProtocol, timeout);
+ SSH_CIRCUIT_BREAKER.recordSuccess(proxyCircuit);
+ return clientSession;
+ } catch (IOException | GeneralSecurityException | RuntimeException
e) {
+ if (isProxyAuthenticationFailure(e)) {
+ SSH_CIRCUIT_BREAKER.recordFailure(proxyCircuit);
+ }
+ throw e;
+ }
+ } finally {
+ identityLock.unlock();
+ }
+ }
+
+ private static ClientSession connectThroughProxyWithIdentity(
+ SshClient sshClient, SshProtocol sshProtocol, int timeout)
throws IOException, GeneralSecurityException {
- CacheIdentifier identifier = CacheIdentifier.builder()
-
.ip(sshProtocol.getHost()).port(sshProtocol.getPort())
-
.username(sshProtocol.getUsername()).password(sshProtocol.getPassword())
- .build();
+ HostConfigEntry proxyConfig = new HostConfigEntry();
+ String proxySpec = String.format("%s@%s:%d",
sshProtocol.getProxyUsername(),
+ sshProtocol.getProxyHost(),
Integer.parseInt(sshProtocol.getProxyPort()));
+ proxyConfig.setHostName(sshProtocol.getHost());
+ proxyConfig.setHost(sshProtocol.getHost());
+ proxyConfig.setPort(Integer.parseInt(sshProtocol.getPort()));
+ proxyConfig.setUsername(sshProtocol.getUsername());
+ proxyConfig.setProxyJump(proxySpec);
+
+ String proxyPassword =
StringUtils.hasText(sshProtocol.getProxyPassword())
+ ? sshProtocol.getProxyPassword()
+ : null;
+ List<KeyPair> proxyKeyPairs = List.of();
ClientSession clientSession = null;
- // When using ProxyJump, force connection reuse:
- // Apache MINA SSHD will pass the proxy password error to the target
host in proxy scenarios, causing the first connection to fail.
- // Reusing connections can skip duplicate authentication and avoid
this problem.
- if (reuseConnection || useProxy) {
- Optional<AbstractConnection<?>> cacheOption =
CONNECTION_COMMON_CACHE.getCache(identifier, true);
- if (cacheOption.isPresent()) {
- SshConnect sshConnect = (SshConnect) cacheOption.get();
- clientSession = sshConnect.getConnection();
+ try {
+ if (proxyPassword != null) {
+ sshClient.addPasswordIdentity(proxyPassword);
+ }
+ if (StringUtils.hasText(sshProtocol.getProxyPrivateKey())) {
try {
- if (clientSession == null || clientSession.isClosed() ||
clientSession.isClosing()) {
- clientSession = null;
- CONNECTION_COMMON_CACHE.removeCache(identifier);
- }
- } catch (Exception e) {
- log.warn(e.getMessage());
- clientSession = null;
- CONNECTION_COMMON_CACHE.removeCache(identifier);
+ proxyKeyPairs = loadKeyPairs(
+ sshProtocol.getProxyHost(),
sshProtocol.getProxyPrivateKey(), null);
+ } catch (IOException | GeneralSecurityException e) {
+ throw new GeneralSecurityException("Failed to load proxy
private key", e);
}
+ proxyKeyPairs.forEach(sshClient::addPublicKeyIdentity);
}
+ clientSession = sshClient.connect(proxyConfig)
+ .verify(timeout, TimeUnit.MILLISECONDS).getSession();
+ if (proxyPassword != null) {
+ clientSession.removePasswordIdentity(proxyPassword);
+ }
+ proxyKeyPairs.forEach(clientSession::removePublicKeyIdentity);
+ return clientSession;
+ } catch (IOException | GeneralSecurityException | RuntimeException e) {
if (clientSession != null) {
- return clientSession;
+ closeSessionOnFailure(clientSession, e);
}
+ throw e;
+ } finally {
+ if (proxyPassword != null) {
+ sshClient.removePasswordIdentity(proxyPassword);
+ }
+ proxyKeyPairs.forEach(sshClient::removePublicKeyIdentity);
}
- SshClient sshClient = CommonSshClient.getSshClient();
- HostConfigEntry proxyConfig = new HostConfigEntry();
- if (useProxy && StringUtils.hasText(sshProtocol.getProxyHost())) {
- String proxySpec = String.format("%s@%s:%d",
sshProtocol.getProxyUsername(), sshProtocol.getProxyHost(),
Integer.parseInt(sshProtocol.getProxyPort()));
- proxyConfig.setHostName(sshProtocol.getHost());
- proxyConfig.setHost(sshProtocol.getHost());
- proxyConfig.setPort(Integer.parseInt(sshProtocol.getPort()));
- proxyConfig.setUsername(sshProtocol.getUsername());
- proxyConfig.setProxyJump(proxySpec);
-
- // Apache SSHD requires the password for the proxy to be preloaded
into the sshClient instance before connecting
- if (StringUtils.hasText(sshProtocol.getProxyPassword())) {
- sshClient.addPasswordIdentity(sshProtocol.getProxyPassword());
- log.debug("Loaded proxy server password authentication:
{}@{}", sshProtocol.getProxyUsername(), sshProtocol.getProxyHost());
+ }
+
+ private static void addIdentity(ClientSession clientSession, String host,
String password, String privateKey,
+ String privateKeyPassphrase) throws
IOException, GeneralSecurityException {
+ if (StringUtils.hasText(password)) {
+ clientSession.addPasswordIdentity(password);
+ return;
+ }
+ if (!StringUtils.hasText(privateKey)) {
+ return;
+ }
+ loadKeyPairs(host, privateKey,
privateKeyPassphrase).forEach(clientSession::addPublicKeyIdentity);
+ }
+
+ private static List<KeyPair> loadKeyPairs(String host, String privateKey,
String privateKeyPassphrase)
+ throws IOException, GeneralSecurityException {
+ String resourceKey = "ssh-key-" + valueOrEmpty(host);
+ FilePasswordProvider passwordProvider = (session, resource, index) ->
+ StringUtils.hasText(privateKeyPassphrase) ?
privateKeyPassphrase : null;
+ try (InputStream keyStream = new
ByteArrayInputStream(privateKey.getBytes(StandardCharsets.UTF_8))) {
+ Iterable<KeyPair> keyPairs = SecurityUtils.loadKeyPairIdentities(
+ null, () -> resourceKey, keyStream, passwordProvider);
+ if (keyPairs == null) {
+ throw new GeneralSecurityException("Failed to load SSH private
key");
}
- if (StringUtils.hasText(sshProtocol.getProxyPrivateKey())) {
-
proxyConfig.setIdentities(List.of(sshProtocol.getProxyPrivateKey()));
- log.debug("Proxy private key loaded into HostConfigEntry");
+ List<KeyPair> loadedKeyPairs = new ArrayList<>();
+ keyPairs.forEach(loadedKeyPairs::add);
+ if (loadedKeyPairs.isEmpty()) {
+ throw new GeneralSecurityException("SSH private key contains
no identities");
}
+ return List.copyOf(loadedKeyPairs);
}
+ }
- if (useProxy && StringUtils.hasText(sshProtocol.getProxyHost())) {
- try {
- clientSession = sshClient.connect(proxyConfig)
- .verify(timeout,
TimeUnit.MILLISECONDS).getSession();
- }
- finally {
-
sshClient.removePasswordIdentity(sshProtocol.getProxyPassword());
+ private static ClientSession getCachedSession(CacheIdentifier identifier) {
+ Optional<AbstractConnection<?>> cacheOption =
CONNECTION_COMMON_CACHE.getCache(identifier, true);
+ if (cacheOption.isEmpty()) {
+ return null;
+ }
+ ClientSession clientSession;
+ try {
+ clientSession = ((SshConnect) cacheOption.get()).getConnection();
+ if (clientSession == null || clientSession.isClosed() ||
clientSession.isClosing()) {
+ CONNECTION_COMMON_CACHE.removeCache(identifier);
+ return null;
}
+ return clientSession;
+ } catch (Exception e) {
+ log.warn("Failed to validate cached SSH session for {}",
identifier, e);
+ CONNECTION_COMMON_CACHE.removeCache(identifier);
+ return null;
+ }
+ }
+
+ private static CacheIdentifier buildCacheIdentifier(
+ String host, String port, String username, List<String>
connectionIdentity) {
+ return CacheIdentifier.builder()
+ .ip(host)
+ .port(port)
+ .username(username)
+ .customArg(fingerprint(connectionIdentity))
+ .build();
+ }
+
+ private static List<String> connectionIdentity(
+ String password, String privateKey, String privateKeyPassphrase) {
+ List<String> identity = new ArrayList<>();
+ addAuthenticationIdentity(identity, password, privateKey,
privateKeyPassphrase);
+ return List.copyOf(identity);
+ }
+
+ private static List<String> connectionIdentity(
+ List<String> targetIdentity, SshProtocol sshProtocol, boolean
proxyEnabled) {
+ List<String> identity = new ArrayList<>(targetIdentity);
+ if (proxyEnabled) {
+ identity.add("proxy");
+ identity.add(valueOrEmpty(sshProtocol.getProxyHost()));
+ identity.add(valueOrEmpty(sshProtocol.getProxyPort()));
+ identity.add(valueOrEmpty(sshProtocol.getProxyUsername()));
+ identity.add(textOrEmpty(sshProtocol.getProxyPassword()));
+ identity.add(textOrEmpty(sshProtocol.getProxyPrivateKey()));
+ }
+ return List.copyOf(identity);
+ }
+
+ private static List<String> proxyIdentity(SshProtocol sshProtocol) {
+ List<String> identity = new ArrayList<>();
+ if (StringUtils.hasText(sshProtocol.getProxyPassword())) {
+ identity.add("proxy-password");
+ identity.add(sshProtocol.getProxyPassword());
+ }
+ if (StringUtils.hasText(sshProtocol.getProxyPrivateKey())) {
+ identity.add("proxy-private-key");
+ identity.add(sshProtocol.getProxyPrivateKey());
+ }
+ if (identity.isEmpty()) {
+ identity.add("proxy-default");
+ }
+ return List.copyOf(identity);
+ }
+
+ private static void addAuthenticationIdentity(
+ List<String> identity, String password, String privateKey, String
privateKeyPassphrase) {
+ if (StringUtils.hasText(password)) {
+ identity.add("target-password");
+ identity.add(password);
+ } else if (StringUtils.hasText(privateKey)) {
+ identity.add("target-private-key");
+ identity.add(privateKey);
+ identity.add(textOrEmpty(privateKeyPassphrase));
} else {
- clientSession = sshClient.connect(sshProtocol.getUsername(),
sshProtocol.getHost(), Integer.parseInt(sshProtocol.getPort()))
- .verify(timeout,
TimeUnit.MILLISECONDS).getSession();
- }
-
- if (StringUtils.hasText(sshProtocol.getPassword())) {
- clientSession.addPasswordIdentity(sshProtocol.getPassword());
- } else if (StringUtils.hasText(sshProtocol.getPrivateKey())) {
- var resourceKey =
PrivateKeyUtils.writePrivateKey(sshProtocol.getHost(),
sshProtocol.getPrivateKey());
- try (InputStream keyStream = new FileInputStream(resourceKey)) {
- FilePasswordProvider passwordProvider = (session, resource,
index) -> {
- if
(StringUtils.hasText(sshProtocol.getPrivateKeyPassphrase())) {
- return sshProtocol.getPrivateKeyPassphrase();
- }
- return null;
- };
- Iterable<KeyPair> keyPairs =
SecurityUtils.loadKeyPairIdentities(null, () -> resourceKey, keyStream,
passwordProvider);
- if (keyPairs != null) {
- keyPairs.forEach(clientSession::addPublicKeyIdentity);
- } else {
- log.error("Failed to load private key pairs from: {}",
resourceKey);
- }
- } catch (IOException e) {
- log.error("Error reading private key file: {}",
e.getMessage());
+ identity.add("target-default");
+ }
+ }
+
+ private static String fingerprint(List<String> connectionIdentity) {
+ try {
+ MessageDigest digest = MessageDigest.getInstance("SHA-256");
+ for (String component : connectionIdentity) {
+ byte[] bytes = component.getBytes(StandardCharsets.UTF_8);
+
digest.update(ByteBuffer.allocate(Integer.BYTES).putInt(bytes.length).array());
+ digest.update(bytes);
}
- } // else auth with localhost private public key certificates
+ return HexFormat.of().formatHex(digest.digest());
+ } catch (NoSuchAlgorithmException e) {
+ throw new IllegalStateException("SHA-256 is unavailable", e);
+ }
+ }
- // auth
- if (!clientSession.auth().verify(timeout,
TimeUnit.MILLISECONDS).isSuccess()) {
- clientSession.close();
- throw new IllegalArgumentException("ssh auth failed.");
+ private static String valueOrEmpty(String value) {
+ return value == null ? "" : value;
+ }
+
+ private static String textOrEmpty(String value) {
+ return StringUtils.hasText(value) ? value : "";
+ }
+
+ private static boolean isCredentialSetupFailure(Exception exception) {
+ return !(exception instanceof InterruptedIOException)
+ && (exception instanceof IOException || exception instanceof
GeneralSecurityException);
+ }
+
+ private static boolean isProxyAuthenticationFailure(Exception exception) {
+ return exception instanceof GeneralSecurityException ||
isAuthenticationRejection(exception);
+ }
+
+ private static boolean isAuthenticationRejection(Exception exception) {
+ Throwable current = exception;
+ while (current != null) {
+ if (current instanceof InterruptedIOException) {
+ return false;
+ }
+ if (current instanceof SshException sshException) {
+ int disconnectCode = sshException.getDisconnectCode();
+ return disconnectCode ==
SshConstants.SSH2_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE
+ || disconnectCode ==
SshConstants.SSH2_DISCONNECT_HOST_AUTHENTICATION_FAILED
+ || disconnectCode ==
SshConstants.SSH2_DISCONNECT_ILLEGAL_USER_NAME;
+ }
+ current = current instanceof RuntimeSshException ?
current.getCause() : null;
}
- if (reuseConnection || useProxy) {
- SshConnect sshConnect = new SshConnect(clientSession);
- CONNECTION_COMMON_CACHE.addCache(identifier, sshConnect);
+ return false;
+ }
+
+ private static void lockInterruptibly(Lock lock, String operation) throws
InterruptedIOException {
+ try {
+ lock.lockInterruptibly();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ InterruptedIOException exception = new
InterruptedIOException(operation + " interrupted");
+ exception.initCause(e);
+ throw exception;
+ }
+ }
+
+ private static void closeSessionOnFailure(ClientSession clientSession,
Exception originalException) {
+ if (clientSession == null) {
+ return;
+ }
+ if (originalException instanceof InterruptedIOException) {
+ try {
+ clientSession.close(true);
+ } catch (RuntimeException closeException) {
+ originalException.addSuppressed(closeException);
+ } finally {
+ Thread.currentThread().interrupt();
+ }
+ return;
+ }
+ try {
+ clientSession.close();
+ } catch (Exception closeException) {
+ if (closeException instanceof InterruptedIOException) {
+ Thread.currentThread().interrupt();
+ }
+ originalException.addSuppressed(closeException);
}
- return clientSession;
}
}
diff --git
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
similarity index 61%
rename from
hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
rename to
hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
index b4ceeedcb3..bfd3eab17c 100644
---
a/hertzbeat-collector/hertzbeat-collector-basic/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/CommonSshClientTest.java
@@ -17,8 +17,12 @@
package org.apache.hertzbeat.collector.collect.common.ssh;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeEach;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.time.Duration;
+import org.apache.sshd.client.SshClient;
+import org.apache.sshd.core.CoreModuleProperties;
import org.junit.jupiter.api.Test;
/**
@@ -26,15 +30,13 @@ import org.junit.jupiter.api.Test;
*/
class CommonSshClientTest {
- @BeforeEach
- void setUp() {
- }
-
- @AfterEach
- void tearDown() {
- }
-
@Test
void getSshClient() {
+ SshClient sshClient = CommonSshClient.getSshClient();
+
+ assertTrue(sshClient.isStarted());
+ assertEquals(Duration.ofSeconds(30),
CoreModuleProperties.HEARTBEAT_INTERVAL.getRequired(sshClient));
+ assertEquals(2,
CoreModuleProperties.HEARTBEAT_NO_REPLY_MAX.getRequired(sshClient));
+
assertTrue(CoreModuleProperties.SOCKET_KEEPALIVE.getRequired(sshClient));
}
}
diff --git
a/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreakerTest.java
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreakerTest.java
new file mode 100644
index 0000000000..e1b46e8b4c
--- /dev/null
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshCircuitBreakerTest.java
@@ -0,0 +1,111 @@
+/*
+ * 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.hertzbeat.collector.collect.common.ssh;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.google.common.base.Ticker;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Tests for {@link SshCircuitBreaker}.
+ */
+class SshCircuitBreakerTest {
+
+ private TestTicker ticker;
+ private SshCircuitBreaker circuitBreaker;
+ private SshCircuitBreaker.Target target;
+
+ @BeforeEach
+ void setUp() {
+ ticker = new TestTicker();
+ circuitBreaker = new SshCircuitBreaker(3, Duration.ofMinutes(5),
ticker);
+ target = circuitBreaker.buildTarget(
+ "127.0.0.1", "22", "root", List.of("target-password",
"bad-password"));
+ }
+
+ @Test
+ void opensAfterThresholdAndExpiresAfterCooldown() {
+ circuitBreaker.recordFailure(target);
+ circuitBreaker.recordFailure(target);
+ assertDoesNotThrow(() -> circuitBreaker.checkOpen(target));
+
+ circuitBreaker.recordFailure(target);
+ assertThrows(IOException.class, () ->
circuitBreaker.checkOpen(target));
+
+ ticker.advance(Duration.ofMinutes(5).plusNanos(1));
+ assertDoesNotThrow(() -> circuitBreaker.checkOpen(target));
+ }
+
+ @Test
+ void successAndCredentialChangeBypassPreviousFailures() {
+ circuitBreaker.recordFailure(target);
+ circuitBreaker.recordFailure(target);
+ circuitBreaker.recordFailure(target);
+
+ SshCircuitBreaker.Target correctedCredentials =
circuitBreaker.buildTarget(
+ "127.0.0.1", "22", "root", List.of("target-password",
"correct-password"));
+ assertNotEquals(target, correctedCredentials);
+ assertDoesNotThrow(() ->
circuitBreaker.checkOpen(correctedCredentials));
+
+ circuitBreaker.recordSuccess(target);
+ assertDoesNotThrow(() -> circuitBreaker.checkOpen(target));
+ }
+
+ @Test
+ void equalTargetsShareConnectionLockWithoutExposingCredentials() {
+ SshCircuitBreaker.Target equalTarget = circuitBreaker.buildTarget(
+ "127.0.0.1", "22", "root", List.of("target-password",
"bad-password"));
+
+ assertSame(circuitBreaker.connectionLock(target),
circuitBreaker.connectionLock(equalTarget));
+ assertEquals("[email protected]:22", target.toString());
+ }
+
+ @Test
+ void rejectsInvalidConfiguration() {
+ assertThrows(IllegalArgumentException.class,
+ () -> new SshCircuitBreaker(0, Duration.ofMinutes(5), ticker));
+ assertThrows(IllegalArgumentException.class,
+ () -> new SshCircuitBreaker(3, Duration.ZERO, ticker));
+ assertThrows(IllegalArgumentException.class,
+ () -> new SshCircuitBreaker(3, Duration.ofSeconds(-1),
ticker));
+ }
+
+ private static final class TestTicker extends Ticker {
+
+ private final AtomicLong nanos = new AtomicLong();
+
+ @Override
+ public long read() {
+ return nanos.get();
+ }
+
+ private void advance(Duration duration) {
+ nanos.addAndGet(duration.toNanos());
+ }
+ }
+}
diff --git
a/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelperTest.java
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelperTest.java
new file mode 100644
index 0000000000..36795d776c
--- /dev/null
+++
b/hertzbeat-collector/hertzbeat-collector-common/src/test/java/org/apache/hertzbeat/collector/collect/common/ssh/SshHelperTest.java
@@ -0,0 +1,551 @@
+/*
+ * 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.hertzbeat.collector.collect.common.ssh;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.io.InterruptedIOException;
+import java.security.KeyPair;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import org.apache.hertzbeat.common.entity.job.protocol.SshProtocol;
+import org.apache.sshd.client.SshClient;
+import org.apache.sshd.client.config.hosts.HostConfigEntry;
+import org.apache.sshd.client.future.AuthFuture;
+import org.apache.sshd.client.future.ConnectFuture;
+import org.apache.sshd.client.session.ClientSession;
+import org.apache.sshd.common.SshConstants;
+import org.apache.sshd.common.SshException;
+import org.apache.sshd.common.config.keys.FilePasswordProvider;
+import org.apache.sshd.common.util.security.SecurityUtils;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.InOrder;
+import org.mockito.MockedStatic;
+
+/**
+ * Tests for {@link SshHelper}.
+ */
+class SshHelperTest {
+
+ private static final int TIMEOUT = 1_000;
+
+ private String host;
+ private SshClient sshClient;
+
+ @BeforeEach
+ void setUp() {
+ host = UUID.randomUUID() + ".example.com";
+ sshClient = mock(SshClient.class);
+ }
+
+ @Test
+ void blocksSameEffectiveCredentialsAndRetriesAfterPasswordChanges() throws
Exception {
+ ConnectAttempt firstFailure = failedAttempt("first auth failure");
+ ConnectAttempt secondFailure = failedAttempt("second auth failure");
+ ConnectAttempt thirdFailure = failedAttempt("third auth failure");
+ ConnectAttempt success = successfulAttempt();
+ when(sshClient.connect("root", host, 22)).thenReturn(
+ firstFailure.connectFuture(), secondFailure.connectFuture(),
+ thirdFailure.connectFuture(), success.connectFuture());
+
+ SshProtocol badProtocol = protocol("bad-password",
"ignored-private-key-a");
+ SshProtocol samePassword = protocol("bad-password",
"ignored-private-key-b");
+ SshProtocol correctedProtocol = protocol("correct-password",
"ignored-private-key-b");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ for (int i = 0; i < SshCircuitBreaker.MAX_FAILURES; i++) {
+ assertThrows(SshException.class,
+ () -> SshHelper.getConnectSession(badProtocol,
TIMEOUT, false, false));
+ }
+ IOException blocked = assertThrows(IOException.class,
+ () -> SshHelper.getConnectSession(samePassword, TIMEOUT,
false, false));
+ assertTrue(blocked.getMessage().contains("temporarily blocked"));
+
+ assertSame(success.clientSession(),
+ SshHelper.getConnectSession(correctedProtocol, TIMEOUT,
false, false));
+ }
+
+ verify(sshClient, times(4)).connect("root", host, 22);
+ verify(firstFailure.clientSession()).close();
+ verify(secondFailure.clientSession()).close();
+ verify(thirdFailure.clientSession()).close();
+ verify(success.clientSession(), never()).close();
+
verify(success.clientSession()).addPasswordIdentity("correct-password");
+ }
+
+ @Test
+ void concurrentReuseCreatesOnePhysicalSession() throws Exception {
+ ConnectAttempt success = successfulAttempt();
+ CountDownLatch authenticationStarted = new CountDownLatch(1);
+ CountDownLatch allowAuthentication = new CountDownLatch(1);
+ when(success.authFuture().verify(TIMEOUT,
TimeUnit.MILLISECONDS)).thenAnswer(invocation -> {
+ authenticationStarted.countDown();
+ assertTrue(allowAuthentication.await(5, TimeUnit.SECONDS));
+ return success.authFuture();
+ });
+ when(success.clientSession().isClosed()).thenReturn(false);
+ when(success.clientSession().isClosing()).thenReturn(false);
+ when(sshClient.connect("root", host,
22)).thenReturn(success.connectFuture());
+
+ CountDownLatch secondStarted = new CountDownLatch(1);
+ try (ExecutorService executor = Executors.newFixedThreadPool(2)) {
+ Future<ClientSession> first = executor.submit(() ->
connectWithThreadLocalClient("password"));
+ assertTrue(authenticationStarted.await(5, TimeUnit.SECONDS));
+ Future<ClientSession> second = executor.submit(() -> {
+ secondStarted.countDown();
+ return connectWithThreadLocalClient("password");
+ });
+ assertTrue(secondStarted.await(5, TimeUnit.SECONDS));
+
+ assertThrows(TimeoutException.class, () -> second.get(200,
TimeUnit.MILLISECONDS));
+ verify(sshClient).connect("root", host, 22);
+ allowAuthentication.countDown();
+ assertSame(success.clientSession(), first.get(5,
TimeUnit.SECONDS));
+ assertSame(success.clientSession(), second.get(5,
TimeUnit.SECONDS));
+ } finally {
+ allowAuthentication.countDown();
+ }
+
+ verify(sshClient).connect("root", host, 22);
+ verify(success.clientSession()).auth();
+ verify(success.authFuture()).verify(TIMEOUT, TimeUnit.MILLISECONDS);
+ }
+
+ @Test
+ void concurrentAuthenticationFailuresCreateAtMostThreePhysicalSessions()
throws Exception {
+ ConnectAttempt firstFailure = failedAttempt("first concurrent
failure");
+ ConnectAttempt secondFailure = failedAttempt("second concurrent
failure");
+ ConnectAttempt thirdFailure = failedAttempt("third concurrent
failure");
+ when(sshClient.connect("root", host, 22)).thenReturn(
+ firstFailure.connectFuture(), secondFailure.connectFuture(),
thirdFailure.connectFuture());
+
+ int taskCount = 8;
+ CountDownLatch ready = new CountDownLatch(taskCount);
+ CountDownLatch start = new CountDownLatch(1);
+ List<Future<Exception>> futures = new ArrayList<>();
+ try (ExecutorService executor =
Executors.newFixedThreadPool(taskCount)) {
+ for (int i = 0; i < taskCount; i++) {
+ futures.add(executor.submit(() -> {
+ ready.countDown();
+ assertTrue(start.await(5, TimeUnit.SECONDS));
+ try {
+
connectWithThreadLocalClient("concurrent-bad-password", false);
+ return null;
+ } catch (Exception e) {
+ return e;
+ }
+ }));
+ }
+ assertTrue(ready.await(5, TimeUnit.SECONDS));
+ start.countDown();
+ }
+
+ int authenticationFailures = 0;
+ int blockedFailures = 0;
+ for (Future<Exception> future : futures) {
+ Exception exception = future.get(5, TimeUnit.SECONDS);
+ if (exception instanceof SshException) {
+ authenticationFailures++;
+ } else {
+ assertTrue(exception instanceof IOException);
+ assertTrue(exception.getMessage().contains("temporarily
blocked"));
+ blockedFailures++;
+ }
+ }
+ assertEquals(SshCircuitBreaker.MAX_FAILURES, authenticationFailures);
+ assertEquals(taskCount - SshCircuitBreaker.MAX_FAILURES,
blockedFailures);
+ verify(sshClient,
times(SshCircuitBreaker.MAX_FAILURES)).connect("root", host, 22);
+ verify(firstFailure.clientSession()).close();
+ verify(secondFailure.clientSession()).close();
+ verify(thirdFailure.clientSession()).close();
+ }
+
+ @Test
+ void proxyFailuresOpenCircuitAndRemoveGlobalPassword() throws Exception {
+ ConnectFuture connectFuture = mock(ConnectFuture.class);
+ ConnectAttempt correctedSuccess = successfulAttempt();
+ when(sshClient.connect(any(HostConfigEntry.class))).thenReturn(
+ connectFuture, connectFuture, connectFuture,
correctedSuccess.connectFuture());
+ when(connectFuture.verify(TIMEOUT, TimeUnit.MILLISECONDS))
+ .thenThrow(new SshException(
+
SshConstants.SSH2_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE, "proxy auth
failed"));
+ SshProtocol protocol = proxyProtocol("bad-proxy-password",
"target-password");
+ SshProtocol otherTarget = proxyProtocol("bad-proxy-password",
"different-target-password");
+ otherTarget.setHost(UUID.randomUUID() + ".example.com");
+ SshProtocol correctedProtocol =
proxyProtocol("correct-proxy-password", "target-password");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ for (int i = 0; i < SshCircuitBreaker.MAX_FAILURES; i++) {
+ assertThrows(SshException.class,
+ () -> SshHelper.getConnectSession(protocol, TIMEOUT,
false, true));
+ }
+ IOException blocked = assertThrows(IOException.class,
+ () -> SshHelper.getConnectSession(protocol, TIMEOUT,
false, true));
+ assertTrue(blocked.getMessage().contains("temporarily blocked"));
+ IOException sharedProxyBlocked = assertThrows(IOException.class,
+ () -> SshHelper.getConnectSession(otherTarget, TIMEOUT,
false, true));
+ assertTrue(sharedProxyBlocked.getMessage().contains("temporarily
blocked"));
+
+ assertSame(correctedSuccess.clientSession(),
+ SshHelper.getConnectSession(correctedProtocol, TIMEOUT,
false, true));
+ }
+
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES +
1)).connect(any(HostConfigEntry.class));
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES))
+ .addPasswordIdentity("bad-proxy-password");
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES))
+ .removePasswordIdentity("bad-proxy-password");
+ verify(sshClient).addPasswordIdentity("correct-proxy-password");
+ verify(sshClient).removePasswordIdentity("correct-proxy-password");
+ }
+
+ @Test
+ void cachedProxySessionDoesNotResetProxyAuthenticationFailures() throws
Exception {
+ ConnectAttempt cachedSuccess = successfulAttempt();
+ ConnectFuture failedProxyConnect = mock(ConnectFuture.class);
+ when(failedProxyConnect.verify(TIMEOUT, TimeUnit.MILLISECONDS))
+ .thenThrow(new SshException(
+
SshConstants.SSH2_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE, "proxy auth
failed"));
+ when(cachedSuccess.clientSession().isClosed()).thenReturn(false);
+ when(cachedSuccess.clientSession().isClosing()).thenReturn(false);
+ when(sshClient.connect(any(HostConfigEntry.class))).thenReturn(
+ cachedSuccess.connectFuture(), failedProxyConnect,
failedProxyConnect, failedProxyConnect);
+ String proxyPassword = UUID.randomUUID() + "-proxy-password";
+ SshProtocol cachedProtocol = proxyProtocol(proxyPassword,
"cached-target-password");
+ SshProtocol failingProtocol = proxyProtocol(proxyPassword,
"failing-target-password");
+ failingProtocol.setHost(UUID.randomUUID() + ".example.com");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ assertSame(cachedSuccess.clientSession(),
+ SshHelper.getConnectSession(cachedProtocol, TIMEOUT,
false, true));
+ for (int i = 0; i < SshCircuitBreaker.MAX_FAILURES; i++) {
+ assertThrows(SshException.class,
+ () -> SshHelper.getConnectSession(failingProtocol,
TIMEOUT, false, true));
+ assertSame(cachedSuccess.clientSession(),
+ SshHelper.getConnectSession(cachedProtocol, TIMEOUT,
false, true));
+ }
+ IOException blocked = assertThrows(IOException.class,
+ () -> SshHelper.getConnectSession(failingProtocol,
TIMEOUT, false, true));
+ assertTrue(blocked.getMessage().contains("temporarily blocked"));
+ }
+
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES + 1))
+ .connect(any(HostConfigEntry.class));
+ }
+
+ @Test
+ void proxyTransportFailuresDoNotOpenAuthenticationCircuit() throws
Exception {
+ ConnectFuture connectFuture = mock(ConnectFuture.class);
+
when(sshClient.connect(any(HostConfigEntry.class))).thenReturn(connectFuture);
+ when(connectFuture.verify(TIMEOUT, TimeUnit.MILLISECONDS))
+ .thenThrow(new SshException(
+ SshConstants.SSH2_DISCONNECT_CONNECTION_LOST, "proxy
transport failed"));
+ String proxyPassword = UUID.randomUUID() + "-proxy-password";
+ SshProtocol protocol = proxyProtocol(proxyPassword, "target-password");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ for (int i = 0; i < SshCircuitBreaker.MAX_FAILURES + 1; i++) {
+ assertThrows(SshException.class,
+ () -> SshHelper.getConnectSession(protocol, TIMEOUT,
false, true));
+ }
+ }
+
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES + 1))
+ .connect(any(HostConfigEntry.class));
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES + 1))
+ .addPasswordIdentity(proxyPassword);
+ verify(sshClient, times(SshCircuitBreaker.MAX_FAILURES + 1))
+ .removePasswordIdentity(proxyPassword);
+ }
+
+ @Test
+ void proxyPasswordIsRemovedBeforeTargetAuthentication() throws Exception {
+ ConnectAttempt success = successfulAttempt();
+
when(sshClient.connect(any(HostConfigEntry.class))).thenReturn(success.connectFuture());
+ SshProtocol protocol = proxyProtocol("proxy-password",
"target-password");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ assertSame(success.clientSession(),
+ SshHelper.getConnectSession(protocol, TIMEOUT, false,
true));
+ }
+
+ verify(sshClient).addPasswordIdentity("proxy-password");
+ verify(sshClient).connect(any(HostConfigEntry.class));
+ verify(success.connectFuture()).verify(TIMEOUT, TimeUnit.MILLISECONDS);
+ verify(sshClient).removePasswordIdentity("proxy-password");
+
+ InOrder order = inOrder(success.clientSession(), success.authFuture());
+
order.verify(success.clientSession()).removePasswordIdentity("proxy-password");
+
order.verify(success.clientSession()).addPasswordIdentity("target-password");
+ order.verify(success.clientSession()).auth();
+ order.verify(success.authFuture()).verify(TIMEOUT,
TimeUnit.MILLISECONDS);
+
+ ArgumentCaptor<HostConfigEntry> configCaptor =
ArgumentCaptor.forClass(HostConfigEntry.class);
+ verify(sshClient).connect(configCaptor.capture());
+ HostConfigEntry proxyConfig = configCaptor.getValue();
+ assertEquals(host, proxyConfig.getHostName());
+ assertEquals("root", proxyConfig.getUsername());
+ assertEquals("[email protected]:2222",
proxyConfig.getProxyJump());
+ }
+
+ @Test
+ void defaultIdentityReconnectsAfterCachedSessionsBecomeUnhealthy() throws
Exception {
+ ConnectAttempt first = successfulAttempt();
+ ConnectAttempt second = successfulAttempt();
+ ConnectAttempt third = successfulAttempt();
+ when(sshClient.connect("root", host, 22)).thenReturn(
+ first.connectFuture(), second.connectFuture(),
third.connectFuture());
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+
+ assertSame(first.clientSession(), SshHelper.getConnectSession(
+ host, "22", "root", null, null, null, TIMEOUT, true));
+ when(first.clientSession().isClosed()).thenReturn(true);
+ assertSame(second.clientSession(), SshHelper.getConnectSession(
+ host, "22", "root", null, null, null, TIMEOUT, true));
+ when(second.clientSession().isClosed()).thenThrow(new
IllegalStateException("health check failed"));
+ assertSame(third.clientSession(), SshHelper.getConnectSession(
+ host, "22", "root", null, null, null, TIMEOUT, true));
+ when(third.clientSession().isClosed()).thenReturn(false);
+ when(third.clientSession().isClosing()).thenReturn(false);
+ assertSame(third.clientSession(), SshHelper.getConnectSession(
+ host, "22", "root", null, null, null, TIMEOUT, true));
+ }
+
+ verify(sshClient, times(3)).connect("root", host, 22);
+ verify(first.clientSession()).close();
+ verify(second.clientSession()).close();
+ verify(third.clientSession(), never()).close();
+ verify(first.clientSession(), never()).addPasswordIdentity(any());
+ verify(first.clientSession(), never()).addPublicKeyIdentity(any());
+ }
+
+ @Test
+ void loadsTargetPrivateKeyDirectlyFromMemory() throws Exception {
+ ConnectAttempt success = successfulAttempt();
+ KeyPair keyPair = mock(KeyPair.class);
+ when(sshClient.connect("root", host,
22)).thenReturn(success.connectFuture());
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class);
+ MockedStatic<SecurityUtils> securityUtils =
mockStatic(SecurityUtils.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+ securityUtils.when(() ->
SecurityUtils.loadKeyPairIdentities(any(), any(), any(), any()))
+ .thenAnswer(invocation -> {
+ FilePasswordProvider provider =
invocation.getArgument(3);
+ assertEquals("passphrase", provider.getPassword(null,
null, 0));
+ return List.of(keyPair);
+ });
+
+ assertSame(success.clientSession(), SshHelper.getConnectSession(
+ host, "22", "root", null, "private-key", "passphrase",
TIMEOUT, false));
+ }
+
+ verify(success.clientSession()).addPublicKeyIdentity(keyPair);
+ verify(success.clientSession(), never()).addPasswordIdentity(any());
+ }
+
+ @Test
+ void proxyPrivateKeyIsRemovedBeforeTargetAuthentication() throws Exception
{
+ ConnectAttempt success = successfulAttempt();
+ KeyPair keyPair = mock(KeyPair.class);
+
when(sshClient.connect(any(HostConfigEntry.class))).thenReturn(success.connectFuture());
+ SshProtocol protocol = proxyProtocol(null, "target-password");
+ protocol.setProxyPrivateKey("proxy-private-key");
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class);
+ MockedStatic<SecurityUtils> securityUtils =
mockStatic(SecurityUtils.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+ securityUtils.when(() ->
SecurityUtils.loadKeyPairIdentities(any(), any(), any(), any()))
+ .thenAnswer(invocation -> {
+ FilePasswordProvider provider =
invocation.getArgument(3);
+ assertNull(provider.getPassword(null, null, 0));
+ return List.of(keyPair);
+ });
+
+ assertSame(success.clientSession(),
+ SshHelper.getConnectSession(protocol, TIMEOUT, false,
true));
+ }
+
+ verify(sshClient).addPublicKeyIdentity(keyPair);
+ verify(sshClient).removePublicKeyIdentity(keyPair);
+ verify(success.clientSession()).removePublicKeyIdentity(keyPair);
+ verify(sshClient, never()).addPasswordIdentity(any());
+ }
+
+ @Test
+ void connectionLockInterruptionIsPropagated() {
+ Thread.currentThread().interrupt();
+ try {
+ InterruptedIOException exception =
assertThrows(InterruptedIOException.class,
+ () -> SshHelper.getConnectSession(
+ host, "22", "root", "password", null, null,
TIMEOUT, true));
+
+ assertTrue(exception.getMessage().contains("SSH connection"));
+ assertTrue(exception.getCause() instanceof InterruptedException);
+ assertTrue(Thread.currentThread().isInterrupted());
+ } finally {
+ Thread.interrupted();
+ }
+ }
+
+ @Test
+ void interruptedAuthenticationForcesSessionCloseAndPreservesInterrupt()
throws Exception {
+ ConnectAttempt attempt = attempt();
+ InterruptedIOException interrupted = new
InterruptedIOException("authentication interrupted");
+ when(attempt.authFuture().verify(TIMEOUT,
TimeUnit.MILLISECONDS)).thenThrow(interrupted);
+ when(sshClient.connect("root", host,
22)).thenReturn(attempt.connectFuture());
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+ try {
+ assertSame(interrupted,
assertThrows(InterruptedIOException.class,
+ () -> SshHelper.getConnectSession(
+ host, "22", "root", "password", null, null,
TIMEOUT, false)));
+ assertTrue(Thread.currentThread().isInterrupted());
+ verify(attempt.clientSession()).close(true);
+ verify(attempt.clientSession(), never()).close();
+ } finally {
+ Thread.interrupted();
+ }
+ }
+ }
+
+ @Test
+ void closeFailureDoesNotHideAuthenticationFailure() throws Exception {
+ ConnectAttempt attempt = attempt();
+ SshException authenticationFailure = new SshException(
+ SshConstants.SSH2_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE,
"auth failed");
+ InterruptedIOException closeFailure = new
InterruptedIOException("close interrupted");
+ when(attempt.authFuture().verify(TIMEOUT,
TimeUnit.MILLISECONDS)).thenThrow(authenticationFailure);
+ doThrow(closeFailure).when(attempt.clientSession()).close();
+ when(sshClient.connect("root", host,
22)).thenReturn(attempt.connectFuture());
+
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+ try {
+ SshException thrown = assertThrows(SshException.class,
+ () -> SshHelper.getConnectSession(
+ host, "22", "root", "password", null, null,
TIMEOUT, false));
+ assertSame(authenticationFailure, thrown);
+ assertEquals(1, thrown.getSuppressed().length);
+ assertSame(closeFailure, thrown.getSuppressed()[0]);
+ assertTrue(Thread.currentThread().isInterrupted());
+ } finally {
+ Thread.interrupted();
+ }
+ }
+ }
+
+ private ClientSession connectWithThreadLocalClient(String password) throws
Exception {
+ return connectWithThreadLocalClient(password, true);
+ }
+
+ private ClientSession connectWithThreadLocalClient(String password,
boolean reuseConnection) throws Exception {
+ try (MockedStatic<CommonSshClient> commonSshClient =
mockStatic(CommonSshClient.class)) {
+
commonSshClient.when(CommonSshClient::getSshClient).thenReturn(sshClient);
+ return SshHelper.getConnectSession(
+ host, "22", "root", password, null, null, TIMEOUT,
reuseConnection);
+ }
+ }
+
+ private SshProtocol protocol(String password, String privateKey) {
+ return SshProtocol.builder()
+ .host(host)
+ .port("22")
+ .username("root")
+ .password(password)
+ .privateKey(privateKey)
+ .build();
+ }
+
+ private SshProtocol proxyProtocol(String proxyPassword, String
targetPassword) {
+ return SshProtocol.builder()
+ .host(host)
+ .port("22")
+ .username("root")
+ .password(targetPassword)
+ .proxyHost("proxy.example.com")
+ .proxyPort("2222")
+ .proxyUsername("proxy-user")
+ .proxyPassword(proxyPassword)
+ .build();
+ }
+
+ private ConnectAttempt failedAttempt(String message) throws IOException {
+ ConnectAttempt attempt = attempt();
+ when(attempt.authFuture().verify(TIMEOUT, TimeUnit.MILLISECONDS))
+ .thenThrow(new SshException(
+
SshConstants.SSH2_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE, message));
+ return attempt;
+ }
+
+ private ConnectAttempt successfulAttempt() throws IOException {
+ ConnectAttempt attempt = attempt();
+ when(attempt.authFuture().verify(TIMEOUT, TimeUnit.MILLISECONDS))
+ .thenReturn(attempt.authFuture());
+ return attempt;
+ }
+
+ private ConnectAttempt attempt() throws IOException {
+ ConnectFuture connectFuture = mock(ConnectFuture.class);
+ ClientSession clientSession = mock(ClientSession.class);
+ AuthFuture authFuture = mock(AuthFuture.class);
+ when(connectFuture.verify(TIMEOUT,
TimeUnit.MILLISECONDS)).thenReturn(connectFuture);
+ when(connectFuture.getSession()).thenReturn(clientSession);
+ when(clientSession.auth()).thenReturn(authFuture);
+ return new ConnectAttempt(connectFuture, clientSession, authFuture);
+ }
+
+ private record ConnectAttempt(
+ ConnectFuture connectFuture, ClientSession clientSession,
AuthFuture authFuture) {
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
