1package main
2
3import (
4	"flag"
5	"fmt"
6	"io/ioutil"
7	"log"
8	"os"
9	"os/exec"
10	"os/user"
11	"strconv"
12	"strings"
13	"time"
14)
15
16type OnFail int
17
18const (
19	IgnoreOnFail OnFail = iota
20	WarnOnFail
21	ExitOnFail
22)
23
24type arrayFlags []string
25
26// Implemented for flag#Value interface
27func (s *arrayFlags) String() string {
28	if s == nil {
29		return ""
30	}
31	return fmt.Sprintf("%v", *s)
32}
33
34// Implemented for flag#Value interface
35func (s *arrayFlags) Set(value string) error {
36	*s = append(*s, value)
37	return nil
38}
39
40// Returns `"foo" "bar"`
41func (s *arrayFlags) AsArgs() string {
42	var result []string
43	for _, value := range *s {
44		result = append(result, fmt.Sprintf("%q", value))
45	}
46	return strings.Join(result, " ")
47}
48
49// Returns `--flag="foo" --flag="bar"`
50func (s *arrayFlags) AsRepeatedFlag(name string) string {
51	var result []string
52	for _, value := range *s {
53		result = append(result, fmt.Sprintf(`--%s="%s"`, name, value))
54	}
55	return strings.Join(result, " ")
56}
57
58var build_instance string
59var build_project string
60var build_zone string
61var dest_image string
62var dest_family string
63var dest_project string
64var launch_instance string
65var arch string
66var source_image_family string
67var source_image_project string
68var repository_url string
69var repository_branch string
70var version string
71var internal_ip_flag string
72var INTERNAL_extra_source string
73var verbose bool
74var username string
75var image_disk_size_gb int
76
77// NOTE: For `gcloud compute ssh` command, `ssh_flags` will be used as SSH_ARGS rather than
78// as `--ssh_flag` repeated flag. Why? because --ssh_flag is not parsed as expected when
79// containing quotes and spaces.
80var ssh_flags arrayFlags
81var host_orchestration_flag bool
82
83func init() {
84	user, err := user.Current()
85	if err != nil {
86		panic(err)
87	}
88	username = user.Username
89
90	flag.StringVar(&build_instance, "build_instance",
91		username+"-build", "Instance name to create for the build")
92	flag.StringVar(&build_project, "build_project",
93		mustShell("gcloud config get-value project"), "Project to use for scratch")
94	// The new get-value output format is different. The result is in 2nd line.
95	str_list := strings.Split(build_project, "\n")
96	if len(str_list) == 2 {
97		build_project = str_list[1]
98	}
99
100	flag.StringVar(&build_zone, "build_zone",
101		mustShell("gcloud config get-value compute/zone"),
102		"Zone to use for scratch resources")
103	// The new get-value output format is different. The result is in 2nd line.
104	str_list = strings.Split(build_zone, "\n")
105	if len(str_list) == 2 {
106		build_zone = str_list[1]
107	}
108
109	flag.StringVar(&dest_image, "dest_image",
110		"vsoc-host-scratch-"+username, "Image to create")
111	flag.StringVar(&dest_family, "dest_family", "",
112		"Image family to add the image to")
113	flag.StringVar(&dest_project, "dest_project",
114		mustShell("gcloud config get-value project"), "Project to use for the new image")
115	// The new get-value output format is different. The result is in 2nd line.
116	str_list = strings.Split(dest_project, "\n")
117	if len(str_list) == 2 {
118		dest_project = str_list[1]
119	}
120
121	flag.StringVar(&launch_instance, "launch_instance", "",
122		"Name of the instance to launch with the new image")
123	flag.StringVar(&arch, "arch", "gce_x86_64",
124		"Which CPU arch, arm/x86_64/gce_x86_64")
125	flag.StringVar(&source_image_family, "source_image_family", "debian-11",
126		"Image familty to use as the base")
127	flag.StringVar(&source_image_project, "source_image_project", "debian-cloud",
128		"Project holding the base image")
129	flag.StringVar(&repository_url, "repository_url",
130		"https://github.com/google/android-cuttlefish.git",
131		"URL to the repository with host changes")
132	flag.StringVar(&repository_branch, "repository_branch",
133		"main", "Branch to check out")
134	flag.StringVar(&version, "version", "", "cuttlefish-common version")
135	flag.StringVar(&internal_ip_flag, "INTERNAL_IP", "",
136		"INTERNAL_IP can be set to --internal-ip run on a GCE instance."+
137			"The instance will need --scope compute-rw.")
138	flag.StringVar(&INTERNAL_extra_source, "INTERNAL_extra_source", "",
139		"INTERNAL_extra_source may be set to a directory containing the source for extra packages to build.")
140	flag.BoolVar(&verbose, "verbose", true, "print commands and output (default: true)")
141	flag.IntVar(&image_disk_size_gb, "image_disk_size_gb", 10, "Image disk size in GB")
142	flag.Var(&ssh_flags, "ssh_flag",
143		"Values for --ssh-flag and --scp_flag for gcloud compute ssh/scp respectively. This flag may be repeated")
144	flag.BoolVar(&host_orchestration_flag, "host_orchestration", false,
145		"assembles image with host orchestration capabilities")
146	flag.Parse()
147}
148
149func shell(cmd string) (string, error) {
150	if verbose {
151		fmt.Println(cmd)
152	}
153	b, err := exec.Command("/bin/sh", "-c", cmd).CombinedOutput()
154	if verbose {
155		fmt.Println(string(b))
156	}
157	if err != nil {
158		return "", err
159	}
160	return strings.TrimSpace(string(b)), nil
161}
162
163func mustShell(cmd string) string {
164	if verbose {
165		fmt.Println(cmd)
166	}
167	out, err := shell(cmd)
168	if err != nil {
169		panic(err)
170	}
171	if verbose {
172		fmt.Println(out)
173	}
174	return strings.TrimSpace(out)
175}
176
177func gce(action OnFail, gceArg string, errorStr ...string) (string, error) {
178	cmd := "gcloud " + gceArg
179	out, err := shell(cmd)
180	if out != "" {
181		fmt.Println(out)
182	}
183	if err != nil && action != IgnoreOnFail {
184		var buf string
185		fmt.Sprintf(buf, "gcloud error occurred: %s", err)
186		if len(errorStr) > 0 {
187			buf += " [" + errorStr[0] + "]"
188		}
189		if action == ExitOnFail {
190			panic(buf)
191		}
192		if action == WarnOnFail {
193			fmt.Println(buf)
194		}
195	}
196	return out, err
197}
198
199func waitForInstance(PZ string) {
200	for {
201		time.Sleep(5 * time.Second)
202		_, err := gce(WarnOnFail, `compute ssh `+internal_ip_flag+` `+PZ+` `+
203			build_instance+` -- `+ssh_flags.AsArgs()+` uptime `)
204		if err == nil {
205			break
206		}
207	}
208}
209
210func packageSource(url string, branch string, subdir string) {
211	repository_dir := url[strings.LastIndex(url, "/")+1:]
212	repository_dir = mustShell(`basename "` + repository_dir + `" .git`)
213	debian_dir := repository_dir
214	if subdir != "" {
215		debian_dir = repository_dir + "/" + subdir
216	}
217	mustShell("git clone " + url + " -b " + branch)
218	mustShell("dpkg-source -b " + debian_dir)
219	mustShell("rm -rf " + repository_dir)
220	mustShell("ls -l")
221	mustShell("pwd")
222}
223
224func createInstance(instance string, arg string) {
225	_, err := gce(WarnOnFail, `compute instances describe "`+instance+`"`)
226	if err != nil {
227		gce(ExitOnFail, `compute instances create `+arg+` "`+instance+`"`)
228	}
229}
230
231func main() {
232	gpu_type := "nvidia-tesla-p100-vws"
233	PZ := "--project=" + build_project + " --zone=" + build_zone
234
235	if arch != "gce_x86_64" {
236		// new path that generate image locally without creating GCE instance
237
238		abt := os.Getenv("ANDROID_BUILD_TOP")
239		cmd := `"` + abt + `/device/google/cuttlefish/tools/create_base_image_combined.sh"`
240		cmd += " " + arch
241		out, err := shell(cmd)
242		if out != "" {
243			fmt.Println(out)
244		}
245		if err != nil {
246			fmt.Println("create_base_image arch %s error occurred: %s", arch, err)
247		}
248
249		// gce operations
250		delete_instances := build_instance + " " + dest_image
251		if launch_instance != "" {
252			delete_instances += " " + launch_instance
253		}
254		zip_file := "disk_" + username + ".raw.tar.gz"
255		gs_file := "gs://cloud-android-testing-esp/" + zip_file
256		cloud_storage_file := "https://storage.googleapis.com/cloud-android-testing-esp/" + zip_file
257		location := "us"
258
259		// delete all previous instances, images and disks
260		gce(WarnOnFail, `compute instances delete -q `+PZ+` `+delete_instances, `Not running`)
261		gce(WarnOnFail, `compute disks delete -q `+PZ+` "`+dest_image+`"`, `No scratch disk`)
262		gce(WarnOnFail, `compute images delete -q --project="`+build_project+`" "`+dest_image+`"`,
263			`Not respinning`)
264		gce(WarnOnFail, `alpha storage rm `+gs_file)
265
266		// upload new local host image into GCE storage
267		gce(WarnOnFail, `alpha storage cp `+abt+`/`+zip_file+` gs://cloud-android-testing-esp`)
268
269		// create GCE image based on new uploaded host image
270		gce(WarnOnFail, `compute images create "`+dest_image+`" --project="`+build_project+
271			`" --family="`+source_image_family+`" --source-uri="`+cloud_storage_file+
272			`" --storage-location="`+location+`" --guest-os-features=UEFI_COMPATIBLE`)
273
274		// find Nvidia GPU and then create GCE instance
275		gce(ExitOnFail, `compute accelerator-types describe "`+gpu_type+`" `+PZ,
276			`Please use a zone with `+gpu_type+` GPUs available.`)
277		createInstance(build_instance, PZ+
278			` --machine-type=n1-standard-16 --network-interface=network-tier=PREMIUM,subnet=default`+
279			` --accelerator="type=`+gpu_type+
280			`,count=1" --maintenance-policy=TERMINATE --provisioning-model=STANDARD`+
281			` --service-account=204446994883-compute@developer.gserviceaccount.com`+
282			` --scopes=https://www.googleapis.com/auth/devstorage.read_only,`+
283			`https://www.googleapis.com/auth/logging.write,`+
284			`https://www.googleapis.com/auth/monitoring.write,`+
285			`https://www.googleapis.com/auth/servicecontrol,`+
286			`https://www.googleapis.com/auth/service.management.readonly,`+
287			`https://www.googleapis.com/auth/trace.append`+
288			` --tags=http-server --create-disk=auto-delete=yes,boot=yes,device-name=`+build_instance+
289			`,image=projects/cloud-android-testing/global/images/`+dest_image+
290			`,mode=rw,size=200,type=projects/cloud-android-testing/zones/`+build_zone+
291			`/diskTypes/pd-balanced --no-shielded-secure-boot --shielded-vtpm`+
292			` --shielded-integrity-monitoring --reservation-affinity=any`)
293
294		// enable serial-port (console)
295		gce(WarnOnFail, `compute instances add-metadata `+build_instance+
296			` --metadata serial-port-enable=TRUE`)
297		return
298	}
299
300	dest_family_flag := ""
301	if dest_family != "" {
302		dest_family_flag = "--family=" + dest_family
303	}
304
305	scratch_dir, err := ioutil.TempDir("", "")
306	if err != nil {
307		log.Fatal(err)
308	}
309
310	oldDir, err := os.Getwd()
311	if err != nil {
312		log.Fatal(err)
313	}
314	os.Chdir(scratch_dir)
315	packageSource(repository_url, repository_branch, "base")
316	packageSource(repository_url, repository_branch, "frontend")
317	os.Chdir(oldDir)
318
319	abt := os.Getenv("ANDROID_BUILD_TOP")
320	source_files := `"` + abt + `/device/google/cuttlefish/tools/create_base_image_gce.sh"`
321	source_files += " " + `"` + abt + `/device/google/cuttlefish/tools/install_nvidia.sh"`
322	source_files += " " + `"` + abt + `/device/google/cuttlefish/tools/update_gce_kernel.sh"`
323	source_files += " " + `"` + abt + `/device/google/cuttlefish/tools/remove_old_gce_kernel.sh"`
324	source_files += " " + scratch_dir + "/*"
325	if INTERNAL_extra_source != "" {
326		source_files += " " + INTERNAL_extra_source + "/*"
327	}
328
329	delete_instances := build_instance + " " + dest_image
330	if launch_instance != "" {
331		delete_instances += " " + launch_instance
332	}
333
334	gce(WarnOnFail, `compute instances delete -q `+PZ+` `+delete_instances,
335		`Not running`)
336	gce(WarnOnFail, `compute disks delete -q `+PZ+` "`+dest_image+
337		`"`, `No scratch disk`)
338	gce(WarnOnFail, `compute images delete -q --project="`+build_project+
339		`" "`+dest_image+`"`, `Not respinning`)
340	gce(WarnOnFail, `compute disks create `+PZ+`  --size=`+strconv.Itoa(image_disk_size_gb)+`G `+
341		`--image-family="`+source_image_family+`" --image-project="`+source_image_project+`" "`+dest_image+`"`)
342	gce(ExitOnFail, `compute accelerator-types describe "`+gpu_type+`" `+PZ,
343		`Please use a zone with `+gpu_type+` GPUs available.`)
344	createInstance(build_instance, PZ+
345		` --machine-type=n1-standard-16 --image-family="`+source_image_family+
346		`" --image-project="`+source_image_project+
347		`" --boot-disk-size=200GiB --accelerator="type=`+gpu_type+
348		`,count=1" --maintenance-policy=TERMINATE --boot-disk-size=200GiB`)
349
350	waitForInstance(PZ)
351
352	// Ubuntu tends to mount the wrong disk as root, so help it by waiting until
353	// it has booted before giving it access to the clean image disk
354	gce(WarnOnFail, `compute instances attach-disk `+PZ+` "`+build_instance+
355		`" --disk="`+dest_image+`"`)
356
357	// beta for the --internal-ip flag that may be passed via internal_ip_flag
358	gce(ExitOnFail, `beta compute scp `+internal_ip_flag+` `+PZ+` `+source_files+
359		` "`+build_instance+`:" `+ssh_flags.AsRepeatedFlag("scp-flag"))
360
361	// Update the host kernel before installing any kernel modules
362	// Needed to guarantee that the modules in the chroot aren't built for the
363	// wrong kernel
364	gce(WarnOnFail, `compute ssh `+internal_ip_flag+` `+PZ+` "`+build_instance+
365		`"`+` -- `+ssh_flags.AsArgs()+` ./update_gce_kernel.sh`)
366	// TODO rammuthiah if the instance is clobbered with ssh commands within
367	// 5 seconds of reboot, it becomes inaccessible. Workaround that by sleeping
368	// 50 seconds.
369	time.Sleep(50 * time.Second)
370	gce(ExitOnFail, `compute ssh `+internal_ip_flag+` `+PZ+` "`+build_instance+
371		`"`+` -- `+ssh_flags.AsArgs()+` ./remove_old_gce_kernel.sh`)
372
373	ho_arg := ""
374	if host_orchestration_flag {
375		ho_arg = "-o"
376	}
377	gce(ExitOnFail, `compute ssh `+internal_ip_flag+` `+PZ+` "`+build_instance+
378		`"`+` -- `+ssh_flags.AsArgs()+` ./create_base_image_gce.sh `+ho_arg)
379
380	// Reboot the instance to force a clean umount of the disk's file system.
381	gce(WarnOnFail, `compute ssh `+internal_ip_flag+` `+PZ+` "`+build_instance+
382		`" -- `+ssh_flags.AsArgs()+` sudo reboot`)
383	waitForInstance(PZ)
384
385	gce(ExitOnFail, `compute instances delete -q `+PZ+` "`+build_instance+`"`)
386	gce(ExitOnFail, `compute images create --project="`+build_project+
387		`" --source-disk="`+dest_image+`" --source-disk-zone="`+build_zone+
388		`" --licenses=https://www.googleapis.com/compute/v1/projects/vm-options/global/licenses/enable-vmx `+
389		dest_family_flag+` "`+dest_image+`"`)
390	gce(ExitOnFail, `compute disks delete -q `+PZ+` "`+dest_image+`"`)
391
392	if launch_instance != "" {
393		createInstance(launch_instance, PZ+
394			` --image-project="`+build_project+`" --image="`+dest_image+
395			`" --machine-type=n1-standard-4 --scopes storage-ro --accelerator="type=`+
396			gpu_type+`,count=1" --maintenance-policy=TERMINATE --boot-disk-size=200GiB`)
397	}
398
399	fmt.Printf("Test and if this looks good, consider releasing it via:\n"+
400		"\n"+
401		"gcloud compute images create \\\n"+
402		"  --project=\"%s\" \\\n"+
403		"  --source-image=\"%s\" \\\n"+
404		"  --source-image-project=\"%s\" \\\n"+
405		"  \"%s\" \\\n"+
406		"  \"%s\"\n",
407		dest_project, dest_image, build_project, dest_family_flag, dest_image)
408}
409