Merge pull request #1039 from ninoseki/fix-parallel-enricher-issue

feat: fix parallel enrichment issue

docs/enrichers/index.md

@@ -16,7 +16,6 @@ options:
   retry_times: ...
   retry_interval: ...
   retry_exponential_backoff: ...
-  parallel: ...
 ```
 
 ### Timeout
@@ -34,7 +33,3 @@ options:
 ### Retry Exponential Backoff
 
 `retry_exponential_backoff` (`bool`) controls whether to do exponential backoff. Optional. Defaults to `true`. Configurable via `RETRY_EXPONENTIAL_BACKOFF` environment variable.
-
-### Parallel
-
-`parallel` (`bool`) controls whether to allow parallel execution or not. Optional. Defaults to `false`. Configurable via `PARALLEL` environment variable.

lib/mihari/actor.rb

@@ -50,13 +50,6 @@ def timeout
       options[:timeout]
     end
 
-    #
-    # @return [Boolean]
-    #
-    def parallel?
-      options[:parallel] || Mihari.config.parallel
-    end
-
     def validate_configuration!
       return if configured?
 

lib/mihari/analyzers/base.rb

@@ -40,6 +40,13 @@ def ignore_error?
         options[:ignore_error] || Mihari.config.ignore_error
       end
 
+      #
+      # @return [Boolean]
+      #
+      def parallel?
+        options[:parallel] || Mihari.config.parallel
+      end
+
       # @return [Array<String>, Array<Mihari::Models::Artifact>]
       def artifacts
         raise NotImplementedError, "You must implement #{self.class}##{__method__}"

lib/mihari/emitters/base.rb

@@ -19,6 +19,13 @@ def initialize(rule:, options: nil)
         @rule = rule
       end
 
+      #
+      # @return [Boolean]
+      #
+      def parallel?
+        options[:parallel] || Mihari.config.parallel
+      end
+
       # A target to emit the data
       #
       # @return [String]

lib/mihari/rule.rb

@@ -174,11 +174,11 @@ def unique_artifacts
     # @return [Array<Mihari::Models::Artifact>]
     #
     def enriched_artifacts
-      @enriched_artifacts ||= unique_artifacts.map do |artifact|
-        serial_enrichers.each { |enricher| enricher.result(artifact) }
-        Parallel.each(parallel_enrichers) { |enricher| enricher.result(artifact) }
-
-        artifact
+      @enriched_artifacts ||= Parallel.map(unique_artifacts) do |artifact|
+        artifact.tap do |tapped|
+          # NOTE: To apply changes correctly, enrichers should be applied to an artifact serially
+          enrichers.each { |enricher| enricher.result(tapped) }
+        end
       end
     end
 
@@ -404,14 +404,6 @@ def enrichers
       end
     end
 
-    def parallel_enrichers
-      enrichers.select(&:parallel?)
-    end
-
-    def serial_enrichers
-      enrichers.reject(&:parallel?)
-    end
-
     #
     # Validate the data format
     #

lib/mihari/schemas/analyzer.rb

@@ -117,6 +117,6 @@ module Analyzers
       end
     end
 
-    Analyzer = Schemas::Analyzers.get_or_composition
+    Analyzer = Schemas::Analyzers.compose_by_or
   end
 end

lib/mihari/schemas/concerns/orrable.rb

@@ -9,7 +9,7 @@ module Concerns
       module Orrable
         extend ActiveSupport::Concern
 
-        def get_or_composition
+        def compose_by_or
           schemas = constants.map { |sym| const_get sym }
           return schemas.first if schemas.length <= 1
 

lib/mihari/schemas/emitter.rb

@@ -10,28 +10,28 @@ module Emitters
 
       Database = Dry::Schema.Params do
         required(:emitter).value(Types::String.enum(*Mihari::Emitters::Database.keys))
-        optional(:options).hash(Options)
+        optional(:options).hash(EmitterOptions)
       end
 
       MISP = Dry::Schema.Params do
         required(:emitter).value(Types::String.enum(*Mihari::Emitters::MISP.keys))
         optional(:url).filled(:string)
         optional(:api_key).filled(:string)
-        optional(:options).hash(Options)
+        optional(:options).hash(EmitterOptions)
       end
 
       TheHive = Dry::Schema.Params do
         required(:emitter).value(Types::String.enum(*Mihari::Emitters::TheHive.keys))
         optional(:url).filled(:string)
         optional(:api_key).filled(:string)
-        optional(:options).hash(Options)
+        optional(:options).hash(EmitterOptions)
       end
 
       Slack = Dry::Schema.Params do
         required(:emitter).value(Types::String.enum(*Mihari::Emitters::Slack.keys))
         optional(:webhook_url).filled(:string)
         optional(:channel).filled(:string)
-        optional(:options).hash(Options)
+        optional(:options).hash(EmitterOptions)
       end
 
       Webhook = Dry::Schema.Params do
@@ -40,10 +40,10 @@ module Emitters
         optional(:method).value(Types::HTTPRequestMethods).default("POST")
         optional(:headers).filled(:hash)
         optional(:template).filled(:string)
-        optional(:options).hash(Options)
+        optional(:options).hash(EmitterOptions)
       end
     end
 
-    Emitter = Schemas::Emitters.get_or_composition
+    Emitter = Schemas::Emitters.compose_by_or
   end
 end

lib/mihari/schemas/enricher.rb

@@ -29,6 +29,6 @@ module Enrichers
       end
     end
 
-    Enricher = Schemas::Enrichers.get_or_composition
+    Enricher = Schemas::Enrichers.compose_by_or
   end
 end

lib/mihari/schemas/options.rb

@@ -3,24 +3,29 @@
 module Mihari
   module Schemas
     Options = Dry::Schema.Params do
-      optional(:retry_times).value(:integer).default(Mihari.config.retry_times)
-      optional(:retry_interval).value(:integer).default(Mihari.config.retry_interval)
-      optional(:retry_exponential_backoff).value(:bool).default(Mihari.config.retry_exponential_backoff)
+      optional(:retry_times).value(:integer)
+      optional(:retry_interval).value(:integer)
+      optional(:retry_exponential_backoff).value(:bool)
       optional(:timeout).value(:integer)
-      optional(:parallel).value(:bool).default(Mihari.config.parallel)
     end
 
-    IgnoreErrorOptions = Dry::Schema.Params do
-      optional(:ignore_error).value(:bool).default(Mihari.config.ignore_error)
+    ParallelOptions = Dry::Schema.Params do
+      optional(:parallel).value(:bool)
     end
 
-    AnalyzerOptions = Options | IgnoreErrorOptions
+    IgnoreErrorOptions = Dry::Schema.Params do
+      optional(:ignore_error).value(:bool)
+    end
 
     PaginationOptions = Dry::Schema.Params do
-      optional(:pagination_interval).value(:integer).default(Mihari.config.pagination_interval)
-      optional(:pagination_limit).value(:integer).default(Mihari.config.pagination_limit)
+      optional(:pagination_interval).value(:integer)
+      optional(:pagination_limit).value(:integer)
     end
 
-    AnalyzerPaginationOptions = AnalyzerOptions | PaginationOptions
+    AnalyzerOptions = Options & IgnoreErrorOptions & ParallelOptions
+
+    AnalyzerPaginationOptions = AnalyzerOptions & PaginationOptions
+
+    EmitterOptions = Options & ParallelOptions
   end
 end

lib/mihari/schemas/rule.rb

@@ -21,9 +21,9 @@ module Schemas
       optional(:created_on).value(:date)
       optional(:updated_on).value(:date)
 
-      required(:queries).value(:array).each { Analyzer } # rubocop:disable Lint/Void
-      optional(:emitters).value(:array).each { Emitter }.default(DEFAULT_EMITTERS) # rubocop:disable Lint/Void
-      optional(:enrichers).value(:array).each { Enricher }.default(DEFAULT_ENRICHERS) # rubocop:disable Lint/Void
+      required(:queries).array { Analyzer }
+      optional(:emitters).array { Emitter }.default(DEFAULT_EMITTERS)
+      optional(:enrichers).array { Enricher }.default(DEFAULT_ENRICHERS)
 
       optional(:data_types).filled(array[Types::DataTypes]).default(Mihari::Types::DataTypes.values)
 

spec/commands/search_spec.rb

@@ -13,6 +13,7 @@ class SearchCLI < Mihari::CLI::Base
 
   before do
     allow(rule).to receive(:enrichers).and_return([])
+    allow(Parallel).to receive(:processor_count).and_return(0)
   end
 
   describe "#search" do