Enhance the constraints for argv and envp (#206)

* Enhance the constraints for `argv` and `envp`

lib/one_gadget/emulators/aarch64.rb

@@ -44,6 +44,24 @@ def argument(idx)
         registers["x#{idx}"]
       end
 
+      # @param [String | Lambda] obj
+      #  A lambda object or its string.
+      # @return [Hash{Integer => Lambda}, nil]
+      #  The corresponding stack (based on sp) that +obj+ used,
+      #  or nil if +obj+ doesn't use the stack.
+      # @example
+      #   get_corresponding_stack('sp+0x10')
+      #   #=> sp_based_stack
+      #   get_corresponding_stack('[sp+0x10]')
+      #   #=> sp_based_stack
+      #   get_corresponding_stack('x21')
+      #   #=> nil
+      def get_corresponding_stack(obj)
+        return nil unless obj.to_s.include?(sp)
+
+        sp_based_stack
+      end
+
       private
 
       def inst_add(dst, src, op2, mode = 'sxtw')
@@ -111,8 +129,8 @@ def inst_stp(reg1, reg2, dst)
         raise_unsupported('stp', reg1, reg2, dst) unless dst_l.obj == sp && dst_l.deref_count.zero?
 
         cur_top = dst_l.evaluate(eval_dict)
-        stack[cur_top] = registers[reg1]
-        stack[cur_top + size_t] = registers[reg2]
+        sp_based_stack[cur_top] = registers[reg1]
+        sp_based_stack[cur_top + size_t] = registers[reg2]
 
         registers[sp] += arg_to_lambda(dst).immi if dst.end_with?('!')
       end
@@ -125,7 +143,7 @@ def inst_str(src, dst, index = 0)
         # Only stores on stack.
         if dst_l.obj == sp && dst_l.deref_count.zero?
           cur_top = dst_l.evaluate(eval_dict)
-          stack[cur_top] = registers[src]
+          sp_based_stack[cur_top] = registers[src]
         else
           # Unlike the stack case, don't know where to save the value.
           # Simply add a constraint.

lib/one_gadget/emulators/amd64.rb

@@ -16,7 +16,7 @@ def bits
 
       # Instantiate an {Amd64} object.
       def initialize
-        super(OneGadget::ABI.amd64, 'rsp', 'rip')
+        super(OneGadget::ABI.amd64, 'rsp', 'rbp', 'rip')
       end
 
       # Return the argument value of calling a function.

lib/one_gadget/emulators/i386.rb

@@ -16,7 +16,7 @@ def bits
 
       # Instantiate an {I386} object.
       def initialize
-        super(OneGadget::ABI.i386, 'esp', 'eip')
+        super(OneGadget::ABI.i386, 'esp', 'ebp', 'eip')
       end
 
       # Get function call arguments.
@@ -29,7 +29,7 @@ def initialize
       # @return [Lambda, Integer]
       def argument(idx)
         cur_top = registers['esp'].evaluate('esp' => 0)
-        stack[cur_top + idx * 4]
+        sp_based_stack[cur_top + idx * 4]
       end
     end
   end

lib/one_gadget/emulators/processor.rb

@@ -9,7 +9,7 @@ module Emulators
     # Base class of a processor.
     class Processor
       attr_reader :registers # @return [Hash{String => OneGadget::Emulators::Lambda}] The current registers' state.
-      attr_reader :stack # @return [Hash{Integer => OneGadget::Emulators::Lambda}] The content on stack.
+      attr_reader :sp_based_stack # @return [Hash{Integer => OneGadget::Emulators::Lambda}] Stack content based on sp.
       attr_reader :sp # @return [String] Stack pointer.
       attr_reader :pc # @return [String] Program counter.
 
@@ -22,7 +22,7 @@ def initialize(registers, sp)
         @registers = registers.map { |reg| [reg, to_lambda(reg)] }.to_h
         @sp = sp
         @constraints = []
-        @stack = Hash.new do |h, k|
+        @sp_based_stack = Hash.new do |h, k|
           h[k] = OneGadget::Emulators::Lambda.new(sp).tap do |lmda|
             lmda.immi = k
             lmda.deref!
@@ -94,6 +94,16 @@ def constraints
         cons.map { |type, obj| type == :writable ? "writable: #{obj}" : obj }.sort
       end
 
+      # Method need to be implemented in inheritors.
+      #
+      # @param [String | Lambda] obj
+      #  A lambda object or its string.
+      # @return [Hash{Integer => Lambda}, nil]
+      #  The corresponding stack that +obj+ used,
+      #  or nil if +obj+ doesn't use the stack.
+      def get_corresponding_stack(obj); raise NotImplementedError
+      end
+
       private
 
       def check_register!(reg)

lib/one_gadget/emulators/x86.rb

@@ -9,10 +9,20 @@ module OneGadget
   module Emulators
     # Super class for amd64 and i386 processor.
     class X86 < Processor
+      attr_reader :bp # @return [String] Stack base register.
+      attr_reader :bp_based_stack # @return [Hash{Integer => OneGadget::Emulators::Lambda}] Stack content based on bp.
+
       # Constructor for a x86 processor.
-      def initialize(registers, sp, pc)
+      def initialize(registers, sp, bp, pc)
         super(registers, sp)
+        @bp = bp
         @pc = pc
+        @bp_based_stack = Hash.new do |h, k|
+          h[k] = OneGadget::Emulators::Lambda.new(bp).tap do |lmda|
+            lmda.immi = k
+            lmda.deref!
+          end
+        end
       end
 
       # Process one command.
@@ -50,23 +60,45 @@ def instructions
         ]
       end
 
