pyrox.dev
1#!/usr/bin/env nix-shell
2#!nix-shell -p awscli -p jq -p qemu -i bash
3# shellcheck shell=bash
4#
5# Future Deprecation?
6# This entire thing should probably be replaced with a generic terraform config
7
8# Uploads and registers NixOS images built from the
9# <nixos/release.nix> amazonImage attribute. Images are uploaded and
10# registered via a home region, and then copied to other regions.
11
12# The home region requires an s3 bucket, and an IAM role named "vmimport"
13# (by default) with access to the S3 bucket. The name can be
14# configured with the "service_role_name" variable. Configuration of the
15# vmimport role is documented in
16# https://docs.aws.amazon.com/vm-import/latest/userguide/vmimport-image-import.html
17
18# set -x
19set -euo pipefail
20
21var () { true; }
22
23# configuration
24var ${state_dir:=$HOME/amis/ec2-images}
25var ${home_region:=eu-west-1}
26var ${bucket:=nixos-amis}
27var ${service_role_name:=vmimport}
28
29# Output of the command:
30# > aws ec2 describe-regions --all-regions --query "Regions[].{Name:RegionName}" --output text | sort
31var ${regions:=
32 af-south-1
33 ap-east-1
34 ap-northeast-1
35 ap-northeast-2
36 ap-northeast-3
37 ap-south-1
38 ap-southeast-1
39 ap-southeast-2
40 ap-southeast-3
41 ca-central-1
42 eu-central-1
43 eu-north-1
44 eu-south-1
45 eu-west-1
46 eu-west-2
47 eu-west-3
48 me-south-1
49 sa-east-1
50 us-east-1
51 us-east-2
52 us-west-1
53 us-west-2
54 }
55
56regions=($regions)
57
58log() {
59 echo "$@" >&2
60}
61
62if [ "$#" -ne 1 ]; then
63 log "Usage: ./upload-amazon-image.sh IMAGE_OUTPUT"
64 exit 1
65fi
66
67# result of the amazon-image from nixos/release.nix
68store_path=$1
69
70if [ ! -e "$store_path" ]; then
71 log "Store path: $store_path does not exist, fetching..."
72 nix-store --realise "$store_path"
73fi
74
75if [ ! -d "$store_path" ]; then
76 log "store_path: $store_path is not a directory. aborting"
77 exit 1
78fi
79
80read_image_info() {
81 if [ ! -e "$store_path/nix-support/image-info.json" ]; then
82 log "Image missing metadata"
83 exit 1
84 fi
85 jq -r "$1" "$store_path/nix-support/image-info.json"
86}
87
88# We handle a single image per invocation, store all attributes in
89# globals for convenience.
90zfs_disks=$(read_image_info .disks)
91is_zfs_image=
92if jq -e .boot <<< "$zfs_disks"; then
93 is_zfs_image=1
94 zfs_boot=".disks.boot"
95fi
96image_label="$(read_image_info .label)${is_zfs_image:+-ZFS}"
97image_system=$(read_image_info .system)
98image_files=( $(read_image_info ".disks.root.file") )
99
100image_logical_bytes=$(read_image_info "${zfs_boot:-.disks.root}.logical_bytes")
101
102if [[ -n "$is_zfs_image" ]]; then
103 image_files+=( $(read_image_info .disks.boot.file) )
104fi
105
106# Derived attributes
107
108image_logical_gigabytes=$(((image_logical_bytes-1)/1024/1024/1024+1)) # Round to the next GB
109
110case "$image_system" in
111 aarch64-linux)
112 amazon_arch=arm64
113 ;;
114 x86_64-linux)
115 amazon_arch=x86_64
116 ;;
117 *)
118 log "Unknown system: $image_system"
119 exit 1
120esac
121
122image_name="NixOS-${image_label}-${image_system}"
123image_description="NixOS ${image_label} ${image_system}"
124
125log "Image Details:"
126log " Name: $image_name"
127log " Description: $image_description"
128log " Size (gigabytes): $image_logical_gigabytes"
129log " System: $image_system"
130log " Amazon Arch: $amazon_arch"
131
132read_state() {
133 local state_key=$1
134 local type=$2
135
136 cat "$state_dir/$state_key.$type" 2>/dev/null || true
137}
138
139write_state() {
140 local state_key=$1
141 local type=$2
142 local val=$3
143
144 mkdir -p "$state_dir"
145 echo "$val" > "$state_dir/$state_key.$type"
146}
147
148wait_for_import() {
149 local region=$1
150 local task_id=$2
151 local state snapshot_id
152 log "Waiting for import task $task_id to be completed"
153 while true; do
154 read -r state message snapshot_id < <(
155 aws ec2 describe-import-snapshot-tasks --region "$region" --import-task-ids "$task_id" | \
156 jq -r '.ImportSnapshotTasks[].SnapshotTaskDetail | "\(.Status) \(.StatusMessage) \(.SnapshotId)"'
157 )
158 log " ... state=$state message=$message snapshot_id=$snapshot_id"
159 case "$state" in
160 active)
161 sleep 10
162 ;;
163 completed)
164 echo "$snapshot_id"
165 return
166 ;;
167 *)
168 log "Unexpected snapshot import state: '${state}'"
169 log "Full response: "
170 aws ec2 describe-import-snapshot-tasks --region "$region" --import-task-ids "$task_id" >&2
171 exit 1
172 ;;
173 esac
174 done
175}
176
177wait_for_image() {
178 local region=$1
179 local ami_id=$2
180 local state
181 log "Waiting for image $ami_id to be available"
182
183 while true; do
184 read -r state < <(
185 aws ec2 describe-images --image-ids "$ami_id" --region "$region" | \
186 jq -r ".Images[].State"
187 )
188 log " ... state=$state"
189 case "$state" in
190 pending)
191 sleep 10
192 ;;
193 available)
194 return
195 ;;
196 *)
197 log "Unexpected AMI state: '${state}'"
198 exit 1
199 ;;
200 esac
201 done
202}
203
204
205make_image_public() {
206 local region=$1
207 local ami_id=$2
208
209 wait_for_image "$region" "$ami_id"
210
211 log "Making image $ami_id public"
212
213 aws ec2 modify-image-attribute \
214 --image-id "$ami_id" --region "$region" --launch-permission 'Add={Group=all}' >&2
215}
216
217upload_image() {
218 local region=$1
219
220 for image_file in "${image_files[@]}"; do
221 local aws_path=${image_file#/}
222
223 if [[ -n "$is_zfs_image" ]]; then
224 local suffix=${image_file%.*}
225 suffix=${suffix##*.}
226 fi
227
228 local state_key="$region.$image_label${suffix:+.${suffix}}.$image_system"
229 local task_id
230 task_id=$(read_state "$state_key" task_id)
231 local snapshot_id
232 snapshot_id=$(read_state "$state_key" snapshot_id)
233 local ami_id
234 ami_id=$(read_state "$state_key" ami_id)
235
236 if [ -z "$task_id" ]; then
237 log "Checking for image on S3"
238 if ! aws s3 ls --region "$region" "s3://${bucket}/${aws_path}" >&2; then
239 log "Image missing from aws, uploading"
240 aws s3 cp --region "$region" "$image_file" "s3://${bucket}/${aws_path}" >&2
241 fi
242
243 log "Importing image from S3 path s3://$bucket/$aws_path"
244
245 task_id=$(aws ec2 import-snapshot --role-name "$service_role_name" --disk-container "{
246 \"Description\": \"nixos-image-${image_label}-${image_system}\",
247 \"Format\": \"vhd\",
248 \"UserBucket\": {
249 \"S3Bucket\": \"$bucket\",
250 \"S3Key\": \"$aws_path\"
251 }
252 }" --region "$region" | jq -r '.ImportTaskId')
253
254 write_state "$state_key" task_id "$task_id"
255 fi
256
257 if [ -z "$snapshot_id" ]; then
258 snapshot_id=$(wait_for_import "$region" "$task_id")
259 write_state "$state_key" snapshot_id "$snapshot_id"
260 fi
261 done
262
263 if [ -z "$ami_id" ]; then
264 log "Registering snapshot $snapshot_id as AMI"
265
266 local block_device_mappings=(
267 "DeviceName=/dev/xvda,Ebs={SnapshotId=$snapshot_id,VolumeSize=$image_logical_gigabytes,DeleteOnTermination=true,VolumeType=gp3}"
268 )
269
270 if [[ -n "$is_zfs_image" ]]; then
271 local root_snapshot_id=$(read_state "$region.$image_label.root.$image_system" snapshot_id)
272
273 local root_image_logical_bytes=$(read_image_info ".disks.root.logical_bytes")
274 local root_image_logical_gigabytes=$(((root_image_logical_bytes-1)/1024/1024/1024+1)) # Round to the next GB
275
276 block_device_mappings+=(
277 "DeviceName=/dev/xvdb,Ebs={SnapshotId=$root_snapshot_id,VolumeSize=$root_image_logical_gigabytes,DeleteOnTermination=true,VolumeType=gp3}"
278 )
279 fi
280
281
282 local extra_flags=(
283 --root-device-name /dev/xvda
284 --sriov-net-support simple
285 --ena-support
286 --virtualization-type hvm
287 )
288
289 block_device_mappings+=("DeviceName=/dev/sdb,VirtualName=ephemeral0")
290 block_device_mappings+=("DeviceName=/dev/sdc,VirtualName=ephemeral1")
291 block_device_mappings+=("DeviceName=/dev/sdd,VirtualName=ephemeral2")
292 block_device_mappings+=("DeviceName=/dev/sde,VirtualName=ephemeral3")
293
294 ami_id=$(
295 aws ec2 register-image \
296 --name "$image_name" \
297 --description "$image_description" \
298 --region "$region" \
299 --architecture $amazon_arch \
300 --block-device-mappings "${block_device_mappings[@]}" \
301 --boot-mode $(read_image_info .boot_mode) \
302 "${extra_flags[@]}" \
303 | jq -r '.ImageId'
304 )
305
306 write_state "$state_key" ami_id "$ami_id"
307 fi
308
309 [[ -v PRIVATE ]] || make_image_public "$region" "$ami_id"
310
311 echo "$ami_id"
312}
313
314copy_to_region() {
315 local region=$1
316 local from_region=$2
317 local from_ami_id=$3
318
319 state_key="$region.$image_label.$image_system"
320 ami_id=$(read_state "$state_key" ami_id)
321
322 if [ -z "$ami_id" ]; then
323 log "Copying $from_ami_id to $region"
324 ami_id=$(
325 aws ec2 copy-image \
326 --region "$region" \
327 --source-region "$from_region" \
328 --source-image-id "$from_ami_id" \
329 --name "$image_name" \
330 --description "$image_description" \
331 | jq -r '.ImageId'
332 )
333
334 write_state "$state_key" ami_id "$ami_id"
335 fi
336
337 [[ -v PRIVATE ]] || make_image_public "$region" "$ami_id"
338
339 echo "$ami_id"
340}
341
342upload_all() {
343 home_image_id=$(upload_image "$home_region")
344 jq -n \
345 --arg key "$home_region.$image_system" \
346 --arg value "$home_image_id" \
347 '$ARGS.named'
348
349 for region in "${regions[@]}"; do
350 if [ "$region" = "$home_region" ]; then
351 continue
352 fi
353 copied_image_id=$(copy_to_region "$region" "$home_region" "$home_image_id")
354
355 jq -n \
356 --arg key "$region.$image_system" \
357 --arg value "$copied_image_id" \
358 '$ARGS.named'
359 done
360}
361
362upload_all | jq --slurp from_entries