fisk 1.0.0 → 2.3.0

data/lib/fisk.rb

@@ -1,18 +1,176 @@
 # frozen_string_literal: true
 
-require "fisk/machine/encoding"
-require "fisk/machine"
+require "stringio"
+require "set"
+
+require "fisk/instructions"
+require "fisk/basic_block"
+require "fisk/errors"
+require "fisk/version"
 
 class Fisk
+  module OperandPredicates
+    def unresolved?;        false; end
+    def register?;          false; end
+    def temp_register?;     false; end
+    def extended_register?; false; end
+    def memory?;            false; end
+    def rip?;               false; end
+    def immediate?;         false; end
+    def absolute_location?; false; end
+  end
+
+  class Operand
+    include OperandPredicates
+
+    def works? op
+      self.type == op.type
+    end
+
+    def rex_value
+      value >> 3
+    end
+
+    def op_value
+      value
+    end
+  end
+
+  class ValueOperand < Operand
+    attr_reader :value
+    alias :type :value
+
+    def initialize value
+      @value = value
+    end
+  end
+
   module Registers
-    class Register < Struct.new(:name, :type, :value)
-      def works? type
+    class Register < Operand
+      attr_reader :name, :type, :value
+
+      def initialize name, type, value
+        @name = name
+        @type = type
+        @value = value
+      end
+
+      def works? op
+        type = op.type
         type == self.name || type == self.type
       end
 
-      def unknown_label?; false; end
+      def op_value
+        value
+      end
+
+      def extended_register?
+        @value > 7 || EXTENDED_R8.include?(self)
+      end
+
+      def register?; true; end
+
+      def to_register
+        self
+      end
+    end
+
+    class Rip < Operand
+      attr_accessor :displacement
+
+      def initialize displacement
+        @displacement = displacement
+      end
+
+      def works? op
+        type = op.type
+        type == "m64" || type == "m"
+      end
+
+      def unresolved?
+        @displacement.is_a?(Fisk::UnknownLabel)
+      end
+
+      def memory?; true; end
+      def rip?; true; end
+
+      def rex_value
+        0x0
+      end
+
+      def op_value
+        0x0
+      end
+    end
+
+    class Temp < Operand
+      attr_reader :name, :type
+      attr_accessor :register, :start_point, :end_point
+
+      def initialize name, type
+        @name        = name
+        @type        = type
+        @start_point = nil
+        @end_point   = nil
+      end
+
+      def temp_register?; true; end
+      def register?; true; end
+
+      def op_value
+        reg.op_value
+      end
+
+      def extended_register?
+        reg.extended_register?
+      end
+
+      def rex_value
+        reg.rex_value
+      end
+
+      def value
+        reg.value
+      end
+
+      def to_register
+        self
+      end
+
+      def reg
+        unless @register
+          raise Errors::UnassignedRegister, "Temporary register #{name} hasn't been assigned a real register"
+        end
+
+        @register
+      end
     end
 
+    AL  = Register.new "al", "r8", 0
+    CL  = Register.new "cl", "r8", 1
+    DL  = Register.new "dl", "r8", 2
+    BL  = Register.new "bl", "r8", 3
+
+    SPL = Register.new "spl", "r8", 4
+    BPL = Register.new "bpl", "r8", 5
+    SIL = Register.new "sil", "r8", 6
+    DIL = Register.new "dil", "r8", 7
+    EXTENDED_R8 = [SPL, BPL, SIL, DIL]
+
+    AH = Register.new "ah", "r8", 4
+    CH = Register.new "ch", "r8", 5
+    DH = Register.new "dh", "r8", 6
+    BH = Register.new "bh", "r8", 7
+
+    AX = Register.new "ax", "r16", 0
+    CX = Register.new "cx", "r16", 1
+    DX = Register.new "dx", "r16", 2
+    BX = Register.new "bx", "r16", 3
+    SP = Register.new "sp", "r16", 4
+    BP = Register.new "bp", "r16", 5
+    SI = Register.new "si", "r16", 6
+    DI = Register.new "di", "r16", 7
+
     EAX = Register.new "eax", "r32", 0
     ECX = Register.new "ecx", "r32", 1
     EDX = Register.new "edx", "r32", 2