+      # @param [String | Lambda] obj
+      #  A lambda object or its string.
+      # @return [Hash{Integer => Lambda}, nil]
+      #  The corresponding stack (based on esp/rsp or ebp/rbp) that +obj+ used,
+      #  or nil if +obj+ doesn't use the stack.
+      # @example
+      #   get_corresponding_stack('rsp+0x10')
+      #   #=> sp_based_stack
+      #   get_corresponding_stack('rbp-0x10')
+      #   #=> bp_based_stack
+      #   get_corresponding_stack('[rbp-0x10]')
+      #   #=> bp_based_stack
+      #   get_corresponding_stack('rax')
+      #   #=> nil
+      def get_corresponding_stack(obj)
+        if obj.to_s.include?(sp)
+          sp_based_stack
+        elsif obj.to_s.include?(bp)
+          bp_based_stack
+        end
+      end
+
       private
 
       def inst_mov(dst, src)
         src = arg_to_lambda(src)
         if register?(dst)
           registers[dst] = src
-        else
-          # Just ignore strange case...
-          # TODO(david942j): #120
-          return add_writable(dst) unless dst.include?(sp)
-
-          dst = arg_to_lambda(dst)
-          return if dst.deref_count != 1 # should not happen
-
-          dst.ref!
-          stack[dst.evaluate(eval_dict)] = src
+          return
         end
+        dst = arg_to_lambda(dst)
+        add_writable(dst.to_s)
+        # TODO: Is it possible that only considering sp and bp is not enough?
+        # If it is, we need to record every memory access
+        stack = get_corresponding_stack(dst)
+        return if stack.nil? || dst.deref_count != 1
+
+        dst.ref!
+        stack[dst.evaluate(eval_dict)] = src
       end
 
       # This instruction moves 128bits.
@@ -76,7 +108,7 @@ def inst_movaps(dst, src)
         off = dst.evaluate(eval_dict)
         @constraints << [:raw, "#{sp} & 0xf == #{0x10 - off & 0xf}"]
         (128 / self.class.bits).times do |i|
-          stack[off + i * size_t] = src[i]
+          sp_based_stack[off + i * size_t] = src[i]
         end
       end
 
@@ -94,7 +126,7 @@ def inst_movq(dst, src)
         dst, src = check_xmm_sp(dst, src) { raise_unsupported('movq', dst, src) }
         off = src.evaluate(eval_dict)
         (64 / self.class.bits).times do |i|
-          dst[i] = stack[off + i * size_t]
+          dst[i] = sp_based_stack[off + i * size_t]
         end
       end
 
