Refactor our polled workers to use a common macro (#709)

apps/core/lib/core/schema/docker_image.ex

@@ -24,6 +24,8 @@ defmodule Core.Schema.DockerImage do
     )
   end
 
+  def search(query \\ __MODULE__, s), do: from(di in query, where: like(di.tag, ^"#{s}%"))
+
   def for_tag(query \\ __MODULE__, tag) do
     from(di in query, where: di.tag == ^tag)
   end

apps/core/lib/core/services/scan.ex

@@ -0,0 +1,80 @@
+defmodule Core.Services.Scan do
+  require Logger
+  alias Core.Services.{Repositories, Versions}
+  alias Core.Schema.{DockerImage, Version, Chart, Terraform}
+  alias Core.Docker.TrivySource
+
+  def scan_image(%DockerImage{} = image) do
+    %{docker_repository: %{repository: repo} = dkr} = img = Core.Repo.preload(image, [docker_repository: :repository])
+    %{publisher: %{owner: owner}} = Core.Repo.preload(repo, [publisher: :owner])
+
+    registry_name = img_name(img)
+    {:ok, registry_token} = Repositories.docker_token([:pull], "#{repo.name}/#{dkr.name}", owner)
+    env = [{"TRIVY_REGISTRY_TOKEN", registry_token} | Core.conf(:docker_env)]
+
+    image = "#{registry_name}:#{image.tag}"
+    Logger.info "Scanning image #{image}"
+    case System.cmd("trivy", ["--quiet", "image", "--format", "json", image, "--timeout", "10m0s"], env: env) do
+      {output, 0} ->
+        case Jason.decode(output) do
+          {:ok, [%{"Vulnerabilities" => vulns} | _]} -> insert_vulns(vulns, img)
+          {:ok, %{"Results" => [_ | _] = res, "SchemaVersion" => 2}} ->
+            Enum.flat_map(res, &Map.get(&1, "Vulnerabilities", []))
+            |> insert_vulns(img)
+            |> log()
+          res ->
+            Logger.info "irregular trivy output #{inspect(res)}"
+            insert_vulns([], img)
+            |> log()
+        end
+      {output, _} ->
+        Logger.info "Trivy failed with: #{output}"
+        handle_trivy_error(output, img)
+    end
+  end
+
+  def terrascan(%Version{} = version) do
+    {type, url} = terrascan_details(version)
+    {output, _} = System.cmd("terrascan", [
+      "scan",
+      "--iac-type", type,
+      "--remote-type", "http",
+      "--remote-url", url,
+      "--output", "json",
+      "--use-colors", "f"
+    ])
+    Versions.record_scan(output, version)
+  end
+
+  defp handle_trivy_error(output, %DockerImage{} = img) do
+    case String.contains?(output, "timeout") do
+      true -> Ecto.Changeset.change(img, %{scan_completed_at: Timex.now()}) |> Core.Repo.update()
+      _ -> :error
+    end
+  end
+
+  defp terrascan_details(%Version{
+    version: v,
+    chart: %Chart{name: chart, repository: %{name: name}}
+  }) do
+    {"helm", "http://chartmuseum:8080/cm/#{name}/charts/#{chart}-#{v}.tgz"}
+  end
+  defp terrascan_details(%Version{terraform: %Terraform{}} = v),
+    do: {"terraform", Core.Storage.url({v.package, v}, :original)}
+
+  defp img_name(%DockerImage{docker_repository: %{repository: repo} = dkr}),
+    do: "#{Core.conf(:registry)}/#{repo.name}/#{dkr.name}"
+
+  defp insert_vulns(nil, img), do: insert_vulns([], img)
+  defp insert_vulns(vulns, img) when is_list(vulns) do
+    Logger.info "found #{length(vulns)} vulnerabilities for #{img_name(img)}"
+    vulns
+    |> Enum.map(&TrivySource.to_vulnerability/1)
+    |> Repositories.add_vulnerabilities(img)
+  end
+
+  defp log({:ok, %DockerImage{id: id, vulnerabilities: vulns}}) when is_list(vulns) do
+    Logger.info "Found #{length(vulns)} vulns for image #{id}"
+  end
+  defp log(_), do: :ok
+end

