diff --git a/app/src/main/java/com/secrm/HostSessions.kt b/app/src/main/java/com/secrm/HostSessions.kt index b633e1f..06addf8 100644 --- a/app/src/main/java/com/secrm/HostSessions.kt +++ b/app/src/main/java/com/secrm/HostSessions.kt @@ -93,7 +93,8 @@ class HostSessions(private val app: SecureRemoteApp) { /** Records the presented key as the one to trust for [change]'s host. The next connect will pin to it. */ fun trustNewKey(change: KeyChange) { - app.connections.save(change.conn.copy(hostKey = change.presented)) + // On the latest record, not the copy taken at connect time, so an edit made since (such as an MCP flag turned off) is kept. + app.connections.update(change.conn.name) { it.copy(hostKey = change.presented) } keyChange.value = null } diff --git a/app/src/main/java/com/secrm/SecureRemoteApp.kt b/app/src/main/java/com/secrm/SecureRemoteApp.kt index c404b6a..68f4410 100644 --- a/app/src/main/java/com/secrm/SecureRemoteApp.kt +++ b/app/src/main/java/com/secrm/SecureRemoteApp.kt @@ -44,6 +44,8 @@ class SecureRemoteApp : Application() { lateinit var agents: AgentTokens; private set lateinit var snippets: SnippetStore; private set lateinit var jobs: JobStore; private set + /** MCP's pooled SSH sessions; closed, with their forwards, when MCP is stopped or a client is revoked. */ + lateinit var mcpSessions: com.secrm.core.ssh.tools.McpSessionPool; private set /** The connected hosts, each with its own shell and file browser. */ val sessions by lazy { HostSessions(this) } @@ -95,7 +97,7 @@ class SecureRemoteApp : Application() { McpServerService.noteActivity() } }) - SshTools.register(dispatcher, connections, newClient()) { + mcpSessions = SshTools.register(dispatcher, connections, newClient()) { snippets.list().map { NamedCommand(it.name, it.command) } } watchNetwork() diff --git a/app/src/main/java/com/secrm/data/JobStore.kt b/app/src/main/java/com/secrm/data/JobStore.kt index 0163f7d..b475195 100644 --- a/app/src/main/java/com/secrm/data/JobStore.kt +++ b/app/src/main/java/com/secrm/data/JobStore.kt @@ -36,6 +36,10 @@ class JobStore(private val store: EncryptedConfigStore) { @Synchronized fun delete(id: String) = write(list().filterNot { it.id == id }) + /** Jobs name their host, so a renamed host takes its jobs with it. */ + @Synchronized + fun renameConnection(from: String, to: String) = write(list().map { if (it.connection == from) it.copy(connection = to) else it }) + @Synchronized fun recordRun(id: String, at: Long, result: String) = write(list().map { if (it.id == id) it.copy(lastRun = at, lastResult = result.take(300)) else it diff --git a/app/src/main/java/com/secrm/data/SnippetStore.kt b/app/src/main/java/com/secrm/data/SnippetStore.kt index e54bc5d..9736009 100644 --- a/app/src/main/java/com/secrm/data/SnippetStore.kt +++ b/app/src/main/java/com/secrm/data/SnippetStore.kt @@ -9,7 +9,10 @@ import java.util.UUID @Serializable data class Snippet(val id: String, val name: String, val command: String) -/** Reusable commands for the terminal, kept in the encrypted config. Inserted, never run: the owner presses Enter. */ +/** + * Reusable commands, kept in the encrypted config. In the terminal they are inserted, not run: the owner presses + * Enter. MCP agents with exec access to a host can also run them there by name (ssh.run_snippet). + */ class SnippetStore(private val store: EncryptedConfigStore) { @Synchronized fun list(): List = store.get(KEY)?.let { diff --git a/app/src/main/java/com/secrm/plugin/RunCommandService.kt b/app/src/main/java/com/secrm/plugin/RunCommandService.kt index 48d0f22..a2a24e8 100644 --- a/app/src/main/java/com/secrm/plugin/RunCommandService.kt +++ b/app/src/main/java/com/secrm/plugin/RunCommandService.kt @@ -25,6 +25,7 @@ import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.cancel import kotlinx.coroutines.launch +import kotlinx.coroutines.withTimeout import java.util.UUID /** @@ -102,7 +103,10 @@ class RunCommandService : Service() { val auth = authFor(conn) ?: return failure("No saved key for ${conn.name}") secrets = Redact.secretsOf(auth) if (!ssh.connectSaved(conn, auth, app.connections)) return failure("Could not connect to ${conn.host}") - val r = ssh.execResult(s.command, live?.sink(red = false), live?.sink(red = true)) + // Capped at the hour the macro host is asked to wait, so a command that never exits doesn't run the service forever. + val r = withTimeout(PluginProtocol.REQUESTED_TIMEOUT_NEVER.toLong()) { + ssh.execResult(s.command, live?.sink(red = false), live?.sink(red = true)) + } PluginProtocol.RESULT_OK to Bundle().apply { putString(PluginProtocol.VAR_STDOUT, r.stdout) putString(PluginProtocol.VAR_STDERR, r.stderr) diff --git a/app/src/main/java/com/secrm/service/McpServerService.kt b/app/src/main/java/com/secrm/service/McpServerService.kt index e707862..cfefde3 100644 --- a/app/src/main/java/com/secrm/service/McpServerService.kt +++ b/app/src/main/java/com/secrm/service/McpServerService.kt @@ -164,6 +164,8 @@ class McpServerService : Service() { if (screenReceiverOn) { runCatching { unregisterReceiver(screenOff) }; screenReceiverOn = false } server?.stop(); server = null _running.value = false + // Otherwise an agent's port forwards would stay open after the owner turned MCP off. + (application as SecureRemoteApp).mcpSessions.closeAll() } override fun onDestroy() { stopServer(); scope.cancel(); super.onDestroy() } diff --git a/app/src/main/java/com/secrm/ui/dashboard/DashboardScreen.kt b/app/src/main/java/com/secrm/ui/dashboard/DashboardScreen.kt index 8507992..e24b08a 100644 --- a/app/src/main/java/com/secrm/ui/dashboard/DashboardScreen.kt +++ b/app/src/main/java/com/secrm/ui/dashboard/DashboardScreen.kt @@ -205,7 +205,17 @@ fun DashboardScreen( onDismiss = { showForm = false; editing = null }, onSave = { conn -> // A corrupt store throws from upsert; report it instead of crashing the app. - val saved = runCatching { app.connections.upsert(conn) } + val old = editing + val saved = runCatching { + // Saved by name: a taken name would silently replace that host and its trusted key. + require(old?.name == conn.name || connections.none { it.name == conn.name }) { "a host named ${conn.name} already exists" } + app.connections.upsert(conn) + if (old != null && old.name != conn.name) { + app.connections.delete(old.name) + app.sessions.close(old.name) + app.jobs.renameConnection(old.name, conn.name) + } + } .onFailure { Toast.makeText(activity, "Couldn't save host: ${it.message}", Toast.LENGTH_LONG).show() } .isSuccess if (saved) { diff --git a/app/src/main/java/com/secrm/ui/server/ServerControlScreen.kt b/app/src/main/java/com/secrm/ui/server/ServerControlScreen.kt index b3d7001..7b28520 100644 --- a/app/src/main/java/com/secrm/ui/server/ServerControlScreen.kt +++ b/app/src/main/java/com/secrm/ui/server/ServerControlScreen.kt @@ -334,7 +334,7 @@ private fun McpSection(app: SecureRemoteApp, onOpenAudit: () -> Unit) { Spacer(Modifier.height(Space.md)) Row(horizontalArrangement = Arrangement.spacedBy(Space.sm)) { PillButton("Pair Claude", PillStyle.Primary, Modifier.weight(1f)) { pairing = app.oauth.newPairingCode() } - PillButton("Disconnect Claude", PillStyle.Ghost, Modifier.weight(1f)) { app.oauth.revokeAll(); pairing = null } + PillButton("Disconnect Claude", PillStyle.Ghost, Modifier.weight(1f)) { app.oauth.revokeAll(); app.mcpSessions.closeAll(); pairing = null } } pairing?.let { Text("Pairing code (10 min, single use)", style = MaterialTheme.typography.bodySmall, color = MaterialTheme.colorScheme.onSurfaceVariant) @@ -523,6 +523,7 @@ private fun AgentTokensSection(app: SecureRemoteApp) { } if (!a.revoked) TextButton(onClick = { app.agents.revoke(a.id) + app.mcpSessions.closeAll() // sessions are shared, so this also ends other clients' forwards agents = app.agents.list() }) { Text("Revoke", color = MaterialTheme.colorScheme.error) } } diff --git a/app/src/main/java/com/secrm/ui/sftp/SftpController.kt b/app/src/main/java/com/secrm/ui/sftp/SftpController.kt index b40fc05..605d81e 100644 --- a/app/src/main/java/com/secrm/ui/sftp/SftpController.kt +++ b/app/src/main/java/com/secrm/ui/sftp/SftpController.kt @@ -302,10 +302,15 @@ class SftpController( } } - /** Remote name for an upload's partial file: the same source and size always map to the same name. */ + /** + * Remote name for an upload's partial file: the same source, size and modified time always map to the same name, + * so a file edited since the last attempt starts over instead of resuming onto old bytes. A provider that + * reports no modified time gives 0, and then only source and size are matched, as before. + */ private fun partialName(u: Uri, size: Long): String { + val modified = runCatching { DocumentFile.fromSingleUri(app, u)?.lastModified() }.getOrNull() ?: 0L val key = java.security.MessageDigest.getInstance("SHA-256") - .digest("$u|$size".toByteArray()) + .digest("$u|$size|$modified".toByteArray()) .take(8).joinToString("") { "%02x".format(it) } return ".secrm-$key.part" } @@ -338,7 +343,12 @@ class SftpController( for (c in src.listFiles()) { if (cancelRequested) return; putTree(c, remote, safeName(c.name ?: continue)) } } else { _state.update { it.copy(transfer = it.transfer?.copy(name = name, done = 0, total = src.length())) } - app.contentResolver.openInputStream(src.uri)!!.use { client.upload(it, remote, { cancelRequested }) { n -> progress(n) } } + try { + app.contentResolver.openInputStream(src.uri)!!.use { client.upload(it, remote, { cancelRequested }) { n -> progress(n) } } + } catch (t: Throwable) { + client.removeQuietly(remote) // no partial file left on the host + throw t + } } } diff --git a/app/src/main/java/com/secrm/ui/terminal/PortForwardSheet.kt b/app/src/main/java/com/secrm/ui/terminal/PortForwardSheet.kt index ba3a26b..112d266 100644 --- a/app/src/main/java/com/secrm/ui/terminal/PortForwardSheet.kt +++ b/app/src/main/java/com/secrm/ui/terminal/PortForwardSheet.kt @@ -76,7 +76,8 @@ fun PortForwardSheet(ssh: SshClientManager, onDismiss: () -> Unit) { lp == null || lp !in 1024..65535 -> error = "Local port must be 1024–65535 (no root on Android)." rp == null || rp !in 1..65535 -> error = "Remote port must be 1–65535." else -> scope.launch { - runCatching { ssh.addLocalForward(lp, remoteHost, rp) } + // The owner's forward is for any app on this phone (a browser, a database client), which cannot send a token. + runCatching { ssh.addLocalForward(lp, remoteHost, rp, requireToken = false) } .onSuccess { error = null; localPort = ""; remotePort = "" } .onFailure { error = "Could not open: ${it.message ?: "port in use"}" } } diff --git a/core/server-http/src/main/kotlin/com/secrm/core/server/http/McpOAuth.kt b/core/server-http/src/main/kotlin/com/secrm/core/server/http/McpOAuth.kt index 0791338..46eee45 100644 --- a/core/server-http/src/main/kotlin/com/secrm/core/server/http/McpOAuth.kt +++ b/core/server-http/src/main/kotlin/com/secrm/core/server/http/McpOAuth.kt @@ -44,13 +44,16 @@ class McpOAuth( private var pairing: String? = null private var pairingExpires = 0L - private var failures = 0 - /** A new single-use pairing code, valid for 10 minutes or 5 wrong tries. */ + /** + * A new single-use pairing code, valid for 10 minutes. Guessing is bounded by the route's rate limit + * (about 100 tries in that time against 31^8 codes), not a wrong-try lockout, which anyone who has + * the public URL could trip to cancel the owner's code. + */ @Synchronized fun newPairingCode(): String = (1..8).map { ALPHABET[random.nextInt(ALPHABET.length)] }.joinToString("") - .also { pairing = it; pairingExpires = now() + 10 * MINUTE; failures = 0 } + .also { pairing = it; pairingExpires = now() + 10 * MINUTE } @Synchronized fun isValid(token: String) = tokens[sha(token)]?.let { it.kind == 'a' && it.expires > now() } == true @@ -64,10 +67,7 @@ class McpOAuth( fun authorize(code: String, challenge: String, redirect: String): String? { codes.values.removeAll { it.expires < now() } val p = pairing - if (p == null || now() > pairingExpires || !MessageDigest.isEqual(code.trim().uppercase().toByteArray(), p.toByteArray())) { - if (++failures >= 5) pairing = null - return null - } + if (p == null || now() > pairingExpires || !MessageDigest.isEqual(code.trim().uppercase().toByteArray(), p.toByteArray())) return null pairing = null return token().also { codes[it] = Code(challenge, redirect, now() + MINUTE) } } @@ -76,6 +76,7 @@ class McpOAuth( @Synchronized fun token(params: Map): Pair? { var family = token() + var refreshedFrom: String? = null when (params["grant_type"]) { "authorization_code" -> { val c = params["code"]?.let(codes::remove) ?: return null @@ -99,12 +100,20 @@ class McpOAuth( persist() return null } + // Only one pair may stay live: the one this retry gets. The first pair's access token goes, and its + // refresh token is marked used, so if both copies were real (one stolen) the other's next refresh + // revokes the family. + reissued.remove(key)?.let { (a, rf) -> + tokens.remove(a) + tokens[rf]?.let { tokens[rf] = it.copy(kind = 'c'); consumedAt.remove(rf) } + } } else { if (r.kind != 'r' || r.expires < now()) return null tokens[key] = r.copy(kind = 'c') consumedAt[key] = now() } family = r.family + refreshedFrom = key } else -> return null } @@ -112,8 +121,10 @@ class McpOAuth( val refresh = token() tokens.values.removeAll { it.expires < now() } consumedAt.values.removeAll { now() - it > REUSE_GRACE_MS } + reissued.keys.retainAll(consumedAt.keys) tokens[sha(access)] = Entry('a', now() + HOUR, family) tokens[sha(refresh)] = Entry('r', now() + 30 * DAY, family) + refreshedFrom?.let { reissued[it] = sha(access) to sha(refresh) } persist() return access to refresh } @@ -121,6 +132,9 @@ class McpOAuth( /** When each consumed refresh token was used. Not persisted: after a restart a reuse is treated as theft. */ private val consumedAt = HashMap() + /** Consumed refresh token hash to the (access, refresh) hashes its use issued, while a retry is still allowed. */ + private val reissued = HashMap>() + private fun persist() = save(tokens.entries.joinToString("\n") { "${it.key} ${it.value.kind} ${it.value.expires} ${it.value.family}" }) private fun token() = b64(ByteArray(32).also(random::nextBytes)) private fun sha(s: String) = sha256(s.toByteArray()).joinToString("") { "%02x".format(it) } diff --git a/core/server-http/src/main/kotlin/com/secrm/core/server/http/OAuthRoutes.kt b/core/server-http/src/main/kotlin/com/secrm/core/server/http/OAuthRoutes.kt index f32c1d1..c91a7bc 100644 --- a/core/server-http/src/main/kotlin/com/secrm/core/server/http/OAuthRoutes.kt +++ b/core/server-http/src/main/kotlin/com/secrm/core/server/http/OAuthRoutes.kt @@ -81,6 +81,7 @@ internal fun Route.oauthRoutes(oauth: McpOAuth, mcpPath: String) { val registerLimit = RateLimit(max = 20) // Per client_id, so one client's burst cannot slow the others' refreshes. val tokenLimit = ClientRateLimits(max = 60) + val pairingLimit = RateLimit(max = 10) // the only bound on guessing a pairing code (see McpOAuth.newPairingCode) get("/.well-known/oauth-protected-resource/{rest...}") { call.json { put("resource", call.base() + mcpPath) @@ -129,6 +130,7 @@ internal fun Route.oauthRoutes(oauth: McpOAuth, mcpPath: String) { } post("/authorize") { val p = call.boundedText()?.let { parseQueryString(it) } ?: return@post call.respond(HttpStatusCode.PayloadTooLarge) + if (!pairingLimit.allow()) return@post call.respondText("Too many tries. Wait a minute and try again.", status = HttpStatusCode.TooManyRequests) val redirect = p["redirect_uri"] if (!McpOAuth.allowedRedirect(redirect)) return@post call.respondText("Only Claude can connect", status = HttpStatusCode.BadRequest) // PKCE is mandatory: without a challenge the code could never be redeemed. diff --git a/core/server-http/src/test/kotlin/com/secrm/core/server/http/McpOAuthTest.kt b/core/server-http/src/test/kotlin/com/secrm/core/server/http/McpOAuthTest.kt index 5f889dc..ab29553 100644 --- a/core/server-http/src/test/kotlin/com/secrm/core/server/http/McpOAuthTest.kt +++ b/core/server-http/src/test/kotlin/com/secrm/core/server/http/McpOAuthTest.kt @@ -41,11 +41,11 @@ class McpOAuthTest { assertNull(pair(oauth, pairing)) } - @Test fun pairingLocksAfterFiveWrongGuesses() { + @Test fun wrongGuessesDoNotCancelThePairingCode() { val oauth = newOAuth() val pairing = oauth.newPairingCode() repeat(5) { assertNull(pair(oauth, "WRONG")) } - assertNull("correct code must fail once locked", pair(oauth, pairing)) + requireNotNull(pair(oauth, pairing)) } @Test fun pairingCodeIsCaseAndWhitespaceTolerant() { @@ -105,6 +105,17 @@ class McpOAuthTest { assertTrue("a benign retry keeps the session", oauth.isValid(access3)) } + @Test fun retryWithinGraceWindowLeavesOnlyOnePairLive() { + val oauth = newOAuth(clock = Clock()) + val (_, refresh1) = requireNotNull(exchange(oauth, requireNotNull(pair(oauth, oauth.newPairingCode())))) + val (access2, refresh2) = requireNotNull(oauth.token(mapOf("grant_type" to "refresh_token", "refresh_token" to refresh1))) + // A stolen copy of refresh1 is replayed within the grace window. + val (access3, _) = requireNotNull(oauth.token(mapOf("grant_type" to "refresh_token", "refresh_token" to refresh1))) + assertFalse("the first pair's access token is gone", oauth.isValid(access2)) + assertNull("its refresh token no longer works", oauth.token(mapOf("grant_type" to "refresh_token", "refresh_token" to refresh2))) + assertFalse("using it is treated as theft and revokes the family", oauth.isValid(access3)) + } + @Test fun omittedRedirectUriIsRejectedWhenAuthorizeSentOne() { val oauth = newOAuth() val code = requireNotNull(pair(oauth, oauth.newPairingCode())) diff --git a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/McpSessionPool.kt b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/McpSessionPool.kt index 7afa875..aa1804a 100644 --- a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/McpSessionPool.kt +++ b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/McpSessionPool.kt @@ -93,18 +93,26 @@ class McpSessionPool( return client } - /** Caller holds [lock]. A session with port forwards open is kept: closing it would close the forwards. */ + /** + * Caller holds [lock]. A session with port forwards open is kept: closing it would close the forwards. + * It is still closed once the owner turns off MCP read access (which forwards need) or deletes the connection. + */ private fun scheduleReap(e: Entry) { e.reaper?.cancel() e.reaper = scope.launch { delay(idleMs) lock.withLock { if (e.users != 0 || entries[e.conn.name] !== e) return@withLock - if (e.client.forwards.value.isEmpty()) evict(e) else scheduleReap(e) + if (e.client.forwards.value.isEmpty() || saved.find(e.conn.name)?.allowMcpRead != true) evict(e) else scheduleReap(e) } } } + /** Closes every shared session and the forwards on them: for when MCP is stopped or a client is revoked. */ + fun closeAll() { + scope.launch { lock.withLock { entries.values.toList().forEach(::evict) } } + } + /** Caller holds [lock]. */ private fun evict(e: Entry) { if (entries[e.conn.name] === e) entries.remove(e.conn.name) diff --git a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/Snippets.kt b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/Snippets.kt index b3eb940..6d2e82f 100644 --- a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/Snippets.kt +++ b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/Snippets.kt @@ -9,9 +9,44 @@ private val PLACEHOLDER = Regex("(? = PLACEHOLDER.findAll(command).map { it.groupValues[1] }.distinct().toList() -/** [command] with each placeholder replaced by its value, single-quoted so a value cannot add commands. */ -fun fillSnippet(command: String, values: Map): String = - PLACEHOLDER.replace(command) { m -> shellQuote(values[m.groupValues[1]].orEmpty()) } +private val NAME = Regex("[A-Za-z_][A-Za-z0-9_]*") + +/** + * [command] with each placeholder replaced by its value, quoted for where the placeholder stands, so a value + * cannot add commands: single-quoted outside quotes, escaped inside '…' or "…". Plain single-quoting would + * end an enclosing '…' (leaving the value unquoted) or leave `$(…)` live inside "…". + */ +fun fillSnippet(command: String, values: Map): String { + val out = StringBuilder() + var quote: Char? = null // the quote the text is inside, if any + var i = 0 + while (i < command.length) { + val c = command[i] + if (c == '{' && (i == 0 || command[i - 1] != '$')) { + val end = command.indexOf('}', i) + val name = if (end > i) command.substring(i + 1, end) else "" + if (NAME.matches(name)) { + val v = values[name].orEmpty() + out.append(when (quote) { + '\'' -> v.replace("'", "'\\''") + '"' -> v.replace(Regex("[\\\\\"$`]")) { "\\" + it.value } + else -> shellQuote(v) + }) + i = end + 1 + continue + } + } + when { + // A backslash outside '…' escapes the next character, which then opens or closes nothing. + c == '\\' && quote != '\'' && i + 1 < command.length -> { out.append(c).append(command[i + 1]); i += 2; continue } + c == '\'' && quote != '"' -> quote = if (quote == '\'') null else '\'' + c == '"' && quote != '\'' -> quote = if (quote == '"') null else '"' + } + out.append(c) + i++ + } + return out.toString() +} /** Single-quotes [s] for a POSIX shell. */ fun shellQuote(s: String): String = "'" + s.replace("'", "'\\''") + "'" diff --git a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/SshTools.kt b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/SshTools.kt index 09786ed..ac2326f 100644 --- a/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/SshTools.kt +++ b/core/ssh-tools/src/main/kotlin/com/secrm/core/ssh/tools/SshTools.kt @@ -15,6 +15,8 @@ object SshTools { /** The read-only tools run fixed, short commands, so they get a short leash. */ private const val READ_TIMEOUT_MS = 15_000L + /** Forwards MCP may keep open on one connection; each holds a listener and up to 16 client threads. */ + private const val MAX_FORWARDS = 8 private const val MAX_READ_BYTES = 1 shl 20 /** [len] bytes, less a trailing UTF-8 character that the cut would split. */ private fun utf8SafeLength(b: ByteArray, len: Int): Int { @@ -46,7 +48,7 @@ object SshTools { saved: SavedConnections, ssh: SshClientManager, commands: () -> List = { emptyList() }, - ) { + ): McpSessionPool { val pool = McpSessionPool(CoroutineScope(Dispatchers.IO + SupervisorJob()), saved) { // A separate client, so this neither replaces nor drops the app's live session. SshClientManager().apply { proxy = ssh.proxy } @@ -155,6 +157,8 @@ object SshTools { "Use a remote port from 1 to 65535.", "Remote port out of range: $port") // A forward needs only read access: it runs no command on the server. return onConnection(saved, pool, args, Need.READ, READ_TIMEOUT_MS) { client -> + if (client.forwards.value.size >= MAX_FORWARDS) return@onConnection ToolResult.error("too_many_forwards", + "Close one with ssh.close_forward first.", "This connection already has $MAX_FORWARDS forwards open.") val token = client.addLocalForward(local, host, port) ToolResult.ok(buildJsonObject { put("local", "127.0.0.1:$local") @@ -181,7 +185,8 @@ object SshTools { object : ToolHandler { override suspend fun call(session: Session, args: JsonObject): ToolResult { val local = args["local_port"]?.jsonPrimitive?.intOrNull ?: return missingArg("local_port") - return onConnection(saved, pool, args, Need.EXEC, READ_TIMEOUT_MS) { client -> + // The same access as opening one, so an agent can always close what it opened. + return onConnection(saved, pool, args, Need.READ, READ_TIMEOUT_MS) { client -> client.removeLocalForward(local) ToolResult.ok("Closed 127.0.0.1:$local") } @@ -347,6 +352,7 @@ object SshTools { } } ) + return pool } private const val CONNECTION_DOC = "Saved connection name, from ssh.list_connections." diff --git a/core/ssh-tools/src/test/kotlin/com/secrm/core/ssh/tools/SnippetsTest.kt b/core/ssh-tools/src/test/kotlin/com/secrm/core/ssh/tools/SnippetsTest.kt index 69e9c28..dd894d6 100644 --- a/core/ssh-tools/src/test/kotlin/com/secrm/core/ssh/tools/SnippetsTest.kt +++ b/core/ssh-tools/src/test/kotlin/com/secrm/core/ssh/tools/SnippetsTest.kt @@ -17,6 +17,20 @@ class SnippetsTest { assertEquals("echo 'it'\\''s'", fillSnippet("echo {msg}", mapOf("msg" to "it's"))) } + @Test fun valueInsideSingleQuotesStaysQuoted() { + // Plain quoting would give grep '''; rm -rf ~; ''' f, which runs the rm. + assertEquals("grep ''\\''; rm -rf ~; '\\''' f", fillSnippet("grep '{p}' f", mapOf("p" to "'; rm -rf ~; '"))) + assertEquals("grep 'a b' f", fillSnippet("grep '{p}' f", mapOf("p" to "a b"))) + } + + @Test fun valueInsideDoubleQuotesCannotExpand() { + assertEquals("echo \"\\\$(id) \\`id\\` \\\"\"", fillSnippet("echo \"{v}\"", mapOf("v" to "\$(id) `id` \""))) + } + + @Test fun escapedQuoteDoesNotChangeContext() { + assertEquals("echo \\' 'x'", fillSnippet("echo \\' {v}", mapOf("v" to "x"))) + } + @Test fun shellExpansionIsLeftAlone() { assertEquals("echo \${HOME}", fillSnippet("echo \${HOME}", emptyMap())) } diff --git a/core/ssh/src/main/kotlin/com/secrm/core/ssh/SshClientManager.kt b/core/ssh/src/main/kotlin/com/secrm/core/ssh/SshClientManager.kt index 48bad36..e3fa2eb 100644 --- a/core/ssh/src/main/kotlin/com/secrm/core/ssh/SshClientManager.kt +++ b/core/ssh/src/main/kotlin/com/secrm/core/ssh/SshClientManager.kt @@ -93,8 +93,8 @@ class SshClientManager { private val _forwards = MutableStateFlow>(emptyMap()) val forwards: StateFlow> = _forwards.asStateFlow() - /** A forward's listener on the phone, and the job that accepts its clients. */ - private class Listener(val server: ServerSocket, val job: Job) + /** A forward's listener on the phone. Closing it ends the thread that accepts its clients. */ + private class Listener(val server: ServerSocket) private companion object { const val MAX_FORWARD_CLIENTS = 16 @@ -104,11 +104,12 @@ class SshClientManager { /** * Forwards [localPort] on this phone's loopback to [host]:[port] beyond the server. Throws if the port is taken. - * Each client must send the returned token as its first line; a client that does not is dropped before any bytes are relayed. + * With [requireToken], each client must send the returned token as its first line; a client that does not is + * dropped before any bytes are relayed. Without it (the owner's own forwards) any client is relayed and null is returned. */ - suspend fun addLocalForward(localPort: Int, host: String, port: Int): String = withContext(Dispatchers.IO) { + suspend fun addLocalForward(localPort: Int, host: String, port: Int, requireToken: Boolean = true): String? = withContext(Dispatchers.IO) { val s = activeSession() ?: error("Not connected") - val token = ByteArray(16).also(random::nextBytes).joinToString("") { "%02x".format(it) } + val token = if (requireToken) ByteArray(16).also(random::nextBytes).joinToString("") { "%02x".format(it) } else null val server = ServerSocket() try { server.reuseAddress = true @@ -119,15 +120,17 @@ class SshClientManager { } // Each client holds a thread for up to its preamble timeout, so the number at once is capped. val clients = java.util.concurrent.atomic.AtomicInteger() - val job = scope.launch { - while (isActive) { + // Own threads, not the shared IO pool: accept and relay block for as long as the forward is open, and + // enough forwards would otherwise take every IO thread the app's SSH, SFTP and MCP work runs on. + kotlin.concurrent.thread(isDaemon = true, name = "ssh-forward-$localPort") { + while (true) { val client = try { server.accept() } catch (e: IOException) { break } if (clients.incrementAndGet() > MAX_FORWARD_CLIENTS) { clients.decrementAndGet() client.close() continue } - launch(Dispatchers.IO) { + kotlin.concurrent.thread(isDaemon = true, name = "ssh-forward-$localPort-client") { try { relay(s, client, host, port, token) } catch (e: Exception) { @@ -140,18 +143,20 @@ class SshClientManager { } } } - synchronized(listeners) { listeners[localPort] = Listener(server, job) } + synchronized(listeners) { listeners[localPort] = Listener(server) } _forwards.update { it + (localPort to "$host:$port") } token } /** Serves one client: checks the token line, then copies bytes both ways through a direct-tcpip channel. */ - private fun relay(s: Session, client: Socket, host: String, port: Int, token: String) { + private fun relay(s: Session, client: Socket, host: String, port: Int, token: String?) { client.use { c -> - c.soTimeout = 10_000 - val given = readLine(c.getInputStream(), token.length + 2) - if (!MessageDigest.isEqual(given.toByteArray(), token.toByteArray())) return - c.soTimeout = 0 + if (token != null) { + c.soTimeout = 10_000 + val given = readLine(c.getInputStream(), token.length + 2) + if (!MessageDigest.isEqual(given.toByteArray(), token.toByteArray())) return + c.soTimeout = 0 + } val ch = s.openChannel("direct-tcpip") as ChannelDirectTCPIP ch.setHost(host) ch.setPort(port) @@ -188,7 +193,6 @@ class SshClientManager { } private fun closeListener(l: Listener) { - l.job.cancel() runCatching { l.server.close() } } diff --git a/core/terminal/src/main/kotlin/com/secrm/core/terminal/TerminalBridge.kt b/core/terminal/src/main/kotlin/com/secrm/core/terminal/TerminalBridge.kt index b79a803..a3e2707 100644 --- a/core/terminal/src/main/kotlin/com/secrm/core/terminal/TerminalBridge.kt +++ b/core/terminal/src/main/kotlin/com/secrm/core/terminal/TerminalBridge.kt @@ -12,9 +12,10 @@ import java.io.InputStream import java.io.OutputStream import org.connectbot.terminal.Terminal import org.connectbot.terminal.TerminalEmulatorFactory -import java.util.concurrent.Executors import java.util.concurrent.LinkedBlockingQueue import java.util.concurrent.Semaphore +import java.util.concurrent.ThreadPoolExecutor +import java.util.concurrent.TimeUnit import kotlin.concurrent.thread /** Pipes an SSH shell channel into ConnectBot termlib's emulator. */ @@ -41,7 +42,8 @@ class TerminalBridge private constructor( /** Called once the output ends (the shell or command closed), from the drawing thread. */ @Volatile var onEnd: (() -> Unit)? = null - private val writer = Executors.newSingleThreadExecutor() + // One thread at most, in order, that exits when idle: a bridge is replaced on every reconnect and never shut down. + private val writer = ThreadPoolExecutor(0, 1, 30, TimeUnit.SECONDS, LinkedBlockingQueue()) @Volatile private var out: OutputStream? = null // Raw session output (last ~256 KB) and where it stood when Enter was last pressed, for copying