com.nerdvision.agent.inst.PluginTransformer Maven / Gradle / Ivy

Go to download
package com.nerdvision.agent.inst;

import com.nerdvision.plugins.datadog.DDTracePlugin;
import com.nerdvision.agent.Settings;
import com.nerdvision.agent.Utils;
import com.nerdvision.agent.api.plugins.IMatcher;
import com.nerdvision.agent.api.plugins.IPlugin;
import com.nerdvision.agent.inst.asm.ClassLoaderAwareClassWriter;
import com.nerdvision.agent.inst.asm.PluginVisitor;
import com.nerdvision.agent.inst.asm.TransformerUtils;
import com.nerdvision.agent.plugins.SentryPlugin;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.tree.ClassNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import static com.nerdvision.agent.inst.InstrumentationBreakpointService.COMPUTE_ON_CLASS_VERSION;

@SuppressWarnings("Convert2Diamond")
public class PluginTransformer implements ClassFileTransformer
{
    private static final Logger LOGGER = LoggerFactory.getLogger( PluginTransformer.class );

    private final Instrumentation inst;
    private final Settings settings;
    private final String disPath;

    private final List plugins = new CopyOnWriteArrayList();


    public PluginTransformer( final Instrumentation inst, final Settings settings )
    {
        this.inst = inst;
        this.settings = settings;
        this.disPath = settings.getSettingAs( "transform.path", String.class );
    }


    /**
     * Register a new set of plugins and scan loaded classes
     *
     * @param plugins the plugins to load
     */
    public void addPlugins( final IPlugin... plugins )
    {
        for( IPlugin plugin : plugins )
        {
            if( plugin.isActive( this.settings ) )
            {
                this.plugins.add( plugin );
            }
        }
        this.scan();
    }


    /**
     * Scan the loaded classes to find any matched that need reloaded
     */
    private void scan()
    {
        final List> matchesClass = new ArrayList<>();
        final Class[] allLoadedClasses = this.inst.getAllLoadedClasses();
        for( Class allLoadedClass : allLoadedClasses )
        {
            if( TransformerUtils.isExcludedClass( allLoadedClass ) || !inst.isModifiableClass( allLoadedClass ) )
            {
                continue;
            }

            for( IPlugin plugin : plugins )
            {
                final List matchers = plugin.getMatchers();
                for( IMatcher matcher : matchers )
                {

                    if( matcher.matchesClass( Utils.internalClass( allLoadedClass ) ) )
                    {
                        matchesClass.add( allLoadedClass );
                    }
                }
            }
        }

        if( !matchesClass.isEmpty() )
        {
            try
            {
                this.inst.retransformClasses( matchesClass.toArray( new Class[0] ) );
            }
            catch( UnmodifiableClassException e )
            {
                LOGGER.error( "Unable to modify class", e );
            }
        }
    }


    @Override
    public byte[] transform( final ClassLoader loader,
                             final String className,
                             final Class classBeingRedefined,
                             final ProtectionDomain protectionDomain,
                             final byte[] classfileBuffer ) throws IllegalClassFormatException
    {
        // no plugins do nothing
        if( plugins.isEmpty() )
        {
            return null;
        }

        // class is excluded do nothing
        if( TransformerUtils.isExcludedClass( Utils.externalClassName( className ) ) )
        {
            return null;
        }

        IMatcher matched = null;
        for( IPlugin plugin : plugins )
        {
            for( IMatcher matcher : plugin.getMatchers() )
            {
                if( matcher.matchesClass( className ) )
                {
                    matched = matcher;
                    break;
                }
            }
        }

        // no matched plugins do nothing
        if( matched == null )
        {
            return null;
        }

        try
        {
            final ClassReader reader = new ClassReader( classfileBuffer );
            final ClassNode cn = new ClassNode();
            // no need to expand frames here as we only need the version out
            reader.accept( cn, ClassReader.SKIP_FRAMES | ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG );

            // if we are not java 1.1 and are greater than COMPUTE_ON_CLASS_VERSION (50 (java 1.6)) then compute frames
            final boolean classVersionNeedsComputeFrames = (cn.version != org.objectweb.asm.Opcodes.V1_1
                    && cn.version >= COMPUTE_ON_CLASS_VERSION);

            final ClassWriter writer = new ClassLoaderAwareClassWriter( reader,
                    // compute the frames if we need to else just maxes as we are a class version that does not have frames
                    classVersionNeedsComputeFrames ? ClassWriter.COMPUTE_FRAMES : ClassWriter.COMPUTE_MAXS, loader );

            final PluginVisitor visitor = new PluginVisitor( writer, className, matched );
            // if we are going to compute frames then we can skip them - as they will be ignored and re computed anyway.
            reader.accept( visitor, classVersionNeedsComputeFrames ? ClassReader.SKIP_FRAMES : ClassReader.EXPAND_FRAMES );

            if( visitor.wasChanged() )
            {
                LOGGER.debug( "Applied {} to class {}.", matched.getClass().getName(), className );
                final byte[] res = writer.toByteArray();
                TransformerUtils.storeUnsafe( this.disPath, classfileBuffer, res, className );
                return res;
            }
            else
            {
                LOGGER.debug( "Class {} not changed.", className );
                return null;
            }
        }
        catch( final Throwable t )
        {
            LOGGER.error( "transform failed for {}", className, t );
        }
        return null;
    }


    public static PluginTransformer init( final Instrumentation inst, final Settings settings )
    {
        final PluginTransformer pluginTransformer = new PluginTransformer( inst, settings );
        inst.addTransformer( pluginTransformer, true );

        pluginTransformer.addPlugins( new SentryPlugin(), new DDTracePlugin() );

        return pluginTransformer;
    }
}