FixDfsTube.java

import java.io.FileOutputStream;
import java.io.InputStream;
import java.net.JarURLConnection;
import java.net.URL;
import java.util.Enumeration;
import java.util.zip.ZipEntry;
import java.util.zip.ZipFile;
import java.util.zip.ZipOutputStream;

import org.apache.commons.lang.StringUtils;

import com.documentum.thirdparty.javassist.ClassPool;
import com.documentum.thirdparty.javassist.CtClass;
import com.documentum.thirdparty.javassist.CtMethod;
import com.emc.documentum.fs.rt.context.impl.ServiceContextAdapter;
import com.emc.documentum.fs.rt.context.impl.ThreadLocalContextStorage;
import com.emc.documentum.fs.rt.impl.content.ContentTransferModeUtil;
import com.emc.documentum.fs.rt.impl.pipe.DfsTube;
import com.sun.xml.ws.api.message.Header;
import com.sun.xml.ws.message.DOMHeader;

/**
 * @author Andrey B. Panfilov <andrey@panfilov.tel>
 */
public class FixDfsTube {

    public static void main(String[] args) throws Exception {
        ClassPool pool = ClassPool.getDefault();
        for (Class cls : new Class[] {ServiceContextAdapter.class,
            ThreadLocalContextStorage.class, Header.class, StringUtils.class,
            DOMHeader.class, ContentTransferModeUtil.class }) {
            pool.importPackage(cls.getPackage().getName());
        }

        String cls = DfsTube.class.getName();
        CtClass cc = pool.get(cls);
        CtMethod m = cc.getDeclaredMethod("processRequest");
        StringBuilder body = new StringBuilder();
        body.append("{\n");
        body.append("ServiceContextAdapter contextAdapter = (ServiceContextAdapter) ThreadLocalContextStorage\n");
        body.append("        .get();\n");
        body.append("String serviceName = $1.endpointAddress.getURL().getFile();\n");
        body.append("if (serviceName.endsWith(CONTEXT_REGISTRY_SERVICE_NAME)\n");
        body.append("        || serviceName.endsWith(AGENT_SERVICE_SERVICE_NAME)\n");
        body.append("        || contextAdapter == null) {\n");
        body.append("    return super.processRequest($1);\n");
        body.append("}\n");
        body.append("Header header;\n");
        body.append("if (StringUtils.isNotBlank(contextAdapter.getConsolidatedContext()\n");
        body.append("        .getToken())) {\n");
        body.append("    header = new DOMHeader(addTokenToHeader(contextAdapter\n");
        body.append("            .getConsolidatedContext().getToken()));\n");
        body.append("    $1.getMessage().getHeaders().add(header);\n");
        body.append("}\n");
        body.append("\n");
        body.append("if (!contextAdapter.isDeltaEmpty()) {\n");
        body.append("    header = new DOMHeader(\n");
        body.append("            createDeltaContextElement(contextAdapter.getDeltaContext()));\n");
        body.append("    $1.getMessage().getHeaders().add(header);\n");
        body.append("}\n");
        body.append("if (ContentTransferModeUtil.isMtomTransfer(contextAdapter)) {\n");
        body.append("    enableMtom($1);\n");
        body.append("} else {\n");
        body.append("    disableMtom($1);\n");
        body.append("}\n");
        body.append("return super.processRequest($1);\n");
        body.append("}");
        m.setBody(body.toString());

        String resourceName = DfsTube.class.getName().replace(".", "/")
                + ".class";
        URL resourceURL = DfsTube.class.getResource("/" + resourceName);
        JarURLConnection connection = (JarURLConnection) resourceURL
                .openConnection();
        URL jarURL = connection.getJarFileURL();
        String fileName = jarURL.getFile();
        ZipFile zipFile = new ZipFile(fileName);
        ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(fileName
                + ".patched"));
        for (Enumeration e = zipFile.entries(); e.hasMoreElements();) {
            ZipEntry entryIn = (ZipEntry) e.nextElement();
            if (!entryIn.getName().equals(resourceName)) {
                zos.putNextEntry(new ZipEntry(entryIn.getName()));
                InputStream is = zipFile.getInputStream(entryIn);
                byte[] buf = new byte[1024];
                int len;
                while ((len = is.read(buf)) > 0) {
                    zos.write(buf, 0, len);
                }
            } else {
                zos.putNextEntry(new ZipEntry(resourceName));
                zos.write(cc.toBytecode());
            }
            zos.closeEntry();
        }
        zos.close();
    }

}