Add more robust startup routine for caddy with tls

ai-dock · Aug 7, 2024 · 29d5c99 · 29d5c99
1 parent f352bcc
commit 29d5c99
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 7 deletions.
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -98,6 +98,9 @@ jobs:
         - {latest: "false", cuda: "12.1.1-base"}
         - {latest: "false", cuda: "12.1.1-cudnn8-runtime"}
         - {latest: "false", cuda: "12.1.1-cudnn8-devel"}
+        - {latest: "false", cuda: "11.8.0-base"}
+        - {latest: "false", cuda: "11.8.0-cudnn8-runtime"}
+        - {latest: "false", cuda: "11.8.0-cudnn8-devel"}
 
     steps:
       -

diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh
@@ -9,7 +9,10 @@ function cleanup() {
 unset -v port
 unset -v url
 
+cert_path=/opt/caddy/tls/container.crt
+key_path=/opt/caddy/tls/container.key
 metrics=""
+
 while getopts l:p: flag
 do
     case "${flag}" in
@@ -23,15 +26,26 @@ if [[ -z $port ]]; then
     exit 1
 fi
 
+function validate_cert_and_key() {
+  if openssl x509 -in "$cert_path" -noout > /dev/null 2>&1 && \
+     openssl rsa -in "$key_path" -check -noout > /dev/null 2>&1; then
+    return 0
+  else
+    return 1
+  fi
+}
+
 function get_scheme() {
-    if [[ ${WEB_ENABLE_HTTPS,,} == true && -f /opt/caddy/tls/container.crt && /opt/caddy/tls/container.key ]]; then
+    if [[ ${WEB_ENABLE_HTTPS,,} == "true" ]] && validate_cert_and_key; then
         echo "https://"
     else
         echo "http://"
     fi
 }
 
-function get_url {
+
+
+function get_url() {
     preset_url=$(jq -r ".service_url" "/run/http_ports/${port}")
     if [[ -n $preset_url ]]; then
         url="$preset_url"

diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh
@@ -8,9 +8,8 @@ function cleanup() {
 
 function start() {
     source /opt/ai-dock/etc/environment.sh
-
     # Give processes time to register their ports
-    sleep 4
+    sleep 2
 
     export SERVICEPORTAL_LOGIN=$(direct-url.sh -p "${SERVICEPORTAL_PORT_HOST:-1111}" -l "/login")
     env-store SERVICEPORTAL_LOGIN
@@ -20,9 +19,31 @@ function start() {
     port_files="/run/http_ports/*"
 
     # Upgrade http to https on the same port
-    if [[ ${WEB_ENABLE_HTTPS,,} == true && -s $(realpath /opt/caddy/tls/container.crt) && -s $(realpath /opt/caddy/tls/container.key) ]]; then
-        export CADDY_TLS_ELEVATION_STRING=$'http_redirect\ntls'
-        export CADDY_TLS_LISTEN_STRING="tls /opt/caddy/tls/container.crt /opt/caddy/tls/container.key"
+
+    if [[ ${WEB_ENABLE_HTTPS,,} == true ]]; then
+        cert_path="/opt/caddy/tls/container.crt"
+        key_path="/opt/caddy/tls/container.key"
+        max_retries=5
+        # Avoid key generation race condition
+        attempts=0
+        while [[ $attempts -lt $max_retries ]]; do
+            if [[ -f $(realpath $cert_path) && -f $(realpath $key_path) ]]; then
+                if validate_cert_and_key; then
+                    echo "Certificate and key are present and valid."
+                    export CADDY_TLS_ELEVATION_STRING=$'http_redirect\ntls'
+                    export CADDY_TLS_LISTEN_STRING="tls /opt/caddy/tls/container.crt /opt/caddy/tls/container.key"
+                    break
+                else
+                    echo "Files are present but invalid, attempt $((attempts + 1)) of $MAX_RETRIES."
+                fi
+            else
+                echo "Waiting for certificate and key to be present, attempt $((attempts + 1)) of $MAX_RETRIES."
+            fi
+            # Increment the retry counter
+            attempts=$((attempts + 1))
+            # Wait before retrying
+            sleep 5
+        done
     fi
 
     cp -f /opt/caddy/share/base_config /opt/caddy/etc/Caddyfile
@@ -49,4 +70,13 @@ function start() {
     caddy run --config /opt/caddy/etc/Caddyfile
 }
 
+function validate_cert_and_key() {
+  if openssl x509 -in "$cert_path" -noout > /dev/null 2>&1 && \
+     openssl rsa -in "$key_path" -check -noout > /dev/null 2>&1; then
+    return 0
+  else
+    return 1
+  fi
+}
+
 start 2>&1