Skip to content

Commit 8aa8e0f

Browse files
ALameLlamaALameLlama
authored
fix(distributed): setup script (#10551)
Assisted-by: OpenCode:GPT-5.5 [Read] [Edit] Signed-off-by: Nicholas Ciechanowski <nicholas@ciech.anow.ski>
1 parent e905639 commit 8aa8e0f

1 file changed

Lines changed: 202 additions & 31 deletions

File tree

‎scripts/nats-auth-setup.sh‎

Lines changed: 202 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,220 @@
11
#!/usr/bin/env bash
2-
# Generate NATS account + service user JWTs for LocalAI distributed mode.
2+
# Generate NATS JWT authentication material and server configuration
3+
# for LocalAI distributed mode.
34
#
45
# Requires: nsc (https://docs.nats.io/running-a-nats-service/configuration/securing_nats/auth_intro/nsc)
56
#
6-
# Usage:
7-
# ./scripts/nats-auth-setup.sh
7+
# Outputs:
8+
# ./nats-keys/localai-nats.env
9+
# ./nats-keys/localai-frontend.creds
10+
# ./nats-keys/nats-auth.conf
11+
# ./nats-keys/nats-server.conf
812
#
9-
# Outputs operator/account seeds and a service user JWT suitable for:
10-
# LOCALAI_NATS_ACCOUNT_SEED
11-
# LOCALAI_NATS_SERVICE_JWT
13+
# Environment overrides:
14+
# NATS_OPERATOR_NAME
15+
# NATS_ACCOUNT_NAME
16+
# NATS_SERVICE_USER
17+
# NATS_KEYS_DIR
18+
1219
#
13-
# Per-node worker JWTs are minted automatically by the frontend at registration
14-
# when LOCALAI_NATS_ACCOUNT_SEED is set.
20+
# LocalAI workers receive their own JWT and user seed when registering
21+
# with the frontend.
1522

1623
set -euo pipefail
1724

25+
# Ensure newly created secret files are private by default.
26+
umask 077
27+
1828
if ! command -v nsc >/dev/null 2>&1; then
1929
echo "nsc is required. Install from https://github.com/nats-io/nsc/releases" >&2
2030
exit 1
2131
fi
2232

2333
OPERATOR="${NATS_OPERATOR_NAME:-localai-operator}"
2434
ACCOUNT="${NATS_ACCOUNT_NAME:-localai}"
35+
SYSTEM_ACCOUNT="${NATS_SYSTEM_ACCOUNT_NAME:-SYS}"
2536
SERVICE_USER="${NATS_SERVICE_USER:-localai-frontend}"
37+
OUTPUT_DIR="${NATS_KEYS_DIR:-./nats-keys}"
38+
39+
CREDS_FILE="$OUTPUT_DIR/${SERVICE_USER}.creds"
40+
ENV_FILE="$OUTPUT_DIR/localai-nats.env"
41+
AUTH_CONFIG_FILE="$OUTPUT_DIR/nats-auth.conf"
42+
SERVER_CONFIG_FILE="$OUTPUT_DIR/nats-server.conf"
43+
44+
mkdir -p "$OUTPUT_DIR"
45+
46+
echo "Configuring NATS operator: $OPERATOR"
47+
48+
# Create the operator if it does not exist, otherwise select it.
49+
if nsc select operator "$OPERATOR" >/dev/null 2>&1; then
50+
echo "[ OK ] using existing operator '$OPERATOR'"
51+
else
52+
nsc add operator \
53+
-n "$OPERATOR" \
54+
--generate-signing-key
55+
56+
nsc select operator "$OPERATOR" >/dev/null
57+
fi
58+
59+
# Create and assign the NATS system account.
60+
if nsc describe account \
61+
-n "$SYSTEM_ACCOUNT" >/dev/null 2>&1; then
62+
echo "[ OK ] using existing system account '$SYSTEM_ACCOUNT'"
63+
else
64+
nsc add account -n "$SYSTEM_ACCOUNT"
65+
fi
66+
67+
nsc edit operator \
68+
--system-account "$SYSTEM_ACCOUNT"
69+
70+
# Create the LocalAI account if it does not exist.
71+
if nsc describe account -n "$ACCOUNT" >/dev/null 2>&1; then
72+
echo "[ OK ] using existing account '$ACCOUNT'"
73+
else
74+
nsc add account -n "$ACCOUNT"
75+
fi
76+
77+
nsc select account "$ACCOUNT" >/dev/null
78+
79+
# Create the frontend service user if it does not exist.
80+
if nsc describe user \
81+
-n "$SERVICE_USER" \
82+
--account "$ACCOUNT" >/dev/null 2>&1; then
83+
echo "[ OK ] using existing user '$SERVICE_USER'"
84+
else
85+
nsc add user \
86+
-n "$SERVICE_USER" \
87+
--account "$ACCOUNT"
88+
fi
89+
90+
# Frontend control-plane permissions.
91+
nsc edit user \
92+
-n "$SERVICE_USER" \
93+
--account "$ACCOUNT" \
94+
--allow-pub "nodes.>,gallery.>,agent.>,staging.>,state.>,jobs.>,mcp.>,cache.>,prefixcache.>,finetune.>" \
95+
--allow-sub "nodes.>,gallery.>,agent.>,staging.>,state.>,jobs.>,mcp.>,cache.>,prefixcache.>,_INBOX.>"
96+
97+
# Generate a credentials file containing the frontend user JWT and seed.
98+
rm -f "$CREDS_FILE"
99+
100+
nsc generate creds \
101+
-a "$ACCOUNT" \
102+
-n "$SERVICE_USER" \
103+
-o "$CREDS_FILE"
104+
105+
# Extract the frontend JWT from the credentials file.
106+
SERVICE_JWT="$(
107+
awk '
108+
/BEGIN NATS USER JWT/ {
109+
capture = 1
110+
next
111+
}
112+
/END NATS USER JWT/ {
113+
capture = 0
114+
}
115+
capture
116+
' "$CREDS_FILE" | tr -d '\r\n'
117+
)"
118+
119+
# Extract the frontend user seed from the credentials file.
120+
SERVICE_SEED="$(
121+
awk '
122+
/BEGIN USER NKEY SEED/ {
123+
capture = 1
124+
next
125+
}
126+
/END USER NKEY SEED/ {
127+
capture = 0
128+
}
129+
capture
130+
' "$CREDS_FILE" | tr -d '\r\n'
131+
)"
132+
133+
# Retrieve the seed belonging to this exact account rather than taking
134+
# the first account key found in the keystore.
135+
ACCOUNT_SEED="$(
136+
nsc list keys \
137+
--account "$ACCOUNT" \
138+
--accounts \
139+
--show-seeds |
140+
awk -F '|' -v expected="$ACCOUNT" '
141+
function trim(value) {
142+
gsub(/^[[:space:]]+|[[:space:]]+$/, "", value)
143+
return value
144+
}
145+
146+
NF >= 3 {
147+
entity = trim($2)
148+
seed = trim($3)
149+
150+
if (entity == expected && seed ~ /^SA[A-Z0-9]+$/) {
151+
print seed
152+
exit
153+
}
154+
}
155+
'
156+
)"
157+
158+
# Validate all extracted values before writing output files.
159+
if [[ ! "$ACCOUNT_SEED" =~ ^SA[A-Z0-9]+$ ]]; then
160+
echo "Unable to extract the account seed for '$ACCOUNT'." >&2
161+
exit 1
162+
fi
163+
164+
if [[ ! "$SERVICE_JWT" =~ ^eyJ ]]; then
165+
echo "Unable to extract the service JWT from '$CREDS_FILE'." >&2
166+
exit 1
167+
fi
168+
169+
if [[ ! "$SERVICE_SEED" =~ ^SU[A-Z0-9]+$ ]]; then
170+
echo "Unable to extract the service seed from '$CREDS_FILE'." >&2
171+
exit 1
172+
fi
173+
174+
# Generate the trusted operator and memory resolver configuration.
175+
# This contains public operator/account JWT claims, not the private seeds.
176+
nsc generate config \
177+
--mem-resolver \
178+
--config-file "$AUTH_CONFIG_FILE" \
179+
--force
180+
181+
# Generate the primary NATS server configuration.
182+
# The include path matches the Docker Compose mounts shown below.
183+
cat >"$SERVER_CONFIG_FILE" <<'NATS_CONFIG'
184+
server_name: localai-nats
185+
port: 4222
186+
http: 8222
187+
188+
jetstream {
189+
store_dir: /data/jetstream
190+
}
191+
192+
include nats-auth.conf
193+
NATS_CONFIG
194+
195+
# Generate the environment file consumed by the LocalAI frontend.
196+
cat >"$ENV_FILE" <<EOF
197+
LOCALAI_NATS_ACCOUNT_SEED=$ACCOUNT_SEED
198+
LOCALAI_NATS_SERVICE_JWT=$SERVICE_JWT
199+
LOCALAI_NATS_SERVICE_SEED=$SERVICE_SEED
200+
EOF
201+
202+
# The environment and credentials files contain private seeds.
203+
chmod 600 "$CREDS_FILE" "$ENV_FILE"
204+
205+
# These contain server configuration and public JWT claims.
206+
chmod 644 "$AUTH_CONFIG_FILE" "$SERVER_CONFIG_FILE"
26207

