setup_client.tftpl.sh

#!/usr/bin/env bash

# Script to bootstrap Nomad as client node.

# This script performs the following tasks:
# - Prepares DNS configuration for exec tasks
# - Renders the Nomad client configuration
# - Optionally, adds Docker configuration to Nomad if the 'enable_docker_plugin' variable is set to true
# - Adds SSH public keys to authorized_keys if provided
# - Starts the Nomad service

set -Eeuo pipefail

declare -r SCRIPT_NAME="$(basename "$0")"
declare -ag AWS_TAGS=()

# Send the log output from this script to user-data.log, syslog, and the console.
exec > >(tee /var/log/user-data.log | logger -t user-data -s 2>/dev/console) 2>&1

# Wrapper to log any outputs from the script to stderr
function log {
  declare -r LVL="$1"
  declare -r MSG="$2"
  declare -r TS=$(date +"%Y-%m-%d %H:%M:%S")
  echo >&2 -e "$TS [$LVL] [$SCRIPT_NAME] $MSG"
}

# Stores AWS tags to use as nomad client meta
# Requires `nomad-cluster` tag to be defined
# within AWS instance tags
store_tags() {
  max_attempts=3
  count=0

  while true; do
    TOKEN=$(curl -s --connect-timeout 1 --retry 3 --retry-delay 3 \
      -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600")

    TAGS=$(curl -s --connect-timeout 1 --retry 3 --retry-delay 3 \
      -H "X-aws-ec2-metadata-token: $TOKEN" http://169.254.169.254/latest/meta-data/tags/instance)

    # If there's no 'nomad-cluster' found in tags, retry.
    if [[ "$${TAGS}" != *"nomad-cluster"* ]]; then
      sleep 1

      count=$((count + 1))

      # If max retries still didn't get the data, fail.
      if [[ $count -eq $max_attempts ]]; then
        log "ERROR" "aborting as max attempts reached"
        exit 1
      fi
      continue
    fi

    readarray -t AWS_TAGS <<<"$TAGS"
    break

  done
}

# Sets hostname for the system
# Replaces `ip` in the hostname with the AWS instance `Name` tag
set_hostname() {
  for t in "$${AWS_TAGS[@]}"; do
    # For servers we'll use the NAME tag of the EC2 instance.
    if [ "$t" == "Name" ]; then
      TAG=$(curl -s --retry 3 --retry-delay 3 --connect-timeout 3 \
        -H "Accept: application/json" -H "X-aws-ec2-metadata-token: $TOKEN" "http://169.254.169.254/latest/meta-data/tags/instance/$t")

      # The original hostname is like `ip-10-x-y-z`
      CURR_HOSTNAME=$(sudo hostnamectl --static)
      # Replace `ip` with tag value.
      HOSTNAME="$${CURR_HOSTNAME//ip/$TAG}"
      log "INFO" "setting hostname as $HOSTNAME"
      sudo hostnamectl set-hostname "$HOSTNAME"
    fi
  done
}

# Ensures the resolv.conf within nomad `exec` jobs
# can access other machines
# see: https://github.com/hashicorp/nomad/issues/11033
prepare_dns_config() {
  cat <<EOF >/etc/nomad.d/route53_resolv.conf
nameserver ${route_53_resolver_address}
search ${aws_region}.compute.internal
EOF
}

# Increase the file limit
modify_nomad_systemd_config() {
  if [ ${nomad_file_limit} > 65536 ]; then
    sudo sed -i '/^LimitNOFILE/s/=.*$/=${nomad_file_limit}/' /lib/systemd/system/nomad.service
  fi
}

# Function to check and load the bridge module
load_bridge() {
  if ! lsmod | grep -q "^bridge"; then
    log "INFO" "Bridge module not loaded. Attempting to load..."
    if ! modprobe bridge; then
      log "ERROR" "Failed to load bridge module. This might affect network functionality."
    else
      log "INFO" "Bridge module loaded successfully."
    fi
  else
    log "INFO" "Bridge module is already loaded."
  fi
}

# Add SSH public keys to authorized_keys file
add_ssh_keys() {
  # Use the configured SSH user or default to "ubuntu"
  SSH_USER="${ssh_user}"

  # Check if the user exists
  if ! id "$SSH_USER" &>/dev/null; then
    log "ERROR" "User $SSH_USER does not exist. Cannot add SSH keys."
    return 1
  fi

  # Get user's home directory from /etc/passwd
  USER_HOME=$(getent passwd "$SSH_USER" | cut -d: -f6)
  if [ -z "$USER_HOME" ]; then
    log "ERROR" "Could not determine home directory for user $SSH_USER"
    return 1
  fi

  SSH_DIR="$USER_HOME/.ssh"
  AUTH_KEYS_FILE="$SSH_DIR/authorized_keys"

  log "INFO" "Adding SSH keys for user $SSH_USER (home: $USER_HOME)"

  # Ensure .ssh directory exists and has correct permissions
  if [ ! -d "$SSH_DIR" ]; then
    mkdir -p "$SSH_DIR"
    log "INFO" "Created $SSH_DIR directory"
  fi
  chmod 700 "$SSH_DIR"
  chown "$SSH_USER:$SSH_USER" "$SSH_DIR"

  # Create authorized_keys file if it doesn't exist
  if [ ! -f "$AUTH_KEYS_FILE" ]; then
    touch "$AUTH_KEYS_FILE"
    log "INFO" "Created $AUTH_KEYS_FILE file"
  fi
  chmod 600 "$AUTH_KEYS_FILE"
  chown "$SSH_USER:$SSH_USER" "$AUTH_KEYS_FILE"

  # Add each provided SSH public key
  %{ for key in ssh_public_keys }
  if ! grep -q "${key}" "$AUTH_KEYS_FILE"; then
    echo "${key}" >> "$AUTH_KEYS_FILE"
    log "INFO" "Added SSH public key to authorized_keys for $SSH_USER"
  else
    log "INFO" "SSH public key already exists in authorized_keys for $SSH_USER"
  fi
  %{ endfor }

  # Ensure the authorized_keys file has the correct ownership and permissions
  # This is crucial since the script runs as root
  chmod 600 "$AUTH_KEYS_FILE"
  chown "$SSH_USER:$SSH_USER" "$AUTH_KEYS_FILE"
}

# Enables nomad systemd service
start_nomad() {
  sudo systemctl daemon-reload
  sudo systemctl enable --now nomad
}

# Restarts nomad systemd service
restart_nomad() {
  sudo systemctl restart nomad
}

# Sets up `/etc/nomad.d`
prepare_nomad_client_config() {
  cat <<EOF >/etc/nomad.d/nomad.hcl
${nomad_client_cfg}
EOF

  cat <<EOF >>/etc/nomad.d/nomad.hcl
client {
  enabled = true
  server_join {
    retry_join = ["provider=aws region=${aws_region} tag_key=${nomad_join_tag_key} tag_value=${nomad_join_tag_value}"]
  }
  meta {
$(for tag in "$${AWS_TAGS[@]}"; do
    key=$${tag//:/_}
    TAG=$(curl -s --connect-timeout 3 --retry 3 --retry-delay 3 \
      -H "Accept: application/json" -H "X-aws-ec2-metadata-token: $TOKEN" "http://169.254.169.254/latest/meta-data/tags/instance/$tag")
    echo -e "\tec2_$key = \"$TAG\""
  done)
  }
  chroot_env {
    # Defaults
    "/bin/"           = "/bin/"
    "/lib"            = "/lib"
    "/lib32"          = "/lib32"
    "/lib64"          = "/lib64"
    "/sbin"           = "/sbin"
    "/usr"            = "/usr"

    "/etc/ld.so.cache"  = "/etc/ld.so.cache"
    "/etc/ld.so.conf"   = "/etc/ld.so.conf"
    "/etc/ld.so.conf.d" = "/etc/ld.so.conf.d"
    "/etc/localtime"    = "/etc/localtime"
    "/etc/passwd"       = "/etc/passwd"
    "/etc/ssl"          = "/etc/ssl"
    "/etc/timezone"     = "/etc/timezone"

    # DNS
    "/etc/nomad.d/route53_resolv.conf" = "/etc/resolv.conf"

  }
}

plugin "exec" {
  config {
    allow_caps = ["audit_write", "chown", "dac_override", "fowner", "fsetid", "kill", "mknod",
    "net_bind_service", "setfcap", "setgid", "setpcap", "setuid", "sys_chroot", "sys_time"]
  }
}

plugin "raw_exec" {
  config {
    enabled = true
  }
}
EOF
}

add_docker_to_nomad() {
  cat <<EOF >>/etc/nomad.d/nomad.hcl
plugin "docker" {
  config {
    auth {
      config = "/etc/docker/config.json"
    }
    allow_privileged = true
    volumes {
      enabled = true
    }
    extra_labels = ["job_name", "job_id", "task_group_name", "task_name", "namespace", "node_name", "node_id"]

    logging {
      type = "json-file"
      config {
        max-size = "10m"
        max-file = 10
      }
    }
  }
}
EOF
}

add_host_volumes_to_nomad() {
  cat <<EOF >>/etc/nomad.d/nomad.hcl
client {
  %{ for name, volume in nomad_client_exec_host_volumes }
  host_volume "${name}" {
    path      = "${volume.path}"
    read_only = ${volume.read_only}
  }
  %{ endfor }
}
EOF

  # Create directories if they don't exist
  %{ for name, volume in nomad_client_exec_host_volumes }
  if [ ! -d "${volume.path}" ]; then
    mkdir -p "${volume.path}"
  fi
  %{ endfor }
}

log "INFO" "Fetching EC2 Tags from AWS"
store_tags

log "INFO" "Setting machine hostname"
set_hostname

log "INFO" "Prepare DNS config for exec tasks"
prepare_dns_config

log "INFO" "Checking and loading bridge module"
load_bridge

log "INFO" "Rendering client config for nomad"
prepare_nomad_client_config

%{ if enable_docker_plugin }
log "INFO" "Adding docker config to Nomad"
add_docker_to_nomad
%{ endif }

%{ if length(nomad_client_exec_host_volumes) > 0 }
log "INFO" "Adding host volumes config to Nomad"
add_host_volumes_to_nomad
%{ else }
log "INFO" "No host volumes configured for Nomad client"
%{ endif }

%{ if length(ssh_public_keys) > 0 }
log "INFO" "Adding SSH public keys to authorized_keys"
add_ssh_keys || log "WARN" "Failed to add SSH keys"
%{ else }
log "INFO" "No SSH public keys provided to add"
%{ endif }

log "INFO" "Modify Nomad systemd config"
modify_nomad_systemd_config

log "INFO" "Starting Nomad service"
start_nomad

log "INFO" "Finished client initializing process! Enjoy Nomad!"