1#!/bin/bash
2# Copyright 2024 Google Inc. All rights reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16color_cyan="\033[0;36m"
17color_plain="\033[0m"
18color_yellow="\033[0;33m"
19
20# validate number of arguments
21if [ "$#" -lt 1 ] || [ "$#" -gt 5 ]; then
22  echo "This script requires 1 mandatory and 4 optional parameters,"
23  echo "server address and optionally cvd instances per docker, and number of " \
24       "docker instances to invoke, vendor_boot image to replace, and config " \
25       "file path for launching configuration."
26  exit 1
27fi
28
29# map arguments to variables
30# $1: ARM server address
31# $2: CVD Instance number per docker (Optional, default is 1)
32# $3: Docker Instance number (Optional, default is 1)
33# $4: Vendor Boot Image path (Optional, default is "")
34server=$1
35
36if [ "$#" -lt 2 ]; then
37 num_instances_per_docker=1
38else
39 num_instances_per_docker=$2
40fi
41
42if [ "$#" -lt 3 ]; then
43 num_dockers=1
44else
45 num_dockers=$3
46fi
47
48if [ "$#" -lt 4 ]; then
49 vendor_boot_image=""
50else
51 vendor_boot_image=$4
52fi
53
54if [ "$#" -lt 5 ]; then
55 config_path=""
56else
57 config_path=$5
58 if [ ! -f $config_path ]; then
59  echo Config file $config_path does not exist
60  exit 1
61 fi
62
63 if ! cat $config_path | jq > /dev/null ; then
64  echo Failed to parse config file $config_path
65  exit 1
66 fi
67fi
68
69
70# set img_dir and cvd_host_tool_dir
71img_dir=${ANDROID_PRODUCT_OUT:-$PWD}
72cvd_host_tool_dir=${ANDROID_HOST_OUT:+"$ANDROID_HOST_OUT/../linux_musl-arm64"}
73cvd_host_tool_dir=${cvd_host_tool_dir:-$PWD}
74
75# upload artifacts into ARM server
76cvd_home_dir=cvd_home
77ssh $server -t "mkdir -p ~/.cvd_artifact; mkdir -p ~/$cvd_home_dir"
78
79# android-info.txt is required for cvd launcher to pick up the correct config file.
80rsync -avch $img_dir/android-info.txt $server:~/$cvd_home_dir --info=progress2
81
82if [ -f $img_dir/required_images ]; then
83  rsync -aSvch --recursive $img_dir --files-from=$img_dir/required_images $server:~/$cvd_home_dir --info=progress2
84  cvd_home_files=($(rsync -rzan --recursive $img_dir --out-format="%n" --files-from=$img_dir/required_images $server:~/$cvd_home_dir --info=name2 | awk '{print $1}'))
85else
86  rsync -aSvch --recursive $img_dir/bootloader $img_dir/*.img $server:~/$cvd_home_dir --info=progress2
87  cvd_home_files=($(rsync -rzan --recursive $img_dir/bootloader --out-format="%n" $img_dir/*.img $server:~/$cvd_home_dir --info=name2 | awk '{print $1}'))
88fi
89
90if [[ $vendor_boot_image != "" ]]; then
91  scp $vendor_boot_image $server:~/$cvd_home_dir/vendor_boot.img
92fi
93
94# upload cvd-host_package.tar.gz into ARM server
95temp_dir=/tmp/cvd_dist
96rm -rf $temp_dir
97mkdir -p $temp_dir
98if [ -d $cvd_host_tool_dir/cvd-host_package ]; then
99  echo "Use contents in cvd-host_package dir"
100  pushd $cvd_host_tool_dir/cvd-host_package > /dev/null
101  tar -cf $temp_dir/cvd-host_package.tar ./*
102  popd > /dev/null
103  pigz -R $temp_dir/cvd-host_package.tar
104elif [ -f $cvd_host_tool_dir/cvd-host_package.tar.gz ]; then
105  echo "Use contents in cvd-host_package.tar.gz"
106  # re-compress with rsyncable option
107  # TODO(b/275312073): remove this if toxbox supports rsyncable
108  pigz -d -c $cvd_host_tool_dir/cvd-host_package.tar.gz | pigz -R > $temp_dir/cvd-host_package.tar.gz
109else
110  echo "There is neither cvd-host_package dir nor cvd-host_package.tar.gz"
111  exit 1
112fi
113rsync -avch $temp_dir/cvd-host_package.tar.gz $server:~/$cvd_home_dir --info=progress2
114cvd_home_files+=("cvd-host_package.tar.gz")
115
116# run root docker instance
117root_container_id=$(ssh $server -t "docker run --privileged -p 2443 -d cuttlefish")
118root_container_id=${root_container_id//$'\r'} # to remove trailing ^M
119echo -e "${color_cyan}Booting root container $root_container_id${color_plain}"
120
121# set trap to stop docker instance
122trap cleanup SIGINT
123cleanup() {
124  echo -e "${color_yellow}SIGINT: stopping the launch instances${color_plain}"
125  ssh $server "docker rm -f $root_container_id ${container_ids[*]} && \
126               docker rmi -f cvd_root_image:$root_container_id && \
127               docker system prune -f"
128  exit 0
129}
130
131# extract Host Orchestrator Port
132docker_inspect=$(ssh $server "docker inspect --format='{{json .NetworkSettings.Ports }}' $root_container_id")
133docker_host_orchestrator_port_parser_script='
134import sys, json;
135json_raw=input()
136data = json.loads(json_raw)
137for k in data:
138  if not data[k]:
139    continue
140
141  original_port = int(k.split("/")[0])
142  assigned_port = int(data[k][0]["HostPort"])
143  if original_port == 2443:
144    print(assigned_port)
145    break
146'
147docker_host_orchestrator_port=$(echo $docker_inspect | python -c "$docker_host_orchestrator_port_parser_script")
148host_orchestrator_url=https://localhost:$docker_host_orchestrator_port
149echo -e "Extracting host orchestrator port in root docker instance"
150
151# create user artifact directory
152create_user_artifacts_dir_script="ssh $server curl -s -k -X POST ${host_orchestrator_url}/userartifacts | jq -r '.name'"
153user_artifacts_dir=$($create_user_artifacts_dir_script)
154while [ -z "$user_artifacts_dir" ]; do
155  echo -e "Failed to create user_artifacts_dir, retrying"
156  sleep 1
157  user_artifacts_dir=$($create_user_artifacts_dir_script)
158done
159echo -e "Succeeded to create user_artifacts_dir"
160
161# upload artifacts and cvd-host_pachage.tar.gz into docker instance
162ssh $server \
163  "for filename in ${cvd_home_files[*]}; do \
164     absolute_path=\$HOME/$cvd_home_dir/\$filename && \
165     size=\$(stat -c%s \$absolute_path) && \
166     echo Uploading \$filename\\(size:\$size\\) ... && \
167     curl -s -k --location -X PUT $host_orchestrator_url/userartifacts/$user_artifacts_dir \
168       -H 'Content-Type: multipart/form-data' \
169       -F chunk_number=1 \
170       -F chunk_total=1 \
171       -F chunk_size_bytes=\$size \
172       -F file=@\$absolute_path; \
173   done"
174echo -e "Done"
175
176# extract cvd-host_package.tar.gz with /:extract API if the API exists
177ssh $server \
178  "job_id=\$(curl -s -k -X POST ${host_orchestrator_url}/userartifacts/$user_artifacts_dir/cvd-host_package.tar.gz/:extract | jq -r '.name' 2>/dev/null) && \
179   if [[ \$job_id != \"\" ]]; then \
180     echo Extracting cvd-host_package.tar.gz ... && \
181     job_done=\"false\" && \
182     while [[ \$job_done == \"false\" ]]; do \
183       sleep 1 && \
184       job_done=\$(curl -s -k ${host_orchestrator_url}/operations/\$job_id | jq -r '.done'); \
185     done && \
186     echo Done; \
187   fi"
188
189echo -e "Creating image from root docker container"
190root_image_id=$(ssh $server -t "docker commit $root_container_id cvd_root_image:$root_container_id")
191root_image_id=${root_image_id//$'\r'} # to remove trailing ^M
192echo -e "${color_cyan}Root image $root_image_id${color_plain}"
193
194echo -e "${color_cyan}Booting containers ... ${color_plain}"
195container_ids=$(ssh $server \
196  "container_ids=() && \
197  for docker_num in \$(seq 1 $num_dockers); do \
198    if [ \$docker_num -eq 1 ]; then \
199      web_port_forward=\"-p 1443 -p 15550-15560 \"; \
200    else \
201      web_port_forward=\"\"; \
202    fi && \
203    adb_port_forward=\"\" && \
204    for instance_num in \$(seq 1 $num_instances_per_docker); do
205      adb_port_forward+=\"-p \$((instance_num + 6520 - 1)) \";
206    done && \
207    container_id=\$(docker run --rm --privileged \$web_port_forward -p 2443 \$adb_port_forward -d $root_image_id) && \
208    container_id=\${container_id//\$'\\r'} && \
209    container_ids+=(\${container_id}); \
210  done && \
211  echo \${container_ids[*]}
212")
213
214echo -e "Extracting host orchestrator ports in docker instances"
215docker_inspects=$(ssh $server \
216  "docker_inspects=() && container_ids=(${container_ids[*]}) &&
217  for container_id in \${container_ids[*]}; do \
218    docker_inspect=\$(docker inspect --format='{{json .NetworkSettings.Ports }}' \$container_id) && \
219    docker_inspects+=(\${docker_inspect}); \
220  done && \
221  echo \${docker_inspects[*]}
222")
223host_orchestrator_ports=()
224for docker_inspect in ${docker_inspects[*]}; do
225  port=$(echo $docker_inspect | python -c "$docker_host_orchestrator_port_parser_script")
226  host_orchestrator_ports+=($port)
227done
228
229if [[ $config_path != "" ]]; then
230  cvd_creation_data=$(cat $config_path | jq -c)
231else
232  cvd_creation_data="{\"cvd\":{\"build_source\": \
233    {\"user_build_source\":{\"artifacts_dir\":\"$user_artifacts_dir\"}}}, \
234    \"additional_instances_num\":$((num_instances_per_docker - 1))}";
235fi
236cvd_creation_data=$(echo $cvd_creation_data | sed s/\$user_artifact_id/$user_artifacts_dir/g)
237
238# start Cuttlefish instance on top of docker instance
239# TODO(b/317942272): support starting the instance with an optional vendor boot debug image.
240echo -e "Starting Cuttlefish"
241ssh $server "job_ids=() && \
242for port in ${host_orchestrator_ports[*]}; do \
243  host_orchestrator_url=https://localhost:\$port && \
244  job_id=\"\" && \
245  while [ -z \"\$job_id\" ]; do \
246    job_id=\$(curl -s -k -X POST \$host_orchestrator_url/cvds \
247      -H 'Content-Type: application/json' \
248      -d '$cvd_creation_data' \
249        | jq -r '.name') && \
250    if [ -z \"\$job_id\" ]; then \
251      echo \"  Failed to request creating Cuttlefish, retrying\" && \
252      sleep 1; \
253    else \
254      echo \"  Succeeded to request: \$job_id\" && \
255      job_ids+=(\${job_id}); \
256    fi; \
257  done; \
258done \
259
260echo \"Waiting Cuttlefish instances to be booted\" && \
261i=0 && \
262for port in ${host_orchestrator_ports[*]}; do \
263  job_id=\${job_ids[\$i]} && \
264  i=\$((i+1)) && \
265  host_orchestrator_url=https://localhost:\$port && \
266  job_done=\"false\" && \
267  while [[ \$job_done == \"false\" ]]; do \
268    sleep 1 && \
269    job_done=\$(curl -s -k \${host_orchestrator_url}/operations/\$job_id | jq -r '.done'); \
270  done && \
271  echo \"  Boot completed: \$job_id\"; \
272done \
273"
274echo -e "Done"
275
276# Web UI port is 3443 instead 1443 because there could be a running operator or host orchestrator in this machine as well.
277web_ui_port=3443
278echo -e "Web UI port: $web_ui_port. ${color_cyan}Please point your browser to https://localhost:$web_ui_port for the UI${color_plain}"
279
280# sets up SSH port forwarding to the remote server for various ports and launch cvd instance
281adb_port=6520
282for docker_num in $(seq 1 $num_dockers); do
283  for instance_num in $(seq 1 $num_instances_per_docker); do
284    device_name="cvd_$instance_num"
285    device_adb_port=$((adb_port + ( (docker_num - 1) * num_instances_per_docker) + instance_num - 1))
286    echo -e "$device_name of docker $docker_num is using adb port $device_adb_port. Try ${color_cyan}adb connect 127.0.0.1:${device_adb_port}${color_plain} if you want to connect to this device"
287  done
288done
289
290docker_port_parser_script='
291import sys, json;
292web_ui_port = int(sys.argv[1])
293adb_port = int(sys.argv[2])
294max_instances = 100
295num_instance = int(sys.argv[3])
296json_raw=input()
297data = json.loads(json_raw)
298for k in data:
299  if not data[k]:
300    continue
301
302  original_port = int(k.split("/")[0])
303  assigned_port = int(data[k][0]["HostPort"])
304  if original_port == 1443:
305    original_port = web_ui_port
306  elif original_port in (1080, 2080, 2443): # Do not expose other operator or host orchestrator port beyond ARM server
307    continue
308  elif original_port >= 6520 and original_port <= 6520 + max_instances:
309    if original_port - 6520 >= num_instance:
310      continue
311    original_port = adb_port + original_port - 6520
312  print(f"-L {original_port}:127.0.0.1:{assigned_port}", end=" ")
313'
314
315ports_forwarding=""
316current_adb_port=$adb_port
317
318for docker_inspect in ${docker_inspects[*]}; do
319  ports_forwarding+=$(echo $docker_inspect | python -c "$docker_port_parser_script" $web_ui_port $current_adb_port $num_instances_per_docker)
320  current_adb_port=$((current_adb_port + num_instances_per_docker))
321done
322
323echo "Set up ssh ports forwarding: $ports_forwarding"
324echo -e "${color_yellow}Please stop the running instances by ctrl+c${color_plain}"
325ssh $server $ports_forwarding "tail -f /dev/null"
326