Merging with the "traced" changeset

Conflicts:
	codepulse/src/main/scala/com/secdec/codepulse/data/trace/TreeBuilder.scala

Author: dylemma

codepulse/src/main/scala/com/secdec/codepulse/data/bytecode/CodeForestBuilder.scala

@@ -31,7 +31,7 @@ object CodeForestBuilder {
 	val JSPGroupName = "JSPs"
 }
 
-class CodeForestBuilder(defaultTracedGroups: List[String]) {
+class CodeForestBuilder {
 	import CodeForestBuilder._
 
 	private val roots = SortedSet.empty[CodeTreeNode]
@@ -45,18 +45,13 @@ class CodeForestBuilder(defaultTracedGroups: List[String]) {
 
 	def result = {
 		val lb = List.newBuilder[TreeNodeData]
-		val tracedGroups = defaultTracedGroups.toSet
 		for (root <- roots) root.visitTree { node =>
 			val id = node.id
 			val name = node.name
 			val parentId = node.parentId
 			val kind = node.kind
 			val size = node.size
-			val traced = kind match {
-				case CodeTreeNodeKind.Grp | CodeTreeNodeKind.Pkg => Some(tracedGroups contains root.name)
-				case _ => None
-			}
-			lb += TreeNodeData(id, parentId, name, kind, size, traced)
+			lb += TreeNodeData(id, parentId, name, kind, size)
 		}
 		lb.result()
 	}

codepulse/src/main/scala/com/secdec/codepulse/data/trace/TreeBuilder.scala

@@ -23,16 +23,23 @@ import com.fasterxml.jackson.core.{ JsonFactory, JsonGenerator }
 
 import com.secdec.codepulse.data.bytecode.CodeTreeNodeKind
 
-case class TreeNode(data: TreeNodeData, children: List[TreeNode])
+case class TreeNode(data: TreeNodeData, children: List[TreeNode])(treeNodeData: TreeNodeDataAccess) {
+	import treeNodeData.ExtendedTreeNodeData
+
+	def traced = data.traced
+}
 
 case class PackageTreeNode(
 	id: Option[Int],
 	kind: CodeTreeNodeKind,
 	label: String,
 	methodCount: Int,
 	otherDescendantIds: List[Int],
-	traced: Option[Boolean],
-	children: List[PackageTreeNode])
+	children: List[PackageTreeNode])(
+	tracedLookup: () => Option[Boolean]) {
+	
+	def traced = tracedLookup()
+}
 
 /** Builds/projects treemap and package tree data as JSON for client.
   * TODO: manage lifetime of cached data internally
@@ -60,7 +67,8 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 
 		def buildNode(id: Int): TreeNode = {
 			val children = childrenFor(id).result.map(buildNode)
-			TreeNode(nodeMap(id), children)
+			val node = nodeMap(id)
+			TreeNode(node, children)(treeNodeData)
 		}
 
 		(roots.result.map(buildNode), nodeMap)
@@ -69,6 +77,7 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 	/** Full package tree, with self nodes */
 	lazy val packageTree = {
 		// build up a package tree with the relevant data
+		import treeNodeData.ExtendedTreeNodeData
 
 		/** A node is eligible for a self node if it is a package node that has at least one
 		  * package child and one non-package child (class/method).
@@ -120,10 +129,10 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 				}
 
 				// build the self node
-				val selfNode = PackageTreeNode(Some(node.data.id), CodeTreeNodeKind.Pkg, "<self>", selfChildren.map(countMethods).sum, otherDescendants, node.data.traced, Nil)
-				PackageTreeNode(None, node.data.kind, node.data.label, countMethods(node), Nil, node.data.traced, selfNode :: filterChildren(children).map(transform))
+				val selfNode = PackageTreeNode(Some(node.data.id), CodeTreeNodeKind.Pkg, "<self>", selfChildren.map(countMethods).sum, otherDescendants, Nil)(() => node.data.traced)
+				PackageTreeNode(None, node.data.kind, node.data.label, countMethods(node), Nil, selfNode :: filterChildren(children).map(transform))(() => node.data.traced)
 			} else {
-				PackageTreeNode(Some(node.data.id), node.data.kind, node.data.label, countMethods(node), otherDescendants, node.data.traced, filterChildren(node.children).map(transform))
+				PackageTreeNode(Some(node.data.id), node.data.kind, node.data.label, countMethods(node), otherDescendants, filterChildren(node.children).map(transform))(() => node.data.traced)
 			}
 		}
 
@@ -165,7 +174,7 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 				}
 
 				if (isSelected || !filteredChildren.isEmpty)
-					Some(TreeNode(node.data, filteredChildren))
+					Some(TreeNode(node.data, filteredChildren)(treeNodeData))
 				else
 					None
 			} else None

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

@@ -29,9 +29,8 @@ import com.secdec.codepulse.data.bytecode.CodeTreeNodeKind
   * @param kind The `CodeTreeNodeKind` of this node
   * @param size A number indicating the size of the node (e.g. lines of code). If unspecified,
   * the size of a node is assumed to be the sum of its childrens' sizes.
-  * @param traced whether or not this treenode is being traced; this value may be unspecified (None)
   */
-case class TreeNodeData(id: Int, parentId: Option[Int], label: String, kind: CodeTreeNodeKind, size: Option[Int], traced: Option[Boolean])
+case class TreeNodeData(id: Int, parentId: Option[Int], label: String, kind: CodeTreeNodeKind, size: Option[Int])
 
 /** Access trait for tree node data.
   *
@@ -64,9 +63,19 @@ trait TreeNodeDataAccess {
 	def storeNode(node: TreeNodeData): Unit
 	def storeNodes(nodes: Iterable[TreeNodeData]) = nodes foreach storeNode
 
+	def isTraced(id: Int): Option[Boolean]
+	def isTraced(node: TreeNodeData): Option[Boolean] = isTraced(node.id)
+
 	def updateTraced(id: Int, traced: Option[Boolean]): Unit
 	def updateTraced(id: Int, traced: Boolean): Unit = updateTraced(id, Some(traced))
 	def updateTraced(node: TreeNodeData, traced: Option[Boolean]): Unit = updateTraced(node.id, traced)
 	def updateTraced(node: TreeNodeData, traced: Boolean): Unit = updateTraced(node.id, Some(traced))
 	def updateTraced(values: Iterable[(Int, Boolean)]): Unit = values foreach { case (id, traced) => updateTraced(id, Some(traced)) }
+
+	implicit class ExtendedTreeNodeData(n: TreeNodeData) {
+		/** whether or not this treenode is being traced; this value may be unspecified (None) */
+		def traced = isTraced(n)
+		def traced_=(newVal: Boolean) = updateTraced(n, newVal)
+		def traced_=(newVal: Option[Boolean]) = updateTraced(n, newVal)
+	}
 }

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

@@ -101,7 +101,18 @@ private[slick] class SlickTreeNodeDataAccess(dao: TreeNodeDataDao, db: Database)
 		dao storeNodes nodes
 	}
 
