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