From 2e3a1ae389fce502359c7296c4762861febb2dbf Mon Sep 17 00:00:00 2001
From: Kevin Baensch <kb@ujamii.com>
Date: Thu, 17 Nov 2022 13:01:07 +0100
Subject: [PATCH] Use SSH Multiplexing

---
 script/myssh           | 79 +++++++++++++++++++++---------------------
 script/sqlproxy_cli.sh | 12 +++----
 sqlproxy_setup.sh      | 21 ++++++++++-
 3 files changed, 66 insertions(+), 46 deletions(-)

diff --git a/script/myssh b/script/myssh
index 9db3322..4ecf047 100755
--- a/script/myssh
+++ b/script/myssh
@@ -1,49 +1,57 @@
 #!/usr/bin/env bash
 
 [ -z $SQL_PROXY_HOST ] && SQL_PROXY_HOST="localhost"
-CACHE_FILE="$HOME/.cache/sqlproxy_$SQL_PROXY_HOST"
+CONNECTION_CACHE="$HOME/.cache/sqlproxy_$SQL_PROXY_HOST"
 
 HELP="Usage: myssh [ls|connect]
 SUBCOMMANDS:
   ls: list available database hosts
   connect: connect to a database host
+  disconnect: Quit existing ssh master session
 
 SYNTAX connect host [-u user] [-p password] [-c client]
 "
 
-ls() {
-    echo $(ssh $SQL_PROXY_HOST ls)
+ssh_status() {
+    ssh -O check -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST > /dev/null 2>&1
+    echo $?
 }
 
-read_cache() {
-    if [ -f $CACHE_FILE ]
+connect() {
+    mkdir -p $HOME/.ssh/controlmasters
+
+    if [ $(ssh_status) -ne 0 ]
     then
-        mapfile -t HOST_LIST < $CACHE_FILE
-        LAST_CHANGED=$(expr $(date +"%s") - "${HOST_LIST[0]}")
-        # Refresh cache if cache is older than a minute
-        if [ $LAST_CHANGED -gt 60 ]
-        then
-            write_cache $(ls)
-            read_cache
-        fi
-    else
-        write_cache $(ls)
-        read_cache
+        echo "" > $CONNECTION_CACHE
+        ssh -o "ControlPersist=10m" -M -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST q
     fi
 }
 
-write_cache() {
-    touch $CACHE_FILE
-    echo -e $(date +"%s")"\n"$1 > $CACHE_FILE
+disconnect() {
+    if [ $(ssh_status) -eq 0 ]
+    then
+        ssh -O stop -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST q
+    fi
+}
+
+port_forward() {
+    ACTIVE_HOST=$(cat $CONNECTION_CACHE)
+    if [ -z $ACTIVE_HOST ] || [ $ACTIVE_HOST != $1 ]
+    then
+        if [ ! -z $ACTIVE_HOST ]
+        then
+            ssh -O cancel -L 3306:$ACTIVE_HOST:3306 -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST q
+        fi
+        ssh -O forward -L 3306:$1:3306 -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST
+    fi
+    echo $1 > $CONNECTION_CACHE
+}
+
+ls_hosts() {
+    echo $(ssh -S $HOME/.ssh/controlmasters/%r@%h:%p $SQL_PROXY_HOST ls)
 }
 
 run_client() {
-    # wait for port to open
-    while ! nc -z localhost 3306 > /dev/null
-    do
-        sleep 0.1
-    done
-
     if [ $(uname -s) = "Linux" ]
     then
         mysql --protocol=TCP -u $1 -p$2 -h localhost -P 3306
@@ -57,21 +65,15 @@ shift
 
 case $MAIN_OPTION in
     ls)
-        RESPONSE=$(ls)
-        write_cache $RESPONSE
-        echo $RESPONSE;;
+        ls_hosts;;
     connect)
-        # Kill open connections on exit
-        # https://stackoverflow.com/questions/360201/how-do-i-kill-background-processes-jobs-when-my-shell-script-exits
-        trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
-
-        # check/update cache
-        read_cache
+        # ensure connection
+        connect
 
         # check if host is valid
         TARGET_HOST=$1
         shift
-        tail -n +2 $CACHE_FILE | grep -qe "^$TARGET_HOST$"
+        ls_hosts | grep -qe "^$TARGET_HOST$"
         GREP_EXIT_CODE=$?
         if [ $GREP_EXIT_CODE -eq 0 ]
         then
@@ -82,18 +84,17 @@ case $MAIN_OPTION in
                     p) MYSQL_PASSWORD="$OPTARG" ;;
                 esac
             done
-            ssh -L 3306:$TARGET_HOST:3306 $SQL_PROXY_HOST > /dev/null 2>&1 &
+            port_forward $TARGET_HOST
             if [ ! -z $MYSQL_USERNAME ] && [ ! -z $MYSQL_PASSWORD ]
             then
                 run_client $MYSQL_USERNAME $MYSQL_PASSWORD
-            else
-                echo 'Press CTRL C to quit this connection'
-                wait
             fi
         else
             echo "Invalid Hostname: $2."
         fi
         ;;
+    disconnect)
+        disconnect;;
     *)
         echo -e "Usage: myssh [ls|connect]\n\n";;
 esac
