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                 try {
123                     in = new FileInputStream(mOriginalConfig);
124                 } catch (FileNotFoundException e) {
125                     throw new RuntimeException(
126                             String.format("Could not find configuration '%s'", mOriginalConfig));
127                 }
128             }
129             doc = builder.parse(in);
130         } catch (ParserConfigurationException | SAXException e) {
131             throw new RuntimeException(e);
132         }
133 
134         if (mPort != null) {
135             // Add subprocess result reporter to a config file.
136             final Node root = doc.getElementsByTagName("configuration").item(0);
137             final Element reporter = doc.createElement(Configuration.RESULT_REPORTER_TYPE_NAME);
138             reporter.setAttribute(ConfigurationUtil.CLASS_NAME, REPORTER_CLASS);
139             final Element options = doc.createElement(ConfigurationUtil.OPTION_NAME);
140             options.setAttribute(ConfigurationUtil.NAME_NAME, OPTION_KEY);
141             options.setAttribute(ConfigurationUtil.VALUE_NAME, mPort);
142             reporter.appendChild(options);
143             root.appendChild(reporter);
144         }
145 
146         File f = new File(mWorkDir, createConfigName(mOriginalConfig));
147         TransformerFactory transformerFactory = TransformerFactory.newInstance();
148         try {
149             Transformer transformer = transformerFactory.newTransformer();
150             transformer.transform(new DOMSource(doc), new StreamResult(f));
151         } catch (TransformerException e) {
152             throw new RuntimeException(e);
153         }
154 
155         return f;
156     }
157 }
158