1 /*
2  * Copyright (C) 2022 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.targetprep;
17 
18 import com.android.tradefed.config.Option;
19 import com.android.tradefed.config.OptionClass;
20 import com.android.tradefed.util.flag.DeviceFeatureFlag;
21 import com.android.tradefed.device.DeviceNotAvailableException;
22 import com.android.tradefed.device.ITestDevice;
23 import com.android.tradefed.invoker.TestInformation;
24 import com.android.tradefed.log.LogUtil.CLog;
25 import com.android.tradefed.result.error.DeviceErrorIdentifier;
26 import com.android.tradefed.result.error.InfraErrorIdentifier;
27 import com.android.tradefed.util.CommandResult;
28 import com.android.tradefed.util.CommandStatus;
29 
30 import com.google.common.base.Strings;
31 
32 import java.io.BufferedReader;
33 import java.io.ByteArrayInputStream;
34 import java.io.File;
35 import java.io.FileInputStream;
36 import java.io.IOException;
37 import java.io.InputStreamReader;
38 import java.util.ArrayList;
39 import java.util.HashMap;
40 import java.util.List;
41 import java.util.Map;
42 import java.util.Objects;
43 import java.util.stream.Collectors;
44 
45 /**
46  * Updates the DeviceConfig (feature flags tuned by a remote service).
47  *
48  * <p>This can be used to reproduce the state of a device (by dumping all flag values to a file
49  * using `adb shell device_config list`) or to bulk enable/disable flags (all-on/all-off testing).
50  *
51  * <p>Example usage:
52  *
53  * <ul>
54  *   <li>To use for all-on/all-off testing, specify the necessary flag file:
55  *       <pre>--flag-file=flag_file_path</pre>
56  *   <li>To override one or more flags, specify their values (can be combined with flag files):
57  *       <pre>--flag-file=flag_file_path --flag-value=namespace/name=value</pre>
58  *   <li>To use for reversibility testing, specify the all-on file followed by the all-off file, and
59  *       enable rebooting between the two files:
60  *       <pre>--flag-file=all_on_file_path --flag-file=all_off_file_path --reboot-between-flag-files
61  *       </pre>
62  * </ul>
63  *
64  * <p>Should be used in combination with {@link DeviceSetup} to disable DeviceConfig syncing during
65  * the test which could overwrite the changes made by this preparer.
66  */
67 @OptionClass(alias = "feature-flags")
68 public class FeatureFlagTargetPreparer extends BaseTargetPreparer {
69 
70     @Option(
71             name = "flag-file",
72             description = "File containing flag values to apply. Can be repeated.")
73     private List<File> mFlagFiles = new ArrayList<>();
74 
75     @Option(name = "flag-value", description = "Additional flag values to apply. Can be repeated.")
76     private List<String> mFlagValues = new ArrayList<>();
77 
78     @Option(
79             name = "reboot-between-flag-files",
80             description = "Enables reboots after each input file. Used for reversibility testing.")
81     private boolean mRebootBetweenFlagFiles = false;
82 
83     private final Map<String, Map<String, String>> mFlagsToRestore = new HashMap<>();
84 
85     @Override
setUp(TestInformation testInformation)86     public void setUp(TestInformation testInformation)
87             throws TargetSetupError, BuildError, DeviceNotAvailableException {
88         ITestDevice device = testInformation.getDevice();
89         if (mFlagFiles.isEmpty() && mFlagValues.isEmpty()) {
90             CLog.i("No flag-file or flag-value option provided, skipping");
91             return;
92         }
93 
94         // Parse input flag files and values.
95         List<Map<String, Map<String, String>>> flagBundles = new ArrayList<>();
96         for (File flagFile : mFlagFiles) {
97             if (flagFile == null || !flagFile.isFile()) {
98                 throw new TargetSetupError(
99                         String.format("Flag file '%s' not found", flagFile),
100                         device.getDeviceDescriptor(),
101                         InfraErrorIdentifier.CONFIGURED_ARTIFACT_NOT_FOUND);
102             }
103             flagBundles.add(parseFlags(device, flagFile));
104         }
105         if (!mFlagValues.isEmpty()) {
106             flagBundles.add(parseFlags(device, mFlagValues, /* throwIfInvalid */ true));
107         }
108 
109         // Keep track of initial flags to be able to restore them during tearDown.
110         Map<String, Map<String, String>> initialFlags = null;
111         boolean flagsUpdated = false;
112 
113         for (Map<String, Map<String, String>> targetFlags : flagBundles) {
114             // Determine current flag values to diff against target flag values.
115             Map<String, Map<String, String>> currentFlags = listFlags(device);
116             if (initialFlags == null) {
117                 initialFlags = currentFlags;
118             }
119 
120             for (String namespace : targetFlags.keySet()) {
121                 // Ignore unchanged flag values.
122                 Map<String, String> currentValues = currentFlags.getOrDefault(namespace, Map.of());
123                 Map<String, String> targetValues = targetFlags.get(namespace);
124                 // target flags - current flags = new flags to apply
125                 targetValues.entrySet().removeAll(currentValues.entrySet());
126                 // new flags to apply - initial flags = new flags to restore
127                 updateFlagsToRestore(
128                         namespace, initialFlags.getOrDefault(namespace, Map.of()), targetValues);
129             }
130 
131             if (targetFlags.values().stream().allMatch(Map::isEmpty)) {
132                 continue; // No flags to update.
133             }
134             updateFlags(device, targetFlags);
135             flagsUpdated = true;
136             if (mRebootBetweenFlagFiles) {
137                 device.reboot();
138             }
139         }
140 
141         if (!mRebootBetweenFlagFiles && flagsUpdated) {
142             device.reboot();
143         }
144     }
145 
146     @Override
tearDown(TestInformation testInformation, Throwable e)147     public void tearDown(TestInformation testInformation, Throwable e)
148             throws DeviceNotAvailableException {
149         if (e instanceof DeviceNotAvailableException
150                 || mFlagsToRestore.isEmpty()
151                 || mFlagsToRestore.values().stream().allMatch(Map::isEmpty)) {
152             return;
153         }
154         try {
155             ITestDevice device = testInformation.getDevice();
156             updateFlags(device, mFlagsToRestore);
157             device.reboot();
158         } catch (TargetSetupError tse) {
159             CLog.e("Failed to restore flags: %s", tse);
160         }
161     }
162 
listFlags(ITestDevice device)163     private Map<String, Map<String, String>> listFlags(ITestDevice device)
164             throws DeviceNotAvailableException, TargetSetupError {
165         String values = runCommand(device, "device_config list");
166         try (ByteArrayInputStream stream = new ByteArrayInputStream(values.getBytes());
167                 BufferedReader reader = new BufferedReader(new InputStreamReader(stream))) {
168             return parseFlags(device, reader.lines().collect(Collectors.toList()), false);
169         } catch (IOException ioe) {
170             throw new TargetSetupError(
171                     "Failed to parse device flags",
172                     ioe,
173                     device.getDeviceDescriptor(),
174                     DeviceErrorIdentifier.DEVICE_UNEXPECTED_RESPONSE);
175         }
176     }
177 
parseFlags(ITestDevice device, File flagFile)178     private Map<String, Map<String, String>> parseFlags(ITestDevice device, File flagFile)
179             throws TargetSetupError {
180         try (FileInputStream stream = new FileInputStream(flagFile);
181                 BufferedReader reader = new BufferedReader(new InputStreamReader(stream))) {
182             return parseFlags(device, reader.lines().collect(Collectors.toList()), false);
183         } catch (IOException ioe) {
184             throw new TargetSetupError(
185                     "Failed to parse flag file",
186                     ioe,
187                     device.getDeviceDescriptor(),
188                     InfraErrorIdentifier.LAB_HOST_FILESYSTEM_ERROR);
189         }
190     }
191 
parseFlags( ITestDevice device, List<String> lines, boolean throwIfInvalid)192     private Map<String, Map<String, String>> parseFlags(
193             ITestDevice device, List<String> lines, boolean throwIfInvalid)
194             throws TargetSetupError {
195         Map<String, Map<String, String>> flags = new HashMap<>();
196         for (String line : lines) {
197             try {
198                 DeviceFeatureFlag deviceFeatureFlag = new DeviceFeatureFlag(line);
199                 flags.computeIfAbsent(deviceFeatureFlag.getNamespace(), ns -> new HashMap<>())
200                         .put(deviceFeatureFlag.getFlagName(), deviceFeatureFlag.getFlagValue());
201             } catch (IllegalArgumentException e) {
202                 if (throwIfInvalid) {
203                     throw new TargetSetupError(
204                             String.format("Failed to parse flag data: %s", line),
205                             e,
206                             device.getDeviceDescriptor(),
207                             InfraErrorIdentifier.OPTION_CONFIGURATION_ERROR);
208                 }
209                 CLog.w("Skipping invalid flag data: %s", line);
210                 continue;
211             }
212         }
213         return flags;
214     }
215 
updateFlagsToRestore( String namespace, Map<String, String> initialValues, Map<String, String> updatedValues)216     private void updateFlagsToRestore(
217             String namespace,
218             Map<String, String> initialValues,
219             Map<String, String> updatedValues) {
220         // Keep track of flag values to restore.
221         for (String name : updatedValues.keySet()) {
222             String initialValue = initialValues.get(name);
223             String updatedValue = updatedValues.get(name);
224             if (Objects.equals(updatedValue, initialValue)) {
225                 // When the first file updates a default flag value and the second file sets the
226                 // flag back to its default, there is no need to restore this flag.
227                 mFlagsToRestore.getOrDefault(namespace, Map.of()).remove(name);
228             } else {
229                 mFlagsToRestore
230                         .computeIfAbsent(namespace, ns -> new HashMap<>())
231                         .put(name, initialValue);
232             }
233         }
234     }
235 
updateFlags(ITestDevice device, Map<String, Map<String, String>> flags)236     private void updateFlags(ITestDevice device, Map<String, Map<String, String>> flags)
237             throws DeviceNotAvailableException, TargetSetupError {
238         for (String namespace : flags.keySet()) {
239             for (Map.Entry<String, String> entry : flags.get(namespace).entrySet()) {
240                 String name = entry.getKey();
241                 String value = entry.getValue();
242                 updateFlag(device, namespace, name, value);
243             }
244         }
245     }
246 
updateFlag(ITestDevice device, String namespace, String name, String value)247     private void updateFlag(ITestDevice device, String namespace, String name, String value)
248             throws DeviceNotAvailableException, TargetSetupError {
249         if (Strings.isNullOrEmpty(value)) { // `device_config put` does not support empty values.
250             runCommand(device, String.format("device_config delete '%s' '%s'", namespace, name));
251         } else {
252             runCommand(
253                     device,
254                     String.format("device_config put '%s' '%s' '%s'", namespace, name, value));
255         }
256     }
257 
runCommand(ITestDevice device, String command)258     private String runCommand(ITestDevice device, String command)
259             throws DeviceNotAvailableException, TargetSetupError {
260         CommandResult result = device.executeShellV2Command(command);
261         if (result.getStatus() != CommandStatus.SUCCESS) {
262             throw new TargetSetupError(
263                     String.format(
264                             "Command %s failed, stdout = [%s], stderr = [%s]",
265                             command, result.getStdout(), result.getStderr()),
266                     device.getDeviceDescriptor(),
267                     DeviceErrorIdentifier.DEVICE_UNEXPECTED_RESPONSE);
268         }
269         return result.getStdout();
270     }
271 }
272