#1367 Fix MISP export

TheHive-Project · May 26, 2020 · 4564863 · 4564863
1 parent fd087c0
commit 4564863
Showing 1 changed file with 16 additions and 9 deletions.
diff --git a/misp/client/src/main/scala/org/thp/misp/client/Base64Flow.scala b/misp/client/src/main/scala/org/thp/misp/client/Base64Flow.scala
@@ -19,7 +19,8 @@ class Base64EncoderFlow extends GraphStage[FlowShape[ByteString, ByteString]] {
   val encoder: Base64.Encoder                           = Base64.getEncoder
 
   override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
-    var remainingBytes: ByteString = ByteString.empty
+    var remainingBytes: ByteString  = ByteString.empty
+    var upstreamIsFinished: Boolean = false
 
     setHandler(
       in,
@@ -35,15 +36,21 @@ class Base64EncoderFlow extends GraphStage[FlowShape[ByteString, ByteString]] {
           }
         }
 
-        override def onUpstreamFinish(): Unit = {
-          push(out, ByteString(encoder.encode(remainingBytes.toArray)))
-          completeStage
-        }
+        override def onUpstreamFinish(): Unit =
+          upstreamIsFinished = true
+      }
+    )
+    setHandler(
+      out,
+      new OutHandler {
+        override def onPull(): Unit =
+          if (!upstreamIsFinished)
+            pull(in)
+          else {
+            push(out, ByteString(encoder.encode(remainingBytes.toArray)))
+            completeStage
+          }
       }
     )
-    setHandler(out, new OutHandler {
-      override def onPull(): Unit =
-        pull(in)
-    })
   }
 }