package com.ctrip.framework.xpipe.redis.instrument;


import com.alibaba.arthas.deps.org.objectweb.asm.ClassReader;
import com.alibaba.arthas.deps.org.objectweb.asm.ClassWriter;
import com.ctrip.framework.xpipe.redis.instrument.adapter.SocketAdapter;
import com.ctrip.framework.xpipe.redis.instrument.adapter.SocketChannelImplAdapter;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class ProxyAgent implements ClassFileTransformer {

    private static Map<String, String> transformedClasses = new ConcurrentHashMap<>();

    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace("/", ".");
        if (className.equals("java.net.Socket") && !transformedClasses.containsKey(className)) {
            ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            SocketAdapter socketAdapter = new SocketAdapter(classWriter);
            ClassReader classReader = new ClassReader(classfileBuffer);
            classReader.accept(socketAdapter, 0);
            transformedClasses.put(className, className);
            return classWriter.toByteArray();
        } else if (className.equals("sun.nio.ch.SocketChannelImpl") && !transformedClasses.containsKey(className)) {
            ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            SocketChannelImplAdapter socketChannelImplAdapter = new SocketChannelImplAdapter(classWriter);
            ClassReader classReader = new ClassReader(classfileBuffer);
            classReader.accept(socketChannelImplAdapter, 0);
            transformedClasses.put(className, className);
            return classWriter.toByteArray();
        }
        return null;
    }

}

@@ -479,223 +479,66 @@
 in;%0A
-import java.util.Map;%0Aimport java.util.concurrent.ConcurrentHashMap;%0A%0Apublic class ProxyAgent implements ClassFileTransformer %7B%0A%0A    private static Map%3CString, String%3E transformedClasses = new ConcurrentHashMap%3C%3E();
+%0Apublic class ProxyAgent implements ClassFileTransformer %7B
 %0A%0A  
@@ -827,54 +827,8 @@
 et%22)
- && !transformedClasses.containsKey(className)
 ) %7B%0A
@@ -1110,66 +1110,8 @@
 0);%0A
-            transformedClasses.put(className, className);%0A
     
@@ -1223,54 +1223,8 @@
 pl%22)
- && !transformedClasses.containsKey(className)
 ) %7B%0A
@@ -1550,66 +1550,8 @@
 0);%0A
-            transformedClasses.put(className, className);%0A
     

remove duplicate loaded class check