Counts source locations for JSP files

agent/src/main/java/com/codedx/codepulse/agent/message/MessageDealer.java

@@ -19,6 +19,7 @@
 
 import java.io.IOException;
 import java.util.BitSet;
+import java.util.HashSet;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -36,9 +37,8 @@
  * {@link BufferService}, then sent via the same BufferService. For events that
  * require mapped ids for the current thread and method signature, those ids
  * (along with the appropriate secondary "map" events) will be automatically
- * generated by an internal {@link MethodId} and {@link ThreadId}. The time of
- * MessageFactory's construction will be saved, and used to calculate the
- * "relative timestamp" for each event that requires one.
+ * generated. The time of MessageFactory's construction will be saved, and used
+ * to calculate the "relative timestamp" for each event that requires one.
  *
  * @author dylanh
  */
@@ -191,6 +191,8 @@ private class MethodIdAdapter
 
 		private final ConcurrentMap<Integer, Boolean> observedIds = new ConcurrentHashMap<Integer, Boolean>();
 
+		private final ConcurrentMap<Integer, Integer> sourceLocationCounts = new ConcurrentHashMap<>();
+
 		private final AtomicInteger nextSourceLocationId = new AtomicInteger();
 		protected final ConcurrentHashMap<String, Integer> sourceLocationMap = new ConcurrentHashMap<>();
 
