From 41543402990b3c7ceae2b944bfadb72c9cb188c6 Mon Sep 17 00:00:00 2001
From: Ng Yat Yan <ng.yatyan@gmail.com>
Date: Mon, 21 Oct 2024 20:03:04 +0800
Subject: [PATCH] V1.4.3 Still display prompt after a command fails

---
 pom.xml                                       |  2 +-
 .../com/example/sshd/config/AppConfig.java    |  7 +++---
 .../sshd/core/EchoSessionListener.java        |  1 +
 .../example/sshd/service/ReplyService.java    | 23 +++++++++++--------
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/pom.xml b/pom.xml
index b307b2e..fa1c157 100644
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>com.example.sshd</groupId>
 	<artifactId>echo-sshd-server</artifactId>
-	<version>1.4.2</version>
+	<version>1.4.3</version>
 	<name>ECHO SSH SERVER</name>
 	<description>Learning Apache Mina SSHD library</description>
 	<parent>
diff --git a/src/main/java/com/example/sshd/config/AppConfig.java b/src/main/java/com/example/sshd/config/AppConfig.java
index 42cf335..be18a88 100644
--- a/src/main/java/com/example/sshd/config/AppConfig.java
+++ b/src/main/java/com/example/sshd/config/AppConfig.java
@@ -1,8 +1,7 @@
 package com.example.sshd.config;
 
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.hc.client5.http.impl.async.CloseableHttpAsyncClient;
 import org.apache.hc.client5.http.impl.async.HttpAsyncClients;
@@ -19,13 +18,13 @@ public class AppConfig {
 	@Bean
 	@Scope(ConfigurableBeanFactory.SCOPE_SINGLETON)
 	public Map<String, Session> remoteSessionMapping() {
-		return Collections.synchronizedMap(new HashMap<>());
+		return new ConcurrentHashMap<>();
 	}
 
 	@Bean
 	@Scope(ConfigurableBeanFactory.SCOPE_SINGLETON)
 	public Map<String, String> ipInfoMapping() {
-		return Collections.synchronizedMap(new HashMap<>());
+		return new ConcurrentHashMap<>();
 	}
 
 	@Bean
diff --git a/src/main/java/com/example/sshd/core/EchoSessionListener.java b/src/main/java/com/example/sshd/core/EchoSessionListener.java
index f842e9a..ab8bb13 100644
--- a/src/main/java/com/example/sshd/core/EchoSessionListener.java
+++ b/src/main/java/com/example/sshd/core/EchoSessionListener.java
@@ -110,6 +110,7 @@ public class EchoSessionListener implements SessionListener {
 			String remoteIpAddress = remoteAddress.getAddress().getHostAddress();
 			logger.info("removing session: {} -> {}", remoteIpAddress, remoteSessionMapping.get(remoteIpAddress));
 			remoteSessionMapping.remove(remoteIpAddress);
+			ipInfoMapping.remove(remoteIpAddress);
 		}
 	}
 
diff --git a/src/main/java/com/example/sshd/service/ReplyService.java b/src/main/java/com/example/sshd/service/ReplyService.java
index a925caf..9aea589 100644
--- a/src/main/java/com/example/sshd/service/ReplyService.java
+++ b/src/main/java/com/example/sshd/service/ReplyService.java
@@ -11,6 +11,7 @@ import org.apache.commons.codec.binary.Base64;
 import org.apache.commons.codec.digest.DigestUtils;
 import org.apache.commons.exec.CommandLine;
 import org.apache.commons.exec.DefaultExecutor;
+import org.apache.commons.exec.ExecuteException;
 import org.apache.commons.exec.PumpStreamHandler;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
@@ -19,9 +20,9 @@ import org.apache.sshd.server.session.ServerSession;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.stereotype.Component;
+import org.springframework.stereotype.Service;
 
-@Component
+@Service
 public class ReplyService {
 
 	private static final Logger logger = LoggerFactory.getLogger(ReplyService.class);
@@ -90,14 +91,18 @@ public class ReplyService {
 						.replace("\\t", "\t");
 				if (reply.isEmpty()) {
 					logger.info("[{}] Execute cmd for real: {} ({})", cmdHash, command.trim(), o.get());
-					CommandLine cmdLine = CommandLine.parse(command.trim());
-					DefaultExecutor executor = DefaultExecutor.builder().get();
 					ByteArrayOutputStream tempOut = new ByteArrayOutputStream();
-					PumpStreamHandler streamHandler = new PumpStreamHandler(tempOut);
-					executor.setStreamHandler(streamHandler);
-					int exitValue = executor.execute(cmdLine);
-					logger.info("[{}] Result: {} ({})", cmdHash, command.trim(), exitValue);
-					reply = new String(tempOut.toByteArray()).replace("\n", "\r\n");
+					try {
+						CommandLine cmdLine = CommandLine.parse(command.trim());
+						DefaultExecutor executor = DefaultExecutor.builder().get();
+						PumpStreamHandler streamHandler = new PumpStreamHandler(tempOut);
+						executor.setStreamHandler(streamHandler);
+						int exitValue = executor.execute(cmdLine);
+						logger.info("[{}] Result: {} ({})", cmdHash, command.trim(), exitValue);
+						reply = new String(tempOut.toByteArray()).replace("\n", "\r\n");
+					} catch (ExecuteException e) {
+						logger.info("[{}] Execute cmd failed: {} ({})", cmdHash, command.trim(), o.get(), e);
+					}
 					out.write(String.format("\r\n%s\r\n%s", reply, prompt).getBytes());
 				} else {
 					logger.info("[{}] Known pattern detected: {} ({})", cmdHash, command.trim(), o.get());