-	def updateTraced(id: Int, traced: Option[Boolean]) = db withTransaction { implicit transaction =>
-		dao.updateTraced(id, traced)
+	private lazy val tracedCache = collection.mutable.Map(db withSession { implicit session => dao.getTraced }: _*)
+
+	def isTraced(id: Int): Option[Boolean] = tracedCache.get(id)
+
+	def updateTraced(id: Int, traced: Option[Boolean]) = {
+		db withTransaction { implicit transaction =>
+			dao.updateTraced(id, traced)
+		}
+
+		traced match {
+			case Some(traced) => tracedCache += id -> traced
+			case None => tracedCache -= id
+		}
 	}
 }

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

@@ -48,11 +48,19 @@ private[slick] class TreeNodeDataDao(val driver: JdbcProfile) {
 		def label = column[String]("label", O.NotNull)
 		def kind = column[CodeTreeNodeKind]("kind", O.NotNull)
 		def size = column[Option[Int]]("size", O.Nullable)
-		def traced = column[Option[Boolean]]("traced", O.Nullable)
-		def * = (id, parentId, label, kind, size, traced) <> (TreeNode.tupled, TreeNode.unapply)
+		def * = (id, parentId, label, kind, size) <> (TreeNode.tupled, TreeNode.unapply)
 	}
 	val treeNodeData = TableQuery[TreeNodeData]
 