@@ -33,177 +191,521 @@ class Fisk
     8.times do |i|
       i += 8
       const_set "R#{i}", Register.new("r#{i}", "r64", i)
+      const_set "R#{i}D", Register.new("r#{i}d", "r32", i)
+      const_set "R#{i}W", Register.new("r#{i}w", "r16", i)
+      const_set "R#{i}B", Register.new("r#{i}b", "r8", i)
     end
-  end
 
-  class Operand < Struct.new(:value)
-    def type; value; end
-    def works? type; self.type == type; end
-    def unknown_label?; false; end
+    # List of *caller* saved registers for the C calling convention
+    CALLER_SAVED = [ RDI, RSI, RDX, RCX, R8, R9, R10, R11 ]
+
+    # List of *callee* saved registers for the C calling convention
+    CALLEE_SAVED = [ RBX, RSP, RBP, R12, R13, R14, R15 ]
   end
 
-  class M64 < Operand
-    def type
-      "m64"
+  class Memory < Operand
+    attr_reader :register, :displacement
+
+    def initialize register, displacement
+      @register     = register
+      @displacement = displacement
+    end
+
+    def extended_register?
+      @register.extended_register?
+    end
+
+    def value
+      @register.value
+    end
+
+    def + displacement
+      self.class.new(register, self.displacement + displacement)
+    end
+
+    def memory?; true; end
+
+    def rex_value
+      @register.rex_value
     end
   end
 
-  def m64 x
-    M64.new x
+  [8, 16, 32, 64].each do |size|
+    class_eval <<~eostr, __FILE__, __LINE__ + 1
+      class M#{size} < Memory
+        def type
+          "m#{size}"
+        end
+
+        def size; #{size}; end
+      end
+
+      def m#{size} val, displacement = 0; M#{size}.new(val, displacement); end
+    eostr
   end
 
-  class Imm8 < Operand
+  class M < Memory
     def type
-      "imm8"
+      "m"
     end
   end
 
-  class Imm32 < Operand
-    def type
-      "imm32"
-    end
+  def rip displacement = 0
+    Registers::Rip.new(displacement)
   end
 
-  class Imm64 < Operand
-    def type
-      "imm64"
+  def m x, displacement = 0
+    M.new x, displacement
+  end
+
+  class ImmediateOperand < ValueOperand
+    def immediate?; true; end
+
+    def works? insn_op
+      insn_op.immediate? && size <= insn_op.size
     end
   end
 
-  class Rel8 < Operand
+  # Define all immediate value methods of different sizes
+  [8, 16, 32, 64].each do |size|
+    class_eval <<~eostr, __FILE__, __LINE__ + 1
+      class Imm#{size} < ImmediateOperand
+        def type
+          "imm#{size}"
+        end
+
+        def size; #{size}; end
+      end
+
+      def imm#{size} val; Imm#{size}.new(val.to_i); end
+    eostr
+  end
+
+  class Relative < ValueOperand; end
+
+  class Rel8 < Relative
     def type
       "rel8"
     end
   end
 
-  class Rel32 < Operand
+  class Rel32 < Relative
     def type
       "rel32"
     end
   end
 
-  class MOffs64 < Operand
+  class MOffs64 < ValueOperand
     def type
       "moffs64"
     end
   end
 
-  class Lit < Operand
+  class Lit < ValueOperand
     def type
       value.to_s
     end
   end
 
-  attr_reader :position
+  class UnknownLabel < Operand
+    attr_reader :name
 
-  def initialize
-    @instructions = []
-    @labels = {}
-    @position = 0
-  end
+    def initialize name
+      @name = name
+    end
 
-  class UnknownLabel < Struct.new(:name, :assembler, :insn)
-    def works? type
+    def works? op
+      type = op.type
       type == "rel32"
     end
 
-    def unknown_label?; true; end
-
-    def value
-      label = assembler.label_for(name)
-      label.position - (insn.position + insn.bytesize)
-    end
+    def unresolved?; true; end
   end
 
