fix(distributed): setup script

scripts/nats-auth-setup.sh

@@ -1,49 +1,220 @@
 #!/usr/bin/env bash
-# Generate NATS account + service user JWTs for LocalAI distributed mode.
+# Generate NATS JWT authentication material and server configuration
+# for LocalAI distributed mode.
 #
 # Requires: nsc (https://docs.nats.io/running-a-nats-service/configuration/securing_nats/auth_intro/nsc)
 #
-# Usage:
-#   ./scripts/nats-auth-setup.sh
+# Outputs:
+#   ./nats-keys/localai-nats.env
+#   ./nats-keys/localai-frontend.creds
+#   ./nats-keys/nats-auth.conf
+#   ./nats-keys/nats-server.conf
 #
-# Outputs operator/account seeds and a service user JWT suitable for:
-#   LOCALAI_NATS_ACCOUNT_SEED
-#   LOCALAI_NATS_SERVICE_JWT
+# Environment overrides:
+#   NATS_OPERATOR_NAME
+#   NATS_ACCOUNT_NAME
+#   NATS_SERVICE_USER
+#   NATS_KEYS_DIR
+
 #
-# Per-node worker JWTs are minted automatically by the frontend at registration
-# when LOCALAI_NATS_ACCOUNT_SEED is set.
+# LocalAI workers receive their own JWT and user seed when registering
+# with the frontend.
 
 set -euo pipefail
 
+# Ensure newly created secret files are private by default.
+umask 077
+
 if ! command -v nsc >/dev/null 2>&1; then
 	echo "nsc is required. Install from https://github.com/nats-io/nsc/releases" >&2
 	exit 1
 fi
 
 OPERATOR="${NATS_OPERATOR_NAME:-localai-operator}"
 ACCOUNT="${NATS_ACCOUNT_NAME:-localai}"
+SYSTEM_ACCOUNT="${NATS_SYSTEM_ACCOUNT_NAME:-SYS}"
 SERVICE_USER="${NATS_SERVICE_USER:-localai-frontend}"
+OUTPUT_DIR="${NATS_KEYS_DIR:-./nats-keys}"
+
+CREDS_FILE="$OUTPUT_DIR/${SERVICE_USER}.creds"
+ENV_FILE="$OUTPUT_DIR/localai-nats.env"
+AUTH_CONFIG_FILE="$OUTPUT_DIR/nats-auth.conf"
+SERVER_CONFIG_FILE="$OUTPUT_DIR/nats-server.conf"
+
+mkdir -p "$OUTPUT_DIR"
+
+echo "Configuring NATS operator: $OPERATOR"
+
+# Create the operator if it does not exist, otherwise select it.
+if nsc select operator "$OPERATOR" >/dev/null 2>&1; then
+	echo "[ OK ] using existing operator '$OPERATOR'"
+else
+	nsc add operator \
+		-n "$OPERATOR" \
+		--generate-signing-key
+
+	nsc select operator "$OPERATOR" >/dev/null
+fi
+
+# Create and assign the NATS system account.
+if nsc describe account \
+	-n "$SYSTEM_ACCOUNT" >/dev/null 2>&1; then
+	echo "[ OK ] using existing system account '$SYSTEM_ACCOUNT'"
+else
+	nsc add account -n "$SYSTEM_ACCOUNT"
+fi
+
+nsc edit operator \
+	--system-account "$SYSTEM_ACCOUNT"
 
-nsc add operator -n "$OPERATOR" --generate-signing-key
-nsc add account -n "$ACCOUNT"
-nsc add user -n "$SERVICE_USER" --account "$ACCOUNT"
+# Create the LocalAI account if it does not exist.
+if nsc describe account -n "$ACCOUNT" >/dev/null 2>&1; then
+	echo "[ OK ] using existing account '$ACCOUNT'"
+else
+	nsc add account -n "$ACCOUNT"
+fi
 
-# Broad publish for frontend control plane (tighten with custom claims in production).
-nsc edit user -n "$SERVICE_USER" --account "$ACCOUNT" \
+nsc select account "$ACCOUNT" >/dev/null
+
+# Create the frontend service user if it does not exist.
+if nsc describe user \
+	-n "$SERVICE_USER" \
+	--account "$ACCOUNT" >/dev/null 2>&1; then
+	echo "[ OK ] using existing user '$SERVICE_USER'"
+else
+	nsc add user \
+		-n "$SERVICE_USER" \
+		--account "$ACCOUNT"
+fi
+
+# Frontend control-plane permissions.
+nsc edit user \
+	-n "$SERVICE_USER" \
+	--account "$ACCOUNT" \
 	--allow-pub "nodes.>,gallery.>,agent.>,jobs.>,mcp.>,cache.>,prefixcache.>,finetune.>" \
 	--allow-sub "nodes.>,gallery.>,agent.>,jobs.>,mcp.>,cache.>,prefixcache.>,_INBOX.>"
 