apps/core/test/services/scan_test.exs

@@ -0,0 +1,77 @@
+defmodule Core.Services.ScanTest do
+  use Core.SchemaCase, async: true
+  use Mimic
+
+  alias Core.Services.Scan
+
+  describe "#scan_image/1" do
+    test "it can execute a trivy command" do
+      image = insert(:docker_image)
+      image_name = "dkr.plural.sh/#{image.docker_repository.repository.name}/#{image.docker_repository.name}:#{image.tag}"
+      vuln = Application.get_env(:core, :vulnerability)
+      expect(System, :cmd, fn
+        "trivy", ["--quiet", "image", "--format", "json", ^image_name, "--timeout", "10m0s"], [env: [{"TRIVY_REGISTRY_TOKEN", _}]] ->
+          {~s([{"Vulnerabilities": [#{vuln}]}]), 0}
+      end)
+
+      {:ok, scanned} = Scan.scan_image(image)
+
+      assert scanned.id == image.id
+      assert scanned.grade == :c
+      assert scanned.scan_completed_at
+
+      [vuln] = scanned.vulnerabilities
+      assert vuln.image_id == scanned.id
+    end
+  end
+
+  describe "terrascan/2" do
+    test "it can scan a terraform version" do
+      result = terrascan_res()
+      version = insert(:version,
+        terraform: insert(:terraform, package: %{file_name: "file.tgz", updated_at: nil}),
+        chart: nil,
+        chart_id: nil
+      )
+      url = Core.Storage.url({version.package, version}, :original)
+      expect(System, :cmd, fn
+        "terrascan", [
+          "scan",
+          "--iac-type", "terraform",
+          "--remote-type", "http",
+          "--remote-url", ^url,
+          "--output", "json",
+          "--use-colors", "f"
+        ] ->
+          {result, 0}
+      end)
+
+      {:ok, scanned} = Scan.terrascan(version)
+
+      assert scanned.grade == :d
+    end
+
+    test "it can scan a chart version" do
+      result = terrascan_res()
+      version = insert(:version, chart: insert(:chart))
+      url = "http://chartmuseum:8080/cm/#{version.chart.repository.name}/charts/#{version.chart.name}-#{version.version}.tgz"
+      expect(System, :cmd, fn
+        "terrascan", [
+          "scan",
+          "--iac-type", "helm",
+          "--remote-type", "http",
+          "--remote-url", ^url,
+          "--output", "json",
+          "--use-colors", "f"
+        ] ->
+          {result, 0}
+      end)
+
+      {:ok, scanned} = Scan.terrascan(version)
+
+      assert scanned.grade == :d
+    end
+  end
+
+  defp terrascan_res(), do: priv_file(:core, "scan.json") |> File.read!()
+end

apps/core/test/test_helper.exs

@@ -28,5 +28,6 @@ Mimic.copy(OAuth2.Client)
 Mimic.copy(Core.OAuth.Github)
 Mimic.copy(Core.Services.Shell.Pods)
 Mimic.copy(Vault)
+Mimic.copy(System)
 
 {:ok, _} = Application.ensure_all_started(:ex_machina)

apps/graphql/lib/graphql/resolvers/docker.ex

@@ -17,6 +17,7 @@ defmodule GraphQl.Resolvers.Docker do
   def list_images(%{docker_repository_id: repo} = args, _) do
     DockerImage.for_repository(repo)
     |> DockerImage.ordered()
+    |> maybe_search(DockerImage, args)
     |> paginate(args)
   end
 

apps/graphql/lib/graphql/schema/docker.ex

@@ -104,6 +104,7 @@ defmodule GraphQl.Schema.Docker do
     connection field :docker_images, node_type: :docker_image do
       middleware Authenticated
       arg :docker_repository_id, non_null(:id)
+      arg :q, :string
 
       resolve &Docker.list_images/2
     end

apps/graphql/test/queries/docker_queries_test.exs

@@ -44,6 +44,27 @@ defmodule GraphQl.DockerQueriesTest do
       assert from_connection(found)
              |> ids_equal(docker_imgs)
     end
+
+    test "it can search by tag" do
+      repo = insert(:docker_repository)
+      insert_list(3, :docker_image, docker_repository: repo)
+      img = insert(:docker_image, docker_repository: repo, tag: "3.0")
+
+      {:ok, %{data: %{"dockerImages" => found}}} = run_query("""
+        query DockerImages($id: ID!) {
+          dockerImages(dockerRepositoryId: $id, q: "3." first: 5) {
+            edges {
+              node {
+                id
+              }
+            }
+          }
+        }
+      """, %{"id" => repo.id}, %{current_user: insert(:user)})
+
+      assert from_connection(found)
+             |> ids_equal([img])
+    end
   end
 
   describe "dockerImage" do

apps/worker/lib/worker/conduit/subscribers/scan.ex

@@ -1,9 +1,6 @@
 defmodule Worker.Conduit.Subscribers.Scan do
   use Conduit.Subscriber
   import Conduit.Message
-  alias Core.Services.Versions
-  alias Core.Schema.{Version, Chart, Terraform}
-  require Logger
 
   def process(%Conduit.Message{body: body} = msg, _) do
     case scan(body) do
@@ -12,26 +9,5 @@ defmodule Worker.Conduit.Subscribers.Scan do
     end
   end
 
-  def scan(version) do
-    {type, url} = scan_details(version)
-    {output, _} = System.cmd("terrascan", [
-      "scan",
-      "--iac-type", type,
-      "--remote-type", "http",
-      "--remote-url", url,
-      "--output", "json",
-      "--use-colors", "f"
-    ])
-    Logger.info "terrascan output: #{output}"
-    Versions.record_scan(output, version)
-  end
-
-  defp scan_details(%Version{
-    version: v,
-    chart: %Chart{name: chart, repository: %{name: name}}
-  }) do
-    {"helm", "http://chartmuseum:8080/cm/#{name}/charts/#{chart}-#{v}.tgz"}
-  end
-  defp scan_details(%Version{terraform: %Terraform{}} = v),
-    do: {"terraform", Core.Storage.url({v.package, v}, :original)}
+  def scan(version), do: Core.Services.Scan.terrascan(version)
 end

apps/worker/lib/worker/demo_projects/producer.ex

@@ -1,13 +1,9 @@
 defmodule Worker.DemoProjects.Producer do
-  use GenStage
-  use Worker.Base
-  require Logger
+  use Worker.PollProducer
   alias Core.Services.Shell.Demo
 
   @max 20
-  @poll :timer.seconds(5)
-
-  defmodule State, do: defstruct [:demand, :draining]
+  @poll Worker.conf(:demo_interval) |>  :timer.seconds()
 
   def start_link(opts \\ []) do
     GenStage.start_link(__MODULE__, opts, name: __MODULE__)
@@ -19,23 +15,11 @@ defmodule Worker.DemoProjects.Producer do
     {:producer, %State{demand: 0}}
   end
 
-  def handle_info(:poll, %State{draining: true} = state), do: empty(state)
-  def handle_info(:poll, %State{demand: demand} = state) do
-    Logger.info "checking for stale demo projects"
-    deliver(demand, %{state | draining: FT.K8S.TrafficDrainHandler.draining?()})
-  end
-
-  def handle_demand(_, %State{draining: true} = state), do: empty(state)
-  def handle_demand(demand, %State{demand: remaining} = state) when demand > 0 do
-    deliver(demand + remaining, state)
-  end
-
-  def handle_demand(_, state), do: empty(state)
-
   defp deliver(demand, state) do
+    Logger.info "checking for expired demo projects"
     case Demo.poll(min(demand, @max)) do
       {:ok, demos} when is_list(demos) ->
-        {:noreply, demos, %{state | demand: demand - length(demos)}}
+        {:noreply, demos, demand(state, demand, demos)}
       _ -> empty(%{state | demand: demand})
     end
   end

apps/worker/lib/worker/docker/pipeline.ex

@@ -8,7 +8,7 @@ defmodule Worker.Docker.Pipeline do
       Logger.info "Scheduling docker scan for #{img.id}"
       img
     end)
-    |> Flow.map(&Worker.Conduit.Subscribers.Docker.scan_image/1)
+    |> Flow.map(&Core.Services.Scan.scan_image/1)
     |> Flow.start_link()
   end