TLA+ safety model

.github/workflows/ci.yml

@@ -65,6 +65,22 @@ jobs:
           dependency-graph: disabled
           validate-wrappers: true
 
+      - name: Download TLA+ tools
+        if: matrix.java == '17'
+        env:
+          TLA2TOOLS_SHA256: 936a262061c914694dfd669a543be24573c45d5aa0ff20a8b96b23d01e050e88
+          TLA2TOOLS_URL: https://github.com/tlaplus/tlaplus/releases/download/v1.7.4/tla2tools.jar
+        run: |
+          tla2tools_jar="${RUNNER_TEMP}/tla2tools-1.7.4.jar"
+          curl --fail --location --proto '=https' --tlsv1.2 \
+            --output "${tla2tools_jar}" "${TLA2TOOLS_URL}"
+          echo "${TLA2TOOLS_SHA256}  ${tla2tools_jar}" | sha256sum --check
+          echo "TLA2TOOLS_JAR=${tla2tools_jar}" >> "${GITHUB_ENV}"
+
+      - name: Check SSH lifecycle TLA+ model
+        if: matrix.java == '17'
+        run: ./gradlew :sshlib:checkSshStateMachineTla
+
       - name: Configure Docker mirror
         run: |
           echo '{"registry-mirrors": ["https://mirror.gcr.io"]}' | sudo tee /etc/docker/daemon.json

sshlib/build.gradle.kts

@@ -67,6 +67,45 @@ tasks.test {
     )
 }
 