-  class Label < Struct.new(:name, :position)
+  module InstructionPredicates
+    def retry?;              false; end
+    def label?;              false; end
+    def jump?;               false; end
+    def has_temp_registers?; false; end
+    def comment?;            false; end
+    def lazy?;               false; end
   end
 
-  def label_for name
-    @labels.fetch name
-  end
+  class Label < Struct.new(:name)
+    include InstructionPredicates
 
-  def label name
-    UnknownLabel.new(name, self)
+    def label?; true; end
   end
 
-  def make_label name
-    @labels[name] = Label.new(name, position)
-  end
+  class Comment < Struct.new(:message)
+    include InstructionPredicates
 
-  Registers.constants.grep(/^[A-Z0-9]*$/).each do |const|
-    val = Registers.const_get const
-    define_method(const.downcase) { val }
+    def comment?; true; end
   end
 
   class Instruction
-    attr_reader :position
+    include InstructionPredicates
 
-    def initialize insn, operands, position
+    def initialize insn, form, operands
       @insn     = insn
+      @form     = form
       @operands = operands
-      @position = position
+    end
+
+    def has_temp_registers?
+      @operands.any?(&:temp_register?)
+    end
+
+    def temp_registers
+      @operands.find_all(&:temp_register?)
     end
 
     def encodings
-      @insn.encodings
+      @form.encodings
     end
 
-    def encode buffer
-      encoding = @insn.encodings.first
+    def encode buffer, labels
+      encoding = @form.encodings.first
       encoding.encode buffer, @operands
     end
+  end
+
+  class UnresolvedRIPInstruction
+    include InstructionPredicates
+
+    def initialize insn, form, operands
+      @insn      = insn
+      @form      = form
+      @operands  = operands
+      @retry     = false
+      @saved_pos = 0
+    end
+
+    def retry?; true; end
+
+    def encode buffer, labels
+      # Use dummy values for any unresolvable operands
+      operands = @operands.map do |op|
+        if op.rip? && op.unresolved?
+          # Try resolving the operands
+          if labels.key?(op.displacement.name)
+            Registers::Rip.new labels[op.displacement.name] - @saved_pos
+          else
+            Registers::Rip.new 0x0CAFE
+          end
+        else
+          op
+        end
+      end
 
-    def bytesize
-      @insn.encodings.first.bytesize
+      @form.encodings.first.encode buffer, operands
+      @saved_pos = buffer.pos
     end
   end
 
-  def gen name, params
-    insns = Machine.instruction_with_name(name)
-    forms = insns.forms.find_all do |insn|
-      if insn.operands.length == params.length
-        params.zip(insn.operands).all? { |want_op, have_op|
-          want_op.works?(have_op.type)
-        }
+  class AbsoluteJumpInstruction
+    include InstructionPredicates
+
+    def initialize insn, form, operand
+      @insn       = insn
+      @form       = form
+      @operand    = operand
+      @must_retry = false
+    end
+
+    def jump?
+      true
+    end
+
+    def target
+      @operand.name
+    end
+
+    def encode buffer, labels
+      form              = find_form "rel32"
+      encoding          = form.encodings.first
+      operand_klass     = Rel32
+
+      pos = buffer.pos
+      rel_jump = 0xCAFE
+      2.times do
+        buffer.seek pos, IO::SEEK_SET
+        encoding.encode buffer, [operand_klass.new(rel_jump)]
+        rel_jump = @operand.value - buffer.address
+      end
+    end
+
+    private
+
+    def find_form form_type
+      @insn.forms.find { |form| form.operands.first.type == form_type }
+    end
+  end
+
+  class UnresolvedJumpInstruction
+    include InstructionPredicates
+
+    def initialize insn, form, operand
+      @insn       = insn
+      @form       = form
+      @operand    = operand
+      @must_retry = false
+    end
+
+    def jump?
+      true
+    end
+
+    def target
+      @operand.name
+    end
+
+    def retry?
+      true
+    end
+
+    def encode buffer, labels
+      # Estimate by using a rel32 offset
+      form              = find_form "rel32"
+      encoding          = form.encodings.first
+      operand_klass     = Rel32
+
+      if labels.key? target
+        unless @must_retry
+          estimated_offset = labels[target] - (buffer.pos + 6) # 6 seems like a good length
+
+          if estimated_offset >= -128 && estimated_offset <= 127
+            # fits in a rel8
+            operand_klass     = Rel8
+            form              = find_form "rel8"
+            encoding          = form.encodings.first
+          end
+        end
+
+        # We don't know what the size of the jump will be, so just write
+        # it out with some dummy data and keep track of the buffer position.
+        # Calculate the real jump offset, then rewind and patch the jump with
+        # the right location. 😩
+        pos = buffer.pos
+        encoding.encode buffer, [operand_klass.new(0xFE)]
+        jump_len = -(buffer.pos - labels[target])
+        buffer.seek pos, IO::SEEK_SET
+        encoding.encode buffer, [operand_klass.new(jump_len)]
       else
