Java tutorial
/* * Copyright 2009-2012 Michael Dalton * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package jtaint; import java.util.ArrayList; import java.util.List; import org.objectweb.asm.ClassAdapter; import org.objectweb.asm.ClassVisitor; import org.objectweb.asm.Label; import org.objectweb.asm.MethodVisitor; import org.objectweb.asm.Opcodes; import org.objectweb.asm.Type; public class XssAdapter extends StubAdapter implements Opcodes { private final List instrumentedMethods; private String className; private int version; private static List xssMethods(String name) { if ("java/io/PrintWriter".equals(name)) return XssContextAdapter.printWriterMethods; else if ("java/io/OutputStream".equals(name)) return XssContextAdapter.outputStreamMethods; throw new RuntimeException("Unknown class " + name); } public XssAdapter(ClassVisitor cv, String className) { super(cv, xssMethods(className)); instrumentedMethods = xssMethods(className); } public XssAdapter(ClassVisitor cv, XssContextAdapter xca) { super(cv, null); instrumentedMethods = xca.instrumentedMethods(); } public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) { super.visit(version, access, name, signature, superName, interfaces); this.className = name; this.version = version; } private void buildGenericWrapper(final MethodVisitor mv, final String methodName, final String desc) { new XssLockBuilder(mv, version, className, methodName, desc) { public void onUnlocked() { mv.visitVarInsn(ALOAD, 0); mv.visitMethodInsn(INVOKEVIRTUAL, className, ByteCodeUtil.internalName("getHtmlValidator"), "()Ljtaint/HtmlValidator;"); Type[] t = Type.getArgumentTypes(methodDesc); for (int i = 0, l = 1; i < t.length; l += t[i].getSize(), i++) mv.visitVarInsn(t[i].getOpcode(ILOAD), l); mv.visitMethodInsn(INVOKEVIRTUAL, "jtaint/HtmlValidator", methodName, Type.getMethodDescriptor(Type.VOID_TYPE, t)); } }.build(); } /* This is called for write(int) methods. This situation is complex * because ServletOutputStream.write(int i) is defined in * java.io.OutputStream to write the value (byte) i to the underlying * buffer. However, PrintWriter.write(int i) is defined in java.io.Writer * to write (char) i to the underlying stream. * * Thus we perform a simple instanceof check against java.io.PrintWriter * to ensure that we perform the correct version (integer-to-byte or * integer-to-character) before calling the underlying HTML validator. * We don't perform the instanceof check using ServletOutputStream because * this class is not included in rt.jar and thus is not guaranteed to * be loaded. */ private void buildWriteIntWrapper(final MethodVisitor mv, final String methodName, final String desc) { new XssLockBuilder(mv, version, className, methodName, desc) { public void onUnlocked() { mv.visitVarInsn(ALOAD, 0); mv.visitMethodInsn(INVOKEVIRTUAL, className, ByteCodeUtil.internalName("getHtmlValidator"), "()Ljtaint/HtmlValidator;"); if ("javax/servlet/ServletOutputStream".equals(className)) { mv.visitVarInsn(ILOAD, 1); mv.visitInsn(I2B); mv.visitMethodInsn(INVOKEVIRTUAL, "jtaint/HtmlValidator", "write", "(B)V"); } else if ("java/io/PrintWriter".equals(className)) { mv.visitVarInsn(ILOAD, 1); mv.visitInsn(I2C); mv.visitMethodInsn(INVOKEVIRTUAL, "jtaint/HtmlValidator", "print", "(C)V"); } else { /* Runtime check required */ Label l0 = new Label(); Label l1 = new Label(); mv.visitVarInsn(ALOAD, 0); mv.visitTypeInsn(INSTANCEOF, "java/io/PrintWriter"); mv.visitJumpInsn(IFNE, l1); /* If-branch, we're a ServletOutputStream. */ mv.visitVarInsn(ILOAD, 1); mv.visitInsn(I2B); mv.visitMethodInsn(INVOKEVIRTUAL, "jtaint/HtmlValidator", "write", "(B)V"); mv.visitJumpInsn(GOTO, l0); /* Else-branch, we're a PrintWriter */ mv.visitLabel(l1); if (version == V1_6) mv.visitFrame(F_SAME1, 0, null, 1, new Object[] { "jtaint/HtmlValidator" }); mv.visitVarInsn(ILOAD, 1); mv.visitInsn(I2C); mv.visitMethodInsn(INVOKEVIRTUAL, "jtaint/HtmlValidator", "print", "(C)V"); mv.visitLabel(l0); if (version == V1_6) { mv.visitFrame(F_SAME, 0, null, 0, null); /* We can't end on a visitFrame because * InstrumentationLockAdapter performs a visitFrame * after this method completes. Two consecutive * visitFrames cause the Java 6 type checker to barf, * so just pad with a single NOP. */ mv.visitInsn(NOP); } } } }.build(); } public MethodVisitor visitMethod(final int access, final String name, final String desc, String signature, String[] exceptions) { final MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions); if (!instrumentedMethods.contains(new MethodDecl(access, name, desc))) return mv; /* Else wrap and rename */ if ("write".equals(name) && "(I)V".equals(desc)) buildWriteIntWrapper(mv, name, desc); else buildGenericWrapper(mv, name, desc); return super.visitMethod(ByteCodeUtil.inheritAccessFlags(access) + ACC_PRIVATE, ByteCodeUtil.internalName(name), desc, signature, exceptions); } public void onEndBuildStubs() { InstrumentationLockBuilder.visitEnd(cv, className); cv.visitField(ACC_PRIVATE + ACC_TRANSIENT, ByteCodeUtil.internalName("htmlValidator"), "Ljtaint/HtmlValidator;", null, null).visitEnd(); ByteCodeUtil.buildGetter(cv, className, ByteCodeUtil.internalName("htmlValidator"), "Ljtaint/HtmlValidator;", 0, ByteCodeUtil.internalName("getHtmlValidator")); ByteCodeUtil.buildSetter(cv, className, ByteCodeUtil.internalName("htmlValidator"), "Ljtaint/HtmlValidator;", 0, ByteCodeUtil.internalName("setHtmlValidator")); } public static InstrumentationBuilder builder(String buildClass) { return Builder.getInstance(buildClass); } private static abstract class XssLockBuilder extends InstrumentationLockBuilder { public XssLockBuilder(MethodVisitor mv, int version, String className, String methodName, String methodDesc) { super(mv, version, className, methodName, methodDesc); } /* if HtmlValidator is null, don't even bother with inc/dec of * lock, just call wrapped method and exit */ public void onMethodEnter() { mv.visitVarInsn(ALOAD, 0); mv.visitMethodInsn(INVOKEVIRTUAL, className, ByteCodeUtil.internalName("getHtmlValidator"), "()Ljtaint/HtmlValidator;"); Label l0 = new Label(); mv.visitJumpInsn(IFNONNULL, l0); mv.visitVarInsn(ALOAD, 0); int l = 1; Type[] t = Type.getArgumentTypes(methodDesc); for (int i = 0; i < t.length; l += t[i].getSize(), i++) mv.visitVarInsn(t[i].getOpcode(ILOAD), l); mv.visitMethodInsn(INVOKEVIRTUAL, className, ByteCodeUtil.internalName(methodName), methodDesc); mv.visitInsn(Type.getReturnType(methodDesc).getOpcode(IRETURN)); mv.visitLabel(l0); if (version == V1_6) mv.visitFrame(F_SAME, 0, null, 0, null); } } private static class Builder implements InstrumentationBuilder { private static final Builder OUTPUTSTREAM; private static final Builder PRINTWRITER; static { OUTPUTSTREAM = new Builder("java/io/OutputStream"); PRINTWRITER = new Builder("java/io/PrintWriter"); } private final String builderClass; public Builder(String builderClass) { this.builderClass = builderClass; } public static InstrumentationBuilder getInstance(String s) { if ("java/io/OutputStream".equals(s)) return OUTPUTSTREAM; else if ("java/io/PrintWriter".equals(s)) return PRINTWRITER; else throw new RuntimeException("Unknown class " + s); } public ClassVisitor build(ClassVisitor cv) { return new XssAdapter(cv, builderClass); } } }