From e5eac347ce477311ec9619c6b6a2caece32eaee2 Mon Sep 17 00:00:00 2001
From: Pasqual Koschmieder <git@derklaro.dev>
Date: Wed, 6 Mar 2024 20:39:14 +0100
Subject: [PATCH] fix: update sshj to 0.38.0 and fix heartbeat issue (#1360)

### Motivation
SSHJ released a new version (0.38.0) in which they improved their
support for strict key exchanges. Unfortunately this also introduced a
bug on their end, as the hearbeater is writing packets to the server
which aren't allowed during key exchange (KEX_INIT must be the first
packet, however due to the hearbeater an IGNORE packet is the first
leading to the error `strict KEX violation: KEXINIT was not the first
packet`).

### Modification
Bump SSHJ to 0.38.0 and write a custom heartbeater which checks if the
key exchange is done before sending heartbeat packets.

### Result
SSHJ is updated to 0.38.0 and everything works as expected again.
---
 gradle/libs.versions.toml                     |  2 +-
 .../ActiveHeartbeatKeepAliveProvider.java     | 26 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index ba61bf4d12..9b7738aa92 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -39,7 +39,7 @@ mysqlConnector = "8.3.0"
 asm = "9.6"
 oshi = "6.4.13"
 vavr = "0.10.4"
-sshj = "0.37.0"
+sshj = "0.38.0"
 jjwt = "0.11.5"
 slf4j = "1.7.36"
 aerogel = "2.1.0"
diff --git a/modules/storage-sftp/src/main/java/eu/cloudnetservice/modules/sftp/sshj/ActiveHeartbeatKeepAliveProvider.java b/modules/storage-sftp/src/main/java/eu/cloudnetservice/modules/sftp/sshj/ActiveHeartbeatKeepAliveProvider.java
index 178175e593..4913bdc9a8 100644
--- a/modules/storage-sftp/src/main/java/eu/cloudnetservice/modules/sftp/sshj/ActiveHeartbeatKeepAliveProvider.java
+++ b/modules/storage-sftp/src/main/java/eu/cloudnetservice/modules/sftp/sshj/ActiveHeartbeatKeepAliveProvider.java
@@ -19,7 +19,10 @@
 import lombok.NonNull;
 import net.schmizz.keepalive.KeepAlive;
 import net.schmizz.keepalive.KeepAliveProvider;
+import net.schmizz.sshj.common.Message;
+import net.schmizz.sshj.common.SSHPacket;
 import net.schmizz.sshj.connection.ConnectionImpl;
+import net.schmizz.sshj.transport.TransportException;
 
 public final class ActiveHeartbeatKeepAliveProvider extends KeepAliveProvider {
 
@@ -28,8 +31,29 @@ public final class ActiveHeartbeatKeepAliveProvider extends KeepAliveProvider {
 
   @Override
   public @NonNull KeepAlive provide(@NonNull ConnectionImpl connection) {
-    var keepAlive = KeepAliveProvider.HEARTBEAT.provide(connection);
+    var keepAlive = new HeartbeatKeepAlive(connection);
     keepAlive.setKeepAliveInterval(HEARTBEAT_DELAY_SECONDS);
     return keepAlive;
   }
+
+  private static final class HeartbeatKeepAlive extends KeepAlive {
+
+    public HeartbeatKeepAlive(@NonNull ConnectionImpl conn) {
+      super(conn, "cloudnet-ssh-heartbeater");
+    }
+
+    @Override
+    protected void doKeepAlive() throws TransportException {
+      // when the server wants a strict key exchange, no other packets are allowed
+      // to be sent in that time interval (KEX_INIT must be the first packet). As it's
+      // very unlikely that a heartbeat is required during the timeframe anyway, we just
+      // don't execute the heartbeat until the key exchange is completed.
+      // Done by ensuring that the service is set. This means that the key
+      // exchange is done and the connection is up.
+      var transport = this.conn.getTransport();
+      if (this.conn.equals(transport.getService())) {
+        transport.write(new SSHPacket(Message.IGNORE));
+      }
+    }
+  }
 }