#2105 Fix injection of case number actor

TheHive-Project · Jul 1, 2021 · eb1bee3 · eb1bee3
1 parent 75d6ff7
commit eb1bee3
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 10 deletions.
diff --git a/migration/src/main/scala/org/thp/thehive/cloner/IntegrityCheckApp.scala b/migration/src/main/scala/org/thp/thehive/cloner/IntegrityCheckApp.scala
@@ -40,6 +40,7 @@ trait IntegrityCheckApp {
               bindActor[DummyActor]("config-actor")
               bindActor[DummyActor]("cortex-actor")
               bindActor[DummyActor]("integrity-check-actor")
+              bindTypedActor(CaseNumberActor.behavior, "case-number-actor")
 
               val integrityCheckOpsBindings = ScalaMultibinder.newSetBinder[GenIntegrityCheckOps](binder)
               integrityCheckOpsBindings.addBinding.to[ProfileIntegrityCheckOps]

diff --git a/thehive/app/org/thp/thehive/TheHiveModule.scala b/thehive/app/org/thp/thehive/TheHiveModule.scala
@@ -107,7 +107,7 @@ class TheHiveModule(environment: Environment, configuration: Configuration) exte
     integrityCheckOpsBindings.addBinding.to[ObservableIntegrityCheckOps]
     integrityCheckOpsBindings.addBinding.to[LogIntegrityCheckOps]
     bind[ActorRef].annotatedWithName("integrity-check-actor").toProvider[IntegrityCheckActorProvider]
-    bind[TypedActorRef[CaseNumberActor.Request]].annotatedWithName("case-number-actor").toProvider[CaseNumberActorProvider]
+    bind[TypedActorRef[CaseNumberActor.Request]].toProvider[CaseNumberActorProvider]
 
     bind[ActorRef].annotatedWithName("flow-actor").toProvider[FlowActorProvider]
 

diff --git a/thehive/app/org/thp/thehive/services/CaseNumber.scala b/thehive/app/org/thp/thehive/services/CaseNumber.scala
@@ -1,6 +1,6 @@
 package org.thp.thehive.services
 
-import akka.actor.typed.scaladsl.Behaviors
+import akka.actor.typed.scaladsl.{ActorContext, Behaviors}
 import akka.actor.typed.scaladsl.adapter.ClassicActorSystemOps
 import akka.actor.typed.{ActorRefResolver, Behavior, ActorRef => TypedActorRef}
 import akka.actor.{ActorSystem, ExtendedActorSystem}
@@ -22,15 +22,25 @@ object CaseNumberActor {
   case class GetNextNumber(replyTo: TypedActorRef[Response]) extends Request
   case class NextNumber(number: Int)                         extends Response
 
-  val behavior: Behavior[Request] = Behaviors.setup[Request] { context =>
+  val behavior: Behavior[Request] = Behaviors.setup[Request](context => waitFirstRequest(context))
+
+  def getNextCaseNumber(context: ActorContext[Request]): Int = {
     val injector = GuiceAkkaExtension(context.system).injector
     val db       = injector.getInstance(classOf[Database])
     val caseSrv  = injector.getInstance(classOf[CaseSrv])
     db.roTransaction { implicit graph =>
-      caseNumberProvider(caseSrv.startTraversal.getLast.headOption.fold(0)(_.number) + 1)
+      caseSrv.startTraversal.getLast.headOption.fold(0)(_.number) + 1
     }
   }
 
+  def waitFirstRequest(context: ActorContext[Request]): Behaviors.Receive[Request] =
+    Behaviors.receiveMessage {
+      case GetNextNumber(replyTo) =>
+        val nextNumber = getNextCaseNumber(context)
+        replyTo ! NextNumber(nextNumber)
+        caseNumberProvider(nextNumber + 1)
+    }
+
   def caseNumberProvider(nextNumber: Int): Behavior[Request] =
     Behaviors.receiveMessage {
       case GetNextNumber(replyTo) =>

diff --git a/thehive/app/org/thp/thehive/services/CaseSrv.scala b/thehive/app/org/thp/thehive/services/CaseSrv.scala
@@ -1,8 +1,9 @@
 package org.thp.thehive.services
 
-import akka.actor.ActorRef
 import akka.actor.typed.scaladsl.AskPattern._
+import akka.actor.typed.scaladsl.adapter.ClassicSchedulerOps
 import akka.actor.typed.{Scheduler, ActorRef => TypedActorRef}
+import akka.actor.{ActorRef, ActorSystem}
 import akka.util.Timeout
 import org.apache.tinkerpop.gremlin.process.traversal.{Order, P}
 import org.apache.tinkerpop.gremlin.structure.Vertex
@@ -34,7 +35,7 @@ import java.lang.{Long => JLong}
 import java.util.{Date, List => JList, Map => JMap}
 import javax.inject.{Inject, Named, Provider, Singleton}
 import scala.concurrent.duration.DurationInt
-import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.concurrent.{Await, ExecutionContextExecutor, Future}
 import scala.util.{Failure, Success, Try}
 
 @Singleton
@@ -53,10 +54,9 @@ class CaseSrv @Inject() (
     userSrv: UserSrv,
     alertSrvProvider: Provider[AlertSrv],
     @Named("integrity-check-actor") integrityCheckActor: ActorRef,
-    @Named("case-number-actor") caseNumberActor: TypedActorRef[CaseNumberActor.Request],
+    caseNumberActor: TypedActorRef[CaseNumberActor.Request],
     cache: SyncCacheApi,
-    implicit val ec: ExecutionContext,
-    implicit val scheduler: Scheduler
+    system: ActorSystem
 ) extends VertexSrv[Case] {
   lazy val alertSrv: AlertSrv = alertSrvProvider.get
 
@@ -133,7 +133,9 @@ class CaseSrv @Inject() (
   }
 
   def nextCaseNumberAsync: Future[Int] = {
-    implicit val timeout: Timeout = Timeout(1.minute)
+    implicit val timeout: Timeout             = Timeout(1.minute)
+    implicit val scheduler: Scheduler         = system.scheduler.toTyped
+    implicit val ec: ExecutionContextExecutor = system.dispatcher
     caseNumberActor.ask[CaseNumberActor.Response](replyTo => CaseNumberActor.GetNextNumber(replyTo)).map {
       case CaseNumberActor.NextNumber(caseNumber) => caseNumber
     }