Merge pull request #17 from codedx/feature/nested-archives

Add support for .ear files and a less naive short-circuit file checker.

Author: Lumifly

codepulse/src/main/scala/bootstrap/liftweb/AppCleanup.scala

@@ -20,11 +20,23 @@
 package bootstrap.liftweb
 
 object AppCleanup {
-	private var hooks: List[() => Unit] = Nil
+	private var preShutdownHooks: List[() => Unit] = Nil
 
-	def add(cleanup: () => Unit) = hooks ::= cleanup
+	private var shutdownHooks: List[() => Unit] = Nil
+
+	def addPreShutdownHook(cleanup: () => Unit) = preShutdownHooks ::= cleanup
+
+	def addShutdownHook(cleanup: () => Unit) = shutdownHooks ::= cleanup
 
 	def runCleanup() = {
-		hooks.reverseIterator.foreach { _() }
+		try {
+			println("Running PreShutdownHooks")
+			preShutdownHooks.reverseIterator.foreach { _() }
+
+			println("Running ShutdownHooks")
+			shutdownHooks.reverseIterator.foreach { _() }
+		} catch {
+			case e: Throwable => e.printStackTrace
+		}
 	}
 }

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

@@ -56,7 +56,7 @@ class ProjectFileUploadHandler(projectManager: ProjectManager) extends RestHelpe
 		case UploadPath("create") Post req => fallbackResponse {
 			for {
 				(inputFile, originalName, cleanup) <- getReqFile(req) ?~! "Creating a new project requires a file"
-				_ <- ProjectUploadData.checkForBinaryZip(inputFile) ?~ {
+				_ <- ProjectUploadData.checkForClassesInNestedArchive(inputFile) ?~ {
 					s"The file you picked doesn't have any compiled Java files."
 				}
 				name <- req.param("name") ?~ "You must specify a name"

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

@@ -44,7 +44,11 @@ import reactive.Observing
 object ProjectManager {
 	lazy val defaultActorSystem = {
 		val sys = ActorSystem("ProjectManagerSystem")
-		AppCleanup.add { () => sys.shutdown() }
+		AppCleanup.addShutdownHook { () =>
+			sys.shutdown()
+			sys.awaitTermination()
+			println("Shutdown ProjectManager's ActorSystem")
+		}
 		sys
 	}
 }
@@ -196,6 +200,9 @@ class ProjectManager(val actorSystem: ActorSystem) extends Observing {
 	}
 
 	// Also make sure any dirty projects are saved when exiting
-	AppCleanup.add { () => flushProjects }
+	AppCleanup.addPreShutdownHook { () =>
+		flushProjects
+		println("Flushed ProjectManager projects")
+	}
 
 }

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

@@ -69,6 +69,12 @@ object ProjectUploadData {
 		}
 	}
 
+	def checkForClassesInNestedArchive(file: File): Boolean = {
+		ZipEntryChecker.findFirstEntry(file) { (filename, entry, contents) =>
+			!entry.isDirectory && FilenameUtils.getExtension(entry.getName) == "class"
+		}
+	}
+
 	/** A preliminary check on a File to see if it looks like an
 	  * exported .pulse file.
 	  */

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

@@ -36,7 +36,7 @@ object TraceServer {
 	private implicit lazy val socketServer = {
 		val ss = SocketServer.default(com.secdec.codepulse.userSettings.tracePort)
 		ss.start()
-		AppCleanup.add { () =>
+		AppCleanup.addPreShutdownHook { () =>
 			ss.shutdown
 			println("Shutdown TracerServer's socketServer")
 		}

codepulse/src/main/scala/com/secdec/codepulse/util/ZipEntryChecker.scala

@@ -19,18 +19,9 @@
 
 package com.secdec.codepulse.util
 
-import java.io.BufferedInputStream
-import java.io.File
-import java.io.FileInputStream
-import java.io.InputStream
-import java.util.zip.ZipEntry
-import java.util.zip.ZipFile
-import java.util.zip.ZipInputStream
-
-import scala.collection.JavaConversions._
-import scala.util.Failure
-import scala.util.Success
-import scala.util.Try
+import java.io.{ BufferedInputStream, File, FileInputStream, InputStream }
+import java.util.zip.{ ZipEntry, ZipFile, ZipInputStream }
+import scala.util.{ Failure, Success, Try }
 
 import org.apache.commons.io.FilenameUtils
 import org.apache.commons.io.input.CloseShieldInputStream
@@ -61,7 +52,7 @@ trait ZipEntryChecker {
 	}
 
 	def isZip(name: String): Boolean = FilenameUtils.getExtension(name) match {
-		case "zip" | "jar" | "war" => true
+		case "zip" | "ear" | "jar" | "war" => true
 		case _ => false
 	}
 
@@ -95,4 +86,32 @@ trait ZipEntryChecker {
 			stream.close
 		}
 	}
+
+	def findFirstEntry(file: File, recursive: Boolean = true)(predicate: (String, ZipEntry, InputStream) => Boolean): Boolean = {
+		val stream = new BufferedInputStream(new FileInputStream(file))
+
+		try findFirstEntry(file.getName, stream, recursive)(predicate) finally stream.close
+	}
+
+	def findFirstEntry(filename: String, stream: InputStream, recursive: Boolean)(predicate: (String, ZipEntry, InputStream) => Boolean): Boolean = {
+		val zipStream = new ZipInputStream(stream)
+
+		try {
+			val entryStream = Stream.continually(Try { zipStream.getNextEntry })
+				.map(_.toOption.flatMap { Option(_) })
+				.takeWhile(_.isDefined)
+				.flatten
+				.filterNot(ZipCleaner.shouldFilter)
+
+			entryStream.exists { entry =>
+				lazy val recurse = recursive &&
+					isZip(entry.getName) &&
+					findFirstEntry(s"$filename/${entry.getName}", new CloseShieldInputStream(zipStream), true)(predicate)
+
+				predicate(filename, entry, zipStream) || recurse
+			}
+		} finally {
+			zipStream.close
+		}
+	}
 }