SshCommand: Handle RequestCancelledException for SSH commands

Catch RequestCancelledException in SshCommand and return proper error
message to the client.

RequestCancelledException is not being thrown yet, this is a preparation
to handle it correctly when we start detecting cancelled requests later
and aborting them by throwing this exception.

Signed-off-by: Edwin Kempin <ekempin@google.com>
Change-Id: I3fc1cee9f182348294ee8e22760d362b54238b9f
diff --git a/java/com/google/gerrit/sshd/SshCommand.java b/java/com/google/gerrit/sshd/SshCommand.java
index c94b25c..93c6c2c 100644
--- a/java/com/google/gerrit/sshd/SshCommand.java
+++ b/java/com/google/gerrit/sshd/SshCommand.java
@@ -19,6 +19,7 @@
 import com.google.gerrit.server.DynamicOptions;
 import com.google.gerrit.server.RequestInfo;
 import com.google.gerrit.server.RequestListener;
+import com.google.gerrit.server.cancellation.RequestCancelledException;
 import com.google.gerrit.server.config.GerritServerConfig;
 import com.google.gerrit.server.logging.PerformanceLogContext;
 import com.google.gerrit.server.logging.PerformanceLogger;
@@ -61,6 +62,12 @@
                   RequestInfo.builder(RequestInfo.RequestType.SSH, user, traceContext).build();
               requestListeners.runEach(l -> l.onRequest(requestInfo));
               SshCommand.this.run();
+            } catch (RequestCancelledException e) {
+              StringBuilder msg = new StringBuilder(e.formatCancellationReason());
+              if (e.getCancellationMessage().isPresent()) {
+                msg.append(String.format(" (%s)", e.getCancellationMessage().get()));
+              }
+              stderr.println(msg.toString());
             } finally {
               stdout.flush();
               stderr.flush();
diff --git a/javatests/com/google/gerrit/acceptance/ssh/SshCancellationIT.java b/javatests/com/google/gerrit/acceptance/ssh/SshCancellationIT.java
new file mode 100644
index 0000000..2cb9637
--- /dev/null
+++ b/javatests/com/google/gerrit/acceptance/ssh/SshCancellationIT.java
@@ -0,0 +1,102 @@
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed 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 com.google.gerrit.acceptance.ssh;
+
+import com.google.gerrit.acceptance.AbstractDaemonTest;
+import com.google.gerrit.acceptance.ExtensionRegistry;
+import com.google.gerrit.acceptance.ExtensionRegistry.Registration;
+import com.google.gerrit.acceptance.UseSsh;
+import com.google.gerrit.server.cancellation.RequestCancelledException;
+import com.google.gerrit.server.cancellation.RequestStateProvider;
+import com.google.gerrit.server.project.CreateProjectArgs;
+import com.google.gerrit.server.validators.ProjectCreationValidationListener;
+import com.google.gerrit.server.validators.ValidationException;
+import com.google.inject.Inject;
+import org.junit.Test;
+
+@UseSsh
+public class SshCancellationIT extends AbstractDaemonTest {
+  @Inject private ExtensionRegistry extensionRegistry;
+
+  @Test
+  public void handleClientDisconnected() throws Exception {
+    ProjectCreationValidationListener projectCreationListener =
+        new ProjectCreationValidationListener() {
+          @Override
+          public void validateNewProject(CreateProjectArgs args) throws ValidationException {
+            throw new RequestCancelledException(
+                RequestStateProvider.Reason.CLIENT_CLOSED_REQUEST, /* cancellationMessage= */ null);
+          }
+        };
+    try (Registration registration =
+        extensionRegistry.newRegistration().add(projectCreationListener)) {
+      adminSshSession.exec("gerrit create-project " + name("new"));
+      adminSshSession.assertFailure("Client Closed Request");
+    }
+  }
+
+  @Test
+  public void handleClientDeadlineExceeded() throws Exception {
+    ProjectCreationValidationListener projectCreationListener =
+        new ProjectCreationValidationListener() {
+          @Override
+          public void validateNewProject(CreateProjectArgs args) throws ValidationException {
+            throw new RequestCancelledException(
+                RequestStateProvider.Reason.CLIENT_PROVIDED_DEADLINE_EXCEEDED,
+                /* cancellationMessage= */ null);
+          }
+        };
+    try (Registration registration =
+        extensionRegistry.newRegistration().add(projectCreationListener)) {
+      adminSshSession.exec("gerrit create-project " + name("new"));
+      adminSshSession.assertFailure("Client Provided Deadline Exceeded");
+    }
+  }
+
+  @Test
+  public void handleServerDeadlineExceeded() throws Exception {
+    ProjectCreationValidationListener projectCreationListener =
+        new ProjectCreationValidationListener() {
+          @Override
+          public void validateNewProject(CreateProjectArgs args) throws ValidationException {
+            throw new RequestCancelledException(
+                RequestStateProvider.Reason.SERVER_DEADLINE_EXCEEDED,
+                /* cancellationMessage= */ null);
+          }
+        };
+    try (Registration registration =
+        extensionRegistry.newRegistration().add(projectCreationListener)) {
+      adminSshSession.exec("gerrit create-project " + name("new"));
+      adminSshSession.assertFailure("Server Deadline Exceeded");
+    }
+  }
+
+  @Test
+  public void handleRequestCancellationWithMessage() throws Exception {
+    ProjectCreationValidationListener projectCreationListener =
+        new ProjectCreationValidationListener() {
+          @Override
+          public void validateNewProject(CreateProjectArgs args) throws ValidationException {
+            throw new RequestCancelledException(
+                RequestStateProvider.Reason.SERVER_DEADLINE_EXCEEDED, "deadline = 10m");
+          }
+        };
+    try (Registration registration =
+        extensionRegistry.newRegistration().add(projectCreationListener)) {
+      adminSshSession.exec("gerrit create-project " + name("new"));
+      adminSshSession.assertFailure("Server Deadline Exceeded (deadline = 10m)");
+    }
+  }
+}