-        false
+        if @must_retry
+          # We retried once, but the label wasn't defined
+          raise Errors::MissingLabel, "Label `#{target}` was used, but not defined"
+        end
+        @must_retry = true
+        encoding.encode buffer, [operand_klass.new(0x0CAFE)]
       end
     end
 
-    raise NotImplementedError, "couldn't find instruction #{name}" if forms.length == 0
+    private
 
-    insn = nil
+    def find_form form_type
+      @insn.forms.find { |form| form.operands.first.type == form_type }
+    end
+  end
 
-    insn = forms.first
+  def initialize
+    @instructions = []
+    @labels = {}
+    # A set of temp registers recorded as we see them (not at allocation time)
+    @temp_registers = Set.new
+    yield self if block_given?
+  end
 
-    insn = Instruction.new(insn, params, position)
-    @position += insn.bytesize
-    @instructions << insn
-    params.each { |param| param.insn = insn if param.unknown_label? }
+  # Mark a temporary register as "done being used" at this point in the
+  # instructions.  Using the register after passing the register to this
+  # method results in undefined behavior.
+  def release_register reg
+    if reg.end_point
+      raise Errors::AlreadyReleasedError, "register #{reg.name} already released at #{reg.end_point}"
+    end
+
+    reg.end_point = (@instructions.length - 1)
+  end
+
+  # Releases all registers that haven't already been released
+  def release_all_registers
+    @temp_registers.each do |reg|
+      next if reg.end_point
+      release_register reg
+    end
+  end
+
+  # Return a list of basic blocks for the instructions
+  def basic_blocks
+    cfg.blocks
+  end
+
+  # Return a cfg for these instructions.  The CFG includes connected basic
+  # blocks as well as any temporary registers
+  def cfg
+    CFG.build @instructions
+  end
+
+  # Create a label to be used with jump instructions.  For example:
+  #
+  #   fisk.jmp(fisk.label(:foo))
+  #   fisk.int(lit(3))
+  #   fisk.put_label(:foo)
+  #
+  def label name
+    UnknownLabel.new(name)
+  end
 
+  Lazy = Struct.new(:block) do # :nodoc:
+    include InstructionPredicates
+
+    def lazy?; true; end
+
+    def encode buffer, labels
+      block.call buffer.pos
+    end
+  end
+
+  # Record a block that will be called during instruction encoding.  The block
+  # will be yielded the current position of the buffer.
+  #
+  # For example:
+  #
+  #   patch_position = nil
+  #
+  #   fisk.nop
+  #   fisk.lazy { |pos| patch_location = pos }
+  #   fisk.jmp(fisk.label(:foo))
+  #   fisk.lazy { |_|
+  #     fisk.mov(fisk.r8, fisk.imm64(patch_location))
+  #   }
+  #   fisk.write_to(buf)
+  #
+  # The first lazy block will be yielded the position of the buffer immediately
+  # after encoding the `nop` instruction and before encoding the `jmp`
+  # instruction.  The second lazy block will be yielded to after the `jmp`
+  # instruction has been encoded.  This gives you a chance to write the buffer
+  # position from certain points in to your assembly
+  def lazy &block
+    @instructions << Lazy.new(block)
     self
   end
 
-  Machine.instructions.keys.each do |insn|
-    define_method(insn.downcase) do |*params|
-      gen insn, params
+  # Insert a label named +name+ at the current position in the instructions.
+  def put_label name
+    if @labels.key? name
+      raise Errors::LabelAlreadyDefined, "Label `#{name}` is already defined"
     end
