diff options
Diffstat (limited to 'src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt')
| -rw-r--r-- | src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt | 187 |
1 files changed, 187 insertions, 0 deletions
diff --git a/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt b/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt new file mode 100644 index 00000000..975f3987 --- /dev/null +++ b/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt @@ -0,0 +1,187 @@ +// Copyright 2021 Code Intelligence GmbH +// +// 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 com.code_intelligence.jazzer.instrumentor + +import com.code_intelligence.jazzer.third_party.org.jacoco.core.analysis.Analyzer +import com.code_intelligence.jazzer.third_party.org.jacoco.core.analysis.ICoverageVisitor +import com.code_intelligence.jazzer.third_party.org.jacoco.core.data.ExecutionDataStore +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.ClassProbesAdapter +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.ClassProbesVisitor +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.IClassProbesAdapterFactory +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.ClassInstrumenter +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.IProbeArrayStrategy +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.IProbeInserterFactory +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.InstrSupport +import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.ProbeInserter +import org.objectweb.asm.ClassReader +import org.objectweb.asm.ClassVisitor +import org.objectweb.asm.ClassWriter +import org.objectweb.asm.MethodVisitor +import java.lang.invoke.MethodHandle +import java.lang.invoke.MethodHandles.publicLookup +import java.lang.invoke.MethodType.methodType +import kotlin.math.max + +/** + * A particular way to instrument bytecode for edge coverage using a coverage map class available to + * hold the collected coverage data at runtime. + */ +interface EdgeCoverageStrategy { + + /** + * Inject bytecode instrumentation on a control flow edge with ID [edgeId], with access to the + * local variable [variable] that is populated at the beginning of each method by the + * instrumentation injected in [loadLocalVariable]. + */ + fun instrumentControlFlowEdge( + mv: MethodVisitor, + edgeId: Int, + variable: Int, + coverageMapInternalClassName: String, + ) + + /** + * The maximal number of stack elements used by [instrumentControlFlowEdge]. + */ + val instrumentControlFlowEdgeStackSize: Int + + /** + * The type of the local variable used by the instrumentation in the format used by + * [MethodVisitor.visitFrame]'s `local` parameter, or `null` if the instrumentation does not use + * one. + * @see https://asm.ow2.io/javadoc/org/objectweb/asm/MethodVisitor.html#visitFrame(int,int,java.lang.Object%5B%5D,int,java.lang.Object%5B%5D) + */ + val localVariableType: Any? + + /** + * Inject bytecode that loads the coverage counters of the coverage map class described by + * [coverageMapInternalClassName] into the local variable [variable]. + */ + fun loadLocalVariable(mv: MethodVisitor, variable: Int, coverageMapInternalClassName: String) + + /** + * The maximal number of stack elements used by [loadLocalVariable]. + */ + val loadLocalVariableStackSize: Int +} + +// An instance of EdgeCoverageInstrumentor should only be used to instrument a single class as it +// internally tracks the edge IDs, which have to be globally unique. +class EdgeCoverageInstrumentor( + private val strategy: EdgeCoverageStrategy, + /** + * The class must have the following public static member + * - method enlargeIfNeeded(int nextEdgeId): Called before a new edge ID is emitted. + */ + coverageMapClass: Class<*>, + private val initialEdgeId: Int, +) : Instrumentor { + private var nextEdgeId = initialEdgeId + + private val coverageMapInternalClassName = coverageMapClass.name.replace('.', '/') + private val enlargeIfNeeded: MethodHandle = + publicLookup().findStatic( + coverageMapClass, + "enlargeIfNeeded", + methodType( + Void::class.javaPrimitiveType, + Int::class.javaPrimitiveType, + ), + ) + + override fun instrument(internalClassName: String, bytecode: ByteArray): ByteArray { + val reader = InstrSupport.classReaderFor(bytecode) + val writer = ClassWriter(reader, 0) + val version = InstrSupport.getMajorVersion(reader) + val visitor = EdgeCoverageClassProbesAdapter( + ClassInstrumenter(edgeCoverageProbeArrayStrategy, edgeCoverageProbeInserterFactory, writer), + InstrSupport.needsFrames(version), + ) + reader.accept(visitor, ClassReader.EXPAND_FRAMES) + return writer.toByteArray() + } + + fun analyze(executionData: ExecutionDataStore, coverageVisitor: ICoverageVisitor, bytecode: ByteArray, internalClassName: String) { + Analyzer(executionData, coverageVisitor, edgeCoverageClassProbesAdapterFactory).run { + analyzeClass(bytecode, internalClassName) + } + } + + val numEdges + get() = nextEdgeId - initialEdgeId + + private fun nextEdgeId(): Int { + enlargeIfNeeded.invokeExact(nextEdgeId) + return nextEdgeId++ + } + + /** + * A [ProbeInserter] that injects bytecode instrumentation at every control flow edge and + * modifies the stack size and number of local variables accordingly. + */ + private inner class EdgeCoverageProbeInserter( + access: Int, + name: String, + desc: String, + mv: MethodVisitor, + arrayStrategy: IProbeArrayStrategy, + ) : ProbeInserter(access, name, desc, mv, arrayStrategy) { + override fun insertProbe(id: Int) { + strategy.instrumentControlFlowEdge(mv, id, variable, coverageMapInternalClassName) + } + + override fun visitMaxs(maxStack: Int, maxLocals: Int) { + val newMaxStack = max(maxStack + strategy.instrumentControlFlowEdgeStackSize, strategy.loadLocalVariableStackSize) + val newMaxLocals = maxLocals + if (strategy.localVariableType != null) 1 else 0 + mv.visitMaxs(newMaxStack, newMaxLocals) + } + + override fun getLocalVariableType() = strategy.localVariableType + } + + private val edgeCoverageProbeInserterFactory = + IProbeInserterFactory { access, name, desc, mv, arrayStrategy -> + EdgeCoverageProbeInserter(access, name, desc, mv, arrayStrategy) + } + + private inner class EdgeCoverageClassProbesAdapter(private val cpv: ClassProbesVisitor, trackFrames: Boolean) : + ClassProbesAdapter(cpv, trackFrames) { + override fun nextId(): Int = nextEdgeId() + + override fun visitEnd() { + cpv.visitTotalProbeCount(numEdges) + // Avoid calling super.visitEnd() as that invokes cpv.visitTotalProbeCount with an + // incorrect value of `count`. + cpv.visitEnd() + } + } + + private val edgeCoverageClassProbesAdapterFactory = IClassProbesAdapterFactory { probesVisitor, trackFrames -> + EdgeCoverageClassProbesAdapter(probesVisitor, trackFrames) + } + + private val edgeCoverageProbeArrayStrategy = object : IProbeArrayStrategy { + override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int { + strategy.loadLocalVariable(mv, variable, coverageMapInternalClassName) + return strategy.loadLocalVariableStackSize + } + + override fun addMembers(cv: ClassVisitor, probeCount: Int) {} + } +} + +fun MethodVisitor.push(value: Int) { + InstrSupport.push(this, value) +} |