diff --git a/script/sqlproxy_cli.sh b/script/sqlproxy_cli.sh
index a27b673..8ee19da 100755
--- a/script/sqlproxy_cli.sh
+++ b/script/sqlproxy_cli.sh
@@ -10,9 +10,9 @@ idle() {
     sleep infinity
 }
 
-if [ -z $SSH_ORIGINAL_COMMAND ]
-then
-    idle
-else
-    ls_hosts
-fi
+case "$SSH_ORIGINAL_COMMAND" in
+    "") idle;;
+    ls) ls_hosts;;
+    q|quit) exit 0;;
+    *) exit 1;;
+esac
diff --git a/sqlproxy_setup.sh b/sqlproxy_setup.sh
index 5923f83..f69bace 100755
--- a/sqlproxy_setup.sh
+++ b/sqlproxy_setup.sh
@@ -1,3 +1,22 @@
+#!/usr/bin/env bash
+
+which myssh > /dev/null 2>&1
+if [ $? -eq 1 ]
+then
+    mkdir -p $HOME/bin
+    cp ./script/myssh $HOME/bin/myssh
+    if [[ ! $PATH =~ $HOME/bin ]]
+    then
+	      if [[ $SHELL =~ bin/bash$ ]]
+        then
+            echo -e 'PATH=$PATH:$HOME/bin' >> $HOME/.bashrc
+        elif [[ $SHELL =~ bin/zsh$ ]]
+        then
+            echo -e 'PATH=$PATH:$HOME/bin' >> $HOME/.zshrc
+        fi
+    fi
+fi
+
 if [ ! -f ./etc/ssh/ssh_host_ed25519_key ]
 then
     echo "Generating sqlproxy SSHD keys"
@@ -14,5 +33,5 @@ case $GEN_KEYS in
         echo -ne "\n\nHost $HOST_NAME\n  User sqlproxy\n  IdentityFile ~/.ssh/$KEY_NAME.key" >> ~/.ssh/config
         echo -e command=\"/sqlproxy_cli.sh\" $(cat ~/.ssh/$KEY_NAME.key.pub) >> ./etc/ssh/.ssh/authorized_keys
         break;;
-       *) echo "Not generating client ssh key.\nPlease put your desired public keys into ./etc/ssh/.ssh/authorized_keys\nAlso add 'command=\"/sqlproxy_cli.sh\" ' in front of your key";;
+       *) echo -e "Not generating client ssh key.\nPlease put your desired public keys into ./etc/ssh/.ssh/authorized_keys\nAlso add 'command=\"/sqlproxy_cli.sh\" ' in front of your key";;
 esac