1#!/usr/bin/python3
2#
3# Copyright (C) 2021 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17#
18# Generates profiles from the set of all methods in a given set of dex/jars and
19# bisects to find minimal repro sets.
20#
21
22import shlex
23import argparse
24import pylibdexfile
25import math
26import subprocess
27from collections import namedtuple
28import sys
29import random
30import os
31
32ApkEntry = namedtuple("ApkEntry", ["file", "location"])
33
34
35def get_parser():
36  parser = argparse.ArgumentParser(
37      description="Bisect profile contents. We will wait while the user runs test"
38  )
39
40  class ApkAction(argparse.Action):
41
42    def __init__(self, option_strings, dest, **kwargs):
43      super(ApkAction, self).__init__(option_strings, dest, **kwargs)
44
45    def __call__(self, parser, namespace, values, option_string=None):
46      lst = getattr(namespace, self.dest)
47      if lst is None:
48        setattr(namespace, self.dest, [])
49        lst = getattr(namespace, self.dest)
50      if len(values) == 1:
51        values = (values[0], values[0])
52      assert len(values) == 2, values
53      lst.append(ApkEntry(*values))
54
55  apks = parser.add_argument_group(title="APK selection")
56  apks.add_argument(
57      "--apk",
58      action=ApkAction,
59      dest="apks",
60      nargs=1,
61      default=[],
62      help="an apk/dex/jar to get methods from. Uses same path as location. " +
63           "Use --apk-and-location if this isn't desired."
64  )
65  apks.add_argument(
66      "--apk-and-location",
67      action=ApkAction,
68      nargs=2,
69      dest="apks",
70      help="an apk/dex/jar + location to get methods from."
71  )
72  profiles = parser.add_argument_group(
73      title="Profile selection").add_mutually_exclusive_group()
74  profiles.add_argument(
75      "--input-text-profile", help="a text profile to use for bisect")
76  profiles.add_argument("--input-profile", help="a profile to use for bisect")
77  parser.add_argument(
78      "--output-source", help="human readable file create the profile from")
79  parser.add_argument("--test-exec", help="file to exec (without arguments) to test a" +
80                                           " candidate. Test should exit 0 if the issue" +
81                                           " is not present and non-zero if the issue is" +
82                                           " present.")
83  parser.add_argument("output_file", help="file we will write the profiles to")
84  return parser
85
86
87def dump_files(meths, args, output):
88  for m in meths:
89    print("HS{}".format(m), file=output)
90  output.flush()
91  profman_args = [
92      "profmand", "--reference-profile-file={}".format(args.output_file),
93      "--create-profile-from={}".format(args.output_source)
94  ]
95  print(" ".join(map(shlex.quote, profman_args)))
96  for apk in args.apks:
97    profman_args += [
98        "--apk={}".format(apk.file), "--dex-location={}".format(apk.location)
99    ]
100  profman = subprocess.run(profman_args)
101  profman.check_returncode()
102
103
104def get_answer(args):
105  if args.test_exec is None:
106    while True:
107      answer = input("Does the file at {} cause the issue (y/n):".format(
108          args.output_file))
109      if len(answer) >= 1 and answer[0].lower() == "y":
110        return "y"
111      elif len(answer) >= 1 and answer[0].lower() == "n":
112        return "n"
113      else:
114        print("Please enter 'y' or 'n' only!")
115  else:
116    test_args = shlex.split(args.test_exec)
117    print(" ".join(map(shlex.quote, test_args)))
118    answer = subprocess.run(test_args)
119    if answer.returncode == 0:
120      return "n"
121    else:
122      return "y"
123
124def run_test(meths, args):
125  with open(args.output_source, "wt") as output:
126    dump_files(meths, args, output)
127    print("Currently testing {} methods. ~{} rounds to go.".format(
128        len(meths), 1 + math.floor(math.log2(len(meths)))))
129  return get_answer(args)
130
131def main():
132  parser = get_parser()
133  args = parser.parse_args()
134  if args.output_source is None:
135    fdnum = os.memfd_create("tempfile_profile")
136    args.output_source = "/proc/{}/fd/{}".format(os.getpid(), fdnum)
137  all_dexs = list()
138  for f in args.apks:
139    try:
140      all_dexs.append(pylibdexfile.FileDexFile(f.file, f.location))
141    except Exception as e1:
142      try:
143        all_dexs += pylibdexfile.OpenJar(f.file)
144      except Exception as e2:
145        parser.error("Failed to open file: {}. errors were {} and {}".format(
146            f.file, e1, e2))
147  if args.input_profile is not None:
148    profman_args = [
149        "profmand", "--dump-classes-and-methods",
150        "--profile-file={}".format(args.input_profile)
151    ]
152    for apk in args.apks:
153      profman_args.append("--apk={}".format(apk.file))
154    print(" ".join(map(shlex.quote, profman_args)))
155    res = subprocess.run(
156        profman_args, capture_output=True, universal_newlines=True)
157    res.check_returncode()
158    meth_list = list(filter(lambda a: a != "", res.stdout.split()))
159  elif args.input_text_profile is not None:
160    with open(args.input_text_profile, "rt") as inp:
161      meth_list = list(filter(lambda a: a != "", inp.readlines()))
162  else:
163    all_methods = set()
164    for d in all_dexs:
165      for m in d.methods:
166        all_methods.add(m.descriptor)
167    meth_list = list(all_methods)
168  print("Found {} methods. Will take ~{} iterations".format(
169      len(meth_list), 1 + math.floor(math.log2(len(meth_list)))))
170  print(
171      "type 'yes' if the behavior you are looking for is present (i.e. the compiled code crashes " +
172      "or something)"
173  )
174  print("Performing single check with all methods")
175  result = run_test(meth_list, args)
176  if result[0].lower() != "y":
177    cont = input(
178        "The behavior you were looking for did not occur when run against all methods. Continue " +
179        "(yes/no)? "
180    )
181    if cont[0].lower() != "y":
182      print("Aborting!")
183      sys.exit(1)
184  needs_dump = False
185  while len(meth_list) > 1:
186    test_methods = list(meth_list[0:len(meth_list) // 2])
187    result = run_test(test_methods, args)
188    if result[0].lower() == "y":
189      meth_list = test_methods
190      needs_dump = False
191    else:
192      meth_list = meth_list[len(meth_list) // 2:]
193      needs_dump = True
194  if needs_dump:
195    with open(args.output_source, "wt") as output:
196      dump_files(meth_list, args, output)
197  print("Found result!")
198  print("{}".format(meth_list[0]))
199  print("Leaving profile at {} and text profile at {}".format(
200      args.output_file, args.output_source))
201
202
203if __name__ == "__main__":
204  main()
205