+val tlaModelDirectory = layout.projectDirectory.dir("src/test/resources/tla")
+val tlaStateDirectory = layout.buildDirectory.dir("tla/states")
+
+tasks.register<JavaExec>("generateSshStateMachineTla") {
+    group = "verification"
+    description = "Regenerates the TLA+ lifecycle model from SshClientStateMachine"
+    dependsOn(tasks.testClasses)
+    classpath = sourceSets.test.get().runtimeClasspath
+    mainClass.set("org.connectbot.sshlib.protocol.SshStateMachineTlaGenerator")
+    args(tlaModelDirectory.file("SshClientStateMachineGenerated.tla").asFile.absolutePath)
+}
+
+val tla2toolsJar = providers.gradleProperty("tla2toolsJar")
+    .orElse(providers.environmentVariable("TLA2TOOLS_JAR"))
+
+tasks.register<JavaExec>("checkSshStateMachineTla") {
+    group = "verification"
+    description = "Checks the generated SSH lifecycle model with TLC"
+    mainClass.set("tlc2.TLC")
+    workingDir(tlaModelDirectory)
+    args(
+        "-workers",
+        "1",
+        "-metadir",
+        tlaStateDirectory.get().asFile.absolutePath,
+        "-config",
+        "SshClientStateMachine.cfg",
+        "SshClientStateMachine.tla",
+    )
+    jvmArgs("-XX:+UseParallelGC")
+    doFirst {
+        val jarPath = tla2toolsJar.orNull
+            ?: throw GradleException(
+                "Set -Ptla2toolsJar=/path/to/tla2tools.jar or TLA2TOOLS_JAR to run TLC",
+            )
+        classpath = files(jarPath)
+    }
+}
+
 java {
     withSourcesJar()
     toolchain {

sshlib/src/main/kotlin/org/connectbot/sshlib/client/AgentChannel.kt

@@ -22,14 +22,18 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.launch
+import org.connectbot.sshlib.SshException
+import org.connectbot.sshlib.protocol.SshChannelEffect
+import org.connectbot.sshlib.protocol.SshChannelState
+import org.connectbot.sshlib.protocol.SshChannelStateMachine
 import org.slf4j.LoggerFactory
 
 internal class AgentChannel(
     private val handler: AgentProtocolHandler,
     private val connection: SshConnection,
     scope: CoroutineScope,
-    private val localChannelNumber: Int,
-    private var remoteChannelNumber: Int,
+    internal val localChannelNumber: Int,
+    internal val remoteChannelNumber: Int,
     private val maxPacketSize: Int,
     remoteWindowSizeInitial: Long,
     initialWindowSize: Int = 64 * 1024,
@@ -38,8 +42,7 @@ internal class AgentChannel(
         private val logger = LoggerFactory.getLogger(AgentChannel::class.java)
     }
 
-    private var _isOpen = true
-    private var closeSent = false
+    private val lifecycle = SshChannelStateMachine(SshChannelState.OPEN)
 
     private val window = LocalChannelWindow(initialWindowSize, remoteInitial = remoteWindowSizeInitial)
     private val windowAvailable = Channel<Unit>(Channel.CONFLATED)
@@ -54,50 +57,66 @@ internal class AgentChannel(
         }
     }
 
-    val isOpen: Boolean get() = _isOpen
+    val isOpen: Boolean get() = lifecycle.isOpen
 
     suspend fun handleData(data: ByteArray) {
-        if (!_isOpen) {
+        if (!lifecycle.receiveData {
+                window.consumeLocal(data.size)
+                logger.debug("Queueing ${data.size} bytes received on agent channel")
+                if (requests.trySend(data).isFailure) {
+                    window.releaseLocal(data.size)
+                    logger.warn("Discarding data received while agent channel is closing")
+                }
+            }
+        ) {
             logger.warn("Received data on closed agent channel")
             return
         }
-        window.consumeLocal(data.size)
-
-        logger.debug("Queueing ${data.size} bytes received on agent channel")
-        if (requests.trySend(data).isFailure) {
-            window.releaseLocal(data.size)
-            logger.warn("Discarding data received while agent channel is closing")
-        }
     }
 
-    fun onWindowAdjust(bytesToAdd: Long) {
-        window.adjustRemote(bytesToAdd)
-        logger.debug("Agent channel window adjust +$bytesToAdd, remote window now ${window.remoteRemaining}")
-        if (window.remoteRemaining > 0) {
-            windowAvailable.trySend(Unit)
+    suspend fun onWindowAdjust(bytesToAdd: Long) {
+        if (!lifecycle.receiveWindowAdjust {
+                window.adjustRemote(bytesToAdd)
+                logger.debug("Agent channel window adjust +$bytesToAdd, remote window now ${window.remoteRemaining}")
+                if (window.remoteRemaining > 0) windowAvailable.trySend(Unit)
+            }
+        ) {
+            throw SshException("Received window adjustment after CLOSE on agent channel $localChannelNumber")
         }
     }
 
-    fun onEof() {
-        logger.debug("Agent channel received EOF")
+    suspend fun onEof() {
+        if (!lifecycle.receiveEof { logger.debug("Agent channel received EOF") }) {
+            throw SshException("Received duplicate EOF or EOF after CLOSE on agent channel $localChannelNumber")
+        }
     }
 
     suspend fun onClose() {
-        logger.debug("Agent channel closed")
-        if (!closeSent) {
-            closeSent = true
-            try {
-                connection.sendChannelClose(remoteChannelNumber)
-            } catch (e: Exception) {
-                logger.debug("Failed to send CHANNEL_CLOSE reply", e)
+        lifecycle.receiveClose { transition ->
+            logger.debug("Agent channel closed")
+            if (SshChannelEffect.SEND_CLOSE in transition.effects) {
+                try {
+                    connection.sendChannelClose(remoteChannelNumber)
+                } catch (e: Exception) {
+                    logger.debug("Failed to send CHANNEL_CLOSE reply", e)
+                }
             }
+            requests.close()
+            requestWorker.cancelAndJoin()
+            windowAvailable.close()
         }
-        _isOpen = false
-        requests.close()
-        requestWorker.cancelAndJoin()
-        windowAvailable.close()
     }
 
+    suspend fun onDisconnected() {
+        lifecycle.disconnect {
+            requests.close()
+            requestWorker.cancelAndJoin()
+            windowAvailable.close()
+        }
+    }
+
+    suspend fun receiveRequest(action: suspend () -> Unit): Boolean = lifecycle.receiveRequest { action() }
+
     private suspend fun sendData(data: ByteArray) {
         var offset = 0
         while (offset < data.size) {
@@ -106,8 +125,13 @@ internal class AgentChannel(
             }
             val chunkSize = window.sendChunkSize(data.size - offset, maxPacketSize)
             val chunk = data.copyOfRange(offset, offset + chunkSize)
-            connection.sendChannelData(remoteChannelNumber, chunk)
-            window.consumeRemote(chunkSize)
+            if (!lifecycle.sendData {
+                    connection.sendChannelData(remoteChannelNumber, chunk)
+                    window.consumeRemote(chunkSize)
+                }
+            ) {
+                throw SshException("Cannot send data after EOF or CLOSE on agent channel $localChannelNumber")
+            }
             offset += chunkSize
         }
     }

sshlib/src/main/kotlin/org/connectbot/sshlib/client/ForwardingChannel.kt

@@ -21,6 +21,10 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.ReceiveChannel
 import kotlinx.coroutines.launch
+import org.connectbot.sshlib.SshException
+import org.connectbot.sshlib.protocol.SshChannelEffect
+import org.connectbot.sshlib.protocol.SshChannelState
+import org.connectbot.sshlib.protocol.SshChannelStateMachine
 import org.slf4j.LoggerFactory
 
 internal class ForwardingChannel(
@@ -31,13 +35,12 @@ internal class ForwardingChannel(
     private val maxPacketSize: Int,
     remoteWindowSizeInitial: Long,
     private val initialWindowSize: Int = 256 * 1024,
+    private val lifecycle: SshChannelStateMachine = SshChannelStateMachine(SshChannelState.OPEN),
 ) {
     companion object {
         private val logger = LoggerFactory.getLogger(ForwardingChannel::class.java)
     }
 
-    private var _isOpen = true
-    private var closeSent = false
     private val window = LocalChannelWindow(
         initialWindowSize,
         remoteInitial = remoteWindowSizeInitial,
@@ -59,45 +62,72 @@ internal class ForwardingChannel(
         }
     }
 
-    val isOpen: Boolean get() = _isOpen
+    val isOpen: Boolean get() = lifecycle.isOpen
 
     internal suspend fun onData(data: ByteArray) {
-        window.consumeLocal(data.size)
-        if (incomingIngress.trySend(data).isFailure) {
-            throw org.connectbot.sshlib.SshException("Received data for a closed forwarding stream")
+        if (!lifecycle.receiveData {
+                window.consumeLocal(data.size)
+                if (incomingIngress.trySend(data).isFailure) {
+                    throw org.connectbot.sshlib.SshException("Received data for a closed forwarding stream")
+                }
+            }
+        ) {
+            throw SshException("Received data after EOF or CLOSE on forwarding channel $localChannelNumber")
         }
     }
 
-    internal fun onWindowAdjust(bytesToAdd: Long) {
-        window.adjustRemote(bytesToAdd)
-        logger.debug("Forwarding channel window adjust +$bytesToAdd, remote window now ${window.remoteRemaining}")
-        if (window.remoteRemaining > 0) {
-            windowAvailable.trySend(Unit)
+    internal suspend fun onWindowAdjust(bytesToAdd: Long) {
+        if (!lifecycle.receiveWindowAdjust {
+                window.adjustRemote(bytesToAdd)
+                logger.debug("Forwarding channel window adjust +$bytesToAdd, remote window now ${window.remoteRemaining}")
+                if (window.remoteRemaining > 0) windowAvailable.trySend(Unit)
+            }
+        ) {
+            throw SshException("Received window adjustment after CLOSE on forwarding channel $localChannelNumber")
         }
     }
 
-    internal fun onEof() {
-        logger.debug("Forwarding channel $localChannelNumber received EOF")
-        incomingIngress.close()
+    internal suspend fun onEof() {
+        if (!lifecycle.receiveEof {
+                logger.debug("Forwarding channel $localChannelNumber received EOF")
+                incomingIngress.close()
+            }
+        ) {
+            throw SshException("Received duplicate EOF or EOF after CLOSE on forwarding channel $localChannelNumber")
+        }
     }
 
     internal suspend fun onClose() {
-        logger.debug("Forwarding channel $localChannelNumber closed")
-        if (!closeSent) {
-            closeSent = true
-            try {
-                connection.sendChannelClose(remoteChannelNumber)
-            } catch (e: Exception) {
-                logger.debug("Failed to send CHANNEL_CLOSE reply", e)
+        if (!lifecycle.receiveClose { transition ->
+                logger.debug("Forwarding channel $localChannelNumber closed")
+                if (SshChannelEffect.SEND_CLOSE in transition.effects) {
+                    try {
+                        connection.sendChannelClose(remoteChannelNumber)
+                    } catch (e: Exception) {
+                        logger.debug("Failed to send CHANNEL_CLOSE reply", e)
+                    }
+                }
+                incomingIngress.close()
+                incomingDeliveryJob.cancel()
+                _incomingData.close()
+                windowAvailable.close()
             }
+        ) {
+            throw SshException("Received duplicate CLOSE on forwarding channel $localChannelNumber")
+        }
+    }
+
+    internal suspend fun onDisconnected() {
+        lifecycle.disconnect {
+            incomingIngress.close()
+            incomingDeliveryJob.cancel()
+            _incomingData.close()
+            windowAvailable.close()
         }
-        _isOpen = false
-        incomingIngress.close()
-        incomingDeliveryJob.cancel()
-        _incomingData.close()
-        windowAvailable.close()
     }
 
+    internal suspend fun receiveRequest(action: suspend () -> Unit): Boolean = lifecycle.receiveRequest { action() }
+
     suspend fun sendData(data: ByteArray) {
         var offset = 0
         while (offset < data.size) {
@@ -106,23 +136,27 @@ internal class ForwardingChannel(
             }
             val chunkSize = window.sendChunkSize(data.size - offset, maxPacketSize)
             val chunk = data.copyOfRange(offset, offset + chunkSize)
-            connection.sendChannelData(remoteChannelNumber, chunk)
-            window.consumeRemote(chunkSize)
+            if (!lifecycle.sendData {
+                    connection.sendChannelData(remoteChannelNumber, chunk)
+                    window.consumeRemote(chunkSize)
+                }
+            ) {
+                throw SshException("Cannot send data after EOF or CLOSE on forwarding channel $localChannelNumber")
+            }
             offset += chunkSize
         }
     }
 
     suspend fun sendEof() {
-        connection.sendChannelEof(remoteChannelNumber)
+        lifecycle.sendEof { connection.sendChannelEof(remoteChannelNumber) }
     }
 
     suspend fun close() {
-        if (!_isOpen) return
-        closeSent = true
-        _isOpen = false
-        incomingIngress.close()
-        incomingDeliveryJob.cancel()
-        _incomingData.close()
-        connection.sendChannelClose(remoteChannelNumber)
+        lifecycle.sendClose {
+            incomingIngress.close()
+            incomingDeliveryJob.cancel()
+            _incomingData.close()
+            connection.sendChannelClose(remoteChannelNumber)
+        }
     }
 }