1/*
2 * Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation.
8 *
9 * This code is distributed in the hope that it will be useful, but WITHOUT
10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
12 * version 2 for more details (a copy is included in the LICENSE file that
13 * accompanied this code).
14 *
15 * You should have received a copy of the GNU General Public License version
16 * 2 along with this work; if not, write to the Free Software Foundation,
17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18 *
19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20 * or visit www.oracle.com if you need additional information or have any
21 * questions.
22 */
23
24package p;
25
26import java.io.File;
27import java.io.IOException;
28import java.io.OutputStream;
29import java.nio.file.Files;
30import java.nio.file.Path;
31import java.nio.file.Paths;
32import java.util.ArrayList;
33import java.util.List;
34import java.util.jar.Attributes;
35import java.util.jar.JarEntry;
36import java.util.jar.JarFile;
37import java.util.jar.JarOutputStream;
38import java.util.jar.Manifest;
39
40import com.sun.tools.attach.VirtualMachine;
41
42public class Main {
43
44    public static void main(String[] args) throws Exception {
45        System.out.println("#modules loaded: " + moduleInfoCont());
46
47        String vmid = "" + ProcessHandle.current().pid();
48        VirtualMachine vm = VirtualMachine.attach(vmid);
49
50        for (String test : args) {
51            switch (test) {
52                case "jmx" :
53                    startJMXAgent(vm);
54                    break;
55                case "javaagent" :
56                    startJavaAgent(vm, createAgentJar());
57                    break;
58            }
59
60            System.out.println("#modules loaded: " + moduleInfoCont());
61        }
62    }
63
64    /**
65     * Locates module-info.class resources to get a count of the module of system
66     * modules.
67     */
68    static long moduleInfoCont() {
69        ClassLoader scl = ClassLoader.getSystemClassLoader();
70        return scl.resources("module-info.class").count();
71    }
72
73    /**
74     * Starts a JMX agent and checks that java.management is loaded.
75     */
76    static void startJMXAgent(VirtualMachine vm) throws Exception {
77        System.out.println("Start JMX agent");
78        vm.startLocalManagementAgent();
79
80        // types in java.management should be visible
81        Class.forName("javax.management.MXBean");
82    }
83
84    /**
85     * Loads a java agent into the VM and checks that java.instrument is loaded.
86     */
87    static void startJavaAgent(VirtualMachine vm, Path agent) throws Exception {
88        System.out.println("Load java agent ...");
89        vm.loadAgent(agent.toString());
90
91        // the Agent class should be visible
92        Class.forName("Agent");
93
94        // types in java.instrument should be visible
95        Class.forName("java.lang.instrument.Instrumentation");
96    }
97
98    /**
99     * Creates a java agent, return the file path to the agent JAR file.
100     */
101    static Path createAgentJar() throws IOException {
102        Manifest man = new Manifest();
103        Attributes attrs = man.getMainAttributes();
104        attrs.put(Attributes.Name.MANIFEST_VERSION, "1.0.0");
105        attrs.put(new Attributes.Name("Agent-Class"), "Agent");
106        Path agent = Paths.get("agent.jar");
107        Path dir = Paths.get(System.getProperty("test.classes"));
108        createJarFile(agent, man, dir, "Agent.class");
109        return agent;
110    }
111
112    /**
113     * Creates a JAR file.
114     *
115     * Equivalent to {@code jar cfm <jarfile> <manifest> -C <dir> file...}
116     *
117     * The input files are resolved against the given directory. Any input
118     * files that are directories are processed recursively.
119     */
120    static void createJarFile(Path jarfile, Manifest man, Path dir, String... files)
121        throws IOException
122    {
123        // create the target directory
124        Path parent = jarfile.getParent();
125        if (parent != null)
126            Files.createDirectories(parent);
127
128        List<Path> entries = new ArrayList<>();
129        for (String file : files) {
130            Files.find(dir.resolve(file), Integer.MAX_VALUE,
131                    (p, attrs) -> attrs.isRegularFile())
132                    .map(e -> dir.relativize(e))
133                    .forEach(entries::add);
134        }
135
136        try (OutputStream out = Files.newOutputStream(jarfile);
137             JarOutputStream jos = new JarOutputStream(out))
138        {
139            if (man != null) {
140                JarEntry je = new JarEntry(JarFile.MANIFEST_NAME);
141                jos.putNextEntry(je);
142                man.write(jos);
143                jos.closeEntry();
144            }
145
146            for (Path entry : entries) {
147                String name = toJarEntryName(entry);
148                jos.putNextEntry(new JarEntry(name));
149                Files.copy(dir.resolve(entry), jos);
150                jos.closeEntry();
151            }
152        }
153    }
154
155    /**
156     * Map a file path to the equivalent name in a JAR file
157     */
158    static String toJarEntryName(Path file) {
159        Path normalized = file.normalize();
160        return normalized.subpath(0, normalized.getNameCount())
161                .toString()
162                .replace(File.separatorChar, '/');
163    }
164}
165