Introduced scoring on constraints (#67)

README.md

@@ -184,10 +184,6 @@ $ one_gadget /lib/x86_64-linux-gnu/libc.so.6 --level 1
 
 ```bash
 $ one_gadget spec/data/aarch64-libc-2.27.so
-# 0x3f15c	execve("/bin/sh", sp+0x70, environ)
-# constraints:
-#   x3+0x7c0 == NULL
-#
 # 0x3f16c	execve("/bin/sh", sp+0x70, environ)
 # constraints:
 #   x3 == NULL
@@ -200,14 +196,6 @@ $ one_gadget spec/data/aarch64-libc-2.27.so
 # constraints:
 #   [x21] == NULL || x21 == NULL
 #
-# 0x63e7c	execl("/bin/sh", "sh", x2+0x7c8)
-# constraints:
-#   x2+0x7c8 == NULL
-#
-# 0x63e88	execl("/bin/sh", x1+0x7c0)
-# constraints:
-#   x1+0x7c0 == NULL
-#
 # 0x63e90	execl("/bin/sh", x1)
 # constraints:
 #   x1 == NULL

lib/one_gadget/abi.rb

@@ -37,5 +37,13 @@ module OneGadget
     def all
       amd64 + aarch64
     end
+
+    # Checks if the register is a stack-related pointer.
+    # @param [String] reg
+    #   Register's name.
+    # @return [Boolean]
+    def stack_register?(reg)
+      %w[esp ebp rsp rbp sp x29].include?(reg)
+    end
   end
 end

lib/one_gadget/emulators/lambda.rb

@@ -109,16 +109,13 @@ module OneGadget
         def parse(argument, predefined: {})
           arg = argument.dup
           return Integer(arg) if OneGadget::Helper.integer?(arg)
+          # nested []
+          return parse(arg[1...arg.rindex(']')], predefined: predefined).deref if arg[0] == '['
 
-          deref_count = 0
-          if arg[0] == '[' # a little hack because there should nerver something like +[[rsp+1]+2]+ to parse.
-            arg = arg[1...arg.rindex(']')]
-            deref_count = 1
-          end
           base, disp = mem_obj(arg)
           obj = predefined[base] || Lambda.new(base)
           obj += disp unless disp.zero?
-          deref_count.zero? ? obj : obj.deref
+          obj
         end
 
         private

lib/one_gadget/gadget.rb

@@ -1,4 +1,5 @@
 require 'one_gadget/abi'
+require 'one_gadget/emulators/lambda'
 require 'one_gadget/error'
 
 module OneGadget
@@ -7,11 +8,11 @@ module OneGadget
     # Information of a gadget.
     class Gadget
       # @return [Integer] The gadget's address offset.
-      attr_accessor :offset
+      attr_reader :offset
       # @return [Array<String>] The constraints need for this gadget.
-      attr_accessor :constraints
+      attr_reader :constraints
       # @return [String] The final result of this gadget.
-      attr_accessor :effect
+      attr_reader :effect
 
       # Initialize method of {Gadget} instance.
       # @param [Integer] offset The relative address offset of this gadget.
@@ -39,6 +40,41 @@ module OneGadget
         end
         str + "\n"
       end
+
+      # @return [Integer]
+      #   The difficulty of constraints.
+      def score
+        @score ||= constraints.reduce(0) { |s, c| s + calculate_score(c) }
+      end
+
+      private
+
+      # REG: OneGadget::ABI.all
+      # IMM: [+-]0x[\da-f]+
+      # Identity: <REG><IMM>?
+      # Identity: [<Identity>]
+      # Expr: <REG> is the GOT address of libc
+      # Expr: <Identity> == NULL
+      # Expr: <Expr> || <Expr>
+      def calculate_score(cons)
+        return cons.split(' || ').map(&method(:calculate_score)).min if cons.include?(' || ')
+        return 1 if cons.include?('GOT address')
+
+        expr = cons.gsub(' == NULL', ' == 0')
+        # raise Error::ArgumentError, cons unless expr.end_with?(' == 0')
+
+        identity = expr.slice(0...expr.rindex(' == 0'))
+        # Thank God we are already able to parse this
+        lmda = OneGadget::Emulators::Lambda.parse(identity)
+        # raise Error::ArgumentError, cons unless OneGadget::ABI.all.include?(lmda.obj)
+        # rax == 0 is easy; rax + 0x10 == 0 is hard.
+        return lmda.immi.zero? ? 1 : 3 if lmda.deref_count.zero?
+
+        # Stack frame registers has difficulty 1
+        # when lmda.deref_count == 1
+        base = OneGadget::ABI.stack_register?(lmda.obj) ? 0 : 1
+        lmda.deref_count + base
+      end
     end
 
     # Define class methods here.

lib/one_gadget/one_gadget.rb

@@ -62,9 +62,9 @@ module OneGadget
       return [] if gadgets.empty?
       return gadgets if level > 0 # currently only supports level > 0 or not
 
-      # remain gadgets with the fewest constraints
-      best = gadgets.map { |g| g.constraints.size }.min
-      gadgets.select { |g| g.constraints.size == best }
+      # remain gadgets with the lowest score
+      best = gadgets.map(&:score).min
+      gadgets.select { |g| g.score == best }
     end
   end
 end

spec/emulators/lambda_spec.rb

@@ -49,5 +49,8 @@ describe OneGadget::Emulators::Lambda do
 
     predefined = { 'rsp' => OneGadget::Emulators::Lambda.new('rsp') + 0x10 }
     expect(OneGadget::Emulators::Lambda.parse('rsp+0x20', predefined: predefined).to_s).to eq 'rsp+0x30'
+
+    # Nested []
+    expect(OneGadget::Emulators::Lambda.parse('[[rsp+0x33]]').to_s).to eq '[[rsp+0x33]]'
   end
 end

spec/gadget_spec.rb

@@ -6,8 +6,8 @@ require 'one_gadget/helper'
 describe OneGadget::Gadget do
   before(:all) do
     @build_id = 'fake_id'
-    OneGadget::Gadget.add(@build_id, 0x1234, constraints: ['[rsp+0x30] == NULL', 'rax == 0'],
-                                             effect: 'execve("/bin/sh", rsp+0x30, rax)')
+    described_class.add(@build_id, 0x1234, constraints: ['[rsp+0x30] == NULL', 'rax == 0'],
+                                           effect: 'execve("/bin/sh", rsp+0x30, rax)')
   end
 
   after(:all) do
@@ -23,6 +23,46 @@ constraints:
     EOS
   end
 
+  context 'score' do
+    def new(cons)
+      OneGadget::Gadget::Gadget.new(0, constraints: cons)
+    end
+
+    it 'empty' do
+      expect(new([]).score).to be_zero
+    end
+
+    it 'level 1' do
+      expect(new(['[rsp+0x30] == NULL']).score).to be 1
+      expect(new(['[esp+0x34] == NULL']).score).to be 1
+      expect(new(['[rbp+0x30] == NULL']).score).to be 1
+      expect(new(['rax == NULL']).score).to be 1
+      expect(new(['x1 == NULL']).score).to be 1
+      expect(new(['[rsi] == NULL || rsi == NULL']).score).to be 1
+      expect(new(['ebx is the GOT address of libc']).score).to be 1
+      expect(new(['[rsi] == NULL || ebx is the GOT address of libc']).score).to be 1
+    end
+
+    it 'level 2' do
+      expect(new(['[[sp+0x38]] == NULL']).score).to be 2
+      expect(new(['[rax] == NULL']).score).to be 2
+      expect(new(['[rsi] == NULL']).score).to be 2
+      expect(new(['[x4+0xad0] == NULL']).score).to be 2
+    end
+
+    it 'level 3' do
+      expect(new(['x4+0xad0 == NULL']).score).to be 3
+      expect(new(['[[x4+0xad0]] == NULL']).score).to be 3
+    end
+
+    it 'more than one' do
+      expect(new([
+                   'rax == NULL',
+                   'rbx+0x333 == NULL'
+                 ]).score).to be 4 # sum
+    end
+  end
+
   it 'remote' do
     id = 'remote_has_this'
     allow(OneGadget::Helper).to receive(:remote_builds).and_return([id])

spec/one_gadget_aarch64_spec.rb

@@ -14,13 +14,12 @@ describe 'one_gadget_aarch64' do
 
     it 'libc-2.24' do
       path = data_path('aarch64-libc-2.24.so')
-      expect(OneGadget.gadgets(file: path, force_file: true)).to eq [0x3c928, 0x3c930, 0x3c970, 0x61484, 0x61488]
+      expect(OneGadget.gadgets(file: path, force_file: true)).to eq [0x3c930, 0x3c970]
     end
 
     it 'libc-2.27' do
       path = data_path('aarch64-libc-2.27.so')
-      expect(OneGadget.gadgets(file: path, force_file: true))
-        .to eq [0x3f15c, 0x3f16c, 0x3f184, 0x3f1a8, 0x63e7c, 0x63e88, 0x63e90]
+      expect(OneGadget.gadgets(file: path, force_file: true)).to eq [0x3f16c, 0x3f184, 0x3f1a8, 0x63e90]
     end
   end