Merge pull request #150 from pomerium/kralicky/ssh-graceful-shutdown

This PR includes refactors to channel lifetime management, and new APIs to allow for graceful interrupt handling.

Using a new ChannelControlAction type, the management server can preconfigure a ChannelData message to be sent to the downstream client in the event that the connection is interrupted. Currently the API is limited in what it allows the server to configure, but the underlying mechanism is much more flexible should we need to support more complex logic in the future. 

One major limitation of the existing code here was that it would treat a management server grpc disconnection as a connection level error, which would send an immediate DisconnectMsg to the downstream, which ignores all other channel state and forcibly ends the whole connection. Sending channel messages on an open channel, then relying on the Disconnect message to clean things up is rather unreliable for several reasons. Instead, we now are able to force a proper channel close handshake for any/all active channels, which will guarantee delivery of the configured interrupt channel data messages (as long as Envoy isn't killed during this, of course). This also allows us to handle server-wide errors, _and_ hook into the builtin envoy drain system to force graceful shutdown of all active connections if pomerium is shutting down.

This is implemented using a new mechanism that allows Envoy to "preempt" an active channel, which kicks off the channel close handshake sequence. This sequence of messages is described in more detail in the code. There are several different states the pair of channels could be in when this happens, so the channel state tracking in ChannelIDManager was made more granular to be able to capture all the edge cases.

Some other smaller changes include:
- Simplified the Channel interface, and added a new `wire::ChannelMessage` type alias which contains only a subset of all message types. This is used in places where only channel-related messages would be received.
- The logic for dealing with grpc errors/disconnects in HandoffChannel has been improved, to work properly with the graceful shutdown system.
- Some connection service tests were split up for compile time reasons

Author: kralicky

WORKSPACE

@@ -54,6 +54,7 @@ http_archive(
         "//patches/envoy:0005-suppress-duplicate-wip-warnings.patch",
         "//patches/envoy:0006-coverage-format.patch",
         "//patches/envoy:0007-user-space-io-handle.patch",
+        "//patches/envoy:0008-fake-upstream.patch",
         "//patches/envoy:tmp-transport-socket-options.patch",
     ],
     sha256 = "bb111b2037e35d8732f12f003ccf82e0d09dfc8a8b7810e849eb081f36d50ddc",

api/extensions/filters/network/ssh/ssh.pb.go

@@ -343,10 +343,20 @@ message SSHChannelControlAction {
     AllowResponse            upstream_auth           = 4;
   }
 
+  message InterruptOptions {
+    // If set, Envoy will send a channel data message to the client just before disconnecting
+    // with the specified contents.
+    bytes send_channel_data = 1;
+  }
+
   oneof action {
     // HandOffUpstream instructs Envoy to end the "hijacked" internal stream, and connect the
     // downstream client to the real upstream server.
     HandOffUpstream hand_off = 2;
+
+    // Configures logic for handling cleanup or other finalization tasks if the connection
+    // is interrupted.
+    InterruptOptions set_interrupt_options = 3;
   }
 }
 

api/extensions/filters/network/ssh/ssh.proto

