1 /*
   2  * Copyright (c) 2014, 2016, 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.  Oracle designates this
   8  * particular file as subject to the "Classpath" exception as provided
   9  * by Oracle in the LICENSE file that accompanied this code.
  10  *
  11  * This code is distributed in the hope that it will be useful, but WITHOUT
  12  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  13  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  14  * version 2 for more details (a copy is included in the LICENSE file that
  15  * accompanied this code).
  16  *
  17  * You should have received a copy of the GNU General Public License version
  18  * 2 along with this work; if not, write to the Free Software Foundation,
  19  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
  20  *
  21  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
  22  * or visit www.oracle.com if you need additional information or have any
  23  * questions.
  24  */
  25 
  26 package jdk.internal.module;
  27 
  28 import java.io.IOException;
  29 import java.io.InputStream;
  30 import java.io.OutputStream;
  31 import java.lang.module.ModuleDescriptor.Version;
  32 import java.util.ArrayList;
  33 import java.util.Collections;
  34 import java.util.HashMap;
  35 import java.util.List;
  36 import java.util.Map;
  37 import java.util.Set;
  38 
  39 import jdk.internal.org.objectweb.asm.Attribute;
  40 import jdk.internal.org.objectweb.asm.ClassReader;
  41 import jdk.internal.org.objectweb.asm.ClassVisitor;
  42 import jdk.internal.org.objectweb.asm.ClassWriter;
  43 import jdk.internal.org.objectweb.asm.Opcodes;
  44 
  45 import static jdk.internal.module.ClassFileAttributes.*;
  46 
  47 /**
  48  * Utility class to extend a module-info.class with additional attributes.
  49  */
  50 
  51 public final class ModuleInfoExtender {
  52 
  53     // the input stream to read the original module-info.class
  54     private final InputStream in;
  55 
  56     // the packages in the ConcealedPackages attribute
  57     private Set<String> conceals;
  58 
  59     // the value of the Version attribute
  60     private Version version;
  61 
  62     // the value of the MainClass attribute
  63     private String mainClass;
  64 
  65     // the values for the TargetPlatform attribute
  66     private String osName;
  67     private String osArch;
  68     private String osVersion;
  69 
  70     // the hashes for the Hashes attribute
  71     private ModuleHashes hashes;
  72 
  73     private ModuleInfoExtender(InputStream in) {
  74         this.in = in;
  75     }
  76 
  77     /**
  78      * Sets the set of packages for the ConcealedPackages attribute
  79      */
  80     public ModuleInfoExtender conceals(Set<String> packages) {
  81         this.conceals = Collections.unmodifiableSet(packages);
  82         return this;
  83     }
  84 
  85     /**
  86      * Sets the value of the Version attribute.
  87      */
  88     public ModuleInfoExtender version(Version version) {
  89         this.version = version;
  90         return this;
  91     }
  92 
  93     /**
  94      * Sets the value of the MainClass attribute.
  95      */
  96     public ModuleInfoExtender mainClass(String mainClass) {
  97         this.mainClass = mainClass;
  98         return this;
  99     }
 100 
 101     /**
 102      * Sets the values for the TargetPlatform attribute.
 103      */
 104     public ModuleInfoExtender targetPlatform(String osName,
 105                                              String osArch,
 106                                              String osVersion) {
 107         this.osName = osName;
 108         this.osArch = osArch;
 109         this.osVersion = osVersion;
 110         return this;
 111     }
 112 
 113     /**
 114      * The Hashes attribute will be emitted to the module-info with
 115      * the hashes encapsulated in the given {@code ModuleHashes}
 116      * object.
 117      */
 118     public ModuleInfoExtender hashes(ModuleHashes hashes) {
 119         this.hashes = hashes;
 120         return this;
 121     }
 122 
 123     /**
 124      * A ClassVisitor that supports adding class file attributes. If an
 125      * attribute already exists then the first occurence of the attribute
 126      * is replaced.
 127      */
 128     private static class AttributeAddingClassVisitor extends ClassVisitor {
 129         private Map<String, Attribute> attrs = new HashMap<>();
 130 
 131         AttributeAddingClassVisitor(int api, ClassVisitor cv) {
 132             super(api, cv);
 133         }
 134 
 135         void addAttribute(Attribute attr) {
 136             attrs.put(attr.type, attr);
 137         }
 138 
 139         @Override
 140         public void visitAttribute(Attribute attr) {
 141             String name = attr.type;
 142             Attribute replacement = attrs.get(name);
 143             if (replacement != null) {
 144                 attr = replacement;
 145                 attrs.remove(name);
 146             }
 147             super.visitAttribute(attr);
 148         }
 149 
 150         /**
 151          * Adds any remaining attributes that weren't replaced to the
 152          * class file.
 153          */
 154         void finish() {
 155             attrs.values().forEach(a -> super.visitAttribute(a));
 156             attrs.clear();
 157         }
 158     }
 159 
 160     /**
 161      * Outputs the modified module-info.class to the given output stream.
 162      * Once this method has been called then the Extender object should
 163      * be discarded.
 164      */
 165     public void write(OutputStream out) throws IOException {
 166         // emit to the output stream
 167         out.write(getBytes());
 168     }
 169 
 170     /**
 171      * Returns the bytes of the modified module-info.class.
 172      * Once this method has been called then the Extender object should
 173      * be discarded.
 174      */
 175     public byte[] getBytes() throws IOException {
 176         ClassWriter cw
 177             = new ClassWriter(ClassWriter.COMPUTE_MAXS + ClassWriter.COMPUTE_FRAMES);
 178 
 179         AttributeAddingClassVisitor cv
 180             = new AttributeAddingClassVisitor(Opcodes.ASM5, cw);
 181 
 182         ClassReader cr = new ClassReader(in);
 183 
 184         if (conceals != null)
 185             cv.addAttribute(new ConcealedPackagesAttribute(conceals));
 186         if (version != null)
 187             cv.addAttribute(new VersionAttribute(version));
 188         if (mainClass != null)
 189             cv.addAttribute(new MainClassAttribute(mainClass));
 190         if (osName != null || osArch != null || osVersion != null)
 191             cv.addAttribute(new TargetPlatformAttribute(osName, osArch, osVersion));
 192         if (hashes != null)
 193             cv.addAttribute(new HashesAttribute(hashes));
 194 
 195         List<Attribute> attrs = new ArrayList<>();
 196 
 197         // prototypes of attributes that should be parsed
 198         attrs.add(new ModuleAttribute());
 199         attrs.add(new ConcealedPackagesAttribute());
 200         attrs.add(new VersionAttribute());
 201         attrs.add(new MainClassAttribute());
 202         attrs.add(new TargetPlatformAttribute());
 203         attrs.add(new HashesAttribute());
 204 
 205         cr.accept(cv, attrs.toArray(new Attribute[0]), 0);
 206 
 207         // add any attributes that didn't replace previous attributes
 208         cv.finish();
 209 
 210         return cw.toByteArray();
 211     }
 212 
 213     /**
 214      * Returns an {@code Extender} that may be used to add additional
 215      * attributes to the module-info.class read from the given input
 216      * stream.
 217      */
 218     public static ModuleInfoExtender newExtender(InputStream in) {
 219         return new ModuleInfoExtender(in);
 220     }
 221 
 222 }