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

"""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

        # --- Resolve instance identity from IMDSv2 for unique hostname ---
        IMDS_TOKEN=$(curl -fsS -X PUT "http://169.254.169.254/latest/api/token" \\
            -H "X-aws-ec2-metadata-token-ttl-seconds: 21600" || true)
        INSTANCE_ID=$(curl -fsS -H "X-aws-ec2-metadata-token: $IMDS_TOKEN" \\
            "http://169.254.169.254/latest/meta-data/instance-id" || true)
        if [ -z "$INSTANCE_ID" ]; then
          INSTANCE_ID="unknown"
        fi

        # --- 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-$INSTANCE_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
    """)