1 /*
2  * Copyright (C) 2019 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 android.device.collectors;
17 
18 import android.app.StatsManager;
19 import android.app.StatsManager.StatsUnavailableException;
20 import android.content.Context;
21 import android.content.res.AssetManager;
22 import android.os.Bundle;
23 import android.os.Environment;
24 import android.os.SystemClock;
25 import android.util.Log;
26 import android.util.StatsLog;
27 
28 import androidx.annotation.VisibleForTesting;
29 import androidx.test.InstrumentationRegistry;
30 
31 import com.android.internal.os.nano.StatsdConfigProto;
32 import com.android.os.nano.AtomsProto;
33 
34 import com.google.protobuf.nano.CodedOutputByteBufferNano;
35 import com.google.protobuf.nano.ExtendableMessageNano;
36 import com.google.protobuf.nano.InvalidProtocolBufferNanoException;
37 
38 import org.junit.runner.Description;
39 import org.junit.runner.Result;
40 
41 import java.io.ByteArrayOutputStream;
42 import java.io.File;
43 import java.io.FileInputStream;
44 import java.io.IOException;
45 import java.io.InputStream;
46 import java.nio.file.Files;
47 import java.nio.file.Path;
48 import java.nio.file.Paths;
49 import java.util.Arrays;
50 import java.util.HashMap;
51 import java.util.List;
52 import java.util.Map;
53 import java.util.UUID;
54 import java.util.concurrent.TimeUnit;
55 import java.util.stream.Collectors;
56 
57 /**
58  * A device-side metric listener that collects statsd-based metrics using bundled config files.
59  *
60  * <p>Statsd configs can either be passed in by name, in which case they must be bundled into the
61  * test APK as assets, or by their absolute on-device path. Comma-separated values are supported.
62  */
63 public class StatsdListener extends BaseMetricListener {
64     private static final String LOG_TAG = StatsdListener.class.getSimpleName();
65 
66     static final String OPTION_CONFIGS_RUN_LEVEL = "statsd-configs-run-level";
67     static final String OPTION_CONFIGS_TEST_LEVEL = "statsd-configs-test-level";
68 
69     // Sub-directory within the test APK's assets/ directory to look for configs.
70     static final String CONFIG_SUB_DIRECTORY = "statsd-configs";
71     // File extension for all statsd configs.
72     static final String PROTO_EXTENSION = ".pb";
73 
74     // Parent directory for all statsd reports.
75     static final String REPORT_PATH_ROOT = "statsd-reports";
76     // Sub-directory for test run reports.
77     static final String REPORT_PATH_RUN_LEVEL = "run-level";
78     // Sub-directory for test-level reports.
79     static final String REPORT_PATH_TEST_LEVEL = "test-level";
80     // Suffix template for test-level metric report files.
81     static final String TEST_SUFFIX_TEMPLATE = "_%s-%d";
82 
83     // Common prefix for the metric key pointing to the report path.
84     static final String REPORT_KEY_PREFIX = "statsd-";
85     // Common prefix for the metric file.
86     static final String REPORT_FILENAME_PREFIX = "statsd-";
87     // Prefix for configs loaded from the device.
88     @VisibleForTesting static final String LOCAL_CONFIG_PREFIX = "local-config-";
89 
90     // Labels used to signify test events to statsd with the AppBreadcrumbReported atom.
91     static final int RUN_EVENT_LABEL = 7;
92     static final int TEST_EVENT_LABEL = 11;
93     // A short delay after pushing the AppBreadcrumbReported event so that metrics can be dumped.
94     static final long METRIC_PULL_DELAY = TimeUnit.SECONDS.toMillis(1);
95 
96     // Configs used for the test run and each test, respectively.
97     private Map<String, StatsdConfigProto.StatsdConfig> mRunLevelConfigs =
98             new HashMap<String, StatsdConfigProto.StatsdConfig>();
99     private Map<String, StatsdConfigProto.StatsdConfig> mTestLevelConfigs =
100             new HashMap<String, StatsdConfigProto.StatsdConfig>();
101 
102     // Map to associate config names with their config Ids.
103     private Map<String, Long> mRunLevelConfigIds = new HashMap<String, Long>();
104     private Map<String, Long> mTestLevelConfigIds = new HashMap<String, Long>();
105 
106     // "Counter" for test iterations, keyed by the display name of each test's description.
107     private Map<String, Integer> mTestIterations = new HashMap<String, Integer>();
108 
109     // Cached stats manager instance.
110     private StatsManager mStatsManager;
111 
112     /** Register the test run configs with {@link StatsManager} before the test run starts. */
113     @Override
onTestRunStart(DataRecord runData, Description description)114     public void onTestRunStart(DataRecord runData, Description description) {
115         // The argument parsing has to be performed here as the instrumentation has not yet been
116         // registered when the constructor of this class is called.
117         mRunLevelConfigs.putAll(getConfigsFromOption(OPTION_CONFIGS_RUN_LEVEL));
118         mTestLevelConfigs.putAll(getConfigsFromOption(OPTION_CONFIGS_TEST_LEVEL));
119 
120         mRunLevelConfigIds = registerConfigsWithStatsManager(mRunLevelConfigs);
121 
122         if (!logStart(RUN_EVENT_LABEL)) {
123             Log.w(LOG_TAG, "Failed to log a test run start event. Metrics might be incomplete.");
124         }
125     }
126 
127     /**
128      * Dump the test run stats reports to the test run subdirectory after the test run ends.
129      *
130      * <p>Dumps the stats regardless of whether all the tests pass.
131      */
132     @Override
onTestRunEnd(DataRecord runData, Result result)133     public void onTestRunEnd(DataRecord runData, Result result) {
134         if (!logStop(RUN_EVENT_LABEL)) {
135             Log.w(LOG_TAG, "Failed to log a test run end event. Metrics might be incomplete.");
136         }
137         SystemClock.sleep(METRIC_PULL_DELAY);
138 
139         Map<String, File> configReports =
140                 pullReportsAndRemoveConfigs(
141                         mRunLevelConfigIds, Paths.get(REPORT_PATH_ROOT, REPORT_PATH_RUN_LEVEL), "");
142         for (String configName : configReports.keySet()) {
143             runData.addFileMetric(REPORT_KEY_PREFIX + configName, configReports.get(configName));
144         }
145     }
146 
147     /** Register the test-level configs with {@link StatsManager} before each test starts. */
148     @Override
onTestStart(DataRecord testData, Description description)149     public void onTestStart(DataRecord testData, Description description) {
150         mTestIterations.computeIfPresent(description.getDisplayName(), (name, count) -> count + 1);
151         mTestIterations.computeIfAbsent(description.getDisplayName(), name -> 1);
152         mTestLevelConfigIds = registerConfigsWithStatsManager(mTestLevelConfigs);
153 
154         if (!logStart(TEST_EVENT_LABEL)) {
155             Log.w(LOG_TAG, "Failed to log a test start event. Metrics might be incomplete.");
156         }
157     }
158 
159     /**
160      * Dump the test-level stats reports to the test-specific subdirectory after the test ends.
161      *
162      * <p>Dumps the stats regardless of whether the test passes.
163      */
164     @Override
onTestEnd(DataRecord testData, Description description)165     public void onTestEnd(DataRecord testData, Description description) {
166         if (!logStop(TEST_EVENT_LABEL)) {
167             Log.w(LOG_TAG, "Failed to log a test end event. Metrics might be incomplete.");
168         }
169         SystemClock.sleep(METRIC_PULL_DELAY);
170 
171         Map<String, File> configReports =
172                 pullReportsAndRemoveConfigs(
173                         mTestLevelConfigIds,
174                         Paths.get(REPORT_PATH_ROOT, REPORT_PATH_TEST_LEVEL),
175                         getTestSuffix(description));
176         for (String configName : configReports.keySet()) {
177             testData.addFileMetric(REPORT_KEY_PREFIX + configName, configReports.get(configName));
178         }
179     }
180 
181     /**
182      * Register a set of statsd configs and return their config IDs in a {@link Map}.
183      *
184      * @param configs Map of (config name, config proto message)
185      * @return Map of (config name, config id)
186      */
registerConfigsWithStatsManager( final Map<String, StatsdConfigProto.StatsdConfig> configs)187     private Map<String, Long> registerConfigsWithStatsManager(
188             final Map<String, StatsdConfigProto.StatsdConfig> configs) {
189         Map<String, Long> configIds = new HashMap<String, Long>();
190         adoptShellPermissionIdentity();
191         for (String configName : configs.keySet()) {
192             try {
193                 long configId = getUniqueIdForConfig(configs.get(configName));
194                 StatsdConfigProto.StatsdConfig newConfig = clone(configs.get(configName));
195                 newConfig.id = configId;
196                 Log.i(LOG_TAG, String.format("Adding config %s with ID %d.", configName, configId));
197                 addStatsConfig(configId, serialize(newConfig));
198                 configIds.put(configName, configId);
199             } catch (IOException | StatsUnavailableException e) {
200                 Log.e(
201                         LOG_TAG,
202                         String.format(
203                                 "Failed to add statsd config %s due to %s.",
204                                 configName, e.toString()));
205             }
206         }
207         dropShellPermissionIdentity();
208         return configIds;
209     }
210 
211     /**
212      * For a set of statsd config ids, retrieve the config reports from {@link StatsManager}, remove
213      * the config and dump the reports into the designated directory on the device's external
214      * storage.
215      *
216      * @param configIds Map of (config name, config Id)
217      * @param directory relative directory on external storage to dump the report in. Each report
218      *     will be named after its config.
219      * @param suffix a suffix to append to the metric report file name, used to differentiate
220      *     between tests and left empty for the test run.
221      * @return Map of (config name, config report file)
222      */
pullReportsAndRemoveConfigs( final Map<String, Long> configIds, Path directory, String suffix)223     private Map<String, File> pullReportsAndRemoveConfigs(
224             final Map<String, Long> configIds, Path directory, String suffix) {
225         File externalStorage = Environment.getExternalStorageDirectory();
226         File saveDirectory = new File(externalStorage, directory.toString());
227         if (!saveDirectory.isDirectory()) {
228             saveDirectory.mkdirs();
229         }
230         Map<String, File> savedConfigFiles = new HashMap<String, File>();
231         adoptShellPermissionIdentity();
232         for (String configName : configIds.keySet()) {
233             // Dump the metric report to external storage.
234             com.android.os.nano.StatsLog.ConfigMetricsReportList reportList;
235             try {
236                 Log.i(
237                         LOG_TAG,
238                         String.format(
239                                 "Pulling metrics for config %s with ID %d.",
240                                 configName, configIds.get(configName)));
241                 reportList =
242                         com.android.os.nano.StatsLog.ConfigMetricsReportList.parseFrom(
243                                 getStatsReports(configIds.get(configName)));
244                 Log.i(
245                         LOG_TAG,
246                         String.format(
247                                 "Found %d metric %s from config %s.",
248                                 reportList.reports.length,
249                                 reportList.reports.length == 1 ? "report" : "reports",
250                                 configName));
251                 File reportFile =
252                         new File(
253                                 saveDirectory,
254                                 REPORT_FILENAME_PREFIX + configName + suffix + PROTO_EXTENSION);
255                 writeToFile(reportFile, serialize(reportList));
256                 savedConfigFiles.put(configName, reportFile);
257             } catch (StatsUnavailableException e) {
258                 Log.e(
259                         LOG_TAG,
260                         String.format(
261                                 "Failed to retrieve metrics for config %s due to %s.",
262                                 configName, e.toString()));
263             } catch (InvalidProtocolBufferNanoException e) {
264                 Log.e(
265                         LOG_TAG,
266                         String.format(
267                                 "Unable to parse report for config %s. Details: %s.",
268                                 configName, e.toString()));
269             } catch (IOException e) {
270                 Log.e(
271                         LOG_TAG,
272                         String.format(
273                                 "Failed to write metric report for config %s to device. "
274                                         + "Details: %s.",
275                                 configName, e.toString()));
276             }
277 
278             // Remove the statsd config.
279             try {
280                 Log.i(
281                         LOG_TAG,
282                         String.format(
283                                 "Removing config %s with ID %d.",
284                                 configName, configIds.get(configName)));
285                 removeStatsConfig(configIds.get(configName));
286             } catch (StatsUnavailableException e) {
287                 Log.e(
288                         LOG_TAG,
289                         String.format(
290                                 "Unable to remove config %s due to %s.", configName, e.toString()));
291             }
292         }
293         dropShellPermissionIdentity();
294         return savedConfigFiles;
295     }
296 
297     /**
298      * Adopt shell permission identity to communicate with {@link StatsManager}.
299      *
300      * @hide
301      */
302     @VisibleForTesting
adoptShellPermissionIdentity()303     protected void adoptShellPermissionIdentity() {
304         InstrumentationRegistry.getInstrumentation()
305                 .getUiAutomation()
306                 .adoptShellPermissionIdentity();
307     }
308 
309     /**
310      * Drop shell permission identity once communication with {@link StatsManager} is done.
311      *
312      * @hide
313      */
314     @VisibleForTesting
dropShellPermissionIdentity()315     protected void dropShellPermissionIdentity() {
316         InstrumentationRegistry.getInstrumentation()
317                 .getUiAutomation()
318                 .dropShellPermissionIdentity();
319     }
320 
321     /** Returns the cached {@link StatsManager} instance; if none exists, request and cache it. */
getStatsManager()322     private StatsManager getStatsManager() {
323         if (mStatsManager == null) {
324             mStatsManager =
325                     (StatsManager)
326                             InstrumentationRegistry.getTargetContext()
327                                     .getSystemService(Context.STATS_MANAGER);
328         }
329         return mStatsManager;
330     }
331 
332     /** Get the suffix for a test + iteration combination to differentiate it from other files. */
333     @VisibleForTesting
getTestSuffix(Description description)334     String getTestSuffix(Description description) {
335         return String.format(
336                 TEST_SUFFIX_TEMPLATE,
337                 formatDescription(description),
338                 mTestIterations.get(description.getDisplayName()));
339     }
340 
341     /** Format a JUnit {@link Description} to a desired string format. */
342     @VisibleForTesting
formatDescription(Description description)343     String formatDescription(Description description) {
344         // Use String.valueOf() to guard agaist a null class name. This normally should not happen
345         // but the Description class does not explicitly guarantee it.
346         String className = String.valueOf(description.getClassName());
347         String methodName = description.getMethodName();
348         return methodName == null ? className : String.join("#", className, methodName);
349     }
350 
351     /**
352      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
353      *
354      * @hide
355      */
356     @VisibleForTesting
addStatsConfig(long configKey, byte[] config)357     protected void addStatsConfig(long configKey, byte[] config) throws StatsUnavailableException {
358         getStatsManager().addConfig(configKey, config);
359     }
360 
361     /**
362      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
363      *
364      * @hide
365      */
366     @VisibleForTesting
removeStatsConfig(long configKey)367     protected void removeStatsConfig(long configKey) throws StatsUnavailableException {
368         mStatsManager.removeConfig(configKey);
369     }
370 
371     /**
372      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
373      *
374      * @hide
375      */
376     @VisibleForTesting
getStatsReports(long configKey)377     protected byte[] getStatsReports(long configKey) throws StatsUnavailableException {
378         return mStatsManager.getReports(configKey);
379     }
380 
381     /**
382      * Allow tests to stub out getting instrumentation arguments.
383      *
384      * @hide
385      */
386     @VisibleForTesting
getArguments()387     protected Bundle getArguments() {
388         return InstrumentationRegistry.getArguments();
389     }
390 
391     /**
392      * Allow tests to stub out file I/O.
393      *
394      * @hide
395      */
396     @VisibleForTesting
writeToFile(File f, byte[] content)397     protected File writeToFile(File f, byte[] content) throws IOException {
398         Files.write(f.toPath(), content);
399         return f;
400     }
401 
402     /**
403      * Allow tests to override the random ID generation. The config is passed in to allow a specific
404      * ID to be associated with a config in the test.
405      *
406      * @hide
407      */
408     @VisibleForTesting
getUniqueIdForConfig(StatsdConfigProto.StatsdConfig config)409     protected long getUniqueIdForConfig(StatsdConfigProto.StatsdConfig config) {
410         return (long) UUID.randomUUID().hashCode();
411     }
412 
413     /**
414      * Allow tests to stub out {@link AssetManager} interactions as that class is final and cannot .
415      * be mocked.
416      *
417      * @hide
418      */
419     @VisibleForTesting
openConfigWithAssetManager(AssetManager manager, String configName)420     protected InputStream openConfigWithAssetManager(AssetManager manager, String configName)
421             throws IOException {
422         String configFilePath =
423                 Paths.get(CONFIG_SUB_DIRECTORY, configName + PROTO_EXTENSION).toString();
424         return manager.open(configFilePath);
425     }
426 
427     /**
428      * Parse a config from its name or on-device path.
429      *
430      * <p>The option name is passed in for better error messaging.
431      */
parseConfig( final AssetManager manager, String optionName, String nameOrPath)432     private StatsdConfigProto.StatsdConfig parseConfig(
433             final AssetManager manager, String optionName, String nameOrPath) {
434         if (new File(nameOrPath).isAbsolute()) {
435             return parseConfigFromPath(optionName, nameOrPath);
436         }
437         return parseConfigFromName(manager, optionName, nameOrPath);
438     }
439 
440     /**
441      * Parse a config from its on-device path.
442      *
443      * <p>The option name is passed in for better error messaging.
444      */
parseConfigFromPath( String optionName, String configPath)445     private StatsdConfigProto.StatsdConfig parseConfigFromPath(
446             String optionName, String configPath) {
447         try (InputStream configStream = new FileInputStream(configPath)) {
448             try {
449                 byte[] serializedConfig = readInputStream(configStream);
450                 return fixPermissions(StatsdConfigProto.StatsdConfig.parseFrom(serializedConfig));
451             } catch (IOException e) {
452                 throw new RuntimeException(
453                         String.format(
454                                 "Cannot parse config %s in option %s.", configPath, optionName),
455                         e);
456             }
457         } catch (IOException e) {
458             throw new IllegalArgumentException(
459                     String.format(
460                             "Config path %s in option %s does not exist", configPath, optionName));
461         }
462     }
463 
464     /**
465      * Parse a config from its name using {@link AssetManager}.
466      *
467      * <p>The option name is passed in for better error messaging.
468      */
parseConfigFromName( final AssetManager manager, String optionName, String configName)469     private StatsdConfigProto.StatsdConfig parseConfigFromName(
470             final AssetManager manager, String optionName, String configName) {
471         try (InputStream configStream = openConfigWithAssetManager(manager, configName)) {
472             try {
473                 byte[] serializedConfig = readInputStream(configStream);
474                 return fixPermissions(StatsdConfigProto.StatsdConfig.parseFrom(serializedConfig));
475             } catch (IOException e) {
476                 throw new RuntimeException(
477                         String.format(
478                                 "Cannot parse config %s in option %s.", configName, optionName),
479                         e);
480             }
481         } catch (IOException e) {
482             throw new IllegalArgumentException(
483                     String.format(
484                             "Config name %s in option %s does not exist", configName, optionName));
485         }
486     }
487 
488     /**
489      * Parse the suppplied option to get a set of statsd configs keyed by their names.
490      *
491      * @hide
492      */
493     @VisibleForTesting
getConfigsFromOption(String optionName)494     protected Map<String, StatsdConfigProto.StatsdConfig> getConfigsFromOption(String optionName) {
495         List<String> configNames =
496                 Arrays.asList(getArguments().getString(optionName, "").split(","))
497                         .stream()
498                         .map(s -> s.trim())
499                         .filter(s -> !s.isEmpty())
500                         .distinct()
501                         .collect(Collectors.toList());
502         // Look inside the APK assets for the configuration file.
503         final AssetManager manager = InstrumentationRegistry.getContext().getAssets();
504         return configNames.stream()
505                 .collect(
506                         Collectors.toMap(
507                                 nameOrPath -> getConfigShortName(nameOrPath),
508                                 nameOrPath -> parseConfig(manager, optionName, nameOrPath)));
509     }
510 
511     /**
512      * Get the "short name" of a statsd config.
513      *
514      * <p>Configs that are bundled into the APK and loaded using the asset manager is used as-is.
515      * Configs that are loaded from an on-device path use their file name, sans the file suffix,
516      * with a prefix specific to local configs.
517      */
getConfigShortName(String nameOrPath)518     private String getConfigShortName(String nameOrPath) {
519         if (new File(nameOrPath).isAbsolute()) {
520             // If the config name/path is an absolute path, it is an on-device local path.
521             return LOCAL_CONFIG_PREFIX
522                     + com.google.common.io.Files.getNameWithoutExtension(nameOrPath);
523         }
524         return nameOrPath;
525     }
526 
527     /**
528      * Log a "start" AppBreadcrumbReported event to statsd. Wraps a static method for testing.
529      *
530      * @hide
531      */
532     @VisibleForTesting
logStart(int label)533     protected boolean logStart(int label) {
534         return StatsLog.logStart(label);
535     }
536 
537     /**
538      * Log a "stop" AppBreadcrumbReported event to statsd. Wraps a static method for testing.
539      *
540      * @hide
541      */
542     @VisibleForTesting
logStop(int label)543     protected boolean logStop(int label) {
544         return StatsLog.logStop(label);
545     }
546 
547     /**
548      * Add a few permission-related options to the statsd config.
549      *
550      * <p>This is related to some new permission restrictions in RVC.
551      */
fixPermissions(StatsdConfigProto.StatsdConfig config)552     private StatsdConfigProto.StatsdConfig fixPermissions(StatsdConfigProto.StatsdConfig config)
553             throws IOException {
554         StatsdConfigProto.StatsdConfig newConfig = clone(config);
555         newConfig.defaultPullPackages =
556                 concat(config.defaultPullPackages, new String[] {"AID_SYSTEM"});
557         newConfig.whitelistedAtomIds =
558                 concat(
559                         config.whitelistedAtomIds,
560                         new int[] {AtomsProto.Atom.APP_BREADCRUMB_REPORTED_FIELD_NUMBER});
561         return newConfig;
562     }
563 
564     // Some utilities for Nano protos.
565 
serialize( ExtendableMessageNano<T> message)566     private static <T extends ExtendableMessageNano<T>> byte[] serialize(
567             ExtendableMessageNano<T> message) throws IOException {
568         byte[] serialized = new byte[message.getSerializedSize()];
569         CodedOutputByteBufferNano buffer = CodedOutputByteBufferNano.newInstance(serialized);
570         message.writeTo(buffer);
571         return serialized;
572     }
573 
clone(StatsdConfigProto.StatsdConfig config)574     private static StatsdConfigProto.StatsdConfig clone(StatsdConfigProto.StatsdConfig config)
575             throws IOException {
576         byte[] output = serialize(config);
577         return StatsdConfigProto.StatsdConfig.parseFrom(output);
578     }
579 
readInputStream(InputStream in)580     private static byte[] readInputStream(InputStream in) throws IOException {
581         ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
582         try {
583             byte[] buffer = new byte[1024];
584             int size = in.read(buffer);
585             while (size > 0) {
586                 outputStream.write(buffer, 0, size);
587                 size = in.read(buffer);
588             }
589             return outputStream.toByteArray();
590         } finally {
591             outputStream.close();
592         }
593     }
594 
595     // Array Concatenation
596 
concat(int[] source, int[] items)597     private static int[] concat(int[] source, int[] items) {
598         int[] concatenated = Arrays.copyOf(source, source.length + items.length);
599         System.arraycopy(items, 0, concatenated, source.length, items.length);
600         return concatenated;
601     }
602 
concat(T[] source, T[] items)603     private static <T> T[] concat(T[] source, T[] items) {
604         T[] concatenated = Arrays.copyOf(source, source.length + items.length);
605         System.arraycopy(items, 0, concatenated, source.length, items.length);
606         return concatenated;
607     }
608 }
609