From 58e923abd8887b4a016358ca5978da42be67e453 Mon Sep 17 00:00:00 2001
From: Sebastian Serth <Sebastian.Serth@hpi.de>
Date: Mon, 11 Oct 2021 23:11:40 +0200
Subject: [PATCH] Add custom websocket header to strategy class

---
 lib/runner/connection.rb                     | 2 +-
 lib/runner/strategy.rb                       | 4 ++++
 lib/runner/strategy/docker_container_pool.rb | 4 ++++
 lib/runner/strategy/poseidon.rb              | 7 +++++++
 4 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/runner/connection.rb b/lib/runner/connection.rb
index 5800045b..5619bf2d 100644
--- a/lib/runner/connection.rb
+++ b/lib/runner/connection.rb
@@ -17,7 +17,7 @@ class Runner::Connection
     # Internally, Faye::WebSocket uses EventMachine and the `ping` value is used to wake the EventMachine thread
     # The `tls` option is used to customize the validation of TLS connections.
     # Passing `nil` as a `root_cert_file` is okay and done so for the DockerContainerPool.
-    @socket = Faye::WebSocket::Client.new(url, [], ping: 0.1, tls: {root_cert_file: Runner.strategy_class.config[:ca_file]})
+    @socket = Faye::WebSocket::Client.new(url, [], strategy.websocket_header.merge(ping: 0.1))
     @strategy = strategy
     @status = :established
     @event_loop = event_loop
diff --git a/lib/runner/strategy.rb b/lib/runner/strategy.rb
index 702311f4..908be168 100644
--- a/lib/runner/strategy.rb
+++ b/lib/runner/strategy.rb
@@ -28,4 +28,8 @@ class Runner::Strategy
   def attach_to_execution(_command)
     raise NotImplementedError
   end
+
+  def websocket_header
+    raise NotImplementedError
+  end
 end
diff --git a/lib/runner/strategy/docker_container_pool.rb b/lib/runner/strategy/docker_container_pool.rb
index 5ae8d62a..5d2fee7d 100644
--- a/lib/runner/strategy/docker_container_pool.rb
+++ b/lib/runner/strategy/docker_container_pool.rb
@@ -76,6 +76,10 @@ class Runner::Strategy::DockerContainerPool < Runner::Strategy
     socket
   end
 
+  def websocket_header
+    {}
+  end
+
   private
 
   def container
diff --git a/lib/runner/strategy/poseidon.rb b/lib/runner/strategy/poseidon.rb
index 7eeef0aa..2c186e1f 100644
--- a/lib/runner/strategy/poseidon.rb
+++ b/lib/runner/strategy/poseidon.rb
@@ -121,6 +121,13 @@ class Runner::Strategy::Poseidon < Runner::Strategy
     raise Runner::Error::FaradayError.new("Request to Poseidon failed: #{e.inspect}")
   end
 
+  def websocket_header
+    {
+      tls: {root_cert_file: self.class.config[:ca_file]},
+      headers: {'Poseidon-Token' => self.class.config[:token]},
+    }
+  end
+
   private
 
   def execute_command(command)