@@ -0,0 +1,78 @@
+diff --git a/test/integration/fake_upstream.cc b/test/integration/fake_upstream.cc
+index 559484e147..fc2d34f653 100644
+--- a/test/integration/fake_upstream.cc
++++ b/test/integration/fake_upstream.cc
+@@ -681,7 +681,8 @@ FakeUpstream::FakeUpstream(Network::DownstreamTransportSocketFactoryPtr&& transp
+       http3_options_(config.http3_options_), quic_options_(config.quic_options_),
+       socket_(Network::SocketSharedPtr(listen_socket.release())),
+       api_(Api::createApiForTest(stats_store_)), time_system_(config.time_system_),
+-      dispatcher_(api_->allocateDispatcher("fake_upstream")),
++      dispatcher_(api_->allocateDispatcher(config.dispatcher_name_ != "" ? config.dispatcher_name_
++                                                                         : "fake_upstream")),
+       handler_(new Server::ConnectionHandlerImpl(*dispatcher_, 0)), config_(config),
+       read_disable_on_new_connection_(true), enable_half_close_(config.enable_half_close_),
+       listener_(*this, http_type_ == Http::CodecType::HTTP3),
+diff --git a/test/integration/fake_upstream.h b/test/integration/fake_upstream.h
+index 42799230cc..382c147ffe 100644
+--- a/test/integration/fake_upstream.h
++++ b/test/integration/fake_upstream.h
+@@ -716,6 +716,7 @@ struct FakeUpstreamConfig {
+   uint32_t max_request_headers_count_ = Http::DEFAULT_MAX_HEADERS_COUNT;
+   envoy::config::core::v3::HttpProtocolOptions::HeadersWithUnderscoresAction
+       headers_with_underscores_action_ = envoy::config::core::v3::HttpProtocolOptions::ALLOW;
++  std::string dispatcher_name_;
+ };
+ 
+ /**
+@@ -765,6 +766,10 @@ public:
+   ABSL_MUST_USE_RESULT
+   testing::AssertionResult assertPendingConnectionsEmpty();
+ 
++  bool hasNewConnections() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_) {
++    return !new_connections_.empty();
++  }
++
+   ABSL_MUST_USE_RESULT
+   testing::AssertionResult
+   waitForRawConnection(FakeRawConnectionPtr& connection,
+@@ -844,9 +849,13 @@ public:
+   const envoy::config::core::v3::Http3ProtocolOptions& http3Options() { return http3_options_; }
+ 
+   Event::DispatcherPtr& dispatcher() { return dispatcher_; }
+-  absl::Mutex& lock() { return lock_; }
++  absl::Mutex& lock() ABSL_LOCK_RETURNED(lock_) { return lock_; }
+ 
+   void runOnDispatcherThread(std::function<void()> cb);
++  AssertionResult
++  runOnDispatcherThreadAndWait(std::function<AssertionResult()> cb,
++                               std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
++  SharedConnectionWrapper& consumeConnection() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_);
+ 
+ protected:
+   const FakeUpstreamConfig& config() const { return config_; }
+@@ -993,11 +1002,7 @@ private:
+   };
+ 
+   void threadRoutine();
+-  SharedConnectionWrapper& consumeConnection() ABSL_EXCLUSIVE_LOCKS_REQUIRED(lock_);
+   Network::FilterStatus onRecvDatagram(Network::UdpRecvData& data);
+-  AssertionResult
+-  runOnDispatcherThreadAndWait(std::function<AssertionResult()> cb,
+-                               std::chrono::milliseconds timeout = TestUtility::DefaultTimeout);
+ 
+   const envoy::config::core::v3::Http2ProtocolOptions http2_options_;
+   const envoy::config::core::v3::Http3ProtocolOptions http3_options_;
+diff --git a/test/integration/server.cc b/test/integration/server.cc
+index 3a75dba9d0..cc5239a4e2 100644
+--- a/test/integration/server.cc
++++ b/test/integration/server.cc
+@@ -282,6 +282,9 @@ void IntegrationTestServerImpl::createAndRunEnvoyServer(
+     server.run();
+   }
+   server_gone_.Notify();
++  server_ = nullptr;
++  admin_address_ = nullptr;
++  stat_store_ = nullptr;
+ }
+ 
+ IntegrationTestServerImpl::~IntegrationTestServerImpl() {

patches/envoy/0008-fake-upstream.patch

@@ -2,6 +2,7 @@
 
 #include "source/common/optref.h"
 #include <source_location>
+#include <type_traits>
 
 // Supplemental type traits
 
@@ -187,6 +188,21 @@ constexpr bool all_types_equal = all_types_equal_to<Rest...>;
 template <typename T, typename... Ts>
 constexpr bool contains_type = index_of_type<T, Ts...>::found;
 
+template <typename T, typename U>
+struct strict_subset : std::false_type {};
+
+template <typename... Ts, typename... Us>
+  requires (sizeof...(Us) > 0 &&
+            sizeof...(Us) < sizeof...(Ts) &&
+            (contains_type<Us, Ts...> && ...))
+struct strict_subset<std::tuple<Ts...>, std::tuple<Us...>>
+    : std::true_type {};
+
+// strict_subset_v is true if T and U are tuples, and the types in U are a strict subset of the
+// types in T, otherwise false. The tuples should not contain duplicate types.
+template <typename T, typename U>
+constexpr bool strict_subset_v = strict_subset<T, U>::value;
+
 // is_vector<T> is true if T is a vector of any type, otherwise false.
 template <typename T>
 struct is_vector : std::false_type {};

source/common/type_traits.h

@@ -10,6 +10,7 @@ licenses(["notice"])
 envoy_cc_library(
     name = "pomerium_ssh",
     srcs = [
+        "channel.cc",
         "client_transport.cc",
         "config.cc",
         "extension_ping.cc",

source/extensions/filters/network/ssh/BUILD

@@ -0,0 +1,43 @@
+#include "source/extensions/filters/network/ssh/channel.h"
+
+namespace Envoy::Extensions::NetworkFilters::GenericProxy::Codec {
+
+Channel::~Channel() {
+  if (callbacks_ != nullptr) {
+    callbacks_->cleanup();
+  }
+};
+
+absl::Status Channel::setChannelCallbacks(ChannelCallbacks& callbacks) {
+  callbacks_ = &callbacks;
+  return absl::OkStatus();
+}
+
+absl::Status PassthroughChannel::readMessage(wire::ChannelMessage&& msg) {
+  return callbacks_->sendMessageRemote(std::move(msg));
+}
+
+absl::Status ForceCloseChannel::readMessage(wire::ChannelMessage&& msg) {
+  return msg.visit(
+    [&](wire::ChannelOpenConfirmationMsg&) {
+      ENVOY_LOG(debug, "channel {}: closing due to peer preemption", callbacks_->channelId());
+      callbacks_->sendMessageLocal(wire::ChannelCloseMsg{
+        .recipient_channel = callbacks_->channelId(),
+      });
+      return absl::OkStatus();
+    },
+    [&](wire::ChannelOpenFailureMsg&) {
+      return absl::OkStatus();
+    },
+    [](wire::ChannelCloseMsg&) {
+      return absl::OkStatus();
+    },
+    [&](auto& msg) {
+      // Ignore any messages received before the reply to our ChannelClose request.
+      // Note: The only way to get here is after receiving a ChannelOpenConfirmation.
+      ENVOY_LOG(debug, "channel {}: dropping message: {}", callbacks_->channelId(), msg.msg_type());
+      return absl::OkStatus();
+    });
+}
+
+} // namespace Envoy::Extensions::NetworkFilters::GenericProxy::Codec

source/extensions/filters/network/ssh/channel.cc

@@ -1,8 +1,10 @@
 #pragma once
 
+#include "source/extensions/filters/network/ssh/transport.h"
 #include "source/extensions/filters/network/ssh/wire/messages.h"
 #pragma clang unsafe_buffer_usage begin
 #include "envoy/stats/scope.h"
+#include "envoy/common/callback.h"
 #include "api/extensions/filters/network/ssh/ssh.pb.h"
 #pragma clang unsafe_buffer_usage end
 
@@ -41,11 +43,34 @@ class ChannelCallbacks {
   // for the channel.
   virtual void setStatsProvider(ChannelStatsProvider& stats_provider) PURE;
 
+  // Adds an interrupt callback to be invoked before a channel is closed to perform graceful
+  // shutdown in the event of an unexpected disconnect or other non-connection-fatal issue.
+  // These callbacks can be invoked using runInterruptCallbacks(). The returned handle can be
+  // deleted to remove the callback.
+  [[nodiscard]]
+  virtual Common::CallbackHandlePtr addInterruptCallback(std::function<void(absl::Status, TransportCallbacks&)> cb) PURE;
+
+  // Invokes all previously added interrupt callbacks, then clears the interrupt callback list.
+  // Deleting a callback handle obtained from addInterruptCallbacks after calling this function
+  // is a no-op; it is safe to let the callback handles go out of scope normally.
+  // This function can be invoked manually from a Channel implementation. It may also be invoked
+  // by the ConnectionService itself. Either way, any added callbacks are only invoked once.
+  virtual void runInterruptCallbacks(absl::Status err) PURE;
+
+  // Terminates the connection with an error. This will send an immediate Disconnect message.
+  virtual void terminate(absl::Status err) PURE;
+
 private:
   friend class Channel;
   virtual void cleanup() PURE;
 };
 
+class ChannelEventCallbacks {
+public:
+  virtual ~ChannelEventCallbacks() = default;
+  virtual void sendChannelEvent(const pomerium::extensions::ssh::ChannelEvent& ev) PURE;
+};
+
 // Channel handles the read path for a single peer (upstream or downstream) for messages on a
 // SSH channel. For channels known to both the upstream server and downstream client, two Channel
 // objects will exist: one managed by the upstream ConnectionService, and one by the downstream
@@ -57,27 +82,12 @@ class ChannelCallbacks {
 // upstream.
 class Channel {
 public:
-  virtual ~Channel() {
-    if (callbacks_ != nullptr) {
-      callbacks_->cleanup();
-    }
-  };
-  virtual absl::Status setChannelCallbacks(ChannelCallbacks& callbacks) {
-    callbacks_ = &callbacks;
-    return absl::OkStatus();
-  }
+  virtual ~Channel();
+  virtual absl::Status setChannelCallbacks(ChannelCallbacks& callbacks);
 
   // Handles a channel message (see concept ChannelMsg) read from the local peer, to be sent to
   // the remote peer.
-  virtual absl::Status readMessage(wire::Message&& msg) PURE;
-
-  // Called when the channel is successfully opened. ChannelOpenConfirmation messages are only
-  // sent here, not to readMessage().
-  virtual absl::Status onChannelOpened(wire::ChannelOpenConfirmationMsg&&) PURE;
-
-  // Called when the channel failed to open. ChannelOpenFailure messages are only sent here,
-  // not to readMessage().
-  virtual absl::Status onChannelOpenFailed(wire::ChannelOpenFailureMsg&&) PURE;
+  virtual absl::Status readMessage(wire::ChannelMessage&& msg) PURE;
 
 protected:
   ChannelCallbacks* callbacks_{};
@@ -87,23 +97,14 @@ class PassthroughChannel : public Channel {
 public:
   PassthroughChannel() = default;
 
-  absl::Status readMessage(wire::Message&& msg) override {
-    return callbacks_->sendMessageRemote(std::move(msg));
-  }
-
-  absl::Status onChannelOpened(wire::ChannelOpenConfirmationMsg&& msg) override {
-    return callbacks_->sendMessageRemote(std::move(msg));
-  }
-
-  absl::Status onChannelOpenFailed(wire::ChannelOpenFailureMsg&& msg) override {
-    return callbacks_->sendMessageRemote(std::move(msg));
-  }
+  absl::Status readMessage(wire::ChannelMessage&& msg) override;
 };
 
-class ChannelEventCallbacks {
+class ForceCloseChannel : public Channel, public Logger::Loggable<Logger::Id::filter> {
 public:
-  virtual ~ChannelEventCallbacks() = default;
-  virtual void sendChannelEvent(const pomerium::extensions::ssh::ChannelEvent& ev) PURE;
+  ForceCloseChannel() = default;
+
+  absl::Status readMessage(wire::ChannelMessage&& msg) override;
 };
 
 } // namespace Envoy::Extensions::NetworkFilters::GenericProxy::Codec

source/extensions/filters/network/ssh/channel.h

@@ -274,34 +274,7 @@ class HandoffChannel : public Channel, public Logger::Loggable<Logger::Id::filte
       : info_(info),
         handoff_callbacks_(callbacks) {}
 
-  absl::Status onChannelOpened(wire::ChannelOpenConfirmationMsg&&) override {
-    if (info_.pty_info == nullptr) {
-      return absl::InvalidArgumentError("session is not interactive");
-    }
-
-    ENVOY_LOG(debug, "handoff started");
-    // 2: PTY open request
-    wire::ChannelRequestMsg channelReq{
-      .recipient_channel = callbacks_->channelId(),
-      .want_reply = true,
-      .request = wire::PtyReqChannelRequestMsg{
-        .term_env = info_.pty_info->term_env(),
-        .width_columns = info_.pty_info->width_columns(),
-        .height_rows = info_.pty_info->height_rows(),
-        .width_px = info_.pty_info->width_px(),
-        .height_px = info_.pty_info->height_px(),
-        .modes = info_.pty_info->modes(),
-      },
-    };
-    callbacks_->sendMessageLocal(std::move(channelReq));
-    return absl::OkStatus();
-  }
-  absl::Status onChannelOpenFailed(wire::ChannelOpenFailureMsg&& msg) override {
-    // this should end the connection
-    return absl::UnavailableError(*msg.description);
-  }
-
-  absl::Status readMessage(wire::Message&& msg) override {
+  absl::Status readMessage(wire::ChannelMessage&& msg) override {
     if (handoff_complete_) {
       return callbacks_->sendMessageRemote(std::move(msg));
     }
@@ -324,12 +297,42 @@ class HandoffChannel : public Channel, public Logger::Loggable<Logger::Id::filte
       [&](const wire::ChannelFailureMsg&) {
         return absl::InternalError("failed to open upstream tty");
       },
+      [&](const wire::ChannelOpenConfirmationMsg& msg) {
+        return onChannelOpened(msg);
+      },
+      [&](const wire::ChannelOpenFailureMsg& msg) {
+        // this should end the connection
+        return absl::UnavailableError(*msg.description);
+      },
       [](const auto& msg) {
         return absl::InternalError(fmt::format("invalid message received during handoff: {}", msg.msg_type()));
       });
   }
 
 private:
+  absl::Status onChannelOpened(const wire::ChannelOpenConfirmationMsg&) {
+    if (info_.pty_info == nullptr) {
+      return absl::InvalidArgumentError("session is not interactive");
+    }
+
+    ENVOY_LOG(debug, "handoff started");
+    // 2: PTY open request
+    wire::ChannelRequestMsg channelReq{
+      .recipient_channel = callbacks_->channelId(),
+      .want_reply = true,
+      .request = wire::PtyReqChannelRequestMsg{
+        .term_env = info_.pty_info->term_env(),
+        .width_columns = info_.pty_info->width_columns(),
+        .height_rows = info_.pty_info->height_rows(),
+        .width_px = info_.pty_info->width_px(),
+        .height_px = info_.pty_info->height_px(),
+        .modes = info_.pty_info->modes(),
+      },
+    };
+    callbacks_->sendMessageLocal(std::move(channelReq));
+    return absl::OkStatus();
+  }
+
   bool handoff_complete_{false};
   const HandoffInfo& info_;
   HandoffChannelCallbacks& handoff_callbacks_;