This is an automated email from the ASF dual-hosted git repository. yuluo-yx pushed a commit to branch 0711-yuluo/fix-ssh in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
commit 4968f75b27f22fa9815a346b281307158b95a9ee Author: yuluo-yx <[email protected]> AuthorDate: Sat Jul 11 19:25:47 2026 +0800 feat: fix ssh connection bug --- .../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]