@@ -104,7 +136,7 @@ def inst_movhps(dst, src)
         dst, src = check_xmm_sp(dst, src) { raise_unsupported('movhps', dst, src) }
         off = src.evaluate(eval_dict)
         (64 / self.class.bits).times do |i|
-          dst[i + 64 / self.class.bits] = stack[off + i * size_t]
+          dst[i + 64 / self.class.bits] = sp_based_stack[off + i * size_t]
         end
       end
 
@@ -147,7 +179,7 @@ def inst_push(val)
         cur_top = registers[sp].evaluate(eval_dict)
         raise Error::InstructionArgumentError, "Corrupted stack pointer: #{cur_top}" unless cur_top.is_a?(Integer)
 
-        stack[cur_top] = val
+        sp_based_stack[cur_top] = val
       end
 
       def inst_xor(dst, src)
@@ -214,6 +246,10 @@ def to_lambda(reg)
           OneGadget::Emulators::Lambda.new(i.zero? ? "#{cast}#{reg}" : "#{cast}(#{reg} >> #{self.class.bits * i})")
         end
       end
+
+      def eval_dict
+        { sp => 0, bp => 0 }
+      end
     end
   end
 end

lib/one_gadget/fetchers/base.rb

@@ -100,7 +100,8 @@ def resolve_execve_args(processor, arg0, arg1, arg2, allow_null_argv: true)
         # arg1 == NULL || [arg1] == NULL
         # arg2 == NULL || [arg2] == NULL || arg[2] == envp
         cons = processor.constraints
-        cons << check_execve_arg(processor, arg1, allow_null_argv)
+        con = check_argv(processor, arg1, allow_null_argv)
+        cons << con unless con.nil?
         return nil unless cons.all?
 
         envp = 'environ'
@@ -112,19 +113,79 @@ def resolve_execve_args(processor, arg0, arg1, arg2, allow_null_argv: true)
         { constraints: cons, envp: envp }
       end
 
-      # arg == NULL || [arg] == NULL
-      def check_execve_arg(processor, arg, allow_null)
-        if arg.start_with?(processor.sp) # arg = sp+<num>
-          # in this case, the only chance is [sp+<num>] == NULL
-          num = Integer(arg[processor.sp.size..-1])
-          slot = processor.stack[num].to_s
-          return if global_var?(slot)
+      def check_argv(processor, arg, allow_null)
+        lmda = OneGadget::Emulators::Lambda.parse(arg)
 
-          "#{slot} == NULL"
-        elsif allow_null
-          "[#{arg}] == NULL || #{arg} == NULL"
+        if lmda.deref_count.zero? && OneGadget::ABI.stack_register?(lmda.obj)
+          return check_stack_argv(processor, lmda, allow_null)
+        end
+
+        check_nonstack_argv(arg, allow_null)
+      end
+
+      def check_stack_argv(processor, lmda, allow_null)
+        stack = processor.get_corresponding_stack(lmda.obj)
+        argv = (0..3).map { |i| stack[lmda.immi + processor.class.bits / 8 * i].to_s }
+
+        # if argv is already valid, no constraints are needed! (but probably won't happen :p)
+        return if argv_already_valid?(argv)
+
+        return generate_argv_with_sh(argv) if global_var?(argv[0])
+
+        generate_argv_without_sh(argv, allow_null)
+      end
+
+      def argv_already_valid?(argv)
+        argv[0] == '0' || (global_var?(argv[0]) && argv[1] == '0')
+      end
+
+      def generate_argv_with_sh(argv)
+        # argv[0] is not controlled by the user, argv[0] probably is "/bin/sh" or "sh" (but actually, the content of
+        # argv[0] doesn't quite matter, just need to make sure it's readable)
+        # So far (I checked glibc 2.37), we can make argv to be {"/bin/sh", sth, NULL} or {"sh", "-c", sth, NULL}
+        # TODO: We need to update this when the above assumption is no longer true
+        if argv[2] == '0' && !global_var?(argv[1])
+          "#{argv[1]} == NULL || {\"/bin/sh\", #{argv[1]}, NULL} is a valid argv"
         else