27-
nsc add operator -n "$OPERATOR" --generate-signing-key
28-
nsc add account -n "$ACCOUNT"
29-
nsc add user -n "$SERVICE_USER" --account "$ACCOUNT"
30-
31-
# Broad publish for frontend control plane (tighten with custom claims in production).
32-
nsc edit user -n "$SERVICE_USER" --account "$ACCOUNT" \
33-
--allow-pub "nodes.>,gallery.>,agent.>,jobs.>,mcp.>,cache.>,prefixcache.>,finetune.>" \
34-
--allow-sub "nodes.>,gallery.>,agent.>,jobs.>,mcp.>,cache.>,prefixcache.>,_INBOX.>"
35-
36-
KEYS_DIR="${NATS_KEYS_DIR:-./nats-keys}"
37-
mkdir -p "$KEYS_DIR"
38-
nsc generate creds -a "$ACCOUNT" -n "$SERVICE_USER" -o "$KEYS_DIR"
39-
40-
ACCOUNT_SEED=$(nsc describe account "$ACCOUNT" -o json | jq -r '.nats.private_key')
41-
SERVICE_JWT=$(cat "$KEYS_DIR/${ACCOUNT}/${SERVICE_USER}.jwt" 2>/dev/null || cat "$KEYS_DIR/${SERVICE_USER}.jwt")
42-
43-
echo ""
44-
echo "=== LocalAI NATS auth material ==="
45-
echo "LOCALAI_NATS_ACCOUNT_SEED=${ACCOUNT_SEED}"
46-
echo "LOCALAI_NATS_SERVICE_JWT=${SERVICE_JWT}"
47-
echo ""
48-
echo "Configure the NATS server with the generated operator/account JWTs under $KEYS_DIR"
49-
echo "and set LOCALAI_NATS_REQUIRE_AUTH=true on frontends and workers in production."
208+
echo
209+
echo "=== LocalAI NATS JWT setup complete ==="
210+
echo
211+
echo "LocalAI environment: $ENV_FILE"
212+
echo "Service credentials: $CREDS_FILE"
213+
echo "NATS server config: $SERVER_CONFIG_FILE"
214+
echo "NATS auth config: $AUTH_CONFIG_FILE"
215+
echo
216+
echo "Keep '$ENV_FILE' and '$CREDS_FILE' secret."
217+
echo "Do not commit them to source control."
218+
echo
219+
echo "=== LocalAI NATS environment ==="
220+
cat "$ENV_FILE"

0 commit comments

Comments
 (0)