module RegSort (#4)

For `shellcraft` using
`regsort` module in [python-pwntools](https://github.com/Gallopsled/pwntools/blob/beta/pwnlib/regsort.py)
Rewrite most part since original version is buggy.

lib/pwnlib/pwn.rb

@@ -7,6 +7,7 @@
 require 'pwnlib/constants/constants'
 require 'pwnlib/context'
 require 'pwnlib/dynelf'
+require 'pwnlib/reg_sort'
 
 require 'pwnlib/util/cyclic'
 require 'pwnlib/util/fiddling'

lib/pwnlib/reg_sort.rb

@@ -0,0 +1,147 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/context'
+
+module Pwnlib
+  # Do topological sort on register assignments.
+  module RegSort
+    # @note Do not create and call instance method here. Instead, call module method on {RegSort}.
+    module ClassMethods
+      # Sorts register dependencies.
+      #
+      # Given a dictionary of registers to desired register contents,
+      # return the optimal order in which to set the registers to
+      # those contents.
+      #
+      # The implementation assumes that it is possible to move from
+      # any register to any other register.
+      #
+      # @param [Hash<Symbol, String => Object>] in_out
+      #   Dictionary of desired register states.
+      #   Keys are registers, values are either registers or any other value.
+      # @param [Array<String>] all_regs
+      #   List of all possible registers.
+      #   Used to determine which values in +in_out+ are registers, versus
+      #   regular values.
+      # @option [Boolean] randomize
+      #   Randomize as much as possible about the order or registers.
+      #
+      # @return [Array]
+      #   Array of instructions, see examples for more details.
+      #
+      # @example
+      #   regs = %w(a b c d x y z)
+      #   regsort({a: 1, b: 2}, regs)
+      #   => [['mov', 'a', 1], ['mov', 'b', 2]]
+      #   regsort({a: 'b', b: 'a'}, regs)
+      #   => [['xchg', 'a', 'b']]
+      #   regsort({a: 1, b: 'a'}, regs)
+      #   => [['mov', 'b', 'a'], ['mov', 'a', 1]]
+      #   regsort({a: 'b', b: 'a', c: 3}, regs)
+      #   => [['mov', 'c', 3], ['xchg', 'a', 'b']]
+      #   regsort({a: 'b', b: 'a', c: 'b'}, regs)
+      #   => [['mov', 'c', 'b'], ['xchg', 'a', 'b']]
+      #   regsort({a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c'}, regs)
+      #   => [['mov', 'x', '1'],
+      #       ['mov', 'y', 'z'],
+      #       ['mov', 'z', 'c'],
+      #       ['xchg', 'a', 'b'],
+      #       ['xchg', 'b', 'c']]
+      #
+      # @note
+      #   Different from python-pwntools, we don't support +tmp+/+xchg+ options
+      #   because there's no such usage at all.
+      def regsort(in_out, all_regs, randomize: nil)
+        # randomize = context.randomize if randomize.nil?
+
+        # TODO(david942j): stringify_keys
+        in_out = in_out.map { |k, v| [k.to_s, v] }.to_h
+        # Drop all registers which will be set to themselves.
+        # Ex. {eax: 'eax'}
+        in_out.reject! { |k, v| k == v }
+
+        # Check input
+        if (in_out.keys - all_regs).any?
+          raise ArgumentError, format('Unknown register! Know: %p.  Got: %p', all_regs, in_out)
+        end
+
+        # Collapse constant values
+        #
+        # Ex. {eax: 1, ebx: 1} can be collapsed to {eax: 1, ebx: 'eax'}.
+        # +post_mov+ are collapsed registers, set their values in the end.
+        post_mov = in_out.group_by { |_, v| v }.each_value.with_object({}) do |list, hash|
+          list.sort!
+          first_reg, val = list.shift
+          # Special case for val.zero? because zeroify registers cost cheaper than mov.
+          next if list.empty? || all_regs.include?(val) || val.zero?
+          list.each do |reg, _|
+            hash[reg] = first_reg
+            in_out.delete(reg)
+          end
+        end
+
+        graph = in_out.dup
+        result = []
+
+        # Let's do the topological sort.
+        # so sad ruby 2.1 doesn't have +itself+...
+        deg = graph.values.group_by { |i| i }.map { |k, v| [k, v.size] }.to_h
+        graph.each_key { |k| deg[k] ||= 0 }
+
+        until deg.empty?
+          min_deg = deg.min_by { |_, v| v }[1]
+          break unless min_deg.zero? # remain are all cycles
+          min_pivs = deg.select { |_, v| v == min_deg }
+          piv = randomize ? min_pivs.sample : min_pivs.first
+          dst = piv.first
+          deg.delete(dst)
+          next unless graph.key?(dst) # Reach an end node.
+          deg[graph[dst]] -= 1
+          result << ['mov', dst, graph[dst]]
+          graph.delete(dst)
+        end
+
+        # Remain must be cycles.
+        graph.each_key do |reg|
+          cycle = check_cycle(reg, graph)
+          cycle.each_cons(2) do |d, s|
+            result << ['xchg', d, s]
+          end
+          cycle.each { |r| graph.delete(r) }
+        end
+
+        # Now assign those collapsed registers.
+        post_mov.sort.each do |dreg, sreg|
+          result << ['mov', dreg, sreg]
+        end
+
+        result
+      end
+
+      private
+
+      # Walk down the assignment list of a register,
+      # return the path walked if it is encountered again.
+      # @example
+      #   check_cycle('a', {'a' => 1}) #=> []
+      #   check_cycle('a', {'a' => 'a'}) #=> ['a']
+      #   check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'b', 'd' => 'a'}) #=> []
+      #   check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'd', 'd' => 'a'})
+      #   #=> ['a', 'b', 'c', 'd']
+      def check_cycle(reg, assignments)
+        check_cycle_(reg, assignments, [])
+      end
+
+      def check_cycle_(reg, assignments, path) # :nodoc:
+        target = assignments[reg]
+        path << reg
+        # No cycle, some other value (e.g. 1)
+        return [] unless assignments.key?(target)
+        # Found a cycle
+        return target == path.first ? path : [] if path.include?(target)
+        check_cycle_(target, assignments, path)
+      end
+    end
+    extend ClassMethods
+  end
+end

pwntools.gemspec

@@ -9,12 +9,12 @@ Gem::Specification.new do |s|
   s.summary       = 'pwntools'
   s.description   = <<-EOS
   Rewrite https://github.com/Gallopsled/pwntools in ruby.
-  Implement useful/easy function first,
+  Implement useful/easy functions first,
   try to be of ruby style and don't follow original pwntools everywhere.
   Would still try to have similar name whenever possible.
   EOS
   s.license       = 'MIT'
-  s.authors       = ['[email protected]']
+  s.authors       = ['[email protected]', '[email protected]']
   s.files         = Dir['lib/**/*.rb'] + %w(README.md Rakefile)
   s.test_files    = Dir['test/**/*']
   s.require_paths = ['lib']

test/reg_sort_test.rb

@@ -0,0 +1,41 @@
+# encoding: ASCII-8BIT
+require 'test_helper'
+require 'pwnlib/reg_sort'
+
+class RegSortTest < MiniTest::Test
+  include ::Pwnlib::RegSort::ClassMethods
+
+  def setup
+    @regs = %w(a b c d x y z)
+  end
+
+  def test_normal
+    assert_equal([['mov', 'a', 1], ['mov', 'b', 2]], regsort({ a: 1, b: 2 }, @regs))
+  end
+
+  def test_post_mov
+    assert_equal([['mov', 'a', 1], %w(mov b a)], regsort({ a: 1, b: 1 }, @regs))
+    assert_equal([%w(mov c a), ['mov', 'a', 1], %w(mov b a)], regsort({ a: 1, b: 1, c: 'a' }, @regs))
+  end
+
+  def test_pseudoforest
+    # only one connected component
+    assert_equal([%w(mov b a), ['mov', 'a', 1]], regsort({ a: 1, b: 'a' }, @regs))
+    assert_equal([['mov', 'c', 3], %w(xchg a b)], regsort({ a: 'b', b: 'a', c: 3 }, @regs))
+    assert_equal([%w(mov c b), %w(xchg a b)], regsort({ a: 'b', b: 'a', c: 'b' }, @regs))
+    assert_equal([%w(mov x 1), %w(mov y z), %w(mov z c), %w(xchg a b), %w(xchg b c)],
+                 regsort({ a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c' }, @regs))
+
+    # more than one connected components
+    assert_equal([%w(xchg a b), %w(xchg c d)], regsort({ a: 'b', b: 'a', c: 'd', d: 'c' }, @regs))
+    assert_equal([%w(mov c b), %w(mov d b), %w(mov z x), %w(xchg a b), %w(xchg x y)],
+                 regsort({ a: 'b', b: 'a', c: 'b', d: 'b', x: 'y', y: 'x', z: 'x' }, @regs))
+  end
+
+  def test_raise
+    err = assert_raises(ArgumentError) do
+      regsort({ a: 1 }, ['b'])
+    end
+    assert_match(/Unknown register!/, err.message)
+  end
+end