From 0a7987541626828f86d8930733280363e0313da0 Mon Sep 17 00:00:00 2001
From: Charles Care <c.p.care@gmail.com>
Date: Wed, 27 Nov 2019 22:59:33 +0000
Subject: [PATCH] Limiting of IO

---
 lib/pipeline/util/external_command.rb       | 83 ++++++++++++++++++---
 lib/pipeline/util/runc_wrapper.rb           |  2 +
 test/pipeline/util/external_command_test.rb | 48 +++++++++++-
 3 files changed, 122 insertions(+), 11 deletions(-)

diff --git a/lib/pipeline/util/external_command.rb b/lib/pipeline/util/external_command.rb
index e5493d1..7677bf1 100644
--- a/lib/pipeline/util/external_command.rb
+++ b/lib/pipeline/util/external_command.rb
@@ -3,10 +3,15 @@ require 'open3'
 module Pipeline::Util
   class ExternalCommand
 
-    attr_accessor :cmd_string, :status, :stdout, :stderr, :suppress_output
+    BLOCK_SIZE = 1024
+
+    attr_accessor :cmd_string, :status, :stdout, :stderr, :suppress_output,
+                  :stderr_limit, :stdout_limit
 
     def initialize(cmd_string)
       @cmd_string = cmd_string
+      @stdout_limit = -1
+      @stderr_limit = -1
     end
 
     def call!
@@ -15,14 +20,9 @@ module Pipeline::Util
     end
 
     def call
-      c = cmd
-      puts "> #{c}"  unless suppress_output
-      @stdout, @stderr, @status = Open3.capture3(c)
-      # Open3.popen3("ls") do |stdout, stderr, status, thread|
-      #   @stdout = stdout.read
-      #   @stderr = stderr.read
-      #   @status = status
-      # end
+      invoke_process
+
+
       puts "status: #{status}" unless suppress_output
       puts "stdout: #{stdout}" unless suppress_output
       puts "stderr: #{stderr}" unless suppress_output
@@ -37,7 +37,11 @@ module Pipeline::Util
     end
 
     def success?
-      status.success?
+      status && status.success?
+    end
+
+    def killed?
+      !! @killed
     end
 
     def exit_status
@@ -56,5 +60,64 @@ module Pipeline::Util
         stderr: stderr
       }
     end
+
+    private
+
+    def invoke_process
+      c = cmd
+      captured_stdout = []
+      captured_stderr = []
+      stdout_size = 0
+      stderr_size = 0
+      puts "> #{c}"  unless suppress_output
+      Open3.popen3(c) do |_stdin, _stdout, _stderr, wait_thr|
+        pid = wait_thr.pid
+        _stdin.close_write
+
+        begin
+          files = [_stdout, _stderr]
+
+          until files.find { |f| !f.eof }.nil? || @killed do
+            ready = IO.select(files)
+
+            if ready
+              readable = ready[0]
+              readable.each do |f|
+                begin
+                  data = f.read_nonblock(BLOCK_SIZE)
+                  if f == _stdout
+                    unless @killed
+                      captured_stdout << data
+                      stdout_size += data.size
+                      if stdout_limit > 0 && stdout_size > stdout_limit
+                        Process.kill("KILL", wait_thr.pid)
+                        @killed = true
+                      end
+                    end
+                  end
+                  if f == _stderr
+                    unless @killed
+                      captured_stderr << data
+                      stderr_size += data.size
+                      if stderr_limit > 0 && stderr_size > stderr_limit
+                        Process.kill("KILL", wait_thr.pid)
+                        @killed = true
+                      end
+                    end
+                  end
+                rescue EOFError => e
+                end
+              end
+            end
+          end
+        rescue IOError => e
+          puts "IOError: #{e}"
+        ensure
+          @stdout = captured_stdout.join
+          @stderr = captured_stderr.join
+          @status = wait_thr.value
+        end
+      end
+    end
   end
 end
diff --git a/lib/pipeline/util/runc_wrapper.rb b/lib/pipeline/util/runc_wrapper.rb
index 3198d48..1aaaff1 100644
--- a/lib/pipeline/util/runc_wrapper.rb
+++ b/lib/pipeline/util/runc_wrapper.rb
@@ -14,6 +14,8 @@ module Pipeline::Util
 
       run_cmd = ExternalCommand.new("bash -x -c 'ulimit -v #{memory_limit}; #{binary_path} --root root-state run #{container_id}'")
       run_cmd.timeout = 5
+      run_cmd.stdout_limit = 1024*1024
+      run_cmd.stderr_limit = 1024*1024
 
       kill_cmd = ExternalCommand.new("#{binary_path} --root root-state kill #{container_id} KILL")
 
diff --git a/test/pipeline/util/external_command_test.rb b/test/pipeline/util/external_command_test.rb
index 895623f..d540ce8 100644
--- a/test/pipeline/util/external_command_test.rb
+++ b/test/pipeline/util/external_command_test.rb
@@ -74,8 +74,54 @@ module Pipeline::Util
       cmd.suppress_output = true
       cmd.call
       refute cmd.success?
-      assert_equal 124, cmd.exit_status
+      assert_nil cmd.exit_status
     end
 
+    def test_stdout_hard_limit
+      cmd = ExternalCommand.new("/bin/sh -c \"echo 'hello'; sleep 4; echo 'world'\"")
+      cmd.suppress_output = true
+      cmd.stdout_limit = 4
+      cmd.call
+      refute cmd.success?
+      assert cmd.killed?
+      assert_equal "hello\n", cmd.stdout
+      assert_equal "", cmd.stderr
+    end
+
+    def test_hard_limits_for_real
+      verbose_command = 'perl -w -e \'$i = 0; while ($i<1000000) { print "x"; $i++ };  print "\n"; $i = 0; while ($i<10000) { print STDERR "y"; $i++ }\''
+
+      # Run without limits
+      cmd = ExternalCommand.new(verbose_command)
+      cmd.suppress_output = true
+      cmd.call
+      assert cmd.success?
+      refute cmd.killed?
+      assert_equal 1000001, cmd.stdout.size
+      assert_equal 10000, cmd.stderr.size
+
+      # limit stdout
+      cmd = ExternalCommand.new(verbose_command)
+      cmd.suppress_output = true
+      cmd.stdout_limit = 1000
+      cmd.call
+      refute cmd.success?
+      assert cmd.killed?
+      # should be truncated between block 1 and 2
+      assert cmd.stdout.size < 2000
+      assert_equal 0, cmd.stderr.size
+
+      # limit stderr
+      cmd = ExternalCommand.new(verbose_command)
+      cmd.suppress_output = true
+      cmd.stderr_limit = 1000
+      cmd.call
+      refute cmd.success?
+      assert cmd.killed?
+      assert_equal 1000001, cmd.stdout.size
+      assert cmd.stderr.size < 2000
+    end
+
+
   end
 end