1 /*
2  * Copyright (C) 2023 The Android Open Source Project
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  */
16 package com.android.tradefed.device.internal;
17 
18 import com.android.tradefed.config.IConfiguration;
19 import com.android.tradefed.config.IConfigurationReceiver;
20 import com.android.tradefed.config.IDeviceConfiguration;
21 import com.android.tradefed.device.DeviceNotAvailableException;
22 import com.android.tradefed.device.ITestDevice;
23 import com.android.tradefed.device.cloud.GceAvdInfo;
24 import com.android.tradefed.device.cloud.RemoteAndroidVirtualDevice;
25 import com.android.tradefed.device.connection.AbstractConnection;
26 import com.android.tradefed.device.connection.AdbSshConnection;
27 import com.android.tradefed.invoker.TestInformation;
28 import com.android.tradefed.result.error.DeviceErrorIdentifier;
29 import com.android.tradefed.service.IRemoteFeature;
30 import com.android.tradefed.targetprep.TargetSetupError;
31 import com.android.tradefed.testtype.ITestInformationReceiver;
32 import com.android.tradefed.util.CommandResult;
33 import com.android.tradefed.util.CommandStatus;
34 import com.android.tradefed.util.SerializationUtil;
35 
36 import com.proto.tradefed.feature.ErrorInfo;
37 import com.proto.tradefed.feature.FeatureRequest;
38 import com.proto.tradefed.feature.FeatureResponse;
39 
40 import java.io.IOException;
41 
42 /** Server side implementation of device snapshot. */
43 public class DeviceSnapshotFeature
44         implements IRemoteFeature, IConfigurationReceiver, ITestInformationReceiver {
45 
46     public static final String DEVICE_SNAPSHOT_FEATURE_NAME = "snapshotDevice";
47     public static final String DEVICE_NAME = "device_name";
48     public static final String SNAPSHOT_ID = "snapshot_id";
49     public static final String RESTORE_FLAG = "restore_flag";
50 
51     private IConfiguration mConfig;
52     private TestInformation mTestInformation;
53 
54     @Override
getName()55     public String getName() {
56         return DEVICE_SNAPSHOT_FEATURE_NAME;
57     }
58 
59     @Override
setConfiguration(IConfiguration configuration)60     public void setConfiguration(IConfiguration configuration) {
61         mConfig = configuration;
62     }
63 
64     @Override
setTestInformation(TestInformation testInformation)65     public void setTestInformation(TestInformation testInformation) {
66         mTestInformation = testInformation;
67     }
68 
69     @Override
getTestInformation()70     public TestInformation getTestInformation() {
71         return mTestInformation;
72     }
73 
74     @Override
execute(FeatureRequest request)75     public FeatureResponse execute(FeatureRequest request) {
76         FeatureResponse.Builder responseBuilder = FeatureResponse.newBuilder();
77         String deviceName = request.getArgsMap().get(DEVICE_NAME);
78         if (deviceName == null) {
79             responseBuilder.setErrorInfo(
80                     ErrorInfo.newBuilder().setErrorTrace("No device_name args specified."));
81             return responseBuilder.build();
82         }
83 
84         IDeviceConfiguration configHolder = mConfig.getDeviceConfigByName(deviceName);
85         int index = 0;
86         for (IDeviceConfiguration deviceConfig : mConfig.getDeviceConfig()) {
87             if (deviceConfig == configHolder) {
88                 break;
89             }
90             index++;
91         }
92 
93         try {
94             mTestInformation.setActiveDeviceIndex(index);
95             AbstractConnection connection = mTestInformation.getDevice().getConnection();
96             // TODO: Support NestedRemoteDevice
97             if ((mTestInformation.getDevice() instanceof RemoteAndroidVirtualDevice)
98                     || (connection instanceof AdbSshConnection)) {
99                 GceAvdInfo info = getAvdInfo(mTestInformation.getDevice(), connection);
100                 if (info == null) {
101                     throw new RuntimeException("GceAvdInfo was null. skipping");
102                 }
103                 Integer offset = info.getDeviceOffset();
104                 String user = info.getInstanceUser();
105 
106                 String snapshotId = request.getArgsMap().get(SNAPSHOT_ID);
107                 boolean restoreFlag = Boolean.parseBoolean(request.getArgsMap().get(RESTORE_FLAG));
108                 if (restoreFlag) {
109                     restoreSnapshot(responseBuilder, connection, user, offset, snapshotId);
110                 } else {
111                     snapshot(responseBuilder, connection, user, offset, snapshotId);
112                 }
113             } else {
114                 String error =
115                         String.format(
116                                 "Device type %s with connection type %s doesn't support"
117                                         + " snapshotting",
118                                 mTestInformation.getDevice().getClass().getSimpleName(),
119                                 connection != null
120                                         ? connection.getClass().getSimpleName()
121                                         : "[null]");
122                 responseBuilder.setErrorInfo(ErrorInfo.newBuilder().setErrorTrace(error));
123                 return responseBuilder.build();
124             }
125         } catch (DeviceNotAvailableException | TargetSetupError e) {
126             String error = "Failed to snapshot device.";
127             try {
128                 error = SerializationUtil.serializeToString(e);
129             } catch (RuntimeException | IOException serializationError) {
130                 // Ignore
131             }
132             responseBuilder.setErrorInfo(ErrorInfo.newBuilder().setErrorTrace(error));
133         } finally {
134             mTestInformation.setActiveDeviceIndex(0);
135         }
136         return responseBuilder.build();
137     }
138 
snapshot( FeatureResponse.Builder responseBuilder, AbstractConnection connection, String user, Integer offset, String snapshotId)139     private void snapshot(
140             FeatureResponse.Builder responseBuilder,
141             AbstractConnection connection,
142             String user,
143             Integer offset,
144             String snapshotId)
145             throws DeviceNotAvailableException, TargetSetupError {
146         String response =
147                 String.format(
148                         "Attempting snapshot device on %s (%s).",
149                         mTestInformation.getDevice().getSerialNumber(),
150                         mTestInformation.getDevice().getClass().getSimpleName());
151         try {
152             long startTime = System.currentTimeMillis();
153             CommandResult result = snapshotGce(connection, user, offset, snapshotId);
154             if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
155                 throw new DeviceNotAvailableException(
156                         String.format(
157                                 "Failed to snapshot device: %s. status:%s\n"
158                                         + "stdout: %s\n"
159                                         + "stderr:%s",
160                                 mTestInformation.getDevice().getSerialNumber(),
161                                 result.getStatus(),
162                                 result.getStdout(),
163                                 result.getStderr()),
164                         mTestInformation.getDevice().getSerialNumber(),
165                         DeviceErrorIdentifier.DEVICE_FAILED_TO_SNAPSHOT);
166             }
167             response +=
168                     String.format(
169                             " Snapshot finished in %d ms.", System.currentTimeMillis() - startTime);
170         } finally {
171             responseBuilder.setResponse(response);
172         }
173     }
174 
restoreSnapshot( FeatureResponse.Builder responseBuilder, AbstractConnection connection, String user, Integer offset, String snapshotId)175     private void restoreSnapshot(
176             FeatureResponse.Builder responseBuilder,
177             AbstractConnection connection,
178             String user,
179             Integer offset,
180             String snapshotId)
181             throws DeviceNotAvailableException, TargetSetupError {
182         String response =
183                 String.format(
184                         "Attempting restore device snapshot on %s (%s) to %s.",
185                         mTestInformation.getDevice().getSerialNumber(),
186                         mTestInformation.getDevice().getClass().getSimpleName(),
187                         snapshotId);
188         try {
189             long startTime = System.currentTimeMillis();
190             CommandResult result = restoreSnapshotGce(connection, user, offset, snapshotId);
191             if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
192                 throw new DeviceNotAvailableException(
193                         String.format(
194                                 "Failed to restore snapshot on device: %s. status:%s\n"
195                                         + "stdout: %s\n"
196                                         + "stderr:%s",
197                                 mTestInformation.getDevice().getSerialNumber(),
198                                 result.getStatus(),
199                                 result.getStdout(),
200                                 result.getStderr()),
201                         mTestInformation.getDevice().getSerialNumber(),
202                         DeviceErrorIdentifier.DEVICE_FAILED_TO_RESTORE_SNAPSHOT);
203             }
204             response +=
205                     String.format(
206                             " Restoring snapshot finished in %d ms.",
207                             System.currentTimeMillis() - startTime);
208         } finally {
209             responseBuilder.setResponse(response);
210         }
211     }
212 
getAvdInfo(ITestDevice device, AbstractConnection connection)213     private GceAvdInfo getAvdInfo(ITestDevice device, AbstractConnection connection) {
214         if (connection instanceof AdbSshConnection) {
215             return ((AdbSshConnection) connection).getAvdInfo();
216         }
217         if (device instanceof RemoteAndroidVirtualDevice) {
218             return ((RemoteAndroidVirtualDevice) device).getAvdInfo();
219         }
220         return null;
221     }
222 
snapshotGce( AbstractConnection connection, String user, Integer offset, String snapshotId)223     private CommandResult snapshotGce(
224             AbstractConnection connection, String user, Integer offset, String snapshotId)
225             throws TargetSetupError {
226         if (connection instanceof AdbSshConnection) {
227             return ((AdbSshConnection) connection).snapshotGce(user, offset, snapshotId);
228         }
229         CommandResult res = new CommandResult(CommandStatus.EXCEPTION);
230         res.setStderr("Incorrect connection type while attempting device snapshot");
231         return res;
232     }
233 
restoreSnapshotGce( AbstractConnection connection, String user, Integer offset, String snapshotId)234     private CommandResult restoreSnapshotGce(
235             AbstractConnection connection, String user, Integer offset, String snapshotId)
236             throws TargetSetupError {
237         if (connection instanceof AdbSshConnection) {
238             return ((AdbSshConnection) connection).restoreSnapshotGce(user, offset, snapshotId);
239         }
240         CommandResult res = new CommandResult(CommandStatus.EXCEPTION);
241         res.setStderr("Incorrect connection type while attempting device restore");
242         return res;
243     }
244 }
245