@@ -214,8 +216,36 @@ public int markSourceLocation(int methodId, int startLine, int endLine, DataBuff
 
 			int classId = methodIdentifier.get(methodId).getClassId();
 
-			LineLevelMapper llm = classIdentifier.get(classId).getLineLevelMapper();
+			ClassIdentifier.ClassInformation classInformation = classIdentifier.get(classId);
+			LineLevelMapper llm = classInformation.getLineLevelMapper();
 			if (llm != null) {
+
+				Boolean sourceLocationCountsMessageSent = sourceLocationCounts.containsKey(classId);
+				if (!sourceLocationCountsMessageSent) {
+					HashSet<String> mappedLocations = new HashSet<>();
+
+					BitSet lineNumbers = classInformation.getLineNumbers();
+					for (int l = lineNumbers.nextSetBit(0); l >= 0; l = lineNumbers.nextSetBit(l + 1)) {
+						BitSet b = new BitSet();
+						b.set(0);
+						LineLevelMapper.MappedCoverage mappedCoverage[] = llm.map(l, b);
+						if (mappedCoverage != null) {
+							StringBuilder s = new StringBuilder();
+							for (LineLevelMapper.MappedCoverage mappedCoverageItem : mappedCoverage) {
+								for (int q = mappedCoverageItem.lines.nextSetBit(0); q >= 0; q = mappedCoverageItem.lines.nextSetBit(q + 1)) {
+									s.append(mappedCoverageItem.startLine + q);
+									s.append("; ");
+								}
+							}
+							mappedLocations.add(s.toString());
+						}
+					}
+
+					int mappingsCount = mappedLocations.size();
+					messageProtocol.writeSourceLocationCount(buffer, methodId, mappingsCount);
+					sourceLocationCounts.put(classId, mappingsCount);
+				}
+
 				BitSet bitSet = new BitSet();
 				bitSet.set(0, endLine - startLine + 1);
 

agent/src/main/java/com/codedx/codepulse/agent/trace/InstrumentationClassVisitor.java

@@ -0,0 +1,31 @@
+package com.codedx.codepulse.agent.trace;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.util.BitSet;
+import java.util.LinkedList;
+
+public class InstrumentationClassVisitor extends ClassVisitor {
+
+    private LinkedList<InstrumentationMethodVisitor> inspectors = new LinkedList<>();
+
+    public InstrumentationClassVisitor() {
+        super(Opcodes.ASM5);
+    }
+
+    public BitSet getLineNumbers() {
+        BitSet lineNumbers = new BitSet();
+        for (InstrumentationMethodVisitor visitor : inspectors) {
+            lineNumbers.or(visitor.getLineNumbers());
+        }
+        return lineNumbers;
+    }
+
+    @Override public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
+        InstrumentationMethodVisitor visitor = new InstrumentationMethodVisitor();
+        inspectors.add(visitor);
+        return visitor;
+    }
+}

agent/src/main/java/com/codedx/codepulse/agent/trace/InstrumentationMethodVisitor.java

@@ -0,0 +1,31 @@
+package com.codedx.codepulse.agent.trace;
+
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.util.BitSet;
+
+public class InstrumentationMethodVisitor extends MethodVisitor {
+
+    private BitSet lineNumbers = new BitSet();
+
+    public InstrumentationMethodVisitor()
+    {
+        super(Opcodes.ASM5);
+    }
+
+    public BitSet getLineNumbers() { return lineNumbers; }
+
+    @Override
+    public void visitCode() {
+        lineNumbers.clear();
+    }
+
+    @Override
+    public void visitLineNumber(int line, Label start)
+    {
+        int lineNumberVisited = line; // line numbers are 1-based
+        lineNumbers.set(lineNumberVisited);
+    }
+}

agent/src/main/java/com/codedx/codepulse/agent/trace/Instrumentor.java

@@ -20,6 +20,7 @@
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.util.BitSet;
 
 import com.codedx.bytefrog.filterinjector.FilterInjector;
 import com.codedx.bytefrog.filterinjector.adapters.Adapter;
@@ -99,7 +100,16 @@ public byte[] instrument(final ClassLoader classLoader, final String className,
 			cr.accept(inspector, 0);
 
 			final ClassInspector.Result inspection = inspector.getResult();
-			final int classId = classIdentifier.record(className, inspection.getFileName(), inspection.getLineLevelMapper());
+
+			BitSet lineNumbers = null;
+			LineLevelMapper lineLevelMapper = inspection.getLineLevelMapper();
+			if (lineLevelMapper != null) {
+				InstrumentationClassVisitor visitor = new InstrumentationClassVisitor();
+				cr.accept(visitor, 0);
+				lineNumbers = visitor.getLineNumbers();
+			}
+
+			final int classId = classIdentifier.record(className, inspection.getFileName(), lineLevelMapper, lineNumbers);
 
 			final ClassInstrumentor ci = new ClassInstrumentor(filterInjectorVisitor != null ? filterInjectorVisitor : cw, methodIdentifier, classId, inspection, handler);
 			cr.accept(ci, ClassReader.EXPAND_FRAMES);

agent/src/test/scala/com/secdec/bytefrog/agent/message/test/MethodIdSpec.scala

@@ -19,6 +19,7 @@ package com.secdec.bytefrog.agent.message.test
 
 import java.io.ByteArrayOutputStream
 import java.io.DataOutputStream
+import java.util
 
 import scala.collection.mutable.ListBuffer
 
@@ -43,7 +44,7 @@ class MethodIdSpec extends FunSpec with Matchers with MockFactory {
 	}
 
 	val classIdentifier = new ClassIdentifier
-	val cId = classIdentifier.record("NA", "NA.source", LineLevelMapper.empty("NA.source"))
+	val cId = classIdentifier.record("NA", "NA.source", LineLevelMapper.empty("NA.source"), new util.BitSet())
 
 	val methodIdentifier = new MethodIdentifier
 	val idA = methodIdentifier.record(cId, 1, "A", "A", 1, 0)

bytefrog/instrumentation/src/main/java/com/codedx/bytefrog/instrumentation/id/ClassIdentifier.java

@@ -18,6 +18,7 @@
 
 package com.codedx.bytefrog.instrumentation.id;
 
+import java.util.BitSet;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -31,9 +32,9 @@ public class ClassIdentifier {
 	private final AtomicInteger nextId = new AtomicInteger();
 	protected final ConcurrentHashMap<Integer, ClassInformation> map = new ConcurrentHashMap<>();
 
-	public int record(String className, String sourceFile, LineLevelMapper lineLevelMapper) {
+	public int record(String className, String sourceFile, LineLevelMapper lineLevelMapper, BitSet lineNumbers) {
 		int id = nextId.getAndIncrement();
-		map.put(id, new ClassInformation(className, sourceFile, lineLevelMapper));
+		map.put(id, new ClassInformation(className, sourceFile, lineLevelMapper, lineNumbers));
 		return id;
 	}
 
@@ -46,11 +47,13 @@ public static class ClassInformation {
 		private final String name;
 		private final String sourceFile;
 		private final LineLevelMapper lineLevelMapper;
+		private final BitSet lineNumbers;
 
-		public ClassInformation(String name, String sourceFile, LineLevelMapper lineLevelMapper) {
+		public ClassInformation(String name, String sourceFile, LineLevelMapper lineLevelMapper, BitSet lineNumbers) {
 			this.name = name;
 			this.sourceFile = sourceFile;
 			this.lineLevelMapper = lineLevelMapper;
+			this.lineNumbers = lineNumbers;
 		}
 
 		/** Gets the name of the class.
@@ -73,5 +76,12 @@ public String getSourceFile() {
 		public LineLevelMapper getLineLevelMapper() {
 			return lineLevelMapper;
 		}
+
+		/** Gets the line numbers for the class.
+		  * @returns the line numbers, or null if line numbers are unavailable for class
+		  */
+		public BitSet getLineNumbers() { 
+			return lineNumbers; 
+		}
 	}
 }

codepulse/src/main/scala/com/secdec/codepulse/data/model/TreeNodeData.scala

@@ -92,6 +92,8 @@ trait TreeNodeDataAccess {
 	def clearFlag(id: Int, flag: TreeNodeFlag): Unit
 	def clearFlag(node: TreeNodeData, flag: TreeNodeFlag): Unit = clearFlag(node.id, flag)
 
+	def updateSourceLocationCount(id: Int, sourceLocationCount: Int)
+
 	implicit class ExtendedTreeNodeData(n: TreeNodeData) {
 		/** whether or not this treenode is being traced; this value may be unspecified (None) */
 		def traced = isTraced(n)

codepulse/src/main/scala/com/secdec/codepulse/data/model/slick/SlickTreeNodeDataAccess.scala

@@ -125,6 +125,12 @@ private[slick] class SlickTreeNodeDataAccess(dao: TreeNodeDataDao, db: Database)
 		}
 	}
 
+	def updateSourceLocationCount(id: Int, sourceLocationCount: Int) = {
+		db withTransaction { implicit transaction =>
+			dao.updateSourceLocationCount(id, sourceLocationCount)
+		}
+	}
+
 	private lazy val flagCache = collection.mutable.Map.empty[Int, List[TreeNodeFlag]]
 
 	def getFlags(id: Int): List[TreeNodeFlag] = flagCache.getOrElse(id, { db withSession { implicit session =>

codepulse/src/main/scala/com/secdec/codepulse/data/model/slick/TreeNodeDataDao.scala

@@ -196,6 +196,11 @@ private[slick] class TreeNodeDataDao(val driver: JdbcProfile, val sourceDataDao:
 		}
 	}
 
+	def updateSourceLocationCount(id: Int, sourceLocationCount: Int) (implicit session: Session): Unit = {
+		val q = for (row <- treeNodeData if row.id === id) yield row.sourceLocationCount
+		q.update(Some(sourceLocationCount))
+	}
+
 	def getFlags(id: Int)(implicit session: Session): List[TreeNodeFlag] = {
 		// wrapped in a try, since older versions may not have the concept of flags
 		Try {

codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceRecorderDataProcessor.scala

@@ -25,12 +25,6 @@ import com.codedx.codepulse.utility.Loggable
 import com.secdec.codepulse.data.model.ProjectData
 import com.secdec.codepulse.data.jsp.JspMapper
 
-case class DeferredMapSourceLocation(sourceLocationId: Int, message: DataMessageContent.MapSourceLocation)
-
-case class DeferredMethodVisit(methodId: Int, message: DataMessageContent.MethodVisit)
-
-case class DeferredMethodEntry(methodId: Int, message: DataMessageContent.MethodEntry)
-
 class TraceRecorderDataProcessor(projectData: ProjectData, transientData: TransientTraceData, jspMapper: Option[JspMapper]) extends DataProcessor with Loggable {
 
 	// list of client-side node IDs that are ignored for tracing because they refer to method signatures missing from the database
@@ -43,6 +37,7 @@ class TraceRecorderDataProcessor(projectData: ProjectData, transientData: Transi
 	val sourceLocationCor = new collection.mutable.HashMap[Int, collection.mutable.HashMap[Int, Option[Int]]]
 
 	val deferredMethodEntries = collection.mutable.Map.empty[Int, collection.mutable.ListBuffer[DataMessageContent.MethodEntry]]
+	val deferredSourceLocationCounts = collection.mutable.Map.empty[Int, collection.mutable.ListBuffer[DataMessageContent.SourceLocationCount]]
 	val deferredMapSourceLocations = collection.mutable.Map.empty[Int, collection.mutable.ListBuffer[DataMessageContent.MapSourceLocation]]
 	val deferredMethodVisits = collection.mutable.Map.empty[Int, collection.mutable.ListBuffer[DataMessageContent.MethodVisit]]
 
@@ -93,11 +88,34 @@ class TraceRecorderDataProcessor(projectData: ProjectData, transientData: Transi
 					logger.info(s"Processing deferred method entry for method ${x.methodId}...")
 					processMessage(x)
 				})
+				deferredSourceLocationCounts.remove(id).getOrElse(collection.mutable.ListBuffer.empty[DataMessageContent.SourceLocationCount]).foreach(x => {
+					logger.info(s"Processing deferred source location count for method $x")
+				})
 				deferredMapSourceLocations.remove(id).getOrElse(collection.mutable.ListBuffer.empty[DataMessageContent.MapSourceLocation]).foreach(x => {
 					logger.info(s"Processing deferred map source location for source location ${x.sourceLocationId} in method ${x.methodId}...")
 					processMessage(x)
 				})
 
+			case sourceLocationCountMessage @ DataMessageContent.SourceLocationCount(methodId, sourceLocationCount) =>
+				val nodeIds = methodCor.get(methodId)
+				if (nodeIds.isEmpty) {
+					if (unknownAndIgnoredMethodCor.contains(methodId)) return
+
+					logger.info(s"Deferring source location count of $sourceLocationCount for method $methodId...")
+
+					var sourceLocationCountMessages = deferredSourceLocationCounts.get(methodId)
+					if (sourceLocationCountMessages.isEmpty) {
+						sourceLocationCountMessages = Option(collection.mutable.ListBuffer.empty[DataMessageContent.SourceLocationCount])
+						deferredSourceLocationCounts.put(methodId, sourceLocationCountMessages.get)
+					}
+					sourceLocationCountMessages.get.append(sourceLocationCountMessage)
+					return
+				}
+
+				nodeIds.get.foreach(nodeId => {
+					projectData.treeNodeData.updateSourceLocationCount(nodeId, sourceLocationCount)
+				})
+
 			case mapSourceLocationMessage @ DataMessageContent.MapSourceLocation(methodId, startLine, endLine, startCharacter, endCharacter, id) =>
 				val nodeIds = methodCor.get(methodId)
 				if (nodeIds.isEmpty) {