FPE sniffer

package paci.samples.bce;

import static org.objectweb.asm.Opcodes.DADD;
import static org.objectweb.asm.Opcodes.DDIV;
import static org.objectweb.asm.Opcodes.DMUL;
import static org.objectweb.asm.Opcodes.DSUB;
import static org.objectweb.asm.Opcodes.IFEQ;
import static org.objectweb.asm.Opcodes.INVOKESTATIC;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;
import java.util.logging.Logger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.commons.GeneratorAdapter;

/*
 * @author [email protected]
 */
public class Agent {
    public static void premain( String agentArgs,Instrumentation instr ) throws Exception {
        String targets = "paci.*";

        if( agentArgs != null )
            targets = agentArgs;

        Set<String> targetSet = new HashSet<String>( );
        for( String str : targets.split( ":" ) )
            targetSet.add( str );
        instr.addTransformer( new TransformerImpl( targetSet ) );
    }
}

class TransformerImpl implements ClassFileTransformer {
    private static final Logger LOGGER = Logger.getLogger(
        TransformerImpl.class.getPackage( ).getName( ) );
    private final Set<Pattern> _targets;

    TransformerImpl( Set<String> targets ) {
        _targets = new HashSet<Pattern>( );
        for( String target : targets )
            _targets.add( Pattern.compile( target ) );
    }

    public byte[ ] transform( ClassLoader loader,
                  final String className,
                  Class<?> classBeingRedefined,
                  ProtectionDomain protectionDomain,
                  byte[ ] classfileBuffer )
        throws IllegalClassFormatException {
        boolean match = false;

        if( className.startsWith( "java/lang" ) || className.startsWith( "java/util" ) )
            return null;

        if( _targets.size( ) > 0 )
        {
            for( Pattern target : _targets )
            {
                Matcher matcher = target.matcher( className );
                if( matcher.matches( ) )
                {
                    match = true;
                    break;
                }
            }
        }
        else
            match = true;

        if( ! match )
            return null;

        try {
            ClassReader cr = new ClassReader( classfileBuffer );
            ClassWriter cw = new ClassWriter( 0 );
            cr.accept(
                    new ClassAdapter( cw ) {
                        public MethodVisitor visitMethod( int access,String name,String desc,String signature,String[ ] exceptions ) {
                            MethodVisitor mv = super.visitMethod( access,name,desc,signature,exceptions );
                            return new BytecodeModifier( access,name,desc,mv );
                        }
                    },ClassReader.EXPAND_FRAMES );

            return cw.toByteArray( );
        }
        catch( Exception e ) {
            // If anything goes wrong, return the class as is and log this fact
            LOGGER.severe( "Instrumentation of class " + className +
                    " failed. Message is " + e.getMessage( ) + "." );

            return null;
        }
    }
}

class BytecodeModifier extends GeneratorAdapter {
    public BytecodeModifier( int access,String name,String desc,MethodVisitor mv ) {
        super( mv,access,name,desc );
    }

    private void generateChecks( String operation ) {
        dup2( );
        // Is it NaN?
        mv.visitMethodInsn( INVOKESTATIC,"java/lang/Double","isNaN","(D)Z" );
        Label label1 = new Label( );
        mv.visitJumpInsn( IFEQ,label1 );
        mv.visitLdcInsn( operation );
        mv.visitMethodInsn(INVOKESTATIC, "paci/samples/bce/FPEListener", "resultIsNan", "(Ljava/lang/String;)V");
        mv.visitLabel( label1 );

        dup2( );
        // Is it infinite?
        mv.visitMethodInsn( INVOKESTATIC,"java/lang/Double","isInfinite","(D)Z" );
        Label label2 = new Label( );
        mv.visitJumpInsn( IFEQ,label2 );
        mv.visitLdcInsn( operation );
        mv.visitMethodInsn(INVOKESTATIC, "paci/samples/bce/FPEListener", "resultIsInfinite", "(Ljava/lang/String;)V");
        mv.visitLabel( label2 );
    }

    @Override
    public void visitInsn( int opcode ) {
        mv.visitInsn( opcode );
        switch( opcode ) {
            case DDIV :
                generateChecks( "DIVIDE" );
                break;
            case DMUL :
                generateChecks( "MULTIPLY" );
                break;
            case DADD :
                generateChecks( "ADD" );
                break;
            case DSUB :
                generateChecks( "SUBSTRACT" );
                break;
        }
    }

    @Override
    public void visitMethodInsn( int opcode,String owner,String name,String desc ) {
        // Slot is an allocation unit on the stack. One slot equals a word (32 bit).
        // An object reference and 32-bit (or shorter) primitives take one slot,
        // 64-bit primitives like long and double take two slots.

        super.visitMethodInsn( opcode,owner,name,desc );
        // Intercept all java.lang.Math calls which are static, take one
        // double as an argument and return double
        if( opcode == INVOKESTATIC && owner.equals( "java/lang/Math" ) && desc.equals( "(D)D" ) ) {
            String operation = name.toUpperCase( );
            generateChecks( operation );
        }
    }

    /**
     * Reports the maximum consumption of stack and local variables in slots.
     */
    @Override
    public void visitMaxs( int maxStack,int maxLocals ) {
        // Parameters are given in slots, not in bytes
        // Our manipulation logic added 2 slots and no local variables
        super.visitMaxs( maxStack + 2,maxLocals );
    }
}