+	class TracedNodes(tag: Tag) extends Table[(Int, Boolean)](tag, "traced_nodes") {
+		def nodeId = column[Int]("node_id", O.PrimaryKey, O.NotNull)
+		def traced = column[Boolean]("traced", O.NotNull)
+		def * = (nodeId, traced)
+
+		def node = foreignKey("tn_node", nodeId, treeNodeData)(_.id, onDelete = ForeignKeyAction.Cascade)
+	}
+	val tracedNodes = TableQuery[TracedNodes]
+
 	class MethodSignatureNodeMap(tag: Tag) extends Table[(String, Int)](tag, "method_signature_node_map") {
 		def signature = column[String]("sig", O.PrimaryKey, O.NotNull)
 		def nodeId = column[Int]("node_id", O.NotNull)
@@ -71,7 +79,7 @@ private[slick] class TreeNodeDataDao(val driver: JdbcProfile) {
 	}
 	val jspNodeMap = TableQuery[JspNodeMap]
 
-	def create(implicit session: Session) = (treeNodeData.ddl ++ methodSignatureNodeMap.ddl ++ jspNodeMap.ddl).create
+	def create(implicit session: Session) = (treeNodeData.ddl ++ tracedNodes.ddl ++ methodSignatureNodeMap.ddl ++ jspNodeMap.ddl).create
 
 	def get(id: Int)(implicit session: Session): Option[TreeNode] = {
 		(for (n <- treeNodeData if n.id === id) yield n).firstOption
@@ -134,8 +142,20 @@ private[slick] class TreeNodeDataDao(val driver: JdbcProfile) {
 		treeNodeData ++= nodes
 	}
 
+	def getTraced()(implicit session: Session) = tracedNodes.list
+
 	def updateTraced(id: Int, traced: Option[Boolean])(implicit session: Session) {
-		val q = for (n <- treeNodeData if n.id === id) yield n.traced
-		q.update(traced)
+		traced match {
+			case Some(traced) =>
+				val q = for (row <- tracedNodes if row.nodeId === id) yield row.traced
+				q.update(traced) match {
+					case 0 => tracedNodes += id -> traced
+					case _ =>
+				}
+
+			case None =>
+				val q = for (row <- tracedNodes if row.nodeId === id) yield row
+				q.delete
+		}
 	}
 }

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

@@ -297,7 +297,7 @@ class TraceAPIServer(manager: TraceManager) extends RestHelper with Loggable {
 				case Full(packages) =>
 					val ids = packages.split(',').flatMap(AsInt.unapply).toSet
 					val tree = target.treeBuilder.projectTree(ids)
-					TreemapDataStreamer.streamTreemapData(tree)
+					TreemapDataStreamer.streamTreemapData(target.traceData.treeNodeData, tree)
 
 				case _ => BadResponse()
 			}

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

@@ -27,7 +27,10 @@ import com.secdec.codepulse.data.trace.TraceData
 object TraceSettingsCreator {
 
 	def generateTraceSettings(traceData: TraceData, jspMapper: Option[JspMapper]): TraceSettings = {
-		val inclusions = traceData.treeNodeData.iterate { it =>
+		val treeNodeData = traceData.treeNodeData
+		import treeNodeData.ExtendedTreeNodeData
+
+		val inclusions = treeNodeData.iterate { it =>
 			(for {
 				treeNode <- it
 				if treeNode.kind == CodeTreeNodeKind.Pkg
@@ -38,7 +41,7 @@ object TraceSettingsCreator {
 				"^" + slashedName + "/[^/]+$"
 			}).toList
 		} ++ jspMapper.map { mapper =>
-			traceData.treeNodeData.iterateJspMappings { it =>
+			treeNodeData.iterateJspMappings { it =>
 				(for ((jspClass, _) <- it) yield mapper getInclusion jspClass).toList
 			}
 		}.getOrElse(Nil)

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

@@ -34,6 +34,8 @@ import com.secdec.codepulse.data.trace.TraceData
 import scala.concurrent.Future
 import scala.concurrent.ExecutionContext.Implicits.global
 import net.liftweb.common.Full
+import com.secdec.codepulse.data.bytecode.CodeTreeNodeKind
+import com.secdec.codepulse.data.trace.TreeNodeData
 
 object TraceUploadData {
 
@@ -105,7 +107,8 @@ object TraceUploadData {
 
 	def handleBinaryZip(file: File): TraceId = createAndLoadTraceData { traceData =>
 		val RootGroupName = "Classes"
-		val builder = new CodeForestBuilder(defaultTracedGroups = RootGroupName :: CodeForestBuilder.JSPGroupName :: Nil)
+		val tracedGroups = (RootGroupName :: CodeForestBuilder.JSPGroupName :: Nil).toSet
+		val builder = new CodeForestBuilder
 		val methodCorrelationsBuilder = collection.mutable.Map.empty[String, Int]
 
 		//TODO: make this configurable somehow
@@ -139,10 +142,22 @@ object TraceUploadData {
 		if (treemapNodes.isEmpty) {
 			throw new NoSuchElementException("No method data found in analyzed upload file.")
 		} else {
+			val treeNodeData = traceData.treeNodeData
+			import treeNodeData.ExtendedTreeNodeData
 
-			traceData.treeNodeData.storeNodes(treemapNodes)
-			traceData.treeNodeData.mapJsps(jspCorrelations)
-			traceData.treeNodeData.mapMethodSignatures(methodCorrelations)
+			val nodeMap = treemapNodes.map(n => n.id -> n).toMap
+			def getRoot(n: TreeNodeData): TreeNodeData = n.parentId match {
+				case Some(parent) => getRoot(nodeMap(parent))
+				case None => n
+			}
+
+			treeNodeData.storeNodes(treemapNodes)
+			for (n <- treemapNodes) n.traced = n.kind match {
+				case CodeTreeNodeKind.Grp | CodeTreeNodeKind.Pkg => Some(tracedGroups contains getRoot(n).label)
+				case _ => None
+			}
+			treeNodeData.mapJsps(jspCorrelations)
+			treeNodeData.mapMethodSignatures(methodCorrelations)
 
 			// The `creationDate` for trace data detected in this manner
 			// should use its default value ('now'), as this is when the data

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

@@ -34,7 +34,9 @@ import net.liftweb.json.Printer
 object TreemapDataStreamer {
 	private val Json = new JsonFactory
 
-	private def writeJson(jg: JsonGenerator)(node: TreeNode) {
+	private def writeJson(treeNodeData: TreeNodeDataAccess, jg: JsonGenerator)(node: TreeNode) {
+		import treeNodeData.ExtendedTreeNodeData
+
 		jg.writeStartObject
 
 		jg.writeNumberField("id", node.data.id)
@@ -46,20 +48,20 @@ object TreemapDataStreamer {
 
 		if (!node.children.isEmpty) {
 			jg writeArrayFieldStart "children"
-			node.children.foreach(writeJson(jg))
+			node.children.foreach(writeJson(treeNodeData, jg))
 			jg.writeEndArray
 		}
 
 		jg.writeEndObject
 	}
 
-	def streamTreemapData(tree: List[TreeNode]): OutputStreamResponse = {
+	def streamTreemapData(treeNodeData: TreeNodeDataAccess, tree: List[TreeNode]): OutputStreamResponse = {
 		def writeData(out: OutputStream) {
 			val jg = Json createGenerator out
 
 			try {
 				jg.writeStartArray
-				tree.foreach(writeJson(jg))
+				tree.foreach(writeJson(treeNodeData, jg))
 				jg.writeEndArray
 			} finally jg.close
 		}

codepulse/src/main/scala/com/secdec/codepulse/tracer/export/TraceExporter.scala

@@ -96,6 +96,8 @@ object TraceExporter extends JsonHelpers {
 	}
 
 	private def writeTreeNodeData(out: OutputStream, treeNodeData: TreeNodeDataAccess) {
+		import treeNodeData.ExtendedTreeNodeData
+
 		streamJson(out) { jg =>
 			jg.writeStartArray