nix-builder-autoscaler/agent/nix_builder_autoscaler/bootstrap/userdata.py

62 lines
2.2 KiB
Python

"""EC2 user-data template rendering for builder instance bootstrap.
The generated script follows the NixOS AMI pattern: write config files
that existing systemd services (tailscale-autoconnect, nix-daemon) consume,
rather than calling ``tailscale up`` directly.
"""
from __future__ import annotations
import textwrap
def render_userdata(slot_id: str, region: str, ssm_param: str = "/nix-builder/ts-authkey") -> str:
"""Render a bash user-data script for builder instance bootstrap.
The returned string is a complete shell script. On NixOS AMIs the script
is executed by ``amazon-init.service``. The caller (EC2Runtime) passes it
to ``run_instances`` as ``UserData``; boto3 base64-encodes automatically.
Args:
slot_id: Autoscaler slot identifier (used as Tailscale hostname suffix).
region: AWS region for SSM parameter lookup.
ssm_param: SSM parameter path containing the Tailscale auth key.
"""
return textwrap.dedent(f"""\
#!/usr/bin/env bash
set -euo pipefail
SLOT_ID="{slot_id}"
REGION="{region}"
SSM_PARAM="{ssm_param}"
# --- Fetch Tailscale auth key from SSM Parameter Store ---
mkdir -p /run/credentials
TS_AUTHKEY=$(aws ssm get-parameter \\
--region "$REGION" \\
--with-decryption \\
--name "$SSM_PARAM" \\
--query 'Parameter.Value' \\
--output text)
printf '%s' "$TS_AUTHKEY" > /run/credentials/tailscale-auth-key
chmod 600 /run/credentials/tailscale-auth-key
# --- Write tailscale-autoconnect config ---
mkdir -p /etc/tailscale
cat > /etc/tailscale/autoconnect.conf <<TSCONF
TS_AUTHKEY_FILE=/run/credentials/tailscale-auth-key
TS_AUTHKEY_EPHEMERAL=true
TS_AUTHKEY_PREAUTHORIZED=true
TS_HOSTNAME=nix-builder-$SLOT_ID
TS_EXTRA_ARGS="--ssh --advertise-tags=tag:nix-builder"
TSCONF
# --- Start/restart tailscale-autoconnect so it picks up the config ---
systemctl restart tailscale-autoconnect.service || true
# --- Ensure nix-daemon is running ---
systemctl start nix-daemon.service || true
# --- Signal readiness ---
echo "ready" > /run/nix-builder-ready
""")