1// Copyright 2020 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Copies all the entries (APKs/APEXes) matching the target configuration from the given
16// APK set into a zip file. Run it without arguments to see usage details.
17package main
18
19import (
20	"flag"
21	"fmt"
22	"io"
23	"log"
24	"math"
25	"os"
26	"regexp"
27	"sort"
28	"strings"
29
30	"github.com/golang/protobuf/proto"
31
32	"android/soong/cmd/extract_apks/bundle_proto"
33	"android/soong/third_party/zip"
34)
35
36type TargetConfig struct {
37	sdkVersion int32
38	screenDpi  map[android_bundle_proto.ScreenDensity_DensityAlias]bool
39	// Map holding <ABI alias>:<its sequence number in the flag> info.
40	abis             map[android_bundle_proto.Abi_AbiAlias]int
41	allowPrereleased bool
42	stem             string
43}
44
45// An APK set is a zip archive. An entry 'toc.pb' describes its contents.
46// It is a protobuf message BuildApkResult.
47type Toc *android_bundle_proto.BuildApksResult
48
49type ApkSet struct {
50	path    string
51	reader  *zip.ReadCloser
52	entries map[string]*zip.File
53}
54
55func newApkSet(path string) (*ApkSet, error) {
56	apkSet := &ApkSet{path: path, entries: make(map[string]*zip.File)}
57	var err error
58	if apkSet.reader, err = zip.OpenReader(apkSet.path); err != nil {
59		return nil, err
60	}
61	for _, f := range apkSet.reader.File {
62		apkSet.entries[f.Name] = f
63	}
64	return apkSet, nil
65}
66
67func (apkSet *ApkSet) getToc() (Toc, error) {
68	var err error
69	tocFile, ok := apkSet.entries["toc.pb"]
70	if !ok {
71		return nil, fmt.Errorf("%s: APK set should have toc.pb entry", apkSet.path)
72	}
73	rc, err := tocFile.Open()
74	if err != nil {
75		return nil, err
76	}
77	bytes := make([]byte, tocFile.FileHeader.UncompressedSize64)
78	if _, err := rc.Read(bytes); err != io.EOF {
79		return nil, err
80	}
81	rc.Close()
82	buildApksResult := new(android_bundle_proto.BuildApksResult)
83	if err = proto.Unmarshal(bytes, buildApksResult); err != nil {
84		return nil, err
85	}
86	return buildApksResult, nil
87}
88
89func (apkSet *ApkSet) close() {
90	apkSet.reader.Close()
91}
92
93// Matchers for selection criteria
94
95type abiTargetingMatcher struct {
96	*android_bundle_proto.AbiTargeting
97}
98
99func (m abiTargetingMatcher) matches(config TargetConfig) bool {
100	if m.AbiTargeting == nil {
101		return true
102	}
103	if _, ok := config.abis[android_bundle_proto.Abi_UNSPECIFIED_CPU_ARCHITECTURE]; ok {
104		return true
105	}
106	// Find the one that appears first in the abis flags.
107	abiIdx := math.MaxInt32
108	for _, v := range m.GetValue() {
109		if i, ok := config.abis[v.Alias]; ok {
110			if i < abiIdx {
111				abiIdx = i
112			}
113		}
114	}
115	if abiIdx == math.MaxInt32 {
116		return false
117	}
118	// See if any alternatives appear before the above one.
119	for _, a := range m.GetAlternatives() {
120		if i, ok := config.abis[a.Alias]; ok {
121			if i < abiIdx {
122				// There is a better alternative. Skip this one.
123				return false
124			}
125		}
126	}
127	return true
128}
129
130type apkDescriptionMatcher struct {
131	*android_bundle_proto.ApkDescription
132}
133
134func (m apkDescriptionMatcher) matches(config TargetConfig) bool {
135	return m.ApkDescription == nil || (apkTargetingMatcher{m.Targeting}).matches(config)
136}
137
138type apkTargetingMatcher struct {
139	*android_bundle_proto.ApkTargeting
140}
141
142func (m apkTargetingMatcher) matches(config TargetConfig) bool {
143	return m.ApkTargeting == nil ||
144		(abiTargetingMatcher{m.AbiTargeting}.matches(config) &&
145			languageTargetingMatcher{m.LanguageTargeting}.matches(config) &&
146			screenDensityTargetingMatcher{m.ScreenDensityTargeting}.matches(config) &&
147			sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
148			multiAbiTargetingMatcher{m.MultiAbiTargeting}.matches(config))
149}
150
151type languageTargetingMatcher struct {
152	*android_bundle_proto.LanguageTargeting
153}
154
155func (m languageTargetingMatcher) matches(_ TargetConfig) bool {
156	if m.LanguageTargeting == nil {
157		return true
158	}
159	log.Fatal("language based entry selection is not implemented")
160	return false
161}
162
163type moduleMetadataMatcher struct {
164	*android_bundle_proto.ModuleMetadata
165}
166
167func (m moduleMetadataMatcher) matches(config TargetConfig) bool {
168	return m.ModuleMetadata == nil ||
169		(m.GetDeliveryType() == android_bundle_proto.DeliveryType_INSTALL_TIME &&
170			moduleTargetingMatcher{m.Targeting}.matches(config) &&
171			!m.IsInstant)
172}
173
174type moduleTargetingMatcher struct {
175	*android_bundle_proto.ModuleTargeting
176}
177
178func (m moduleTargetingMatcher) matches(config TargetConfig) bool {
179	return m.ModuleTargeting == nil ||
180		(sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
181			userCountriesTargetingMatcher{m.UserCountriesTargeting}.matches(config))
182}
183
184// A higher number means a higher priority.
185// This order must be kept identical to bundletool's.
186var multiAbiPriorities = map[android_bundle_proto.Abi_AbiAlias]int{
187	android_bundle_proto.Abi_ARMEABI:     1,
188	android_bundle_proto.Abi_ARMEABI_V7A: 2,
189	android_bundle_proto.Abi_ARM64_V8A:   3,
190	android_bundle_proto.Abi_X86:         4,
191	android_bundle_proto.Abi_X86_64:      5,
192	android_bundle_proto.Abi_MIPS:        6,
193	android_bundle_proto.Abi_MIPS64:      7,
194}
195
196type multiAbiTargetingMatcher struct {
197	*android_bundle_proto.MultiAbiTargeting
198}
199
200func (t multiAbiTargetingMatcher) matches(config TargetConfig) bool {
201	if t.MultiAbiTargeting == nil {
202		return true
203	}
204	if _, ok := config.abis[android_bundle_proto.Abi_UNSPECIFIED_CPU_ARCHITECTURE]; ok {
205		return true
206	}
207	// Find the one with the highest priority.
208	highestPriority := 0
209	for _, v := range t.GetValue() {
210		for _, a := range v.GetAbi() {
211			if _, ok := config.abis[a.Alias]; ok {
212				if highestPriority < multiAbiPriorities[a.Alias] {
213					highestPriority = multiAbiPriorities[a.Alias]
214				}
215			}
216		}
217	}
218	if highestPriority == 0 {
219		return false
220	}
221	// See if there are any matching alternatives with a higher priority.
222	for _, v := range t.GetAlternatives() {
223		for _, a := range v.GetAbi() {
224			if _, ok := config.abis[a.Alias]; ok {
225				if highestPriority < multiAbiPriorities[a.Alias] {
226					// There's a better one. Skip this one.
227					return false
228				}
229			}
230		}
231	}
232	return true
233}
234
235type screenDensityTargetingMatcher struct {
236	*android_bundle_proto.ScreenDensityTargeting
237}
238
239func (m screenDensityTargetingMatcher) matches(config TargetConfig) bool {
240	if m.ScreenDensityTargeting == nil {
241		return true
242	}
243	if _, ok := config.screenDpi[android_bundle_proto.ScreenDensity_DENSITY_UNSPECIFIED]; ok {
244		return true
245	}
246	for _, v := range m.GetValue() {
247		switch x := v.GetDensityOneof().(type) {
248		case *android_bundle_proto.ScreenDensity_DensityAlias_:
249			if _, ok := config.screenDpi[x.DensityAlias]; ok {
250				return true
251			}
252		default:
253			log.Fatal("For screen density, only DPI name based entry selection (e.g. HDPI, XHDPI) is implemented")
254		}
255	}
256	return false
257}
258
259type sdkVersionTargetingMatcher struct {
260	*android_bundle_proto.SdkVersionTargeting
261}
262
263func (m sdkVersionTargetingMatcher) matches(config TargetConfig) bool {
264	const preReleaseVersion = 10000
265	if m.SdkVersionTargeting == nil {
266		return true
267	}
268	if len(m.Value) > 1 {
269		log.Fatal(fmt.Sprintf("sdk_version_targeting should not have multiple values:%#v", m.Value))
270	}
271	// Inspect only sdkVersionTargeting.Value.
272	// Even though one of the SdkVersionTargeting.Alternatives values may be
273	// better matching, we will select all of them
274	return m.Value[0].Min == nil ||
275		m.Value[0].Min.Value <= config.sdkVersion ||
276		(config.allowPrereleased && m.Value[0].Min.Value == preReleaseVersion)
277}
278
279type textureCompressionFormatTargetingMatcher struct {
280	*android_bundle_proto.TextureCompressionFormatTargeting
281}
282
283func (m textureCompressionFormatTargetingMatcher) matches(_ TargetConfig) bool {
284	if m.TextureCompressionFormatTargeting == nil {
285		return true
286	}
287	log.Fatal("texture based entry selection is not implemented")
288	return false
289}
290
291type userCountriesTargetingMatcher struct {
292	*android_bundle_proto.UserCountriesTargeting
293}
294
295func (m userCountriesTargetingMatcher) matches(_ TargetConfig) bool {
296	if m.UserCountriesTargeting == nil {
297		return true
298	}
299	log.Fatal("country based entry selection is not implemented")
300	return false
301}
302
303type variantTargetingMatcher struct {
304	*android_bundle_proto.VariantTargeting
305}
306
307func (m variantTargetingMatcher) matches(config TargetConfig) bool {
308	if m.VariantTargeting == nil {
309		return true
310	}
311	return sdkVersionTargetingMatcher{m.SdkVersionTargeting}.matches(config) &&
312		abiTargetingMatcher{m.AbiTargeting}.matches(config) &&
313		multiAbiTargetingMatcher{m.MultiAbiTargeting}.matches(config) &&
314		screenDensityTargetingMatcher{m.ScreenDensityTargeting}.matches(config) &&
315		textureCompressionFormatTargetingMatcher{m.TextureCompressionFormatTargeting}.matches(config)
316}
317
318type SelectionResult struct {
319	moduleName string
320	entries    []string
321}
322
323// Return all entries matching target configuration
324func selectApks(toc Toc, targetConfig TargetConfig) SelectionResult {
325	var result SelectionResult
326	for _, variant := range (*toc).GetVariant() {
327		if !(variantTargetingMatcher{variant.GetTargeting()}.matches(targetConfig)) {
328			continue
329		}
330		for _, as := range variant.GetApkSet() {
331			if !(moduleMetadataMatcher{as.ModuleMetadata}.matches(targetConfig)) {
332				continue
333			}
334			for _, apkdesc := range as.GetApkDescription() {
335				if (apkDescriptionMatcher{apkdesc}).matches(targetConfig) {
336					result.entries = append(result.entries, apkdesc.GetPath())
337					// TODO(asmundak): As it turns out, moduleName which we get from
338					// the ModuleMetadata matches the module names of the generated
339					// entry paths just by coincidence, only for the split APKs. We
340					// need to discuss this with bundletool folks.
341					result.moduleName = as.GetModuleMetadata().GetName()
342				}
343			}
344			// we allow only a single module, so bail out here if we found one
345			if result.moduleName != "" {
346				return result
347			}
348		}
349	}
350	return result
351}
352
353type Zip2ZipWriter interface {
354	CopyFrom(file *zip.File, name string) error
355}
356
357// Writes out selected entries, renaming them as needed
358func (apkSet *ApkSet) writeApks(selected SelectionResult, config TargetConfig,
359	writer Zip2ZipWriter, partition string) ([]string, error) {
360	// Renaming rules:
361	//  splits/MODULE-master.apk to STEM.apk
362	// else
363	//  splits/MODULE-*.apk to STEM>-$1.apk
364	// TODO(asmundak):
365	//  add more rules, for .apex files
366	renameRules := []struct {
367		rex  *regexp.Regexp
368		repl string
369	}{
370		{
371			regexp.MustCompile(`^.*/` + selected.moduleName + `-master\.apk$`),
372			config.stem + `.apk`,
373		},
374		{
375			regexp.MustCompile(`^.*/` + selected.moduleName + `(-.*\.apk)$`),
376			config.stem + `$1`,
377		},
378		{
379			regexp.MustCompile(`^universal\.apk$`),
380			config.stem + ".apk",
381		},
382	}
383	renamer := func(path string) (string, bool) {
384		for _, rr := range renameRules {
385			if rr.rex.MatchString(path) {
386				return rr.rex.ReplaceAllString(path, rr.repl), true
387			}
388		}
389		return "", false
390	}
391
392	entryOrigin := make(map[string]string) // output entry to input entry
393	var apkcerts []string
394	for _, apk := range selected.entries {
395		apkFile, ok := apkSet.entries[apk]
396		if !ok {
397			return nil, fmt.Errorf("TOC refers to an entry %s which does not exist", apk)
398		}
399		inName := apkFile.Name
400		outName, ok := renamer(inName)
401		if !ok {
402			log.Fatalf("selected an entry with unexpected name %s", inName)
403		}
404		if origin, ok := entryOrigin[inName]; ok {
405			log.Fatalf("selected entries %s and %s will have the same output name %s",
406				origin, inName, outName)
407		}
408		entryOrigin[outName] = inName
409		if err := writer.CopyFrom(apkFile, outName); err != nil {
410			return nil, err
411		}
412		if partition != "" {
413			apkcerts = append(apkcerts, fmt.Sprintf(
414				`name="%s" certificate="PRESIGNED" private_key="" partition="%s"`, outName, partition))
415		}
416	}
417	sort.Strings(apkcerts)
418	return apkcerts, nil
419}
420
421func (apkSet *ApkSet) extractAndCopySingle(selected SelectionResult, outFile *os.File) error {
422	if len(selected.entries) != 1 {
423		return fmt.Errorf("Too many matching entries for extract-single:\n%v", selected.entries)
424	}
425	apk, ok := apkSet.entries[selected.entries[0]]
426	if !ok {
427		return fmt.Errorf("Couldn't find apk path %s", selected.entries[0])
428	}
429	inputReader, _ := apk.Open()
430	_, err := io.Copy(outFile, inputReader)
431	return err
432}
433
434// Arguments parsing
435var (
436	outputFile   = flag.String("o", "", "output file containing extracted entries")
437	targetConfig = TargetConfig{
438		screenDpi: map[android_bundle_proto.ScreenDensity_DensityAlias]bool{},
439		abis:      map[android_bundle_proto.Abi_AbiAlias]int{},
440	}
441	extractSingle = flag.Bool("extract-single", false,
442		"extract a single target and output it uncompressed. only available for standalone apks and apexes.")
443	apkcertsOutput = flag.String("apkcerts", "",
444		"optional apkcerts.txt output file containing signing info of all outputted apks")
445	partition = flag.String("partition", "", "partition string. required when -apkcerts is used.")
446)
447
448// Parse abi values
449type abiFlagValue struct {
450	targetConfig *TargetConfig
451}
452
453func (a abiFlagValue) String() string {
454	return "all"
455}
456
457func (a abiFlagValue) Set(abiList string) error {
458	for i, abi := range strings.Split(abiList, ",") {
459		v, ok := android_bundle_proto.Abi_AbiAlias_value[abi]
460		if !ok {
461			return fmt.Errorf("bad ABI value: %q", abi)
462		}
463		targetConfig.abis[android_bundle_proto.Abi_AbiAlias(v)] = i
464	}
465	return nil
466}
467
468// Parse screen density values
469type screenDensityFlagValue struct {
470	targetConfig *TargetConfig
471}
472
473func (s screenDensityFlagValue) String() string {
474	return "none"
475}
476
477func (s screenDensityFlagValue) Set(densityList string) error {
478	if densityList == "none" {
479		return nil
480	}
481	if densityList == "all" {
482		targetConfig.screenDpi[android_bundle_proto.ScreenDensity_DENSITY_UNSPECIFIED] = true
483		return nil
484	}
485	for _, density := range strings.Split(densityList, ",") {
486		v, found := android_bundle_proto.ScreenDensity_DensityAlias_value[density]
487		if !found {
488			return fmt.Errorf("bad screen density value: %q", density)
489		}
490		targetConfig.screenDpi[android_bundle_proto.ScreenDensity_DensityAlias(v)] = true
491	}
492	return nil
493}
494
495func processArgs() {
496	flag.Usage = func() {
497		fmt.Fprintln(os.Stderr, `usage: extract_apks -o <output-file> -sdk-version value -abis value `+
498			`-screen-densities value {-stem value | -extract-single} [-allow-prereleased] `+
499			`[-apkcerts <apkcerts output file> -partition <partition>] <APK set>`)
500		flag.PrintDefaults()
501		os.Exit(2)
502	}
503	version := flag.Uint("sdk-version", 0, "SDK version")
504	flag.Var(abiFlagValue{&targetConfig}, "abis",
505		"comma-separated ABIs list of ARMEABI ARMEABI_V7A ARM64_V8A X86 X86_64 MIPS MIPS64")
506	flag.Var(screenDensityFlagValue{&targetConfig}, "screen-densities",
507		"'all' or comma-separated list of screen density names (NODPI LDPI MDPI TVDPI HDPI XHDPI XXHDPI XXXHDPI)")
508	flag.BoolVar(&targetConfig.allowPrereleased, "allow-prereleased", false,
509		"allow prereleased")
510	flag.StringVar(&targetConfig.stem, "stem", "", "output entries base name in the output zip file")
511	flag.Parse()
512	if (*outputFile == "") || len(flag.Args()) != 1 || *version == 0 ||
513		(targetConfig.stem == "" && !*extractSingle) || (*apkcertsOutput != "" && *partition == "") {
514		flag.Usage()
515	}
516	targetConfig.sdkVersion = int32(*version)
517
518}
519
520func main() {
521	processArgs()
522	var toc Toc
523	apkSet, err := newApkSet(flag.Arg(0))
524	if err == nil {
525		defer apkSet.close()
526		toc, err = apkSet.getToc()
527	}
528	if err != nil {
529		log.Fatal(err)
530	}
531	sel := selectApks(toc, targetConfig)
532	if len(sel.entries) == 0 {
533		log.Fatalf("there are no entries for the target configuration: %#v", targetConfig)
534	}
535
536	outFile, err := os.Create(*outputFile)
537	if err != nil {
538		log.Fatal(err)
539	}
540	defer outFile.Close()
541
542	if *extractSingle {
543		err = apkSet.extractAndCopySingle(sel, outFile)
544	} else {
545		writer := zip.NewWriter(outFile)
546		defer func() {
547			if err := writer.Close(); err != nil {
548				log.Fatal(err)
549			}
550		}()
551		apkcerts, err := apkSet.writeApks(sel, targetConfig, writer, *partition)
552		if err == nil && *apkcertsOutput != "" {
553			apkcertsFile, err := os.Create(*apkcertsOutput)
554			if err != nil {
555				log.Fatal(err)
556			}
557			defer apkcertsFile.Close()
558			for _, a := range apkcerts {
559				_, err = apkcertsFile.WriteString(a + "\n")
560				if err != nil {
561					log.Fatal(err)
562				}
563			}
564		}
565	}
566	if err != nil {
567		log.Fatal(err)
568	}
569}
570