+    @labels[name] = nil
+    @instructions << Label.new(name)
+    self
   end
+  alias :make_label :put_label
 
-  def moffs64 val
-    MOffs64.new val
+  # Insert a comment at the current position in the instructions.
+  def comment message
+    @instructions << Comment.new(message)
+    self
   end
 
-  def imm8 val
-    Imm8.new val
+  # Allocate and return a new register.  These registers will be replaced with
+  # real registers when `assign_registers` is called.
+  def register name = "temp"
+    Registers::Temp.new name, "r64"
   end
 
-  def imm32 val
-    Imm32.new val
+  # Allocate a register and yield it to the block.  The register will be
+  # automatically released after the block is finished.
+  def with_register name = "temp"
+    reg = register(name)
+    yield reg
+    release_register reg
   end
 
-  def imm64 val
-    Imm64.new val
+  # Assign registers to any temporary registers.  Only registers in +list+
+  # will be used when selecting register assignments
+  def assign_registers list, local: false
+    temp_registers = @temp_registers
+
+    # This mutates the temp registers, setting their end_point based on
+    # the CFG
+    self.cfg unless local
+
+    temp_registers.each do |reg|
+      unless reg.end_point
+        raise Errors::UnreleasedRegisterError, "Register #{reg.name} hasn't been released"
+      end
+    end
+
+    temp_registers = temp_registers.sort_by(&:start_point)
+
+    active         = []
+    free_registers = list.reverse
+    register_count = list.length
+
+    temp_registers.each do |temp_reg|
+      # expire old intervals
+      active, dead = active.sort_by(&:end_point).partition do |j|
+        j.end_point >= temp_reg.start_point
+      end
+
+      # Add unused registers back to the free register list
+      dead.each { |tr| free_registers << tr.register }
+
+      if active.length == register_count
+        raise NotImplementedError, "Register spilled"
+      end
+
+      temp_reg.register = free_registers.pop
+      active << temp_reg
+    end
+  end
+
+  Registers.constants.grep(/^[A-Z0-9]*$/).each do |const|
+    val = Registers.const_get const
+    define_method(const.downcase) { val }
+  end
+
+  include Fisk::Instructions::DSLMethods
+
+  def moffs64 val
+    MOffs64.new val
+  end
+
+  # Create a signed immediate value of the right width
+  def imm val
+    val = val.to_i
+
+    if val >= -0x7F - 1 && val <= 0x7F
+      imm8 val
+    elsif val >=  -0x7FFF - 1 && val <= 0x7FFF
+      imm16 val
+    elsif val >=  -0x7FFFFFFF - 1 && val <= 0x7FFFFFFF
+      imm32 val
+    elsif val >=  -0x7FFFFFFFFFFFFFFF - 1 && val <= 0x7FFFFFFFFFFFFFFF
+      imm64 val
+    else
+      raise ArgumentError, "#{val} is larger than a 64 bit int"
+    end
+  end
+
+  # Create an unsigned immediate value of the right width
+  def uimm val
+    val = val.to_i
+
+    if val < 0
+      raise ArgumentError, "#{val} is negative"
+    elsif val <= 0xFF
+      imm8 val
+    elsif val <= 0xFFFF
+      imm16 val
+    elsif val <= 0xFFFFFFFF
+      imm32 val
+    elsif val <= 0xFFFFFFFFFFFFFFFF
+      imm64 val
+    else
+      raise ArgumentError, "#{val} is too large for a 64 bit int"
+    end
   end
 
   def rel8 val
@@ -214,27 +716,157 @@ class Fisk
     Rel32.new val
   end
 
+  class AbsoluteLocation < Rel32
+    def absolute_location?;  true; end
+  end
+
+  # Creates an "absolute" location.  Use this for JUMP instructions where you
+  # know the absolute location but want a relative location calculated later.
+  #
+  # For example:
+  #
+  #   fisk.jmp fisk.absolute(0xCAFE)
+  #
+  # The emitted JMP instruction will use a calculated relative address (this is
+  # because x64 doesn't have a jump to absolute position).
+  def absolute val
+    AbsoluteLocation.new val
+  end
+
   def lit val
