Add option to generate coverage report

The new --coverage_report option triggers a coverage report to be
written on fuzzer exit.

The report is generated with the JaCoCo analyzer. The information about
observed coverage IDs is obtained from libFuzzer and combined with the
coverage obtained during fuzzerInitialize as well as the current run.

3 files changed, 193 insertions, 0 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
index ad31db7c..91b796a3 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/agent/RuntimeInstrumentor.kt
@@ -15,6 +15,7 @@
 package com.code_intelligence.jazzer.agent
 
 import com.code_intelligence.jazzer.instrumentor.ClassInstrumentor
+import com.code_intelligence.jazzer.instrumentor.CoverageRecorder
 import com.code_intelligence.jazzer.instrumentor.Hook
 import com.code_intelligence.jazzer.instrumentor.InstrumentationType
 import com.code_intelligence.jazzer.instrumentor.loadHooks
@@ -159,6 +160,7 @@ internal class RuntimeInstrumentor(
                 } finally {
                     coverageIdSynchronizer.commitIdCount(actualNumEdgeIds)
                 }
+                CoverageRecorder.recordInstrumentedClass(internalClassName, bytecode, firstId, firstId + actualNumEdgeIds)
             } else {
                 hooks(customHooks)
             }
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/CoverageRecorder.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/CoverageRecorder.kt
new file mode 100644
index 00000000..a6d5e406
--- /dev/null
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/CoverageRecorder.kt
@@ -0,0 +1,177 @@
+// 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.runtime.CoverageMap
+import com.code_intelligence.jazzer.third_party.jacoco.core.analysis.CoverageBuilder
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.ExecutionData
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.ExecutionDataReader
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.ExecutionDataStore
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.ExecutionDataWriter
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.SessionInfo
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.SessionInfoStore
+import com.code_intelligence.jazzer.third_party.jacoco.core.internal.data.CRC64
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import java.time.Instant
+import java.util.UUID
+
+private data class InstrumentedClassInfo(
+    val classId: Long,
+    val initialEdgeId: Int,
+    val nextEdgeId: Int,
+    val bytecode: ByteArray
+)
+
+object CoverageRecorder {
+    private val instrumentedClassInfo = mutableMapOf<String, InstrumentedClassInfo>()
+    private var startTimestamp: Instant? = null
+    private val additionalCoverage = mutableSetOf<Int>()
+
+    fun recordInstrumentedClass(internalClassName: String, bytecode: ByteArray, firstId: Int, numIds: Int) {
+        if (startTimestamp == null)
+            startTimestamp = Instant.now()
+        instrumentedClassInfo[internalClassName] = InstrumentedClassInfo(
+            CRC64.classId(bytecode), firstId, firstId + numIds, bytecode
+        )
+    }
+
+    /**
+     * Manually records coverage IDs based on the current state of [CoverageMap.mem].
+     * Should be called after static initializers have run.
+     */
+    @JvmStatic
+    fun updateCoveredIdsWithCoverageMap() {
+        val mem = CoverageMap.mem
+        val size = mem.capacity()
+        additionalCoverage.addAll((0 until size).filter { mem[it] > 0 })
+    }
+
+    @JvmStatic
+    fun computeFileCoverage(coveredIds: IntArray): String {
+        val coverage = analyzeCoverage(coveredIds.toSet()) ?: return "No classes were instrumented"
+        return coverage.sourceFiles.joinToString(
+            "\n",
+            prefix = "Branch coverage:\n",
+            postfix = "\n\n"
+        ) { fileCoverage ->
+            val counter = fileCoverage.branchCounter
+            val percentage = 100 * counter.coveredRatio
+            "${fileCoverage.name}: ${counter.coveredCount}/${counter.totalCount} (${percentage.format(2)}%)"
+        } + coverage.sourceFiles.joinToString(
+            "\n",
+            prefix = "Line coverage:\n",
+            postfix = "\n\n"
+        ) { fileCoverage ->
+            val counter = fileCoverage.lineCounter
+            val percentage = 100 * counter.coveredRatio
+            "${fileCoverage.name}: ${counter.coveredCount}/${counter.totalCount} (${percentage.format(2)}%)"
+        } + coverage.sourceFiles.joinToString(
+            "\n",
+            prefix = "Incompletely covered lines:\n",
+            postfix = "\n\n"
+        ) { fileCoverage ->
+            "${fileCoverage.name}: " + (fileCoverage.firstLine..fileCoverage.lastLine).filter {
+                val instructions = fileCoverage.getLine(it).instructionCounter
+                instructions.coveredCount in 1 until instructions.totalCount
+            }.toString()
+        } + coverage.sourceFiles.joinToString(
+            "\n",
+            prefix = "Missed lines:\n",
+        ) { fileCoverage ->
+            "${fileCoverage.name}: " + (fileCoverage.firstLine..fileCoverage.lastLine).filter {
+                val instructions = fileCoverage.getLine(it).instructionCounter
+                instructions.coveredCount == 0 && instructions.totalCount > 0
+            }.toString()
+        }
+    }
+
+    private fun Double.format(digits: Int) = "%.${digits}f".format(this)
+
+    private fun dumpJacocoCoverage(coveredIds: Set<Int>): ByteArray? {
+        // Update the list of covered IDs with the coverage information for the current run.
+        updateCoveredIdsWithCoverageMap()
+
+        val dumpTimestamp = Instant.now()
+        val outStream = ByteArrayOutputStream()
+        val outWriter = ExecutionDataWriter(outStream)
+        // Return null if no class has been instrumented.
+        val startTimestamp = startTimestamp ?: return null
+        outWriter.visitSessionInfo(
+            SessionInfo(UUID.randomUUID().toString(), startTimestamp.epochSecond, dumpTimestamp.epochSecond)
+        )
+
+        val sortedCoveredIds = (additionalCoverage + coveredIds).sorted().toIntArray()
+        for ((internalClassName, info) in instrumentedClassInfo) {
+            // Determine the subarray of coverage IDs in sortedCoveredIds that contains the IDs generated while
+            // instrumenting the current class. Since the ID array is sorted, use binary search.
+            var coveredIdsStart = sortedCoveredIds.binarySearch(info.initialEdgeId)
+            if (coveredIdsStart < 0) {
+                coveredIdsStart = -(coveredIdsStart + 1)
+            }
+            var coveredIdsEnd = sortedCoveredIds.binarySearch(info.nextEdgeId)
+            if (coveredIdsEnd < 0) {
+                coveredIdsEnd = -(coveredIdsEnd + 1)
+            }
+            if (coveredIdsStart == coveredIdsEnd) {
+                // No coverage data for the class.
+                continue
+            }
+            check(coveredIdsStart in 0 until coveredIdsEnd && coveredIdsEnd <= sortedCoveredIds.size) {
+                "Invalid range [$coveredIdsStart, $coveredIdsEnd) with coveredIds.size=${sortedCoveredIds.size}"
+            }
+            // Generate a probes array for the current class only, i.e., mapping info.initialEdgeId to 0.
+            val probes = BooleanArray(info.nextEdgeId - info.initialEdgeId)
+            (coveredIdsStart until coveredIdsEnd).asSequence()
+                .map {
+                    val globalEdgeId = sortedCoveredIds[it]
+                    globalEdgeId - info.initialEdgeId
+                }
+                .forEach { classLocalEdgeId ->
+                    probes[classLocalEdgeId] = true
+                }
+            outWriter.visitClassExecution(ExecutionData(info.classId, internalClassName, probes))
+        }
+        return outStream.toByteArray()
+    }
+
+    private fun analyzeCoverage(coveredIds: Set<Int>): CoverageBuilder? {
+        return try {
+            val coverage = CoverageBuilder()
+            val rawExecutionData = dumpJacocoCoverage(coveredIds) ?: return null
+            val executionDataStore = ExecutionDataStore()
+            val sessionInfoStore = SessionInfoStore()
+            ByteArrayInputStream(rawExecutionData).use { stream ->
+                ExecutionDataReader(stream).run {
+                    setExecutionDataVisitor(executionDataStore)
+                    setSessionInfoVisitor(sessionInfoStore)
+                    read()
+                }
+            }
+            for ((internalClassName, info) in instrumentedClassInfo) {
+                EdgeCoverageInstrumentor(0).analyze(
+                    executionDataStore,
+                    coverage,
+                    info.bytecode,
+                    internalClassName
+                )
+            }
+            coverage
+        } catch (e: Exception) {
+            e.printStackTrace()
+            null
+        }
+    }
+}
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
index e22cd2c0..2e2a0b11 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
@@ -16,8 +16,12 @@ package com.code_intelligence.jazzer.instrumentor
 
 import com.code_intelligence.jazzer.generated.JAVA_NO_THROW_METHODS
 import com.code_intelligence.jazzer.runtime.CoverageMap
+import com.code_intelligence.jazzer.third_party.jacoco.core.analysis.Analyzer
+import com.code_intelligence.jazzer.third_party.jacoco.core.analysis.ICoverageVisitor
+import com.code_intelligence.jazzer.third_party.jacoco.core.data.ExecutionDataStore
 import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.ClassProbesAdapter
 import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.ClassProbesVisitor
+import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.IClassProbesAdapterFactory
 import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.IMethodProbesAdapterFactory
 import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.IProbeIdGenerator
 import com.code_intelligence.jazzer.third_party.jacoco.core.internal.flow.MethodProbesAdapter
@@ -53,6 +57,12 @@ class EdgeCoverageInstrumentor(
         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
 
@@ -232,6 +242,10 @@ class EdgeCoverageInstrumentor(
         override fun nextId(): Int = nextEdgeId()
     }
 
+    private val edgeCoverageClassProbesAdapterFactory = IClassProbesAdapterFactory { probesVisitor, trackFrames ->
+        EdgeCoverageClassProbesAdapter(probesVisitor, trackFrames)
+    }
+
     private val edgeCoverageProbeArrayStrategy = object : IProbeArrayStrategy {
         override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
             loadCoverageMap(mv, variable)