// Licensed to the Software Freedom Conservancy (SFC) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The SFC licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

package org.openqa.selenium.grid;

import java.io.BufferedInputStream;

import java.io.ByteArrayOutputStream;

import java.io.File;

import java.io.IOException;

import java.io.InputStream;

import java.io.UncheckedIOException;

import java.lang.reflect.Method;

import java.net.MalformedURLException;

import java.net.URL;

import java.net.URLClassLoader;

import java.security.AccessController;

import java.security.PrivilegedAction;

import java.util.ArrayList;

import java.util.HashSet;

import java.util.List;

import java.util.Set;

import java.util.logging.Level;

import java.util.logging.Logger;

public class Bootstrap {

private static final String MAIN_CLASS = Bootstrap.class.getPackage().getName() + ".Main";

private static final Logger LOG = Logger.getLogger(Bootstrap.class.getName());

public static void main(String[] args) {

ClassLoader classLoader = Bootstrap.class.getClassLoader();

if (args.length == 0) {

runMain(classLoader, args);

return;

}

if ("--ext".equals(args[0])) {

if (args.length < 2) {

runMain(classLoader, args);

return;

}

ClassLoader parent = createExtendedClassLoader(args[1]);

String[] remainingArgs = new String[args.length - 2];

System.arraycopy(args, 2, remainingArgs, 0, args.length - 2);

args = remainingArgs;

classLoader = new PossessiveClassLoader(parent);

// Ensure that we use our freshly minted classloader by default.

Thread.currentThread().setContextClassLoader(classLoader);

}

runMain(classLoader, args);

}

private static void runMain(ClassLoader loader, String[] args) {

try {

Class<?> clazz = loader.loadClass(MAIN_CLASS);

Method main = clazz.getMethod("main", String[].class);

main.invoke(null, new Object[] {args});

} catch (ReflectiveOperationException e) {

e.printStackTrace();

System.exit(1);

}

}

private static ClassLoader createExtendedClassLoader(String ext) {

List<File> jars = new ArrayList<>();

for (String part : ext.split(File.pathSeparator)) {

File file = new File(part);

if (!file.exists()) {

LOG.warning("Extension file or directory does not exist: " + file);

continue;

}

if (file.isDirectory()) {

File[] files = file.listFiles();

if (files == null) {

LOG.warning("Cannot list files in directory: " + file);

} else {

for (File subdirFile : files) {

if (subdirFile.isFile() && subdirFile.getName().endsWith(".jar")) {

jars.add(subdirFile);

}

}

}

} else {

jars.add(file);

}

}

URL[] jarUrls = jars.stream()

.map(file -> {

try {

return file.toURI().toURL();

} catch (MalformedURLException e) {

LOG.log(Level.SEVERE, "Unable to find JAR file " + file, e);

throw new UncheckedIOException(e);

}

})

.toArray(URL[]::new);

return AccessController.doPrivileged((PrivilegedAction<URLClassLoader>) () ->

new URLClassLoader(jarUrls, Bootstrap.class.getClassLoader()));

}

private static class PossessiveClassLoader extends ClassLoader {

private final ClassLoader delegate;

private final Set<String> blessed;

PossessiveClassLoader(ClassLoader delegate) {

super(delegate);

this.delegate = delegate;

blessed = new HashSet<>();

blessed.add("java.");

blessed.add("javax.");

blessed.add("sun.");

blessed.add("jdk.");

}

@Override

protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {

for (String prefix : blessed) {

if (name.startsWith(prefix)) {

return super.loadClass(name, resolve);

}

}

Class<?> clazz = findLoadedClass(name);

if (clazz != null) {

return clazz;

}

URL resource = delegate.getResource(name.replaceAll("\\.", "/") + ".class");

if (resource == null) {

throw new ClassNotFoundException("Unable to find " + name);

}

ByteArrayOutputStream bos = new ByteArrayOutputStream();

try (InputStream is = resource.openStream();

BufferedInputStream bis = new BufferedInputStream(is)) {

int read;

byte[] bytes = new byte[4096];

while ((read = bis.read(bytes, 0, bytes.length)) != -1) {

bos.write(bytes, 0, read);

}

} catch (IOException e) {

throw new ClassNotFoundException("Could not load " + name, e);

}

byte[] rawClass = bos.toByteArray();

Class<?> defined = defineClass(name, rawClass, 0, rawClass.length);

if (resolve) {

resolveClass(defined);

}

return defined;

}

}

}