-    # to_s because we're getting the value from JSON as a string
     Lit.new val
   end
 
-  def asm buf = StringIO.new(''.b), &block
-    @position = 0
+  # Instance eval's a given block and writes encoded instructions to +buf+.
+  # For example:
+  #
+  #   fisk = Fisk.new
+  #   fisk.asm do
+  #     mov r9, imm64(32)
+  #   end
+  #
+  def asm buf = StringIO.new(''.b), metadata: {}, &block
     instance_eval(&block)
-    write_to_buffer buf
+    write_to buf, metadata: metadata
     buf
   end
 
-  def to_binary
+  # Encodes all instructions and returns a binary string with the encoded
+  # instructions.
+  def to_binary(metadata: {})
     io = StringIO.new ''.b
-    write_to_buffer io
+    write_to io, metadata: metadata
     io.string
   end
 
+  RetryRequest = Struct.new(:insn, :io_seek_pos)
+
+  # Encode all instructions and write them to +buffer+.  +buffer+ should be an
+  # IO object.
+  def write_to buffer, metadata: {}
+    labels = {}
+    comments = {}
+    unresolved = []
+    instructions = @instructions.dup
+    backup = @instructions.dup
+    @instructions.clear
+
+    while insn = instructions.shift
+      if insn.label?
+        labels[insn.name] = buffer.pos
+      elsif insn.comment?
+        comments.update({buffer.pos => insn.message}) { |_, *lines| lines.join($/) }
+      elsif insn.lazy?
+        insn.encode buffer, labels
+        instructions.unshift(*@instructions)
+        @instructions.clear
+      else
+        if insn.retry?
+          retry_req = RetryRequest.new(insn, buffer.pos)
+          unresolved << retry_req
+        end
+
+        write_instruction insn, buffer, labels
+      end
+    end
+
+    metadata[:comments] = comments
+    @instructions = backup
+
+    return if unresolved.empty?
+
+    pos = buffer.pos
+    unresolved.each do |req|
+      insn = req.insn
+      buffer.seek req.io_seek_pos, IO::SEEK_SET
+      insn.encode buffer, labels
+    end
+    buffer.seek pos, IO::SEEK_SET
+
+    buffer
+  end
+
+  def gen_with_insn insns, params
+    forms = insns.forms.find_all do |insn|
+      if insn.operands.length == params.length
+        params.zip(insn.operands).all? { |want_op, have_op|
+          want_op.works?(have_op)
+        }
+      else
+        false
+      end
+    end
+
+    if forms.length == 0
+      valid_forms = insns.forms.map { |form|
+        "    #{insns.name} #{form.operands.map(&:type).join(", ")}"
+      }.join "\n"
+      msg = <<~eostr
+      Couldn't find instruction #{insns.name} #{params.map(&:type).join(", ")}
+      Valid forms:
+      #{valid_forms}
+      eostr
+      raise NotImplementedError, msg
+    end
+
+    form = forms.first
+
+    insn = nil
+
+    params.each do |param|
+      if param.unresolved?
+        if insns.name =~ /^J/ # I hope all jump instructions start with J!
+          insn = UnresolvedJumpInstruction.new(insns, form, params.first)
+        else
+          # If it's not a jump instruction, assume unresolved RIP relative 😬
+          insn = UnresolvedRIPInstruction.new(insns, form, params)
+        end
+      end
+
+      if param.absolute_location? && insns.name =~ /^(?:J|CALL)/
+        insn = AbsoluteJumpInstruction.new(insns, form, params.first)
+      end
+
+      if param.temp_register?
+        if param.end_point
+          raise Errors::UseAfterInvalidationError, "Register #{param.name} used after release"
+        end
+        @temp_registers << param
+
+        param.start_point ||= @instructions.length
+      end
+    end
+
+    insn ||= Instruction.new(insns, form, params)
+
+    @instructions << insn
+
+    self
+  end
+
   private
 
-  def write_to_buffer buffer
-    @instructions.each { |insn| insn.encode buffer }
+  def write_instruction insn, buffer, labels
+    insn.encode buffer, labels
   end
 end