From 5b5bd1dcfa0c2fc250e079e1ebcd643b51f735eb Mon Sep 17 00:00:00 2001
From: Lyor Goldstein <lgoldstein@apache.org>
Date: Fri, 26 Feb 2021 06:54:43 +0200
Subject: [PATCH] [SSHD-1125] Added option to require immediate close of
 channel in command ExitCallback invocation

---
 .../org/apache/sshd/server/ExitCallback.java  | 27 +++++++++++++++++--
 .../sshd/server/channel/ChannelSession.java   | 10 +++----
 .../sshd/util/test/BogusExitCallback.java     | 12 ++++++---
 .../server/subsystem/sftp/SftpSubsystem.java  |  6 +++--
 4 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/sshd-core/src/main/java/org/apache/sshd/server/ExitCallback.java b/sshd-core/src/main/java/org/apache/sshd/server/ExitCallback.java
index dfa55be..eb1bca1 100644
--- a/sshd-core/src/main/java/org/apache/sshd/server/ExitCallback.java
+++ b/sshd-core/src/main/java/org/apache/sshd/server/ExitCallback.java
@@ -30,7 +30,17 @@ public interface ExitCallback {
      * @param exitValue the exit value
      */
     default void onExit(int exitValue) {
-        onExit(exitValue, "");
+        onExit(exitValue, false);
+    }
+
+    /**
+     * Informs the SSH server that the shell has exited
+     *
+     * @param exitValue        the exit value
+     * @param closeImmediately whether to also terminate the channel immediately or do a graceful close.
+     */
+    default void onExit(int exitValue, boolean closeImmediately) {
+        onExit(exitValue, "", closeImmediately);
     }
 
     /**
@@ -39,5 +49,18 @@ public interface ExitCallback {
      * @param exitValue   the exit value
      * @param exitMessage exit value description
      */
-    void onExit(int exitValue, String exitMessage);
+    default void onExit(int exitValue, String exitMessage) {
+        onExit(exitValue, exitMessage, false);
+    }
+
+    /**
+     *
+     * Informs the SSH client/server that the shell has exited
+     *
+     * @param exitValue        the exit value
+     * @param exitMessage      exit value description
+     * @param closeImmediately whether to also terminate the channel immediately or do a graceful close.
+     */
+    void onExit(int exitValue, String exitMessage, boolean closeImmediately);
+
 }
diff --git a/sshd-core/src/main/java/org/apache/sshd/server/channel/ChannelSession.java b/sshd-core/src/main/java/org/apache/sshd/server/channel/ChannelSession.java
index e54e0e4..484b4f0 100644
--- a/sshd-core/src/main/java/org/apache/sshd/server/channel/ChannelSession.java
+++ b/sshd-core/src/main/java/org/apache/sshd/server/channel/ChannelSession.java
@@ -684,9 +684,9 @@ public class ChannelSession extends AbstractServerChannel {
             tempBuffer = null;
             doWriteData(buffer.array(), buffer.rpos(), buffer.available());
         }
-        command.setExitCallback((exitValue, exitMessage) -> {
+        command.setExitCallback((exitValue, exitMessage, closeImmediately) -> {
             try {
-                closeShell(exitValue);
+                closeShell(exitValue, closeImmediately);
                 if (log.isDebugEnabled()) {
                     log.debug("onExit({}) code={} message='{}' shell closed", ChannelSession.this, exitValue, exitMessage);
                 }
@@ -798,9 +798,9 @@ public class ChannelSession extends AbstractServerChannel {
         return env;
     }
 
-    protected void closeShell(int exitValue) throws IOException {
+    protected void closeShell(int exitValue, boolean closeImmediately) throws IOException {
         if (log.isDebugEnabled()) {
-            log.debug("closeShell({}) exit code={}", this, exitValue);
+            log.debug("closeShell({}) exit code={}, immediate={}", this, exitValue, closeImmediately);
         }
 
         if (!isClosing()) {
@@ -810,7 +810,7 @@ public class ChannelSession extends AbstractServerChannel {
             sendEof();
             sendExitStatus(exitValue);
             commandExitFuture.setClosed();
-            close(false);
+            close(closeImmediately);
         } else {
             commandExitFuture.setClosed();
         }
diff --git a/sshd-core/src/test/java/org/apache/sshd/util/test/BogusExitCallback.java b/sshd-core/src/test/java/org/apache/sshd/util/test/BogusExitCallback.java
index a4f1ff6..9de602a 100644
--- a/sshd-core/src/test/java/org/apache/sshd/util/test/BogusExitCallback.java
+++ b/sshd-core/src/test/java/org/apache/sshd/util/test/BogusExitCallback.java
@@ -25,21 +25,23 @@ public class BogusExitCallback implements ExitCallback {
     private boolean exited;
     private int exitValue;
     private String exitMessage;
+    private boolean closeImmediately;
 
     public BogusExitCallback() {
         super();
     }
 
     @Override
-    public void onExit(int exitValue) {
-        onExit(exitValue, String.valueOf(exitValue));
+    public void onExit(int exitValue, boolean closeImmediately) {
+        onExit(exitValue, String.valueOf(exitValue), closeImmediately);
     }
 
     @Override
-    public void onExit(int exitValue, String exitMessage) {
+    public void onExit(int exitValue, String exitMessage, boolean closeImmediately) {
         this.exited = true;
         this.exitValue = exitValue;
         this.exitMessage = exitMessage;
+        this.closeImmediately = closeImmediately;
     }
 
     public boolean isExited() {
@@ -53,4 +55,8 @@ public class BogusExitCallback implements ExitCallback {
     public String getExitMessage() {
         return exitMessage;
     }
+
+    public boolean isCloseImmediately() {
+        return closeImmediately;
+    }
 }
diff --git a/sshd-sftp/src/main/java/org/apache/sshd/server/subsystem/sftp/SftpSubsystem.java b/sshd-sftp/src/main/java/org/apache/sshd/server/subsystem/sftp/SftpSubsystem.java
index 66a0ced..c18b55e 100644
--- a/sshd-sftp/src/main/java/org/apache/sshd/server/subsystem/sftp/SftpSubsystem.java
+++ b/sshd-sftp/src/main/java/org/apache/sshd/server/subsystem/sftp/SftpSubsystem.java
@@ -299,6 +299,7 @@ public class SftpSubsystem
 
     @Override
     public void run() {
+        int exitCode = 0;
         try {
             while (true) {
                 Buffer buffer = requests.take();
@@ -318,10 +319,11 @@ public class SftpSubsystem
                 if (log.isDebugEnabled()) {
                     log.debug("run(" + session + ") caught exception details", t);
                 }
+		exitCode = -1;
             }
         } finally {
             closeAllHandles();
-            callback.onExit(0);
+	    callback.onExit(exitCode, exitCode != 0);
         }
     }
 
-- 
2.27.0

