    001    package cpw.mods.fml.common.asm.transformers;
    002
    003    import java.io.BufferedOutputStream;
    004    import java.io.BufferedReader;
    005    import java.io.ByteArrayOutputStream;
    006    import java.io.DataInputStream;
    007    import java.io.File;
    008    import java.io.FileInputStream;
    009    import java.io.FileNotFoundException;
    010    import java.io.FileOutputStream;
    011    import java.io.IOException;
    012    import java.io.InputStream;
    013    import java.io.InputStreamReader;
    014    import java.util.ArrayList;
    015    import java.util.Collections;
    016    import java.util.Enumeration;
    017    import java.util.HashSet;
    018    import java.util.Hashtable;
    019    import java.util.LinkedHashSet;
    020    import java.util.List;
    021    import java.util.Map.Entry;
    022    import java.util.zip.ZipEntry;
    023    import java.util.zip.ZipFile;
    024    import java.util.zip.ZipOutputStream;
    025
    026    import org.objectweb.asm.ClassReader;
    027    import org.objectweb.asm.ClassWriter;
    028    import org.objectweb.asm.Type;
    029    import org.objectweb.asm.tree.AnnotationNode;
    030    import org.objectweb.asm.tree.ClassNode;
    031    import org.objectweb.asm.tree.FieldNode;
    032    import org.objectweb.asm.tree.MethodNode;
    033
    034    import com.google.common.base.Objects;
    035    import com.google.common.collect.Lists;
    036    import com.google.common.collect.Sets;
    037
    038    import cpw.mods.fml.relauncher.Side;
    039    import cpw.mods.fml.relauncher.SideOnly;
    040
    041    public class MCPMerger
    042    {
    043        private static Hashtable<String, ClassInfo> clients = new Hashtable<String, ClassInfo>();
    044        private static Hashtable<String, ClassInfo> shared  = new Hashtable<String, ClassInfo>();
    045        private static Hashtable<String, ClassInfo> servers = new Hashtable<String, ClassInfo>();
    046        private static HashSet<String> copyToServer = new HashSet<String>();
    047        private static HashSet<String> copyToClient = new HashSet<String>();
    048        private static HashSet<String> dontAnnotate = new HashSet<String>();
    049        private static final boolean DEBUG = false;
    050
    051        public static void main(String[] args)
    052        {
    053            if (args.length != 3)
    054            {
    055                System.out.println("Usage: MCPMerger <MapFile> <minecraft.jar> <minecraft_server.jar>");
    056                System.exit(1);
    057            }
    058
    059            File map_file = new File(args[0]);
    060            File client_jar = new File(args[1]);
    061            File server_jar = new File(args[2]);
    062            File client_jar_tmp = new File(args[1] + ".MergeBack");
    063            File server_jar_tmp = new File(args[2] + ".MergeBack");
    064
    065
    066            if (client_jar_tmp.exists() && !client_jar_tmp.delete())
    067            {
    068                System.out.println("Could not delete temp file: " + client_jar_tmp);
    069            }
    070
    071            if (server_jar_tmp.exists() && !server_jar_tmp.delete())
    072            {
    073                System.out.println("Could not delete temp file: " + server_jar_tmp);
    074            }
    075
    076            if (!client_jar.exists())
    077            {
    078                System.out.println("Could not find minecraft.jar: " + client_jar);
    079                System.exit(1);
    080            }
    081
    082            if (!server_jar.exists())
    083            {
    084                System.out.println("Could not find minecraft_server.jar: " + server_jar);
    085                System.exit(1);
    086            }
    087
    088            if (!client_jar.renameTo(client_jar_tmp))
    089            {
    090                System.out.println("Could not rename file: " + client_jar + " -> " + client_jar_tmp);
    091                System.exit(1);
    092            }
    093
    094            if (!server_jar.renameTo(server_jar_tmp))
    095            {
    096                System.out.println("Could not rename file: " + server_jar + " -> " + server_jar_tmp);
    097                System.exit(1);
    098            }
    099
    100            if (!readMapFile(map_file))
    101            {
    102                System.out.println("Could not read map file: " + map_file);
    103                System.exit(1);
    104            }
    105
    106            try
    107            {
    108                processJar(client_jar_tmp, server_jar_tmp, client_jar, server_jar);
    109            }
    110            catch (IOException e)
    111            {
    112                e.printStackTrace();
    113                System.exit(1);
    114            }
    115
    116            if (!client_jar_tmp.delete())
    117            {
    118                System.out.println("Could not delete temp file: " + client_jar_tmp);
    119            }
    120
    121            if (!server_jar_tmp.delete())
    122            {
    123                System.out.println("Could not delete temp file: " + server_jar_tmp);
    124            }
    125        }
    126
    127        private static boolean readMapFile(File mapFile)
    128        {
    129            try
    130            {
    131                FileInputStream fstream = new FileInputStream(mapFile);
    132                DataInputStream in = new DataInputStream(fstream);
    133                BufferedReader br = new BufferedReader(new InputStreamReader(in));
    134
    135                String line;
    136                while ((line = br.readLine()) != null)
    137                {
    138                    line = line.split("#")[0];
    139                    char cmd = line.charAt(0);
    140                    line = line.substring(1).trim();
    141
    142                    switch (cmd)
    143                    {
    144                        case '!': dontAnnotate.add(line); break;
    145                        case '<': copyToClient.add(line); break;
    146                        case '>': copyToServer.add(line); break;
    147                    }
    148                }
    149
    150                in.close();
    151                return true;
    152            }
    153            catch (Exception e)
    154            {
    155                System.err.println("Error: " + e.getMessage());
    156                return false;
    157            }
    158        }
    159
    160        public static void processJar(File clientInFile, File serverInFile, File clientOutFile, File serverOutFile) throws IOException
    161        {
    162            ZipFile cInJar = null;
    163            ZipFile sInJar = null;
    164            ZipOutputStream cOutJar = null;
    165            ZipOutputStream sOutJar = null;
    166
    167            try
    168            {
    169                try
    170                {
    171                    cInJar = new ZipFile(clientInFile);
    172                    sInJar = new ZipFile(serverInFile);
    173                }
    174                catch (FileNotFoundException e)
    175                {
    176                    throw new FileNotFoundException("Could not open input file: " + e.getMessage());
    177                }
    178                try
    179                {
    180                    cOutJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(clientOutFile)));
    181                    sOutJar = new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(serverOutFile)));
    182                }
    183                catch (FileNotFoundException e)
    184                {
    185                    throw new FileNotFoundException("Could not open output file: " + e.getMessage());
    186                }
    187                Hashtable<String, ZipEntry> cClasses = getClassEntries(cInJar, cOutJar);
    188                Hashtable<String, ZipEntry> sClasses = getClassEntries(sInJar, sOutJar);
    189                HashSet<String> cAdded = new HashSet<String>();
    190                HashSet<String> sAdded = new HashSet<String>();
    191
    192                for (Entry<String, ZipEntry> entry : cClasses.entrySet())
    193                {
    194                    String name = entry.getKey();
    195                    ZipEntry cEntry = entry.getValue();
    196                    ZipEntry sEntry = sClasses.get(name);
    197
    198                    if (sEntry == null)
    199                    {
    200                        if (!copyToServer.contains(name))
    201                        {
    202                            copyClass(cInJar, cEntry, cOutJar, null, true);
    203                            cAdded.add(name);
    204                        }
    205                        else
    206                        {
    207                            if (DEBUG)
    208                            {
    209                                System.out.println("Copy class c->s : " + name);
    210                            }
    211                            copyClass(cInJar, cEntry, cOutJar, sOutJar, true);
    212                            cAdded.add(name);
    213                            sAdded.add(name);
    214                        }
    215                        continue;
    216                    }
    217
    218                    sClasses.remove(name);
    219                    ClassInfo info = new ClassInfo(name);
    220                    shared.put(name, info);
    221
    222                    byte[] cData = readEntry(cInJar, entry.getValue());
    223                    byte[] sData = readEntry(sInJar, sEntry);
    224                    byte[] data = processClass(cData, sData, info);
    225
    226                    ZipEntry newEntry = new ZipEntry(cEntry.getName());
    227                    cOutJar.putNextEntry(newEntry);
    228                    cOutJar.write(data);
    229                    sOutJar.putNextEntry(newEntry);
    230                    sOutJar.write(data);
    231                    cAdded.add(name);
    232                    sAdded.add(name);
    233                }
    234
    235                for (Entry<String, ZipEntry> entry : sClasses.entrySet())
    236                {
    237                    if (DEBUG)
    238                    {
    239                        System.out.println("Copy class s->c : " + entry.getKey());
    240                    }
    241                    copyClass(sInJar, entry.getValue(), cOutJar, sOutJar, false);
    242                }
    243
    244                for (String name : new String[]{SideOnly.class.getName(), Side.class.getName()})
    245                {
    246                    String eName = name.replace(".", "/");
    247                    byte[] data = getClassBytes(name);
    248                    ZipEntry newEntry = new ZipEntry(name.replace(".", "/").concat(".class"));
    249                    if (!cAdded.contains(eName))
    250                    {
    251                        cOutJar.putNextEntry(newEntry);
    252                        cOutJar.write(data);
    253                    }
    254                    if (!sAdded.contains(eName))
    255                    {
    256                        sOutJar.putNextEntry(newEntry);
    257                        sOutJar.write(data);
    258                    }
    259                }
    260
    261            }
    262            finally
    263            {
    264                if (cInJar != null)
    265                {
    266                    try { cInJar.close(); } catch (IOException e){}
    267                }
    268
    269                if (sInJar != null)
    270                {
    271                    try { sInJar.close(); } catch (IOException e) {}
    272                }
    273                if (cOutJar != null)
    274                {
    275                    try { cOutJar.close(); } catch (IOException e){}
    276                }
    277
    278                if (sOutJar != null)
    279                {
    280                    try { sOutJar.close(); } catch (IOException e) {}
    281                }
    282            }
    283        }
    284
    285        private static void copyClass(ZipFile inJar, ZipEntry entry, ZipOutputStream outJar, ZipOutputStream outJar2, boolean isClientOnly) throws IOException
    286        {
    287            ClassReader reader = new ClassReader(readEntry(inJar, entry));
    288            ClassNode classNode = new ClassNode();
    289
    290            reader.accept(classNode, 0);
    291
    292            if (!dontAnnotate.contains(classNode.name))
    293            {
    294                if (classNode.visibleAnnotations == null) classNode.visibleAnnotations = new ArrayList<AnnotationNode>();
    295                classNode.visibleAnnotations.add(getSideAnn(isClientOnly));
    296            }
    297
    298            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    299            classNode.accept(writer);
    300            byte[] data = writer.toByteArray();
    301
    302            ZipEntry newEntry = new ZipEntry(entry.getName());
    303            if (outJar != null)
    304            {
    305                outJar.putNextEntry(newEntry);
    306                outJar.write(data);
    307            }
    308            if (outJar2 != null)
    309            {
    310                outJar2.putNextEntry(newEntry);
    311                outJar2.write(data);
    312            }
    313        }
    314
    315        private static AnnotationNode getSideAnn(boolean isClientOnly)
    316        {
    317            AnnotationNode ann = new AnnotationNode(Type.getDescriptor(SideOnly.class));
    318            ann.values = new ArrayList<Object>();
    319            ann.values.add("value");
    320            ann.values.add(new String[]{ Type.getDescriptor(Side.class), (isClientOnly ? "CLIENT" : "SERVER")});
    321            return ann;
    322        }
    323
    324        @SuppressWarnings("unchecked")
    325        private static Hashtable<String, ZipEntry> getClassEntries(ZipFile inFile, ZipOutputStream outFile) throws IOException
    326        {
    327            Hashtable<String, ZipEntry> ret = new Hashtable<String, ZipEntry>();
    328            for (ZipEntry entry : Collections.list((Enumeration<ZipEntry>)inFile.entries()))
    329            {
    330                if (entry.isDirectory())
    331                {
    332                    outFile.putNextEntry(entry);
    333                    continue;
    334                }
    335                String entryName = entry.getName();
    336                if (!entryName.endsWith(".class") || entryName.startsWith("."))
    337                {
    338                    ZipEntry newEntry = new ZipEntry(entry.getName());
    339                    outFile.putNextEntry(newEntry);
    340                    outFile.write(readEntry(inFile, entry));
    341                }
    342                else
    343                {
    344                    ret.put(entryName.replace(".class", ""), entry);
    345                }
    346            }
    347            return ret;
    348        }
    349        private static byte[] readEntry(ZipFile inFile, ZipEntry entry) throws IOException
    350        {
    351            return readFully(inFile.getInputStream(entry));
    352        }
    353        private static byte[] readFully(InputStream stream) throws IOException
    354        {
    355            byte[] data = new byte[4096];
    356            ByteArrayOutputStream entryBuffer = new ByteArrayOutputStream();
    357            int len;
    358            do
    359            {
    360                len = stream.read(data);
    361                if (len > 0)
    362                {
    363                    entryBuffer.write(data, 0, len);
    364                }
    365            } while (len != -1);
    366
    367            return entryBuffer.toByteArray();
    368        }
    369        private static class ClassInfo
    370        {
    371            public String name;
    372            public ArrayList<FieldNode> cField = new ArrayList<FieldNode>();
    373            public ArrayList<FieldNode> sField = new ArrayList<FieldNode>();
    374            public ArrayList<MethodNode> cMethods = new ArrayList<MethodNode>();
    375            public ArrayList<MethodNode> sMethods = new ArrayList<MethodNode>();
    376            public ClassInfo(String name){ this.name = name; }
    377            public boolean isSame() { return (cField.size() == 0 && sField.size() == 0 && cMethods.size() == 0 && sMethods.size() == 0); }
    378        }
    379
    380        public static byte[] processClass(byte[] cIn, byte[] sIn, ClassInfo info)
    381        {
    382            ClassNode cClassNode = getClassNode(cIn);
    383            ClassNode sClassNode = getClassNode(sIn);
    384
    385            processFields(cClassNode, sClassNode, info);
    386            processMethods(cClassNode, sClassNode, info);
    387
    388            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    389            cClassNode.accept(writer);
    390            return writer.toByteArray();
    391        }
    392
    393        private static ClassNode getClassNode(byte[] data)
    394        {
    395            ClassReader reader = new ClassReader(data);
    396            ClassNode classNode = new ClassNode();
    397            reader.accept(classNode, 0);
    398            return classNode;
    399        }
    400
    401        @SuppressWarnings("unchecked")
    402        private static void processFields(ClassNode cClass, ClassNode sClass, ClassInfo info)
    403        {
    404            List<FieldNode> cFields = cClass.fields;
    405            List<FieldNode> sFields = sClass.fields;
    406
    407            int sI = 0;
    408            for (int x = 0; x < cFields.size(); x++)
    409            {
    410                FieldNode cF = cFields.get(x);
    411                if (sI < sFields.size())
    412                {
    413                    if (!cF.name.equals(sFields.get(sI).name))
    414                    {
    415                        boolean serverHas = false;
    416                        for (int y = sI + 1; y < sFields.size(); y++)
    417                        {
    418                            if (cF.name.equals(sFields.get(y).name))
    419                            {
    420                                serverHas = true;
    421                                break;
    422                            }
    423                        }
    424                        if (serverHas)
    425                        {
    426                            boolean clientHas = false;
    427                            FieldNode sF = sFields.get(sI);
    428                            for (int y = x + 1; y < cFields.size(); y++)
    429                            {
    430                                if (sF.name.equals(cFields.get(y).name))
    431                                {
    432                                    clientHas = true;
    433                                    break;
    434                                }
    435                            }
    436                            if (!clientHas)
    437                            {
    438                                if  (sF.visibleAnnotations == null) sF.visibleAnnotations = new ArrayList<AnnotationNode>();
    439                                sF.visibleAnnotations.add(getSideAnn(false));
    440                                cFields.add(x++, sF);
    441                                info.sField.add(sF);
    442                            }
    443                        }
    444                        else
    445                        {
    446                            if  (cF.visibleAnnotations == null) cF.visibleAnnotations = new ArrayList<AnnotationNode>();
    447                            cF.visibleAnnotations.add(getSideAnn(true));
    448                            sFields.add(sI, cF);
    449                            info.cField.add(cF);
    450                        }
    451                    }
    452                }
    453                else
    454                {
    455                    if  (cF.visibleAnnotations == null) cF.visibleAnnotations = new ArrayList<AnnotationNode>();
    456                    cF.visibleAnnotations.add(getSideAnn(true));
    457                    sFields.add(sI, cF);
    458                    info.cField.add(cF);
    459                }
    460                sI++;
    461            }
    462            if (sFields.size() != cFields.size())
    463            {
    464                for (int x = cFields.size(); x < sFields.size(); x++)
    465                {
    466                    FieldNode sF = sFields.get(x);
    467                    if  (sF.visibleAnnotations == null) sF.visibleAnnotations = new ArrayList<AnnotationNode>();
    468                    sF.visibleAnnotations.add(getSideAnn(true));
    469                    cFields.add(x++, sF);
    470                    info.sField.add(sF);
    471                }
    472            }
    473        }
    474
    475        private static class MethodWrapper
    476        {
    477            private MethodNode node;
    478            public boolean client;
    479            public boolean server;
    480            public MethodWrapper(MethodNode node)
    481            {
    482                this.node = node;
    483            }
    484            @Override
    485            public boolean equals(Object obj)
    486            {
    487                if (obj == null || !(obj instanceof MethodWrapper)) return false;
    488                MethodWrapper mw = (MethodWrapper) obj;
    489                boolean eq = Objects.equal(node.name, mw.node.name) && Objects.equal(node.desc, mw.node.desc);
    490                if (eq)
    491                {
    492                    mw.client = this.client | mw.client;
    493                    mw.server = this.server | mw.server;
    494                    this.client = this.client | mw.client;
    495                    this.server = this.server | mw.server;
    496                    if (DEBUG)
    497                    {
    498                        System.out.printf(" eq: %s %s\n", this, mw);
    499                    }
    500                }
    501                return eq;
    502            }
    503
    504            @Override
    505            public int hashCode()
    506            {
    507                return Objects.hashCode(node.name, node.desc);
    508            }
    509            @Override
    510            public String toString()
    511            {
    512                return Objects.toStringHelper(this).add("name", node.name).add("desc",node.desc).add("server",server).add("client",client).toString();
    513            }
    514        }
    515        @SuppressWarnings("unchecked")
    516        private static void processMethods(ClassNode cClass, ClassNode sClass, ClassInfo info)
    517        {
    518            List<MethodNode> cMethods = (List<MethodNode>)cClass.methods;
    519            List<MethodNode> sMethods = (List<MethodNode>)sClass.methods;
    520            LinkedHashSet<MethodWrapper> allMethods = Sets.newLinkedHashSet();
    521
    522            int cPos = 0;
    523            int sPos = 0;
    524            int cLen = cMethods.size();
    525            int sLen = sMethods.size();
    526            String clientName = "";
    527            String lastName = clientName;
    528            String serverName = "";
    529            while (cPos < cLen || sPos < sLen)
    530            {
    531                do
    532                {
    533                    if (sPos>=sLen)
    534                    {
    535                        break;
    536                    }
    537                    MethodNode sM = sMethods.get(sPos);
    538                    serverName = sM.name;
    539                    if (!serverName.equals(lastName) && cPos != cLen)
    540                    {
    541                        if (DEBUG)
    542                        {
    543                            System.out.printf("Server -skip : %s %s %d (%s %d) %d [%s]\n", sClass.name, clientName, cLen - cPos, serverName, sLen - sPos, allMethods.size(), lastName);
    544                        }
    545                        break;
    546                    }
    547                    MethodWrapper mw = new MethodWrapper(sM);
    548                    mw.server = true;
    549                    allMethods.add(mw);
    550                    if (DEBUG)
    551                    {
    552                        System.out.printf("Server *add* : %s %s %d (%s %d) %d [%s]\n", sClass.name, clientName, cLen - cPos, serverName, sLen - sPos, allMethods.size(), lastName);
    553                    }
    554                    sPos++;
    555                }
    556                while (sPos < sLen);
    557                do
    558                {
    559                    if (cPos>=cLen)
    560                    {
    561                        break;
    562                    }
    563                    MethodNode cM = cMethods.get(cPos);
    564                    lastName = clientName;
    565                    clientName = cM.name;
    566                    if (!clientName.equals(lastName) && sPos != sLen)
    567                    {
    568                        if (DEBUG)
    569                        {
    570                            System.out.printf("Client -skip : %s %s %d (%s %d) %d [%s]\n", cClass.name, clientName, cLen - cPos, serverName, sLen - sPos, allMethods.size(), lastName);
    571                        }
    572                        break;
    573                    }
    574                    MethodWrapper mw = new MethodWrapper(cM);
    575                    mw.client = true;
    576                    allMethods.add(mw);
    577                    if (DEBUG)
    578                    {
    579                        System.out.printf("Client *add* : %s %s %d (%s %d) %d [%s]\n", cClass.name, clientName, cLen - cPos, serverName, sLen - sPos, allMethods.size(), lastName);
    580                    }
    581                    cPos++;
    582                }
    583                while (cPos < cLen);
    584            }
    585
    586            cMethods.clear();
    587            sMethods.clear();
    588
    589            for (MethodWrapper mw : allMethods)
    590            {
    591                if (DEBUG)
    592                {
    593                    System.out.println(mw);
    594                }
    595                cMethods.add(mw.node);
    596                sMethods.add(mw.node);
    597                if (mw.server && mw.client)
    598                {
    599                    // no op
    600                }
    601                else
    602                {
    603                    if (mw.node.visibleAnnotations == null) mw.node.visibleAnnotations = Lists.newArrayListWithExpectedSize(1);
    604                    mw.node.visibleAnnotations.add(getSideAnn(mw.client));
    605                    if (mw.client)
    606                    {
    607                        info.sMethods.add(mw.node);
    608                    }
    609                    else
    610                    {
    611                        info.cMethods.add(mw.node);
    612                    }
    613                }
    614            }
    615        }
    616
    617        public static byte[] getClassBytes(String name) throws IOException
    618        {
    619            InputStream classStream = null;
    620            try
    621            {
    622                classStream = MCPMerger.class.getResourceAsStream("/" + name.replace('.', '/').concat(".class"));
    623                return readFully(classStream);
    624            }
    625            finally
    626            {
    627                if (classStream != null)
    628                {
    629                    try
    630                    {
    631                        classStream.close();
    632                    }
    633                    catch (IOException e){}
    634                }
    635            }
    636        }
    637    }