<lambda>null1 // Copyright 2021 Code Intelligence GmbH
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 package com.code_intelligence.jazzer.instrumentor
16 
17 import com.code_intelligence.jazzer.third_party.org.jacoco.core.analysis.Analyzer
18 import com.code_intelligence.jazzer.third_party.org.jacoco.core.analysis.ICoverageVisitor
19 import com.code_intelligence.jazzer.third_party.org.jacoco.core.data.ExecutionDataStore
20 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.ClassProbesAdapter
21 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.ClassProbesVisitor
22 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.flow.IClassProbesAdapterFactory
23 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.ClassInstrumenter
24 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.IProbeArrayStrategy
25 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.IProbeInserterFactory
26 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.InstrSupport
27 import com.code_intelligence.jazzer.third_party.org.jacoco.core.internal.instr.ProbeInserter
28 import org.objectweb.asm.ClassReader
29 import org.objectweb.asm.ClassVisitor
30 import org.objectweb.asm.ClassWriter
31 import org.objectweb.asm.MethodVisitor
32 import java.lang.invoke.MethodHandle
33 import java.lang.invoke.MethodHandles.publicLookup
34 import java.lang.invoke.MethodType.methodType
35 import kotlin.math.max
36 
37 /**
38  * A particular way to instrument bytecode for edge coverage using a coverage map class available to
39  * hold the collected coverage data at runtime.
40  */
41 interface EdgeCoverageStrategy {
42 
43     /**
44      * Inject bytecode instrumentation on a control flow edge with ID [edgeId], with access to the
45      * local variable [variable] that is populated at the beginning of each method by the
46      * instrumentation injected in [loadLocalVariable].
47      */
48     fun instrumentControlFlowEdge(
49         mv: MethodVisitor,
50         edgeId: Int,
51         variable: Int,
52         coverageMapInternalClassName: String,
53     )
54 
55     /**
56      * The maximal number of stack elements used by [instrumentControlFlowEdge].
57      */
58     val instrumentControlFlowEdgeStackSize: Int
59 
60     /**
61      * The type of the local variable used by the instrumentation in the format used by
62      * [MethodVisitor.visitFrame]'s `local` parameter, or `null` if the instrumentation does not use
63      * one.
64      * @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)
65      */
66     val localVariableType: Any?
67 
68     /**
69      * Inject bytecode that loads the coverage counters of the coverage map class described by
70      * [coverageMapInternalClassName] into the local variable [variable].
71      */
72     fun loadLocalVariable(mv: MethodVisitor, variable: Int, coverageMapInternalClassName: String)
73 
74     /**
75      * The maximal number of stack elements used by [loadLocalVariable].
76      */
77     val loadLocalVariableStackSize: Int
78 }
79 
80 // An instance of EdgeCoverageInstrumentor should only be used to instrument a single class as it
81 // internally tracks the edge IDs, which have to be globally unique.
82 class EdgeCoverageInstrumentor(
83     private val strategy: EdgeCoverageStrategy,
84     /**
85      * The class must have the following public static member
86      *  - method enlargeIfNeeded(int nextEdgeId): Called before a new edge ID is emitted.
87      */
88     coverageMapClass: Class<*>,
89     private val initialEdgeId: Int,
90 ) : Instrumentor {
91     private var nextEdgeId = initialEdgeId
92 
93     private val coverageMapInternalClassName = coverageMapClass.name.replace('.', '/')
94     private val enlargeIfNeeded: MethodHandle =
95         publicLookup().findStatic(
96             coverageMapClass,
97             "enlargeIfNeeded",
98             methodType(
99                 Void::class.javaPrimitiveType,
100                 Int::class.javaPrimitiveType,
101             ),
102         )
103 
instrumentnull104     override fun instrument(internalClassName: String, bytecode: ByteArray): ByteArray {
105         val reader = InstrSupport.classReaderFor(bytecode)
106         val writer = ClassWriter(reader, 0)
107         val version = InstrSupport.getMajorVersion(reader)
108         val visitor = EdgeCoverageClassProbesAdapter(
109             ClassInstrumenter(edgeCoverageProbeArrayStrategy, edgeCoverageProbeInserterFactory, writer),
110             InstrSupport.needsFrames(version),
111         )
112         reader.accept(visitor, ClassReader.EXPAND_FRAMES)
113         return writer.toByteArray()
114     }
115 
analyzenull116     fun analyze(executionData: ExecutionDataStore, coverageVisitor: ICoverageVisitor, bytecode: ByteArray, internalClassName: String) {
117         Analyzer(executionData, coverageVisitor, edgeCoverageClassProbesAdapterFactory).run {
118             analyzeClass(bytecode, internalClassName)
119         }
120     }
121 
122     val numEdges
123         get() = nextEdgeId - initialEdgeId
124 
nextEdgeIdnull125     private fun nextEdgeId(): Int {
126         enlargeIfNeeded.invokeExact(nextEdgeId)
127         return nextEdgeId++
128     }
129 
130     /**
131      * A [ProbeInserter] that injects bytecode instrumentation at every control flow edge and
132      * modifies the stack size and number of local variables accordingly.
133      */
134     private inner class EdgeCoverageProbeInserter(
135         access: Int,
136         name: String,
137         desc: String,
138         mv: MethodVisitor,
139         arrayStrategy: IProbeArrayStrategy,
140     ) : ProbeInserter(access, name, desc, mv, arrayStrategy) {
insertProbenull141         override fun insertProbe(id: Int) {
142             strategy.instrumentControlFlowEdge(mv, id, variable, coverageMapInternalClassName)
143         }
144 
visitMaxsnull145         override fun visitMaxs(maxStack: Int, maxLocals: Int) {
146             val newMaxStack = max(maxStack + strategy.instrumentControlFlowEdgeStackSize, strategy.loadLocalVariableStackSize)
147             val newMaxLocals = maxLocals + if (strategy.localVariableType != null) 1 else 0
148             mv.visitMaxs(newMaxStack, newMaxLocals)
149         }
150 
getLocalVariableTypenull151         override fun getLocalVariableType() = strategy.localVariableType
152     }
153 
154     private val edgeCoverageProbeInserterFactory =
155         IProbeInserterFactory { access, name, desc, mv, arrayStrategy ->
156             EdgeCoverageProbeInserter(access, name, desc, mv, arrayStrategy)
157         }
158 
159     private inner class EdgeCoverageClassProbesAdapter(private val cpv: ClassProbesVisitor, trackFrames: Boolean) :
160         ClassProbesAdapter(cpv, trackFrames) {
nextIdnull161         override fun nextId(): Int = nextEdgeId()
162 
163         override fun visitEnd() {
164             cpv.visitTotalProbeCount(numEdges)
165             // Avoid calling super.visitEnd() as that invokes cpv.visitTotalProbeCount with an
166             // incorrect value of `count`.
167             cpv.visitEnd()
168         }
169     }
170 
trackFramesnull171     private val edgeCoverageClassProbesAdapterFactory = IClassProbesAdapterFactory { probesVisitor, trackFrames ->
172         EdgeCoverageClassProbesAdapter(probesVisitor, trackFrames)
173     }
174 
175     private val edgeCoverageProbeArrayStrategy = object : IProbeArrayStrategy {
storeInstancenull176         override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
177             strategy.loadLocalVariable(mv, variable, coverageMapInternalClassName)
178             return strategy.loadLocalVariableStackSize
179         }
180 
addMembersnull181         override fun addMembers(cv: ClassVisitor, probeCount: Int) {}
182     }
183 }
184 
pushnull185 fun MethodVisitor.push(value: Int) {
186     InstrSupport.push(this, value)
187 }
188