From 6ba2b15ab2bbed8421a04bfb0792c19d0c8d1bc5 Mon Sep 17 00:00:00 2001
From: Spencer McIntyre <zeroSteiner@gmail.com>
Date: Fri, 13 May 2022 09:06:51 -0400
Subject: [PATCH] Overhaul retry_until_true specs

Co-authored-by: adfoster-r7 <60357436+adfoster-r7@users.noreply.github.com>
---
 spec/lib/msf/core/exploit/retry_spec.rb | 118 +++++++++++++++++++++---
 1 file changed, 107 insertions(+), 11 deletions(-)

diff --git a/spec/lib/msf/core/exploit/retry_spec.rb b/spec/lib/msf/core/exploit/retry_spec.rb
index d5f6c06bd1..5f3bfb62c4 100644
--- a/spec/lib/msf/core/exploit/retry_spec.rb
+++ b/spec/lib/msf/core/exploit/retry_spec.rb
@@ -10,27 +10,123 @@ RSpec.describe Msf::Exploit::Retry do
   end
 
   describe '#retry_until_true' do
-    let(:timeout) { Proc.new { false } }
-
     subject do
       create_exploit
     end
 
-    it 'does not call the block if the timeout is negative' do
-      expect { |b| subject.retry_until_true(timeout: -1, &b) }.to_not yield_control
+    # Quick workaround for Timecop not supporting Process.clock_gettime
+    # Timecop feature request: https://github.com/travisjeffery/timecop/issues/220
+    let(:mock_clock) do
+      clazz = Class.new do
+        def initialize
+          @current_time = 0
+        end
+
+        def now
+          @current_time
+        end
+
+        def time_travel(new_time)
+          @current_time = new_time
+        end
+      end
+
+      clazz.new
     end
 
-    it 'does call the block if the timeout is positive' do
-      expect { |b| subject.retry_until_true(timeout: 1, &b) }.to yield_with_no_args
+    def increment_time_by(seconds)
+        mock_clock.time_travel(mock_clock.now + seconds)
     end
 
-    it 'returns false when the timeout elapses' do
-      expect(subject.retry_until_true(timeout: 1, &timeout)).to eq false
+    def expect_sleep_calls(calls)
+      expect(subject).to have_received(:sleep).exactly(calls.length).times
+      calls.each do |value|
+        expect(subject).to have_received(:sleep).with(value)
+      end
     end
 
-    it 'returns the block result' do
-      result = Object.new
-      expect(subject.retry_until_true(timeout: 1, &Proc.new { result })).to eq result
+    before(:each) do
+      allow(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC, :second) { |*_args| mock_clock.now }
+      allow(subject).to receive(:sleep) { |seconds| increment_time_by(seconds) }
+    end
+
+    context 'when the timeout is negative' do
+      it 'does not yield control' do
+        result = nil
+        expect do |block|
+          result = subject.retry_until_true(timeout: -1, &block)
+        end.to_not yield_control
+        expect_sleep_calls([])
+        expect(result).to be false
+      end
+    end
+
+    context 'when the timeout is 0' do
+      it 'does not yield control' do
+        result = nil
+        expect do |block|
+          result = subject.retry_until_true(timeout: -1, &block)
+        end.to_not yield_control
+        expect_sleep_calls([])
+        expect(result).to be false
+      end
+    end
+
+    context 'when the timeout is 40' do
+      let(:timeout) { 40 }
+
+      it 'only yields once if the block takes longer than the allocated time' do
+        expect do |block|
+          subject.retry_until_true(timeout: timeout) do
+            block.to_proc.call
+            increment_time_by(timeout + 1)
+          end
+        end.to yield_control.once
+        expect_sleep_calls([])
+      end
+
+      it 'yields exponentially until the timeout has surpassed' do
+        result = nil
+        expect do |block|
+          result = subject.retry_until_true(timeout: timeout, &block)
+        end.to yield_control.exactly(5).times
+        expect_sleep_calls([2, 4, 8, 16, 10])
+        expect(result).to eq(false)
+      end
+
+      it 'returns the yielded value if it is immediately truthy' do
+        result = nil
+        expect do |block|
+          result = subject.retry_until_true(timeout: timeout) do
+            block.to_proc.call
+            :success
+          end
+        end.to yield_control.exactly(1).times
+        expect_sleep_calls([])
+        expect(result).to eq(:success)
+      end
+
+      it 'returns the yielded value if it is eventually truthy' do
+        result = nil
+        expect do |block|
+          result = subject.retry_until_true(timeout: timeout) do
+            block.to_proc.call
+            mock_clock.now >= 8
+          end
+        end.to yield_control.exactly(4).times
+        expect_sleep_calls([2, 4, 8])
+        expect(result).to eq(true)
+      end
+
+      it 'raises any unhandled exceptions' do
+        error = StandardError.new
+        expect do
+          subject.retry_until_true(timeout: timeout) do
+            raise error
+          end
+        end.to raise_error(error)
+        expect(subject).to_not have_received(:sleep)
+      end
     end
   end
 end