diff --git a/codepulse/src/main/scala/bootstrap/liftweb/AppCleanup.scala b/codepulse/src/main/scala/bootstrap/liftweb/AppCleanup.scala index 350640a..bdad742 100644 --- a/codepulse/src/main/scala/bootstrap/liftweb/AppCleanup.scala +++ b/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 + } } } \ No newline at end of file diff --git a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectFileUploadHandler.scala b/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectFileUploadHandler.scala index 8d21efe..4869bf0 100644 --- a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectFileUploadHandler.scala +++ b/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" diff --git a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectManager.scala b/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectManager.scala index 6668dad..08837ae 100644 --- a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectManager.scala +++ b/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") + } } \ No newline at end of file diff --git a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectUploadData.scala b/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectUploadData.scala index d3f4cdb..1ee5bae 100644 --- a/codepulse/src/main/scala/com/secdec/codepulse/tracer/ProjectUploadData.scala +++ b/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. */ diff --git a/codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceServer.scala b/codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceServer.scala index 49b6d9a..387ea1c 100644 --- a/codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceServer.scala +++ b/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") } diff --git a/codepulse/src/main/scala/com/secdec/codepulse/util/ZipEntryChecker.scala b/codepulse/src/main/scala/com/secdec/codepulse/util/ZipEntryChecker.scala index c0c298e..4a1f449 100644 --- a/codepulse/src/main/scala/com/secdec/codepulse/util/ZipEntryChecker.scala +++ b/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 + } + } } \ No newline at end of file