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 com.android.tradefed.cluster;
17 
18 import com.android.tradefed.config.Configuration;
19 import com.android.tradefed.config.ConfigurationUtil;
20 import com.android.tradefed.result.LegacySubprocessResultsReporter;
21 import com.android.tradefed.util.FileUtil;
22 
23 import com.google.common.base.Strings;
24 
25 import org.w3c.dom.Document;
26 import org.w3c.dom.Element;
27 import org.w3c.dom.Node;
28 import org.xml.sax.SAXException;
29 
30 import java.io.File;
31 import java.io.FileInputStream;
32 import java.io.FileNotFoundException;
33 import java.io.IOException;
34 import java.io.InputStream;
35 import java.net.URL;
36 import java.net.URLClassLoader;
37 import java.util.ArrayList;
38 import java.util.List;
39 import javax.xml.parsers.DocumentBuilder;
40 import javax.xml.parsers.DocumentBuilderFactory;
41 import javax.xml.parsers.ParserConfigurationException;
42 import javax.xml.transform.Transformer;
43 import javax.xml.transform.TransformerException;
44 import javax.xml.transform.TransformerFactory;
45 import javax.xml.transform.dom.DOMSource;
46 import javax.xml.transform.stream.StreamResult;
47 
48 /**
49  * Build a wrapper TF config XML for an existing TF config.
50  *
51  * <p>A wrapper XML allows to enable subprocess reporting on an existing TF config.
52  */
53 public class SubprocessConfigBuilder {
54     private static final String REPORTER_CLASS = LegacySubprocessResultsReporter.class.getName();
55     private static final String OPTION_KEY = "subprocess-report-port";
56     private String mClasspath;
57 
58     private File mWorkDir;
59 
60     private String mOriginalConfig;
61 
62     private String mPort;
63 
setClasspath(String classpath)64     public SubprocessConfigBuilder setClasspath(String classpath) {
65         mClasspath = classpath;
66         return this;
67     }
68 
setWorkingDir(File dir)69     public SubprocessConfigBuilder setWorkingDir(File dir) {
70         mWorkDir = dir;
71         return this;
72     }
73 
setOriginalConfig(String config)74     public SubprocessConfigBuilder setOriginalConfig(String config) {
75         mOriginalConfig = config;
76         return this;
77     }
78 
setPort(String port)79     public SubprocessConfigBuilder setPort(String port) {
80         mPort = port;
81         return this;
82     }
83 
84     /**
85      * Current handling of ATS for the naming of injected config. Exposed so it can be used to align
86      * the test harness side.
87      */
createConfigName(String originalConfigName)88     public static String createConfigName(String originalConfigName) {
89         return "_" + originalConfigName.replace("/", "$") + ".xml";
90     }
91 
build()92     public File build() throws IOException {
93         final List<URL> urls = new ArrayList<>();
94         for (final String path : mClasspath.split(File.pathSeparator)) {
95             if (path.endsWith("*")) {
96                 final File dir = new File(path.substring(0, path.length() - 1));
97                 if (!dir.exists()) {
98                     continue;
99                 }
100                 for (final File file :
101                         dir.listFiles((parent, name) -> name.toLowerCase().endsWith(".jar"))) {
102                     urls.add(file.toURI().toURL());
103                 }
104             } else {
105                 urls.add(new File(path).toURI().toURL());
106             }
107         }
108 
109         // Read the original config file.
110         final DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
111         Document doc = null;
112         try (URLClassLoader loader = new URLClassLoader(urls.toArray(new URL[urls.size()]), null)) {
113             final DocumentBuilder builder = factory.newDocumentBuilder();
114             final String ext = FileUtil.getExtension(mOriginalConfig);
115             InputStream in = null;
116             if (Strings.isNullOrEmpty(ext)) {
117                 in = loader.getResourceAsStream(String.format("config/%s.xml", mOriginalConfig));
118             } else {
119                 in = loader.getResourceAsStream(String.format("config/%s", mOriginalConfig));
120             }
121             if (in == null) {
122                 File f = new File(mOriginalConfig);
123                 if (!f.isAbsolute()) {
124                     f = new File(mWorkDir, mOriginalConfig);
125                 }
126                 try {
127                     in = new FileInputStream(f);
128                 } catch (FileNotFoundException e) {
129                     throw new RuntimeException(
130                             String.format("Could not find configuration '%s'", mOriginalConfig));
131                 }
132             }
133             doc = builder.parse(in);
134         } catch (ParserConfigurationException | SAXException e) {
135             throw new RuntimeException(e);
136         }
137 
138         if (mPort != null) {
139             // Add subprocess result reporter to a config file.
140             final Node root = doc.getElementsByTagName("configuration").item(0);
141             final Element reporter = doc.createElement(Configuration.RESULT_REPORTER_TYPE_NAME);
142             reporter.setAttribute(ConfigurationUtil.CLASS_NAME, REPORTER_CLASS);
143             final Element options = doc.createElement(ConfigurationUtil.OPTION_NAME);
144             options.setAttribute(ConfigurationUtil.NAME_NAME, OPTION_KEY);
145             options.setAttribute(ConfigurationUtil.VALUE_NAME, mPort);
146             reporter.appendChild(options);
147             root.appendChild(reporter);
148         }
149 
150         File f = new File(mWorkDir, mOriginalConfig);
151         if (!f.exists() || !f.isFile()) {
152             // If the original config is an existing file, we need to update it since some old TFs
153             // check the file system first before bundled configs when loading configs.
154             // If the original config is not an existing file, we can use any name since the
155             // original config name will be assigned when creating a injection jar.
156             f = File.createTempFile("subprocess_config_", ".xml", mWorkDir);
157         }
158         TransformerFactory transformerFactory = TransformerFactory.newInstance();
159         try {
160             Transformer transformer = transformerFactory.newTransformer();
161             transformer.transform(new DOMSource(doc), new StreamResult(f));
162         } catch (TransformerException e) {
163             throw new RuntimeException(e);
164         }
165 
166         return f;
167     }
168 }
169