-KEYS_DIR="${NATS_KEYS_DIR:-./nats-keys}"
-mkdir -p "$KEYS_DIR"
-nsc generate creds -a "$ACCOUNT" -n "$SERVICE_USER" -o "$KEYS_DIR"
+# Generate a credentials file containing the frontend user JWT and seed.
+rm -f "$CREDS_FILE"
+
+nsc generate creds \
+	-a "$ACCOUNT" \
+	-n "$SERVICE_USER" \
+	-o "$CREDS_FILE"
+
+# Extract the frontend JWT from the credentials file.
+SERVICE_JWT="$(
+	awk '
+		/BEGIN NATS USER JWT/ {
+			capture = 1
+			next
+		}
+		/END NATS USER JWT/ {
+			capture = 0
+		}
+		capture
+	' "$CREDS_FILE" | tr -d '\r\n'
+)"
+
+# Extract the frontend user seed from the credentials file.
+SERVICE_SEED="$(
+	awk '
+		/BEGIN USER NKEY SEED/ {
+			capture = 1
+			next
+		}
+		/END USER NKEY SEED/ {
+			capture = 0
+		}
+		capture
+	' "$CREDS_FILE" | tr -d '\r\n'
+)"
+
+# Retrieve the seed belonging to this exact account rather than taking
+# the first account key found in the keystore.
+ACCOUNT_SEED="$(
+	nsc list keys \
+		--account "$ACCOUNT" \
+		--accounts \
+		--show-seeds |
+	awk -F '|' -v expected="$ACCOUNT" '
+		function trim(value) {
+			gsub(/^[[:space:]]+|[[:space:]]+$/, "", value)
+			return value
+		}
+
+		NF >= 3 {
+			entity = trim($2)
+			seed = trim($3)
+
+			if (entity == expected && seed ~ /^SA[A-Z0-9]+$/) {
+				print seed
+				exit
+			}
+		}
+	'
+)"
+
+# Validate all extracted values before writing output files.
+if [[ ! "$ACCOUNT_SEED" =~ ^SA[A-Z0-9]+$ ]]; then
+	echo "Unable to extract the account seed for '$ACCOUNT'." >&2
+	exit 1
+fi
+
+if [[ ! "$SERVICE_JWT" =~ ^eyJ ]]; then
+	echo "Unable to extract the service JWT from '$CREDS_FILE'." >&2
+	exit 1
+fi
+
+if [[ ! "$SERVICE_SEED" =~ ^SU[A-Z0-9]+$ ]]; then
+	echo "Unable to extract the service seed from '$CREDS_FILE'." >&2
+	exit 1
+fi
+
+# Generate the trusted operator and memory resolver configuration.
+# This contains public operator/account JWT claims, not the private seeds.
+nsc generate config \
+	--mem-resolver \
+	--config-file "$AUTH_CONFIG_FILE" \
+	--force
+
+# Generate the primary NATS server configuration.
+# The include path matches the Docker Compose mounts shown below.
+cat >"$SERVER_CONFIG_FILE" <<'NATS_CONFIG'
+server_name: localai-nats
+port: 4222
+http: 8222
+
+jetstream {
+	store_dir: /data/jetstream
+}
+
+include nats-auth.conf
+NATS_CONFIG
+
+# Generate the environment file consumed by the LocalAI frontend.
+cat >"$ENV_FILE" <<EOF
+LOCALAI_NATS_ACCOUNT_SEED=$ACCOUNT_SEED
+LOCALAI_NATS_SERVICE_JWT=$SERVICE_JWT
+LOCALAI_NATS_SERVICE_SEED=$SERVICE_SEED
+EOF
+
+# The environment and credentials files contain private seeds.
+chmod 600 "$CREDS_FILE" "$ENV_FILE"
 
-ACCOUNT_SEED=$(nsc describe account "$ACCOUNT" -o json | jq -r '.nats.private_key')
-SERVICE_JWT=$(cat "$KEYS_DIR/${ACCOUNT}/${SERVICE_USER}.jwt" 2>/dev/null || cat "$KEYS_DIR/${SERVICE_USER}.jwt")
+# These contain server configuration and public JWT claims.
+chmod 644 "$AUTH_CONFIG_FILE" "$SERVER_CONFIG_FILE"
 
-echo ""
-echo "=== LocalAI NATS auth material ==="
-echo "LOCALAI_NATS_ACCOUNT_SEED=${ACCOUNT_SEED}"
-echo "LOCALAI_NATS_SERVICE_JWT=${SERVICE_JWT}"
-echo ""
-echo "Configure the NATS server with the generated operator/account JWTs under $KEYS_DIR"
-echo "and set LOCALAI_NATS_REQUIRE_AUTH=true on frontends and workers in production."
+echo
+echo "=== LocalAI NATS JWT setup complete ==="
+echo
+echo "LocalAI environment: $ENV_FILE"
+echo "Service credentials: $CREDS_FILE"
+echo "NATS server config:   $SERVER_CONFIG_FILE"
+echo "NATS auth config:     $AUTH_CONFIG_FILE"
+echo
+echo "Keep '$ENV_FILE' and '$CREDS_FILE' secret."
+echo "Do not commit them to source control."
+echo
+echo "=== LocalAI NATS environment ==="
+cat "$ENV_FILE"