#1708 Fix integrity check serializer

TheHive-Project · Dec 17, 2020 · b96bd97 · b96bd97
1 parent eb0230c
commit b96bd97
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/thehive/app/org/thp/thehive/services/IntegrityCheckActor.scala b/thehive/app/org/thp/thehive/services/IntegrityCheckActor.scala
@@ -20,10 +20,10 @@ import scala.util.Success
 
 sealed trait IntegrityCheckMessage
 case class EntityAdded(name: String) extends IntegrityCheckMessage
+case class NeedCheck(name: String)   extends IntegrityCheckMessage
+case class Check(name: String)       extends IntegrityCheckMessage
 
 class IntegrityCheckActor() extends Actor {
-  case class NeedCheck(name: String)
-  case class Check(name: String)
 
   lazy val logger: Logger               = Logger(getClass)
   lazy val injector: Injector           = GuiceAkkaExtension(context.system).injector

diff --git a/thehive/app/org/thp/thehive/services/IntegrityCheckSerializer.scala b/thehive/app/org/thp/thehive/services/IntegrityCheckSerializer.scala
@@ -12,12 +12,16 @@ class IntegrityCheckSerializer extends Serializer {
   override def toBinary(o: AnyRef): Array[Byte] =
     o match {
       case EntityAdded(name) => 0.toByte +: name.getBytes
+      case NeedCheck(name)   => 1.toByte +: name.getBytes
+      case Check(name)       => 2.toByte +: name.getBytes
       case _                 => throw new NotSerializableException
     }
 
   override def fromBinary(bytes: Array[Byte], manifest: Option[Class[_]]): AnyRef =
     bytes(0) match {
       case 0 => EntityAdded(new String(bytes.tail))
+      case 1 => NeedCheck(new String(bytes.tail))
+      case 2 => Check(new String(bytes.tail))
       case _ => throw new NotSerializableException
     }
 }
diff --git a/thehive/conf/play/reference-overrides.conf b/thehive/conf/play/reference-overrides.conf
@@ -24,12 +24,14 @@ akka.actor {
     notification = "org.thp.thehive.services.notification.NotificationSerializer"
     thehive-schema-updater = "org.thp.thehive.models.SchemaUpdaterSerializer"
     flow = "org.thp.thehive.services.FlowSerializer"
+    integrity = "org.thp.thehive.services.IntegrityCheckSerializer"
   }
 
   serialization-bindings {
     "org.thp.thehive.services.StreamMessage" = stream
     "org.thp.thehive.services.notification.NotificationMessage" = notification
     "org.thp.thehive.models.SchemaUpdaterMessage" = thehive-schema-updater
     "org.thp.thehive.services.FlowMessage" = flow
+    "org.thp.thehive.services.IntegrityCheckMessage" = integrity
   }
 }