-          "[#{arg}] == NULL"
+          argv_gte3 = argv[3] == '0' ? 'NULL' : "#{argv[3]}, ..."
+          if global_var?(argv[1])
+            "{\"sh\", \"-c\", #{argv[2]}, #{argv_gte3}} is a valid argv"
+          else
+            "#{argv[1]} == NULL || {\"sh\", #{argv[1]}, #{argv[2]}, #{argv_gte3}} is a valid argv"
+          end
+        end
+      end
+
+      def generate_argv_without_sh(argv, allow_null)
+        argv_cons = "{#{argv[0]}"
+        (1..argv.length - 1).each do |i|
+          if argv[i] == '0'
+            argv_cons += ', NULL'
+            break
+          elsif i == 1 && global_var?(argv[i])
+            # TODO: We probably need to get the true content of the global variable for a more accurate result
+            argv_cons += ', "-c"'
+          else
+            argv_cons += ", #{argv[i]}"
+          end
+        end
+        argv_cons += ', ...' unless argv_cons.end_with?('NULL')
+        argv_cons += '} is a valid argv'
+
+        if allow_null && argv.all? { |a| OneGadget::ABI.stack_register?(a) }
+          # If libc writes something into the stack, arg cannot be NULL.
+          # TODO: Find a better way to check can arg be NULL
+          "#{arg} == NULL || #{argv[0]} == NULL || #{argv_cons}"
+        else
+          "#{argv[0]} == NULL || #{argv_cons}"
+        end
+      end
+
+      def check_nonstack_argv(arg, allow_null)
+        if allow_null
+          "[#{arg}] == NULL || #{arg} == NULL || #{arg} is a valid argv"
+        else
+          "[#{arg}] == NULL || #{arg} is a valid argv"
         end
       end
 
@@ -134,8 +195,16 @@ def check_envp(processor, arg)
         # If it starts with [[ but not a global var, drop it.
         return global_var?(arg) if arg.start_with?('[[')
 
-        # normal
-        cons = check_execve_arg(processor, arg, true)
+        lmda = OneGadget::Emulators::Lambda.parse(arg)
+        if lmda.deref_count.zero? && OneGadget::ABI.stack_register?(lmda.obj)
+          # I haven't see this case after some tests, but just in case :)
+          stack = processor.get_corresponding_stack(lmda.obj)
+          envp = (0..3).map { |i| stack[lmda.immi + processor.class.bits / 8 * i].to_s }
+          # TODO: Handle the case when libc will write something into envp
+          cons = global_var?(envp[0]) ? nil : "#{arg} == NULL || {#{envp.join(', ')}, ...} is a valid envp"
+        else
+          cons = "[#{arg}] == NULL || #{arg} == NULL || #{arg} is a valid envp"
+        end
         return nil if cons.nil?
 
         yield cons

lib/one_gadget/gadget.rb

@@ -71,6 +71,8 @@ def score
       # Expr: <CAST>?<Identity> == NULL
       # Expr: <REG> & 0xf == <IMM>
       # Expr: (s32)[<Identity>] <= 0
+      # Expr: .+ is a valid argv
+      # Expr: .+ is a valid envp
       # Expr: <Expr> || <Expr>
       def calculate_score(expr)
         return expr.split(' || ').map(&method(:calculate_score)).max if expr.include?(' || ')
@@ -81,6 +83,7 @@ def calculate_score(expr)
         when /^writable/ then calculate_writable_score(expr.sub('writable: ', ''))
         when / == NULL$/ then calculate_null_score(expr.sub(' == NULL', ''))
         when / <= 0$/ then calculate_null_score(expr.sub(' <= 0', ''))
+        when / is a valid (argv|envp)$/ then 0.2 # This usually means the register has to be a readable pointer.
         end
       end
 

spec/cli_spec.rb

@@ -80,7 +80,7 @@
     skip_unless_objdump
 
     expect { described_class.work(%w[--force --raw --level 1] + [libc_file]) }.to output(<<-EOS).to_stdout
-324293 324386 939679 940120 940127 940131 1090444 1090456
+324279 324286 324293 324386 939679 940120 940127 940131 1090444 1090456
     EOS
   end