From caf2938f8d61faa299d785f8221e9763bd410551 Mon Sep 17 00:00:00 2001 From: Vasil Vasilev Date: Fri, 31 Jul 2026 10:27:49 +0200 Subject: [PATCH] Instrument @NotNull assertions into the compiled JVM bytecode #SCL-22085 Add a NotNull/Nullable bytecode instrumentation feature matching what IntelliJ's JPS build and the IntelliJ Platform Gradle Plugin do: null checks are generated for @NotNull-annotated method parameters (throwing IllegalArgumentException) and return values (IllegalStateException). The instrumenter itself (NotNullVerifyingInstrumenter, AuxiliaryMethodGenerator, FailSafeClassReader, FailSafeMethodVisitor) is copied verbatim from IntelliJ IDEA Community Edition (Apache 2.0, commit 7a35a7d7fe64), with only the package, the ASM package, and Opcodes.API_VERSION -> Opcodes.ASM9 changed. It is wired into the existing manipulateBytecode flow next to the threading-annotations instrumentation and controlled by two new settings: instrumentNotNullAnnotations (default false) and notNullAnnotations (default org.jetbrains.annotations.NotNull). Co-Authored-By: Claude Fable 5 --- README.md | 20 + build.sbt | 3 + .../AuxiliaryMethodGenerator.java | 282 +++++++++ .../FailSafeClassReader.java | 43 ++ .../FailSafeMethodVisitor.java | 35 ++ .../NotNullVerifyingInstrumenter.java | 557 ++++++++++++++++++ .../scala/org/jetbrains/sbtidea/Init.scala | 2 + .../scala/org/jetbrains/sbtidea/Keys.scala | 8 + .../instrumentation/ManipulateBytecode.scala | 40 +- .../instrumentation/NotNullInstrumenter.scala | 32 + .../NotNullInstrumenterTest.scala | 230 ++++++++ 11 files changed, 1246 insertions(+), 6 deletions(-) create mode 100644 ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/AuxiliaryMethodGenerator.java create mode 100644 ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeClassReader.java create mode 100644 ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeMethodVisitor.java create mode 100644 ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/NotNullVerifyingInstrumenter.java create mode 100644 ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenter.scala create mode 100644 ideaSupport/src/test/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenterTest.scala diff --git a/README.md b/README.md index c89d32e1..b55599a5 100644 --- a/README.md +++ b/README.md @@ -439,6 +439,26 @@ Generate JVM bytecode to assert that a method is called on the correct IDEA thre See: [IntelliJ IDEA ThreadingAssertions.java](https://github.com/JetBrains/intellij-community/blob/5758eb99b4a1971ebe75cda755693cc930949465/platform/core-api/src/com/intellij/util/concurrency/ThreadingAssertions.java) +#### `instrumentNotNullAnnotations :: SettingKey[Boolean]` + +**Default**: `false` + +Generate JVM bytecode to assert that values of `@NotNull`-annotated method parameters and return values are not null. +A null value passed for an annotated parameter throws an `IllegalArgumentException`; a null value returned from an +annotated method throws an `IllegalStateException`. The set of annotation classes is configured with the +`notNullAnnotations` setting. + +This is the same instrumentation that IntelliJ IDEA applies to its own codebase and that the +[IntelliJ Platform Gradle Plugin](https://plugins.jetbrains.com/docs/intellij/tools-intellij-platform-gradle-plugin-tasks.html#instrumentCode) +performs as part of its `instrumentCode` task. Kotlin classes are skipped, as kotlinc generates its own +nullability assertions. + +#### `notNullAnnotations :: SettingKey[Seq[String]]` + +**Default**: `Seq("org.jetbrains.annotations.NotNull")` + +Fully qualified names of the annotation classes instrumented by `instrumentNotNullAnnotations`. + #### `packageOutputDir :: SettingKey[File]` **Default**: `target.value / "plugin" / intellijPluginName.in(ThisBuild).value.removeSpaces` diff --git a/build.sbt b/build.sbt index 0b4bd8d2..79890b5f 100644 --- a/build.sbt +++ b/build.sbt @@ -133,6 +133,9 @@ lazy val ideaSupport = (project in file("ideaSupport")) "org.ow2.asm" % "asm" % "9.10.1", "io.get-coursier" %% "coursier" % "2.1.24", "commons-io" % "commons-io" % "2.22.0", + + // Used to compile the test fixtures of the @NotNull instrumentation tests + "org.jetbrains" % "annotations" % "26.1.0" % Test ), ) diff --git a/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/AuxiliaryMethodGenerator.java b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/AuxiliaryMethodGenerator.java new file mode 100644 index 00000000..bed5a684 --- /dev/null +++ b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/AuxiliaryMethodGenerator.java @@ -0,0 +1,282 @@ +// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.notNullVerification. +// The only changes are the package name, the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm) +// and Opcodes.API_VERSION (a JetBrains ASM addition) -> Opcodes.ASM9. +package org.jetbrains.sbtidea.instrumentation.notNullVerification; + +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassVisitor; +import org.objectweb.asm.Label; +import org.objectweb.asm.MethodVisitor; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.objectweb.asm.Opcodes.AASTORE; +import static org.objectweb.asm.Opcodes.ACC_INTERFACE; +import static org.objectweb.asm.Opcodes.ACC_PRIVATE; +import static org.objectweb.asm.Opcodes.ACC_STATIC; +import static org.objectweb.asm.Opcodes.ACC_SYNTHETIC; +import static org.objectweb.asm.Opcodes.ANEWARRAY; +import static org.objectweb.asm.Opcodes.ASM9; +import static org.objectweb.asm.Opcodes.ATHROW; +import static org.objectweb.asm.Opcodes.BIPUSH; +import static org.objectweb.asm.Opcodes.DUP; +import static org.objectweb.asm.Opcodes.DUP_X1; +import static org.objectweb.asm.Opcodes.GOTO; +import static org.objectweb.asm.Opcodes.ICONST_0; +import static org.objectweb.asm.Opcodes.ICONST_1; +import static org.objectweb.asm.Opcodes.ICONST_2; +import static org.objectweb.asm.Opcodes.ICONST_3; +import static org.objectweb.asm.Opcodes.ICONST_4; +import static org.objectweb.asm.Opcodes.ICONST_5; +import static org.objectweb.asm.Opcodes.ILOAD; +import static org.objectweb.asm.Opcodes.INVOKESPECIAL; +import static org.objectweb.asm.Opcodes.INVOKESTATIC; +import static org.objectweb.asm.Opcodes.NEW; +import static org.objectweb.asm.Opcodes.SIPUSH; +import static org.objectweb.asm.Opcodes.SWAP; + +final class AuxiliaryMethodGenerator { + private static final String STRING_CLASS_NAME = "java/lang/String"; + private static final String OBJECT_CLASS_NAME = "java/lang/Object"; + private static final String CONSTRUCTOR_NAME = ""; + private static final String EXCEPTION_INIT_SIGNATURE = "(L" + STRING_CLASS_NAME + ";)V"; + private static final String REPORTING_METHOD_DESC = "(I)V"; + + private final ClassReader myOriginalClass; + private final boolean myIsInterface; + private final List myReportingPlaces = new ArrayList<>(); + private String myReportingMethod; + private int myMaxArgCount; + + AuxiliaryMethodGenerator(ClassReader originalClass) { + myOriginalClass = originalClass; + myIsInterface = (myOriginalClass.getAccess() & ACC_INTERFACE) == ACC_INTERFACE; + } + + private String getReportingMethodName() { + if (myReportingMethod == null) { + myReportingMethod = suggestUniqueName(); + } + return myReportingMethod; + } + + private String suggestUniqueName() { + Set existingMethods = populateExistingMethods(); + for (int i = 0;; i++) { + String name = "$$$reportNull$$$" + i; + if (!existingMethods.contains(name)) { + return name; + } + } + } + + private Set populateExistingMethods() { + final Set existingMethods = new HashSet<>(); + myOriginalClass.accept(new ClassVisitor(ASM9) { + @Override + public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) { + existingMethods.add(name); + return null; + } + }, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES); + return existingMethods; + } + + void reportError(MethodVisitor mv, String className, String exceptionClass, String descrPattern, String[] args) { + myMaxArgCount = Math.max(myMaxArgCount, args.length); + + int index = myReportingPlaces.size(); + myReportingPlaces.add(new ReportingPlace(exceptionClass, descrPattern, args)); + pushIntConstant(mv, index); + + mv.visitMethodInsn(INVOKESTATIC, className, getReportingMethodName(), REPORTING_METHOD_DESC, myIsInterface); + } + + private static void pushIntConstant(MethodVisitor mv, int i) { + if (i <= 5) { + mv.visitInsn(getSmallIntConstantInstruction(i)); + } + else if (i <= Byte.MAX_VALUE) { + mv.visitIntInsn(BIPUSH, i); + } + else { + mv.visitIntInsn(SIPUSH, i); + } + } + + private static int getSmallIntConstantInstruction(int i) { + switch (i) { + case 0: return ICONST_0; + case 1: return ICONST_1; + case 2: return ICONST_2; + case 3: return ICONST_3; + case 4: return ICONST_4; + case 5: return ICONST_5; + default: throw new AssertionError(i); + } + } + + void generateReportingMethod(ClassVisitor cw) { + if (myReportingPlaces.isEmpty()) return; + + MethodVisitor mv = cw.visitMethod(ACC_PRIVATE | ACC_SYNTHETIC | ACC_STATIC, getReportingMethodName(), REPORTING_METHOD_DESC, null, null); + pushExceptionMessage(mv); + createExceptionObject(mv); + mv.visitInsn(ATHROW); + mv.visitMaxs(0, 0); + } + + private void createExceptionObject(final MethodVisitor mv) { + new SwitchGenerator() { + @Override + void generateCaseBody(String exceptionClass) { + mv.visitTypeInsn(NEW, exceptionClass); + mv.visitInsn(DUP_X1); + mv.visitInsn(SWAP); + + mv.visitMethodInsn(INVOKESPECIAL, exceptionClass, CONSTRUCTOR_NAME, EXCEPTION_INIT_SIGNATURE, false); + } + + @Override + String getSwitchedValue(ReportingPlace place) { + return place.exceptionClass; + } + }.generateSwitch(mv); + } + + private void pushExceptionMessage(MethodVisitor mv) { + pushFormatPattern(mv); + + createFormatArgArray(mv); + for (int i = 0; i < myMaxArgCount; i++) { + pushFormatArg(mv, i); + } + + //noinspection SpellCheckingInspection + mv.visitMethodInsn(INVOKESTATIC, STRING_CLASS_NAME, "format", "(Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;", false); + } + + private void createFormatArgArray(final MethodVisitor mv) { + new SwitchGenerator(){ + @Override + void generateCaseBody(Integer argCount) { + pushIntConstant(mv, argCount); + } + + @Override + Integer getSwitchedValue(ReportingPlace place) { + return place.args.length; + } + }.generateSwitch(mv); + mv.visitTypeInsn(ANEWARRAY, OBJECT_CLASS_NAME); + } + + private void pushFormatArg(final MethodVisitor mv, final int index) { + new SwitchGenerator() { + @Override + protected String getSwitchedValue(ReportingPlace place) { + return index < place.args.length ? place.args[index] : null; + } + + @Override + void generateCaseBody(String value) { + if (value != null) { + mv.visitInsn(DUP); + pushIntConstant(mv, index); + mv.visitLdcInsn(value); + mv.visitInsn(AASTORE); + } + } + }.generateSwitch(mv); + } + + private void pushFormatPattern(final MethodVisitor mv) { + new SwitchGenerator() { + @Override + protected String getSwitchedValue(ReportingPlace place) { + return place.descrPattern; + } + + @Override + void generateCaseBody(String descrPattern) { + mv.visitLdcInsn(descrPattern); + } + }.generateSwitch(mv); + } + + private static class ReportingPlace { + final String exceptionClass; + final String descrPattern; + final String[] args; + + ReportingPlace(String exceptionClass, String descrPattern, String[] args) { + this.exceptionClass = exceptionClass; + this.descrPattern = descrPattern; + this.args = args; + } + } + + private abstract class SwitchGenerator { + void generateSwitch(MethodVisitor mv) { + Label[] labels = getCaseLabels(); + if (labels == null) { + // all places behave in a same way, don't bother with switch + generateCaseBody(getSwitchedValue(myReportingPlaces.get(0))); + } else { + reallyGenerateSwitch(mv, labels, deduplicateLabels(labels)); + } + } + + private void reallyGenerateSwitch(MethodVisitor mv, Label[] labels, Map label2Place) { + Label afterSwitch = new Label(); + mv.visitVarInsn(ILOAD, 0); + mv.visitTableSwitchInsn(0, labels.length - 1 , labels[0], labels); + + for (Map.Entry entry : label2Place.entrySet()) { + mv.visitLabel(entry.getKey()); + generateCaseBody(getSwitchedValue(entry.getValue())); + mv.visitJumpInsn(GOTO, afterSwitch); + } + + mv.visitLabel(afterSwitch); + } + + private Map deduplicateLabels(Label[] labels) { + Map label2Place = new LinkedHashMap<>(); + for (int i = 0; i < labels.length; i++) { + if (!label2Place.containsKey(labels[i])) { + label2Place.put(labels[i], myReportingPlaces.get(i)); + } + } + return label2Place; + } + + private Label[] getCaseLabels() { + Map labelsByValue = new HashMap<>(); + Label[] labels = new Label[myReportingPlaces.size()]; + for (int i = 0; i < myReportingPlaces.size(); i++) { + labels[i] = getOrCreateLabel(labelsByValue, getSwitchedValue(myReportingPlaces.get(i))); + } + return labelsByValue.size() == 1 ? null : labels; + } + + private Label getOrCreateLabel(Map labelsByValue, T key) { + Label label = labelsByValue.get(key); + if (label == null) { + labelsByValue.put(key, label = new Label()); + } + return label; + } + + abstract void generateCaseBody(T switchedValue); + + abstract T getSwitchedValue(ReportingPlace place); + } +} diff --git a/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeClassReader.java b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeClassReader.java new file mode 100644 index 00000000..88fbd341 --- /dev/null +++ b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeClassReader.java @@ -0,0 +1,43 @@ +// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.instrumentation. +// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm). +package org.jetbrains.sbtidea.instrumentation.notNullVerification; + +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.Label; + +import java.io.IOException; +import java.io.InputStream; + +public final class FailSafeClassReader extends ClassReader { + private static final Label INVALID = new Label(); + + public FailSafeClassReader(byte[] b) { + super(b); + } + + public FailSafeClassReader(byte[] b, int off, int len) { + super(b, off, len); + } + + public FailSafeClassReader(InputStream is) throws IOException { + super(is); + } + + public FailSafeClassReader(String name) throws IOException { + super(name); + } + + @Override + protected Label readLabel(int offset, Label[] labels) { + // attempt to workaround javac bug: + // annotation table from original method is duplicated for synthetic bridge methods. + // All offsets in the duplicated table is taken from original annotations table and obviously are not relevant for the bridge method + if (offset >= 0 && offset < labels.length) { + return super.readLabel(offset, labels); + } + else { + return INVALID; + } + } +} diff --git a/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeMethodVisitor.java b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeMethodVisitor.java new file mode 100644 index 00000000..6209c9ff --- /dev/null +++ b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeMethodVisitor.java @@ -0,0 +1,35 @@ +// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.instrumentation. +// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm). +package org.jetbrains.sbtidea.instrumentation.notNullVerification; + +import org.objectweb.asm.AnnotationVisitor; +import org.objectweb.asm.Label; +import org.objectweb.asm.MethodVisitor; +import org.objectweb.asm.TypePath; + +/** + * To be used together with FailSafeClassReader: adds null checks for labels describing annotation visibility range. + * For incorrectly generated annotations FailSafeClassReader returns null labels. Local variables annotations with null labels + * will be ignored by this visitor. + */ +public class FailSafeMethodVisitor extends MethodVisitor { + public FailSafeMethodVisitor(int api, MethodVisitor mv) { + super(api, mv); + } + + @Override + public AnnotationVisitor visitLocalVariableAnnotation(int typeRef, TypePath typePath, Label[] start, Label[] end, int[] index, String desc, boolean visible) { + for (Label aStart : start) { + if (aStart == null) { + return null; + } + } + for (Label anEnd : end) { + if (anEnd == null) { + return null; + } + } + return super.visitLocalVariableAnnotation(typeRef, typePath, start, end, index, desc, visible); + } +} diff --git a/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/NotNullVerifyingInstrumenter.java b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/NotNullVerifyingInstrumenter.java new file mode 100644 index 00000000..fadea399 --- /dev/null +++ b/ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/NotNullVerifyingInstrumenter.java @@ -0,0 +1,557 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.notNullVerification. +// The only changes are the package name, the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm) +// and Opcodes.API_VERSION (a JetBrains ASM addition) -> Opcodes.ASM9. +package org.jetbrains.sbtidea.instrumentation.notNullVerification; + +import org.objectweb.asm.AnnotationVisitor; +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassVisitor; +import org.objectweb.asm.Handle; +import org.objectweb.asm.Label; +import org.objectweb.asm.MethodVisitor; +import org.objectweb.asm.Type; +import org.objectweb.asm.TypePath; +import org.objectweb.asm.TypeReference; + +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +import static org.objectweb.asm.Opcodes.ACC_BRIDGE; +import static org.objectweb.asm.Opcodes.ACC_ENUM; +import static org.objectweb.asm.Opcodes.ACC_FINAL; +import static org.objectweb.asm.Opcodes.ACC_PRIVATE; +import static org.objectweb.asm.Opcodes.ACC_STATIC; +import static org.objectweb.asm.Opcodes.ALOAD; +import static org.objectweb.asm.Opcodes.ANEWARRAY; +import static org.objectweb.asm.Opcodes.ARETURN; +import static org.objectweb.asm.Opcodes.ASM9; +import static org.objectweb.asm.Opcodes.CHECKCAST; +import static org.objectweb.asm.Opcodes.DUP; +import static org.objectweb.asm.Opcodes.DUP2; +import static org.objectweb.asm.Opcodes.DUP2_X1; +import static org.objectweb.asm.Opcodes.DUP2_X2; +import static org.objectweb.asm.Opcodes.DUP_X1; +import static org.objectweb.asm.Opcodes.DUP_X2; +import static org.objectweb.asm.Opcodes.GOTO; +import static org.objectweb.asm.Opcodes.IFNONNULL; +import static org.objectweb.asm.Opcodes.IINC; +import static org.objectweb.asm.Opcodes.INVOKEDYNAMIC; +import static org.objectweb.asm.Opcodes.INVOKESPECIAL; +import static org.objectweb.asm.Opcodes.INVOKESTATIC; +import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL; +import static org.objectweb.asm.Opcodes.JSR; +import static org.objectweb.asm.Opcodes.LDC; +import static org.objectweb.asm.Opcodes.LOOKUPSWITCH; +import static org.objectweb.asm.Opcodes.MULTIANEWARRAY; +import static org.objectweb.asm.Opcodes.NEW; +import static org.objectweb.asm.Opcodes.NEWARRAY; +import static org.objectweb.asm.Opcodes.NOP; +import static org.objectweb.asm.Opcodes.RET; +import static org.objectweb.asm.Opcodes.TABLESWITCH; + +public final class NotNullVerifyingInstrumenter extends ClassVisitor { + private static final String IAE_CLASS_NAME = "java/lang/IllegalArgumentException"; + private static final String ISE_CLASS_NAME = "java/lang/IllegalStateException"; + private static final String KOTLIN_METADATA_ANNOTATION_CLASS_DESCRIPTOR = "Lkotlin/Metadata;"; + + private static final String ANNOTATION_DEFAULT_METHOD = "value"; + + @SuppressWarnings("SSBasedInspection") + private static final String[] EMPTY_STRING_ARRAY = new String[0]; + + private final MethodData myMethodData; + private boolean myIsModification = false; + private RuntimeException myPostponedError; + private final AuxiliaryMethodGenerator myAuxGenerator; + + private NotNullVerifyingInstrumenter(ClassVisitor classVisitor, ClassReader reader, String[] notNullAnnotations) { + super(ASM9, classVisitor); + Set annoSet = new HashSet<>(); + for (String annotation : notNullAnnotations) { + annoSet.add('L' + annotation.replace('.', '/') + ';'); + } + myMethodData = collectMethodData(reader, annoSet); + myAuxGenerator = new AuxiliaryMethodGenerator(reader); + } + + public static boolean processClassFile(ClassReader reader, ClassVisitor writer, String[] notNullAnnotations) { + NotNullVerifyingInstrumenter instrumenter = new NotNullVerifyingInstrumenter(writer, reader, notNullAnnotations); + if (instrumenter.myMethodData.myIsKotlinBytecode) { + // skip Kotlin-generated bytecode, as nullability assertions are handled on compiler level by kotlinc + return false; + } + reader.accept(instrumenter, 0); + return instrumenter.myIsModification; + } + + private static class MethodInfo { + final NotNullState nullability = new NotNullState(); + final Map paramNames = new HashMap<>(); + final Map paramNullability = new LinkedHashMap<>(); + boolean isStable; + int paramAnnotationOffset; + + NotNullState obtainParameterNullability(int index) { + NotNullState state = paramNullability.get(index); + if (state == null) { + state = new NotNullState(); + paramNullability.put(index, state); + } + return state; + } + } + + private static final class MethodData { + private String myClassName; + private boolean myIsKotlinBytecode; + private final Map myMethodInfos = new HashMap<>(); + + static String key(String methodName, String desc) { + return methodName + desc; + } + + String lookupParamName(String methodName, String desc, Integer num) { + MethodInfo info = myMethodInfos.get(key(methodName, desc)); + Map names = info == null ? null : info.paramNames; + return names != null ? names.get(num) : null; + } + + boolean isAlwaysNotNull(String className, String methodName, String desc) { + if (myClassName.equals(className)) { + MethodInfo info = myMethodInfos.get(key(methodName, desc)); + return info != null && info.isStable && info.nullability.isNotNull(); + } + return false; + } + } + + private static MethodData collectMethodData(ClassReader reader, final Set notNullAnnotations) { + final MethodData result = new MethodData(); + reader.accept(new ClassVisitor(ASM9) { + private boolean myEnum, myInner; + + @Override + public AnnotationVisitor visitAnnotation(String desc, boolean visible) { + if (KOTLIN_METADATA_ANNOTATION_CLASS_DESCRIPTOR.equals(desc)) { + result.myIsKotlinBytecode = true; + } + return super.visitAnnotation(desc, visible); + } + + @Override + public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) { + super.visit(version, access, name, signature, superName, interfaces); + result.myClassName = name; + myEnum = (access & ACC_ENUM) != 0; + } + + @Override + public void visitInnerClass(String name, String outerName, String innerName, int access) { + super.visitInnerClass(name, outerName, innerName, access); + if (result.myClassName.equals(name)) { + myInner = (access & ACC_STATIC) == 0; + } + } + + @Override + public MethodVisitor visitMethod(int access, final String name, final String desc, String signature, String[] exceptions) { + final Type[] args = Type.getArgumentTypes(desc); + final boolean methodCanHaveNullability = isReferenceType(Type.getReturnType(desc)); + + final Map paramSlots = new LinkedHashMap<>(); // map: localVariableSlot -> methodParameterIndex + int slotIndex = isStatic(access) ? 0 : 1; + for (int paramIndex = 0; paramIndex < args.length; paramIndex++) { + Type arg = args[paramIndex]; + paramSlots.put(slotIndex, paramIndex); + slotIndex += arg.getSize(); + } + + final MethodInfo methodInfo = new MethodInfo(); + methodInfo.isStable = (access & (ACC_FINAL | ACC_STATIC | ACC_PRIVATE)) != 0; + methodInfo.paramAnnotationOffset = !"".equals(name) ? 0 : myEnum ? 2 : myInner ? 1 : 0; + result.myMethodInfos.put(MethodData.key(name, desc), methodInfo); + + return new MethodVisitor(api) { + private int myParamAnnotationOffset = methodInfo.paramAnnotationOffset; + + @Override + public void visitAnnotableParameterCount(int parameterCount, boolean visible) { + if (myParamAnnotationOffset != 0 && parameterCount == args.length) { + myParamAnnotationOffset = 0; + } + super.visitAnnotableParameterCount(parameterCount, visible); + } + + @Override + public AnnotationVisitor visitParameterAnnotation(int parameter, String anno, boolean visible) { + AnnotationVisitor base = super.visitParameterAnnotation(parameter, anno, visible); + return checkParameterNullability(parameter + myParamAnnotationOffset, anno, base, false); + } + + @Override + public AnnotationVisitor visitAnnotation(String anno, boolean isRuntime) { + AnnotationVisitor base = super.visitAnnotation(anno, isRuntime); + if (methodCanHaveNullability && notNullAnnotations.contains(anno)) { + return collectNotNullArgs(base, methodInfo.nullability.withNotNull(anno, ISE_CLASS_NAME)); + } + return base; + } + + @Override + public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String anno, boolean visible) { + AnnotationVisitor base = super.visitTypeAnnotation(typeRef, typePath, anno, visible); + if (typePath != null) return base; + + TypeReference ref = new TypeReference(typeRef); + if (methodCanHaveNullability && ref.getSort() == TypeReference.METHOD_RETURN) { + if (notNullAnnotations.contains(anno)) { + return collectNotNullArgs(base, methodInfo.nullability.withNotNull(anno, ISE_CLASS_NAME)); + } + else if (seemsNullable(anno)) { + methodInfo.nullability.hasTypeUseNullable = true; + } + } + else if (ref.getSort() == TypeReference.METHOD_FORMAL_PARAMETER) { + return checkParameterNullability(ref.getFormalParameterIndex() + methodInfo.paramAnnotationOffset, anno, base, true); + } + + return base; + } + + private boolean seemsNullable(String anno) { + String shortName = getAnnoShortName(anno); + // use hardcoded short names until it causes trouble + // this is to avoid cumbersome passing of configured nullable names from the IDE + return shortName.contains("Nullable") || shortName.equals("CheckForNull"); + } + + private AnnotationVisitor collectNotNullArgs(AnnotationVisitor base, final NotNullState state) { + return new AnnotationVisitor(ASM9, base) { + @Override + public void visit(String methodName, Object o) { + if (ANNOTATION_DEFAULT_METHOD.equals(methodName) && !((String) o).isEmpty()) { + state.message = (String) o; + } + else if ("exception".equals(methodName) && o instanceof Type && !((Type)o).getClassName().equals(Exception.class.getName())) { + state.exceptionType = ((Type)o).getInternalName(); + } + super.visit(methodName, o); + } + }; + } + + private AnnotationVisitor checkParameterNullability(int parameter, String anno, AnnotationVisitor av, boolean typeUse) { + if (parameter >= 0 && parameter < args.length && isReferenceType(args[parameter])) { + if (notNullAnnotations.contains(anno)) { + return collectNotNullArgs(av, methodInfo.obtainParameterNullability(parameter).withNotNull(anno, IAE_CLASS_NAME)); + } + else if (typeUse && seemsNullable(anno)) { + methodInfo.obtainParameterNullability(parameter).hasTypeUseNullable = true; + } + } + + return av; + } + + @Override + public void visitLocalVariable(String name2, String desc, String signature, Label start, Label end, int slotIndex) { + Integer paramIndex = paramSlots.get(slotIndex); + if (paramIndex != null) { + methodInfo.paramNames.put(paramIndex, name2); + } + } + }; + } + }, ClassReader.SKIP_FRAMES); + return result; + } + + private static class NotNullState { + String message; + String exceptionType; + String notNullAnno; + boolean hasTypeUseNullable; + + NotNullState withNotNull(String notNullAnno, String exceptionType) { + this.notNullAnno = notNullAnno; + this.exceptionType = exceptionType; + return this; + } + + boolean isNotNull() { + return notNullAnno != null && !hasTypeUseNullable; + } + + String getNullParamMessage(String paramName) { + if (message != null) return message; + String shortName = getAnnoShortName(notNullAnno); + if (paramName != null) return "Argument for @" + shortName + " parameter '%s' of %s.%s must not be null"; + return "Argument %s for @" + shortName + " parameter of %s.%s must not be null"; + } + + String getNullResultMessage() { + if (message != null) return message; + String shortName = getAnnoShortName(notNullAnno); + return "@" + shortName + " method %s.%s must not return null"; + } + } + + private static String getAnnoShortName(String anno) { + String fullName = anno.substring(1, anno.length() - 1); // "Lpk/name;" -> "pk/name" + return fullName.substring(fullName.lastIndexOf('/') + 1); + } + + @Override + public MethodVisitor visitMethod(int access, final String name, final String desc, String signature, String[] exceptions) { + final MethodInfo info = myMethodData.myMethodInfos.get(MethodData.key(name, desc)); + if ((access & ACC_BRIDGE) != 0 || info == null) { + return new FailSafeMethodVisitor(ASM9, super.visitMethod(access, name, desc, signature, exceptions)); + } + + final boolean isStatic = isStatic(access); + final Type[] args = Type.getArgumentTypes(desc); + final NotNullInstructionTracker instrTracker = new NotNullInstructionTracker(cv.visitMethod(access, name, desc, signature, exceptions)); + return new FailSafeMethodVisitor(ASM9, instrTracker) { + private Label myStartGeneratedCodeLabel; + + @Override + public void visitCode() { + for (Iterator iterator = info.paramNullability.values().iterator(); iterator.hasNext(); ) { + if (!iterator.next().isNotNull()) { + iterator.remove(); + } + } + if (!info.paramNullability.isEmpty()) { + myStartGeneratedCodeLabel = new Label(); + mv.visitLabel(myStartGeneratedCodeLabel); + } + for (Map.Entry entry : info.paramNullability.entrySet()) { + Integer param = entry.getKey(); + int var = isStatic ? 0 : 1; + for (int i = 0; i < param; ++i) { + var += args[i].getSize(); + } + mv.visitVarInsn(ALOAD, var); + + Label end = new Label(); + mv.visitJumpInsn(IFNONNULL, end); + + NotNullState state = entry.getValue(); + String paramName = myMethodData.lookupParamName(name, desc, param); + String descrPattern = state.getNullParamMessage(paramName); + String[] args = state.message != null + ? EMPTY_STRING_ARRAY + : new String[]{paramName != null ? paramName : String.valueOf(param - info.paramAnnotationOffset), myMethodData.myClassName, name}; + reportError(state.exceptionType, end, descrPattern, args); + } + } + + @Override + public void visitLocalVariable(String name, String desc, String signature, Label start, Label end, int index) { + boolean isParameterOrThisRef = isStatic ? index < args.length : index <= args.length; + Label label = (isParameterOrThisRef && myStartGeneratedCodeLabel != null) ? myStartGeneratedCodeLabel : start; + mv.visitLocalVariable(name, desc, signature, label, end, index); + } + + @Override + public void visitInsn(int opcode) { + if (opcode == ARETURN && instrTracker.canBeNull() && info.nullability.isNotNull()) { + mv.visitInsn(DUP); + Label skipLabel = new Label(); + mv.visitJumpInsn(IFNONNULL, skipLabel); + String descrPattern = info.nullability.getNullResultMessage(); + String[] args = info.nullability.message != null ? EMPTY_STRING_ARRAY : new String[]{myMethodData.myClassName, name}; + reportError(info.nullability.exceptionType, skipLabel, descrPattern, args); + } + + mv.visitInsn(opcode); + } + + private void reportError(String exceptionClass, Label end, String descrPattern, String[] args) { + myAuxGenerator.reportError(mv, myMethodData.myClassName, exceptionClass, descrPattern, args); + mv.visitLabel(end); + myIsModification = true; + processPostponedErrors(); + } + + @Override + @SuppressWarnings("SpellCheckingInspection") + public void visitMaxs(int maxStack, int maxLocals) { + try { + super.visitMaxs(maxStack, maxLocals); + } + catch (Throwable e) { + registerError(name, "visitMaxs", e); + } + } + }; + } + + @Override + public void visitEnd() { + myAuxGenerator.generateReportingMethod(cv); + super.visitEnd(); + } + + private static boolean isStatic(int access) { + return (access & ACC_STATIC) != 0; + } + + private static boolean isReferenceType(Type type) { + return type.getSort() == Type.OBJECT || type.getSort() == Type.ARRAY; + } + + private void registerError(String methodName, @SuppressWarnings("SameParameterValue") String operationName, Throwable t) { + if (myPostponedError == null) { + // throw the first error that occurred + Throwable cause = t.getCause(); + if (cause != null) t = cause; + + String message = t.getMessage(); + + StringWriter writer = new StringWriter(); + t.printStackTrace(new PrintWriter(writer)); + + StringBuilder text = new StringBuilder(); + text.append("Operation '").append(operationName).append("' failed for ").append(myMethodData.myClassName).append(".").append(methodName).append("(): "); + if (message != null) text.append(message); + text.append('\n').append(writer.getBuffer()); + myPostponedError = new RuntimeException(text.toString(), cause); + } + if (myIsModification) { + processPostponedErrors(); + } + } + + private void processPostponedErrors() { + RuntimeException error = myPostponedError; + if (error != null) { + throw error; + } + } + + private final class NotNullInstructionTracker extends MethodVisitor { + private boolean myCanBeNull = true; // initially assume the value can be null + + NotNullInstructionTracker(MethodVisitor delegate) { + super(ASM9, delegate); + } + + public boolean canBeNull() { + return myCanBeNull; + } + + @Override + public void visitIntInsn(int opcode, int operand) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitIntInsn(opcode, operand); + } + + @Override + public void visitVarInsn(int opcode, int var) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitVarInsn(opcode, var); + } + + @Override + public void visitTypeInsn(int opcode, String type) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitTypeInsn(opcode, type); + } + + @Override + public void visitFieldInsn(int opcode, String owner, String name, String descriptor) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitFieldInsn(opcode, owner, name, descriptor); + } + + @Override + public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) { + myCanBeNull = nextCanBeNullValue(opcode, owner, name, descriptor); /*is not a constructor call*/ + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + } + + @Override + public void visitInvokeDynamicInsn(String name, String descriptor, Handle bootstrapMethodHandle, Object... bootstrapMethodArguments) { + myCanBeNull = nextCanBeNullValue(INVOKEDYNAMIC); + super.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); + } + + @Override + public void visitJumpInsn(int opcode, Label label) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitJumpInsn(opcode, label); + } + + @Override + public void visitLdcInsn(Object value) { + myCanBeNull = nextCanBeNullValue(LDC); + super.visitLdcInsn(value); + } + + @Override + public void visitIincInsn(int var, int increment) { + myCanBeNull = nextCanBeNullValue(IINC); + super.visitIincInsn(var, increment); + } + + @Override + public void visitTableSwitchInsn(int min, int max, Label defaultLabel, Label... labels) { + myCanBeNull = nextCanBeNullValue(TABLESWITCH); + super.visitTableSwitchInsn(min, max, defaultLabel, labels); + } + + @Override + public void visitLookupSwitchInsn(Label defaultLabel, int[] keys, Label[] labels) { + myCanBeNull = nextCanBeNullValue(LOOKUPSWITCH); + super.visitLookupSwitchInsn(defaultLabel, keys, labels); + } + + @Override + public void visitMultiANewArrayInsn(String descriptor, int numDimensions) { + myCanBeNull = nextCanBeNullValue(MULTIANEWARRAY); + super.visitMultiANewArrayInsn(descriptor, numDimensions); + } + + @Override + public void visitInsn(int opcode) { + myCanBeNull = nextCanBeNullValue(opcode); + super.visitInsn(opcode); + } + + private boolean nextCanBeNullValue(int nextMethodCallOpcode, String owner, String name, String descriptor) { + if (nextMethodCallOpcode == INVOKESPECIAL && ("".equals(name) || myMethodData.isAlwaysNotNull(owner, name, descriptor))) { + // a constructor call or a NotNull marked own method + return false; + } + if ((nextMethodCallOpcode == INVOKESTATIC || nextMethodCallOpcode == INVOKEVIRTUAL) && + myMethodData.isAlwaysNotNull(owner, name, descriptor)) { + return false; + } + return true; + } + + private boolean nextCanBeNullValue(int nextOpcode) { + // if instruction guaranteed produces non-null stack value + if (nextOpcode == LDC || nextOpcode == NEW || nextOpcode == ANEWARRAY || nextOpcode == NEWARRAY || nextOpcode == MULTIANEWARRAY) { + return false; + } + // for some instructions, it is safe not to change the previously calculated flag value + if (nextOpcode == DUP || nextOpcode == DUP_X1 || nextOpcode == DUP_X2 || + nextOpcode == DUP2 || nextOpcode == DUP2_X1 || nextOpcode == DUP2_X2 || + nextOpcode == JSR || nextOpcode == GOTO || nextOpcode == NOP || + nextOpcode == RET || nextOpcode == CHECKCAST) { + return myCanBeNull; + } + // by default assume nullable + return true; + } + } +} diff --git a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Init.scala b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Init.scala index 21a74ed1..2fd8815a 100644 --- a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Init.scala +++ b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Init.scala @@ -217,6 +217,8 @@ trait Init { this: Keys.type => unmanagedResourceDirectories in Test += baseDirectory.value / "testResources", instrumentThreadingAnnotations := false, + instrumentNotNullAnnotations := false, + notNullAnnotations := Seq("org.jetbrains.annotations.NotNull"), Compile / manipulateBytecode := ManipulateBytecode.manipulateBytecodeTask(Compile).value, Test / manipulateBytecode := ManipulateBytecode.manipulateBytecodeTask(Test).value, diff --git a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Keys.scala b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Keys.scala index 93049935..0bbed703 100644 --- a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Keys.scala +++ b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/Keys.scala @@ -133,6 +133,14 @@ object Keys extends Defns with Init with Utils with Quirks { "Generate JVM bytecode to assert that a method is called on the correct IDEA thread " ++ "(supported method annotations: @RequiresBackgroundThread, @RequiresEdt, @RequiresReadLock, @RequiresReadLockAbsence, @RequiresWriteLock)") + lazy val instrumentNotNullAnnotations = settingKey[Boolean]( + "Generate JVM bytecode to assert that values of @NotNull-annotated method parameters and return values are not null " ++ + "(throws IllegalArgumentException for parameters and IllegalStateException for return values; " ++ + "the annotation classes are configured with the notNullAnnotations setting)") + + lazy val notNullAnnotations = settingKey[Seq[String]]( + "Fully qualified names of the annotation classes instrumented by instrumentNotNullAnnotations") + /* Deprecated task aliases */ lazy val buildIntellijOptionsIndex = taskKey[Unit]( diff --git a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/ManipulateBytecode.scala b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/ManipulateBytecode.scala index f933fcf6..843f29ee 100644 --- a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/ManipulateBytecode.scala +++ b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/ManipulateBytecode.scala @@ -1,6 +1,6 @@ package org.jetbrains.sbtidea.instrumentation -import org.jetbrains.sbtidea.Keys.instrumentThreadingAnnotations +import org.jetbrains.sbtidea.Keys.{instrumentNotNullAnnotations, instrumentThreadingAnnotations, notNullAnnotations} import sbt.* import sbt.Keys.* import sbt.internal.inc.{Analysis, Stamps} @@ -8,28 +8,56 @@ import xsbti.compile.CompileResult import xsbti.compile.analysis.Stamp import xsbti.{FileConverter, VirtualFileRef} +import java.net.URLClassLoader import java.nio.file.Path +import scala.util.control.NonFatal object ManipulateBytecode { def manipulateBytecodeTask(config: Configuration): Def.Initialize[Task[CompileResult]] = Def.taskDyn { - val doInstrument = instrumentThreadingAnnotations.value + val instrumentThreading = instrumentThreadingAnnotations.value + val instrumentNotNull = instrumentNotNullAnnotations.value val currentResult = (config / manipulateBytecode).value - if (doInstrument) { - instrumentTask(config, currentResult) + if (instrumentThreading || instrumentNotNull) { + instrumentTask(config, currentResult, instrumentThreading, instrumentNotNull) } else { Def.task(currentResult) } } - private def instrumentTask(config: Configuration, currentResult: CompileResult): Def.Initialize[Task[CompileResult]] = Def.task { + private def instrumentTask( + config: Configuration, + currentResult: CompileResult, + instrumentThreading: Boolean, + instrumentNotNull: Boolean + ): Def.Initialize[Task[CompileResult]] = Def.task { val previousResult = (config / previousCompile).value val converter = fileConverter.value + val annotations = notNullAnnotations.value + // fullClasspath cannot be used here: it depends on the compile task, which would create a cycle with manipulateBytecode + val classpath = (config / classDirectory).value +: (config / dependencyClasspath).value.map(_.data) val previousAnalysis = previousResult.analysis().asScala.collect { case a: Analysis => a }.getOrElse(Analysis.empty) val currentAnalysis = currentResult.analysis() match { case a: Analysis => a } val changed = changedClasses(currentAnalysis.stamps, previousAnalysis.stamps, converter) - changed.foreach(ThreadingAnnotationInstrumenter.instrument) + if (instrumentThreading) { + changed.foreach(ThreadingAnnotationInstrumenter.instrument) + } + if (instrumentNotNull) { + val classpathLoader = new URLClassLoader(classpath.map(_.toURI.toURL).toArray) + try { + changed.foreach { classFile => + try { + NotNullInstrumenter.instrument(classFile, annotations, classpathLoader) + } catch { + case NonFatal(e) => + throw new MessageOnlyException(s"Failed to instrument @NotNull assertions into $classFile: ${e.getMessage}") + } + } + } finally { + classpathLoader.close() + } + } val stamper = Stamps.timeWrapBinaryStamps(converter) diff --git a/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenter.scala b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenter.scala new file mode 100644 index 00000000..5274d7af --- /dev/null +++ b/ideaSupport/src/main/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenter.scala @@ -0,0 +1,32 @@ +package org.jetbrains.sbtidea.instrumentation + +import org.jetbrains.sbtidea.instrumentation.notNullVerification.{FailSafeClassReader, NotNullVerifyingInstrumenter} +import org.objectweb.asm.{ClassWriter, Opcodes} + +import java.nio.file.{Files, Path} + +private object NotNullInstrumenter { + + def instrument(classFile: Path, notNullAnnotations: Seq[String], classpathLoader: ClassLoader): Unit = { + val bytes = Files.readAllBytes(classFile) + val reader = new FailSafeClassReader(bytes) + val version = classFileVersion(reader) + if (reader.getClassName != "module-info" && (version & 0xFFFF) >= Opcodes.V1_5) { + // COMPUTE_FRAMES resolves common superclasses of the classes referenced in the rewritten methods, + // which must happen against the project classpath rather than this plugin's own classpath. + val writer = new ClassWriter(reader, asmClassWriterFlags(version)) { + override def getClassLoader: ClassLoader = classpathLoader + } + if (NotNullVerifyingInstrumenter.processClassFile(reader, writer, notNullAnnotations.toArray)) { + Files.write(classFile, writer.toByteArray) + } + } + } + + /** Class file version in the `minor << 16 | major` format (see `com.intellij.compiler.instrumentation.InstrumenterClassWriter`). */ + private def classFileVersion(reader: FailSafeClassReader): Int = + reader.readInt(4) + + private def asmClassWriterFlags(version: Int): Int = + if ((version & 0xFFFF) >= Opcodes.V1_6) ClassWriter.COMPUTE_FRAMES else ClassWriter.COMPUTE_MAXS +} diff --git a/ideaSupport/src/test/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenterTest.scala b/ideaSupport/src/test/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenterTest.scala new file mode 100644 index 00000000..6e4cc3a8 --- /dev/null +++ b/ideaSupport/src/test/scala/org/jetbrains/sbtidea/instrumentation/NotNullInstrumenterTest.scala @@ -0,0 +1,230 @@ +package org.jetbrains.sbtidea.instrumentation + +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers + +import java.lang.reflect.InvocationTargetException +import java.net.URLClassLoader +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Path, Paths} +import javax.tools.ToolProvider +import scala.collection.JavaConverters.asScalaIteratorConverter + +/** + * Modeled on IntelliJ IDEA's `com.intellij.java.compiler.notNullVerification.NotNullVerifyingInstrumenterTest`: + * compile small annotated Java fixtures at test runtime, instrument the class files, load them + * in a throwaway classloader and assert on the thrown exceptions. Loading and invoking the + * instrumented classes also runs the JVM bytecode verifier over the rewritten methods. + */ +class NotNullInstrumenterTest extends AnyFunSuite with Matchers { + + private val NotNullFqn = "org.jetbrains.annotations.NotNull" + + test("@NotNull parameter check without debug info uses the parameter index") { + val dir = compileAndInstrument(Map( + "SimpleParam.java" -> + """import org.jetbrains.annotations.NotNull; + |public class SimpleParam { + | public static void test(@NotNull Object o) {} + |}""".stripMargin + )) + val clazz = loadClass(dir, "SimpleParam") + invokeStatic(clazz, "test", new Object) // non-null argument must pass + val cause = interceptCause(clazz, "test", null) + cause shouldBe an[IllegalArgumentException] + cause.getMessage shouldBe "Argument 0 for @NotNull parameter of SimpleParam.test must not be null" + } + + test("@NotNull parameter check with debug info uses the parameter name") { + val dir = compileAndInstrument(Map( + "SimpleParam.java" -> + """import org.jetbrains.annotations.NotNull; + |public class SimpleParam { + | public static void test(@NotNull Object important) {} + |}""".stripMargin + ), debugInfo = true) + val clazz = loadClass(dir, "SimpleParam") + val cause = interceptCause(clazz, "test", null) + cause shouldBe an[IllegalArgumentException] + cause.getMessage shouldBe "Argument for @NotNull parameter 'important' of SimpleParam.test must not be null" + } + + test("@NotNull return value check throws IllegalStateException") { + val dir = compileAndInstrument(Map( + "SimpleReturn.java" -> + """import org.jetbrains.annotations.NotNull; + |public class SimpleReturn { + | @NotNull + | public static Object test(Object o) { return o; } + |}""".stripMargin + )) + val clazz = loadClass(dir, "SimpleReturn") + invokeStatic(clazz, "test", new Object) // non-null return value must pass + val cause = interceptCause(clazz, "test", null) + cause shouldBe an[IllegalStateException] + cause.getMessage shouldBe "@NotNull method SimpleReturn.test must not return null" + } + + test("custom message given as the annotation value is used verbatim") { + val dir = compileAndInstrument(Map( + "CustomMessage.java" -> + """import org.jetbrains.annotations.NotNull; + |public class CustomMessage { + | public static void test(@NotNull("null is not allowed here") Object o) {} + |}""".stripMargin + )) + val cause = interceptCause(loadClass(dir, "CustomMessage"), "test", null) + cause shouldBe an[IllegalArgumentException] + cause.getMessage shouldBe "null is not allowed here" + } + + test("custom annotation from the configured list with a custom exception class") { + val dir = compileAndInstrument(Map( + "MyNotNull.java" -> + """import java.lang.annotation.*; + |@Retention(RetentionPolicy.CLASS) + |@Target({ElementType.METHOD, ElementType.PARAMETER}) + |@interface MyNotNull { + | String value() default ""; + | Class exception() default Exception.class; + |}""".stripMargin, + "MyException.java" -> + """public class MyException extends RuntimeException { + | public MyException(String message) { super(message); } + |}""".stripMargin, + "CustomException.java" -> + """public class CustomException { + | public static void test(@MyNotNull(exception = MyException.class) Object o) {} + |}""".stripMargin + ), annotations = Seq("MyNotNull")) + val cause = interceptCause(loadClass(dir, "CustomException"), "test", null) + cause.getClass.getName shouldBe "MyException" + cause.getMessage shouldBe "Argument 0 for @MyNotNull parameter of CustomException.test must not be null" + } + + test("enum constructor parameters are checked despite the synthetic name/ordinal parameters") { + val dir = compileAndInstrument(Map( + "TestEnum.java" -> + """import org.jetbrains.annotations.NotNull; + |public enum TestEnum { + | OK("ok"), BAD(null); + | TestEnum(@NotNull String s) {} + |}""".stripMargin + )) + val error = intercept[ExceptionInInitializerError](Class.forName("TestEnum", true, classLoader(dir))) + error.getCause shouldBe an[IllegalArgumentException] + error.getCause.getMessage should fullyMatch regex "Argument \\d+ for @NotNull parameter of TestEnum\\. must not be null" + } + + test("non-static inner class constructor parameters are checked despite the synthetic outer-instance parameter") { + val dir = compileAndInstrument(Map( + "Outer.java" -> + """import org.jetbrains.annotations.NotNull; + |public class Outer { + | public class Inner { + | public Inner(@NotNull String s) {} + | } + | public static void create(String s) { new Outer().new Inner(s); } + |}""".stripMargin + )) + val clazz = loadClass(dir, "Outer") + invokeStatic(clazz, "create", "not null") + val cause = interceptCause(clazz, "create", null) + cause shouldBe an[IllegalArgumentException] + cause.getMessage should fullyMatch regex "Argument \\d+ for @NotNull parameter of Outer\\$Inner\\. must not be null" + } + + test("no check is generated when the returned value is provably non-null") { + val dir = compileFixture(Map( + "NewObject.java" -> + """import org.jetbrains.annotations.NotNull; + |public class NewObject { + | @NotNull + | public static Object test() { return new Object(); } + |}""".stripMargin + )) + val classFile = dir.resolve("NewObject.class") + val before = Files.readAllBytes(classFile) + instrumentAll(dir, Seq(NotNullFqn)) + val after = Files.readAllBytes(classFile) + assert(java.util.Arrays.equals(before, after), "class file must not be modified") + invokeStatic(loadClass(dir, "NewObject"), "test") + } + + test("Kotlin bytecode is skipped (kotlinc generates its own nullability assertions)") { + val dir = compileFixture(Map( + "kotlin/Metadata.java" -> + """package kotlin; + |import java.lang.annotation.*; + |@Retention(RetentionPolicy.RUNTIME) + |@Target(ElementType.TYPE) + |public @interface Metadata {}""".stripMargin, + "KotlinLike.java" -> + """@kotlin.Metadata + |public class KotlinLike { + | public static void test(@org.jetbrains.annotations.NotNull Object o) {} + |}""".stripMargin + )) + val classFile = dir.resolve("KotlinLike.class") + val before = Files.readAllBytes(classFile) + instrumentAll(dir, Seq(NotNullFqn)) + val after = Files.readAllBytes(classFile) + assert(java.util.Arrays.equals(before, after), "class file must not be modified") + invokeStatic(loadClass(dir, "KotlinLike"), "test", null) // no assertion is generated + } + + private def annotationsJar: Path = + Paths.get(classOf[org.jetbrains.annotations.NotNull].getProtectionDomain.getCodeSource.getLocation.toURI) + + private def compileFixture(sources: Map[String, String], debugInfo: Boolean = false): Path = { + val dir = Files.createTempDirectory("notnull-instrumentation-test") + val javaFiles = sources.map { case (relativePath, content) => + val file = dir.resolve(relativePath) + Files.createDirectories(file.getParent) + Files.write(file, content.getBytes(StandardCharsets.UTF_8)) + file.toString + }.toSeq + val args = Seq("-classpath", annotationsJar.toString, "-d", dir.toString) ++ + (if (debugInfo) Seq("-g") else Seq.empty) ++ + javaFiles + val exitCode = ToolProvider.getSystemJavaCompiler.run(null, null, null, args: _*) + assert(exitCode == 0, s"javac exited with code $exitCode") + dir + } + + private def instrumentAll(dir: Path, annotations: Seq[String]): Unit = { + val loader = new URLClassLoader(Array(dir.toUri.toURL, annotationsJar.toUri.toURL)) + try { + val stream = Files.walk(dir) + val classFiles = try stream.iterator().asScala.filter(_.toString.endsWith(".class")).toList finally stream.close() + classFiles.foreach(NotNullInstrumenter.instrument(_, annotations, loader)) + } finally { + loader.close() + } + } + + private def compileAndInstrument( + sources: Map[String, String], + annotations: Seq[String] = Seq(NotNullFqn), + debugInfo: Boolean = false + ): Path = { + val dir = compileFixture(sources, debugInfo) + instrumentAll(dir, annotations) + dir + } + + private def classLoader(dir: Path): ClassLoader = + new URLClassLoader(Array(dir.toUri.toURL), getClass.getClassLoader) + + private def loadClass(dir: Path, name: String): Class[_] = + classLoader(dir).loadClass(name) + + private def invokeStatic(clazz: Class[_], methodName: String, args: AnyRef*): AnyRef = { + val method = clazz.getMethods.find(_.getName == methodName) + .getOrElse(fail(s"Method $methodName not found in ${clazz.getName}")) + method.invoke(null, args: _*) + } + + private def interceptCause(clazz: Class[_], methodName: String, args: AnyRef*): Throwable = + intercept[InvocationTargetException](invokeStatic(clazz, methodName, args: _*)).getCause +}