Time machine

package paci.samples.bce;

import static org.objectweb.asm.Opcodes.INVOKESTATIC;
import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
import static org.objectweb.asm.Opcodes.LADD;

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.concurrent.TimeUnit;
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.MethodVisitor;
import org.objectweb.asm.commons.GeneratorAdapter;

/*
 * @author paci.salminen@gmail.com
 */
public class TimeMachineAgent {
    public static void premain( String agentArgs,Instrumentation instr ) throws Exception {
        String targets = "paci.*:org/joda.*";
        long timeOffset = TimeUnit.HOURS.toMillis( 6 );
        if( agentArgs != null )
            targets = agentArgs;

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

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

    TransformerImpl2( Set<String> targets,long timeOffset ) {
        _targets = new HashSet<Pattern>( );
        for( String target : targets )
            _targets.add( Pattern.compile( target ) );
        this.timeOffset = timeOffset;
    }

    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" ) || className.startsWith( "com/ibm" ) || className.startsWith( "sun/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;

        //System.out.println( "Codesource: " + protectionDomain.getCodeSource( ).getLocation( ) );

        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 TimeMachineAdapter( access,name,desc,mv,timeOffset );
                        }
                    },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 TimeMachineAdapter extends GeneratorAdapter {
    private long timeOffset;

    public TimeMachineAdapter( int access,String name,String desc,MethodVisitor mv,long offset ) {
        super( mv,access,name,desc );
        timeOffset = offset;
    }

    /**
     * Modifies calls to methods which may be used get current time.
     * Modification is accomplished using bytecode manipulation, so that e.g
     * an instantiation of java.util.Date with a default constructor is intercepted
     * and offset (positive or negative) is added to current time. This method is not used
     * in any way during execution of the code; bytecode manipulation happens during
     * class loading time.
     */
    @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.

        // Intercept java.util.Date default constructor call
        if( opcode == INVOKESPECIAL && owner.equals( "java/util/Date" ) &&
            name.equals( "<init>" ) && desc.equals( "()V" ) ) {
            // Bytecode NEW is already executed
            /*
             *  Duplicate reference before constructor call
             *  FIXES BUG when no assignment is done
             *  private static void foo( ) {
             *      new Date( );
             *  }
            */
            // Stack contains now a reference to java.util.Date (slots reserved 1)
            // Note that constructor call is not made yet.
            // Dup the reference (slots used +1)
            dup( );

            // Now generate a constructor call (consumes one slot so slots used is 1 after the call)
            super.visitMethodInsn( opcode,owner,name,desc );
            // Dup the reference (slots used is 2) ...
            dup( );
            // Invoke getTime method (slots used after this call is 3)
            visitMethodInsn( INVOKEVIRTUAL,"java/util/Date","getTime", "()J");
            // Push offset to the stack (slots used after this call is 5)
            visitLdcInsn( new Long( timeOffset ) );
            // Add the operands (slots used after this call is 3)
            visitInsn( LADD );
            // And finally call setTime method (slots used after this call is 1)
            visitMethodInsn( INVOKEVIRTUAL,"java/util/Date","setTime", "(J)V");
            // One reference to java.util.Date remains on the stack
        }
        //Intercept java.lang.System.currentTimeMills call
        else if(    opcode == INVOKESTATIC &&
                    owner.equals( "java/lang/System") &&
                    name.equals( "currentTimeMillis" ) &&
                    desc.equals( "()J" ) ) {
            // Generate a currentTimeMillis call (after this call slots used is 2)
            super.visitMethodInsn( opcode,owner,name,desc );
            // Push offset to the stack (slots used after this call is 4)
            visitLdcInsn( new Long( timeOffset ) );
            // Add the operands (slots used after this call is 2)
            visitInsn( LADD );
            // Only current time remains on the stack
        }
        //Intercept java.util.Calendar getInstance call
        else if( opcode == INVOKESTATIC &&
                owner.equals( "java/util/Calendar" ) &&
                name.equals( "getInstance" ) &&
                desc.equals( "()Ljava/util/Calendar;" ) ) {
            // Generate a getInstance call (after this call slots used is 1)
            super.visitMethodInsn( opcode,owner,name,desc );
            // Dup the reference (slots used after this is 2)
            dup( );
            // Dup the reference (slots used after this is 3)
            dup( );
            // Call getTimeInMillis method (slots used after this call is 4)
            visitMethodInsn( INVOKEVIRTUAL,"java/util/Calendar","getTimeInMillis", "()J");
            // Push offset to the stack (slots used after this call is 6)
            visitLdcInsn( new Long( timeOffset ) );
            // Add the operands (slots used after this call is 4)
            visitInsn( LADD );
            // And finally call setTimeInMillis method (slots used after this call is 1)
            visitMethodInsn( INVOKEVIRTUAL,"java/util/Calendar","setTimeInMillis", "(J)V");
            // One reference to java.util.Calendar remains on the stack
        }
        else
            // Otherwise generate whatever instruction was encountered
            super.visitMethodInsn( opcode,owner,name,desc );
    }

    /**
     * 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 date manipulation logic added 5 slots and no local variables
        super.visitMaxs( maxStack + 4,maxLocals );
    }
}