|
| 1 | +# encoding: ASCII-8BIT |
| 2 | + |
| 3 | +require 'pwnlib/context' |
| 4 | + |
| 5 | +module Pwnlib |
| 6 | + # Do topological sort on register assignments. |
| 7 | + module RegSort |
| 8 | + # @note Do not create and call instance method here. Instead, call module method on {RegSort}. |
| 9 | + module ClassMethods |
| 10 | + # Sorts register dependencies. |
| 11 | + # |
| 12 | + # Given a dictionary of registers to desired register contents, |
| 13 | + # return the optimal order in which to set the registers to |
| 14 | + # those contents. |
| 15 | + # |
| 16 | + # The implementation assumes that it is possible to move from |
| 17 | + # any register to any other register. |
| 18 | + # |
| 19 | + # @param [Hash<Symbol, String => Object>] in_out |
| 20 | + # Dictionary of desired register states. |
| 21 | + # Keys are registers, values are either registers or any other value. |
| 22 | + # @param [Array<String>] all_regs |
| 23 | + # List of all possible registers. |
| 24 | + # Used to determine which values in +in_out+ are registers, versus |
| 25 | + # regular values. |
| 26 | + # @option [Boolean] randomize |
| 27 | + # Randomize as much as possible about the order or registers. |
| 28 | + # |
| 29 | + # @return [Array] |
| 30 | + # Array of instructions, see examples for more details. |
| 31 | + # |
| 32 | + # @example |
| 33 | + # regs = %w(a b c d x y z) |
| 34 | + # regsort({a: 1, b: 2}, regs) |
| 35 | + # => [['mov', 'a', 1], ['mov', 'b', 2]] |
| 36 | + # regsort({a: 'b', b: 'a'}, regs) |
| 37 | + # => [['xchg', 'a', 'b']] |
| 38 | + # regsort({a: 1, b: 'a'}, regs) |
| 39 | + # => [['mov', 'b', 'a'], ['mov', 'a', 1]] |
| 40 | + # regsort({a: 'b', b: 'a', c: 3}, regs) |
| 41 | + # => [['mov', 'c', 3], ['xchg', 'a', 'b']] |
| 42 | + # regsort({a: 'b', b: 'a', c: 'b'}, regs) |
| 43 | + # => [['mov', 'c', 'b'], ['xchg', 'a', 'b']] |
| 44 | + # regsort({a: 'b', b: 'c', c: 'a', x: '1', y: 'z', z: 'c'}, regs) |
| 45 | + # => [['mov', 'x', '1'], |
| 46 | + # ['mov', 'y', 'z'], |
| 47 | + # ['mov', 'z', 'c'], |
| 48 | + # ['xchg', 'a', 'b'], |
| 49 | + # ['xchg', 'b', 'c']] |
| 50 | + # |
| 51 | + # @note |
| 52 | + # Different from python-pwntools, we don't support +tmp+/+xchg+ options |
| 53 | + # because there's no such usage at all. |
| 54 | + def regsort(in_out, all_regs, randomize: nil) |
| 55 | + # randomize = context.randomize if randomize.nil? |
| 56 | + |
| 57 | + # TODO(david942j): stringify_keys |
| 58 | + in_out = in_out.map { |k, v| [k.to_s, v] }.to_h |
| 59 | + # Drop all registers which will be set to themselves. |
| 60 | + # Ex. {eax: 'eax'} |
| 61 | + in_out.reject! { |k, v| k == v } |
| 62 | + |
| 63 | + # Check input |
| 64 | + if (in_out.keys - all_regs).any? |
| 65 | + raise ArgumentError, format('Unknown register! Know: %p. Got: %p', all_regs, in_out) |
| 66 | + end |
| 67 | + |
| 68 | + # Collapse constant values |
| 69 | + # |
| 70 | + # Ex. {eax: 1, ebx: 1} can be collapsed to {eax: 1, ebx: 'eax'}. |
| 71 | + # +post_mov+ are collapsed registers, set their values in the end. |
| 72 | + post_mov = in_out.group_by { |_, v| v }.each_value.with_object({}) do |list, hash| |
| 73 | + list.sort! |
| 74 | + first_reg, val = list.shift |
| 75 | + # Special case for val.zero? because zeroify registers cost cheaper than mov. |
| 76 | + next if list.empty? || all_regs.include?(val) || val.zero? |
| 77 | + list.each do |reg, _| |
| 78 | + hash[reg] = first_reg |
| 79 | + in_out.delete(reg) |
| 80 | + end |
| 81 | + end |
| 82 | + |
| 83 | + graph = in_out.dup |
| 84 | + result = [] |
| 85 | + |
| 86 | + # Let's do the topological sort. |
| 87 | + # so sad ruby 2.1 doesn't have +itself+... |
| 88 | + deg = graph.values.group_by { |i| i }.map { |k, v| [k, v.size] }.to_h |
| 89 | + graph.each_key { |k| deg[k] ||= 0 } |
| 90 | + |
| 91 | + until deg.empty? |
| 92 | + min_deg = deg.min_by { |_, v| v }[1] |
| 93 | + break unless min_deg.zero? # remain are all cycles |
| 94 | + min_pivs = deg.select { |_, v| v == min_deg } |
| 95 | + piv = randomize ? min_pivs.sample : min_pivs.first |
| 96 | + dst = piv.first |
| 97 | + deg.delete(dst) |
| 98 | + next unless graph.key?(dst) # Reach an end node. |
| 99 | + deg[graph[dst]] -= 1 |
| 100 | + result << ['mov', dst, graph[dst]] |
| 101 | + graph.delete(dst) |
| 102 | + end |
| 103 | + |
| 104 | + # Remain must be cycles. |
| 105 | + graph.each_key do |reg| |
| 106 | + cycle = check_cycle(reg, graph) |
| 107 | + cycle.each_cons(2) do |d, s| |
| 108 | + result << ['xchg', d, s] |
| 109 | + end |
| 110 | + cycle.each { |r| graph.delete(r) } |
| 111 | + end |
| 112 | + |
| 113 | + # Now assign those collapsed registers. |
| 114 | + post_mov.sort.each do |dreg, sreg| |
| 115 | + result << ['mov', dreg, sreg] |
| 116 | + end |
| 117 | + |
| 118 | + result |
| 119 | + end |
| 120 | + |
| 121 | + private |
| 122 | + |
| 123 | + # Walk down the assignment list of a register, |
| 124 | + # return the path walked if it is encountered again. |
| 125 | + # @example |
| 126 | + # check_cycle('a', {'a' => 1}) #=> [] |
| 127 | + # check_cycle('a', {'a' => 'a'}) #=> ['a'] |
| 128 | + # check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'b', 'd' => 'a'}) #=> [] |
| 129 | + # check_cycle('a', {'a' => 'b', 'b' => 'c', 'c' => 'd', 'd' => 'a'}) |
| 130 | + # #=> ['a', 'b', 'c', 'd'] |
| 131 | + def check_cycle(reg, assignments) |
| 132 | + check_cycle_(reg, assignments, []) |
| 133 | + end |
| 134 | + |
| 135 | + def check_cycle_(reg, assignments, path) # :nodoc: |
| 136 | + target = assignments[reg] |
| 137 | + path << reg |
| 138 | + # No cycle, some other value (e.g. 1) |
| 139 | + return [] unless assignments.key?(target) |
| 140 | + # Found a cycle |
| 141 | + return target == path.first ? path : [] if path.include?(target) |
| 142 | + check_cycle_(target, assignments, path) |
| 143 | + end |
| 144 | + end |
| 145 | + extend ClassMethods |
| 146 | + end |
| 147 | +end |
0 commit comments