OSNCL/lib/pecoff-utilities.lisp

;;;; PMSF-Lib --- PMSF Common Lisp Utility Library
;;;; This is copyrighted software.  See documentation for terms.
;;;; 
;;;; pecoff-utilities.lisp --- Utilities concerned with PE/COFF
;;;; 
;;;; $Id$

(cl:in-package #:pmsf-pecoff)

;;;; %File Description:
;;;; 
;;;; This file contains utilities for PE/COFF file handling
;;;; 

(defvar *pecoff-processing-pathname* nil)

(define-condition pecoff-processing-error (simple-error)
  ((pathname :initarg :pathname :initform *pecoff-processing-pathname*
             :reader pecoff-processing-error-pathname))
  (:report 
   (lambda (c s)
     (with-standard-io-syntax
       (let ((*print-readably* nil))
         (format s 
                 "Error processing PE/COFF file~@[ for ~A~]: ~?"
                 (pecoff-processing-error-pathname c)
                 (simple-condition-format-control c)
                 (simple-condition-format-arguments c)))))))

(defun pecoff-processing-error (datum &rest arguments)
  (error 'pecoff-processing-error :format-control datum
         :format-arguments arguments))

;;;
;;; PE/COFF Structures
;;;

(defclass file ()
  ((mz-stub :initarg :mz-stub :accessor file-mz-stub)
   (signature :initarg :signature :accessor file-signature)
   (header :initarg :header :accessor file-header)
   (optional-header :initarg :optional-header 
                    :accessor file-optional-header)
   (sections :initarg :sections :accessor file-sections)
   (raw-data :initarg :raw-data :accessor file-raw-data)))

(defclass mz-stub ()
  ((raw-magic :initarg :raw-magic :accessor mz-stub-raw-magic)
   (pecoff-offset :initarg :pecoff-offset
                  :accessor mz-stub-pecoff-offset)
   (raw-data :initarg :raw-data :accessor mz-stub-raw-data)))

(defconstant +header-machine-type-i386+ #x14c)
(defconstant +header-machine-type-amd64+ #x8664)
(defconstant +header-machine-type-arm+ #x1c0)
(defconstant +header-machine-type-thumb+ #x1c2)
(defconstant +header-machine-type-armnt+ #x1c4)
(defconstant +header-machine-type-arm64+ #xaa64)

(defclass header ()
  ((machine-type :initarg :machine-type
                 :accessor header-machine-type)
   (number-of-sections :initarg :number-of-sections
                       :accessor header-number-of-sections)
   (time-date-stamp :initarg :time-date-stamp
                    :accessor header-time-date-stamp)
   (offset-to-symbol-table :initarg :offset-to-symbol-table
                           :accessor header-offset-to-symbol-table)
   (number-of-symbols :initarg :number-of-symbols
                      :accessor header-number-of-symbols)
   (size-of-optional-header :initarg :size-of-optional-header 
                            :accessor header-size-of-optional-header)
   (characteristics :initarg :characteristics 
                    :accessor header-characteristics)))

(defclass optional-header ()
  ((linker-major-version :initarg :linker-major-version
                         :accessor optional-header-linker-major-version)
   (linker-minor-version :initarg :linker-minor-version
                         :accessor optional-header-linker-minor-version)
   (size-of-code :initarg :size-of-code
                 :accessor optional-header-size-of-code)
   (size-of-initialized-data :initarg :size-of-initialized-data
                             :accessor optional-header-size-of-initialized-data)
   (size-of-uninitialized-data :initarg :size-of-uninitialized-data
                               :accessor optional-header-size-of-uninitialized-data)
   (address-of-entry-point :initarg :address-of-entry-point
                           :accessor optional-header-address-of-entry-point)
   (base-of-code :initarg :base-of-code
                 :accessor optional-header-base-of-code)
   ;; Windows-specific part
   (image-base :initarg :image-base 
               :accessor optional-header-image-base)
   (section-alignment :initarg :section-alignment
                      :accessor optional-header-section-alignment)
   (file-alignment :initarg :file-alignment
                   :accessor optional-header-file-alignment)
   (major-os-version :initarg :major-os-version
                     :accessor optional-header-major-os-version)
   (minor-os-version :initarg :minor-os-version
                     :accessor optional-header-minor-os-version)
   (major-image-version :initarg :major-image-version
                        :accessor optional-header-major-image-version)
   (minor-image-version :initarg :minor-image-version
                        :accessor optional-header-minor-image-version)
   (major-subsystem-version :initarg :major-subsystem-version
                            :accessor optional-header-major-subsystem-version)
   (minor-subsystem-version :initarg :minor-subsystem-version
                            :accessor optional-header-minor-subsystem-version)
   (win32-version-value :initarg :win32-version-value
                        :accessor optional-header-win32-version-value)
   (size-of-image :initarg :size-of-image
                  :accessor optional-header-size-of-image)
   (size-of-headers :initarg :size-of-headers
                    :accessor optional-header-size-of-headers)
   (checksum :initarg :checksum
             :accessor optional-header-checksum)
   (subsystem :initarg :subsystem
              :accessor optional-header-subsystem)
   (dll-characteristics :initarg :dll-characteristics
                        :accessor optional-header-dll-characteristics)
   (size-of-stack-reserve :initarg :size-of-stack-reserve
                          :accessor optional-header-size-of-stack-reserve)
   (size-of-stack-commit :initarg :size-of-stack-commit
                         :accessor optional-header-size-of-stack-commit)
   (size-of-heap-reserve :initarg :size-of-heap-reserve
                         :accessor optional-header-size-of-heap-reserve)
   (size-of-heap-commit :initarg :size-of-heap-commit
                        :accessor optional-header-size-of-heap-commit)
   (loader-flags :initarg :loader-flags
                 :accessor optional-header-loader-flags)
   (number-of-rva-and-sizes :initarg :number-of-rva-and-sizes
                            :accessor optional-header-number-of-rva-and-sizes)
   (data-directories :initarg :data-directories
                     :accessor optional-header-data-directories)))

(defclass data-directory ()
  ((virtual-address :initarg :virtual-address
                    :accessor data-directory-virtual-address)
   (size :initarg :size :accessor data-directory-size)))

(defclass optional-header-pe32 (optional-header)
  ((base-of-data :initarg :base-of-data
                 :accessor optional-header-base-of-data)))

(defclass optional-header-pe32+ (optional-header)
  ())

(defclass section ()
  ((name :initarg :name :accessor section-name)
   (virtual-size :initarg :virtual-size :accessor section-virtual-size)
   (virtual-address :initarg :virtual-address :accessor section-virtual-address)
   (size-of-raw-data :initarg :size-of-raw-data :accessor section-size-of-raw-data)
   (pointer-to-raw-data :initarg :pointer-to-raw-data :accessor section-pointer-to-raw-data)
   (pointer-to-relocations :initarg :pointer-to-relocations :accessor section-pointer-to-relocations)
   (pointer-to-line-numbers :initarg :pointer-to-line-numbers :accessor section-pointer-to-line-numbers)
   (number-of-relocations :initarg :number-of-relocations :accessor section-number-of-relocations)
   (number-of-line-numbers :initarg :number-of-line-numbers :accessor section-number-of-line-numbers)
   (characteristics :initarg :characteristics :accessor section-characteristics)
   (raw-data :initarg :raw-data :accessor section-raw-data)))

(defclass export-directory ()
  ((export-flags :initarg :export-flags 
                 :accessor export-directory-export-flags)
   (time-date-stamp :initarg :time-date-stamp
                    :accessor export-directory-time-date-stamp)
   (major-version :initarg :major-version
                  :accessor export-directory-major-version)
   (minor-version :initarg :minor-version
                  :accessor export-directory-minor-version)
   (name-rva :initarg :name-rva
             :accessor export-directory-name-rva)
   (name :initarg :name
         :accessor export-directory-name)
   (ordinal-base :initarg :ordinal-base
                 :accessor export-directory-ordinal-name)
   (address-table-entries :initarg :address-table-entries
                          :accessor export-directory-address-table-entries)
   (name-table-entries :initarg :name-table-entries
                       :accessor export-directory-name-table-entries)
   (address-table-rva :initarg :address-table-rva
                      :accessor export-directory-address-table-rva)
   (name-table-rva :initarg :name-table-rva
                   :accessor export-directory-name-table-rva)
   (ordinal-table-rva :initarg :ordinal-table-rva
                      :accessor export-directory-ordinal-table-rva)
   (address-table :initarg :address-table
                  :accessor export-directory-address-table)
   (name-table :initarg :name-table
                  :accessor export-directory-name-table)
   (ordinal-table :initarg :ordinal-table
                  :accessor export-directory-ordinal-table)
   (raw-data :initarg :raw-data
             :accessor export-directory-raw-data)))

;;;
;;; Utility Accessors
;;;

(defun get-string-from-rva (file address)
  (loop for section across (file-sections file)
        for start = (section-virtual-address section)
        for stop = (1- (+ start (section-virtual-size section)))
        for raw-size = (section-size-of-raw-data section)
        for raw-data = (section-raw-data section)
        thereis
        (when (and (<= start address stop)
                   (<= (- address start) raw-size))
          (let ((end (position 0 raw-data :start (- address start))))
            (make-array (list (1+ (- end (- address start))))
                        :element-type '(unsigned-byte 8)
                        :displaced-to raw-data
                        :displaced-index-offset (- address start))))))

(defun access-data-directory-contents (file index)
  (let* ((opt-header (file-optional-header file))
         (entry (aref (optional-header-data-directories opt-header)
                      index)))
    (loop with address = (data-directory-virtual-address entry)
          with size = (data-directory-size entry)
          for section across (file-sections file)
          for start = (section-virtual-address section)
          for stop = (1- (+ start (section-virtual-size section)))
          for raw-size = (section-size-of-raw-data section)
          for raw-data = (section-raw-data section)
          thereis
          (when (and (<= start address stop)
                     (<= start (1- (+ address size)) stop)
                     (<= (- address start) raw-size))
            (make-array (list size)
                        :element-type '(unsigned-byte 8)
                        :displaced-to raw-data
                        :displaced-index-offset (- address start))))))

;;;
;;; Basic COFF I/O
;;;

(defvar *current-source*)
(defvar *current-offset*)
(defvar *current-kind*)

(defmacro with-pecoff-reading (source &body body)
  `(let ((*current-source* ,source)
         (*current-offset* 0))
     ,@body))

(defmacro with-pecoff-reading-kind (kind &body body)
  `(let ((*current-kind* ,kind))
     ,@body))

(defun read-coff-byte (&optional offset)
  (when offset (setq *current-offset* offset))
  (prog1 (aref *current-source* *current-offset*)
    (incf *current-offset*)))

(defun read-coff-half (&optional offset)
  (let ((lsb (read-coff-byte offset)))
    (dpb (read-coff-byte) (byte 8 8) lsb)))

(defun read-coff-word (&optional offset)
  (let* ((lsb (read-coff-byte offset))
         (lsb2 (read-coff-byte))
         (lsb3 (read-coff-byte))
         (msb (read-coff-byte)))
    (dpb msb (byte 8 24) (dpb lsb3 (byte 8 16) (dpb lsb2 (byte 8 8) lsb)))))

(defun read-coff-xword (&optional offset)
  (let* ((lsb (read-coff-byte offset))
         (lsb2 (read-coff-byte))
         (lsb3 (read-coff-byte))
         (lsb4 (read-coff-byte))
         (lsb5 (read-coff-byte))
         (lsb6 (read-coff-byte))
         (lsb7 (read-coff-byte))
         (msb (read-coff-byte)))
    (dpb msb (byte 8 56) 
         (dpb lsb7 (byte 8 48) 
              (dpb lsb6 (byte 8 40) 
                   (dpb lsb5 (byte 8 32)
                        (dpb lsb4 (byte 8 24)
                             (dpb lsb3 (byte 8 16)
                                  (dpb lsb2 (byte 8 8) lsb)))))))))

(defun read-coff-word/xword (&optional offset)
  (ecase *current-kind*
    (:pe32 (read-coff-word offset))
    (:pe32+ (read-coff-xword offset))))

(defun read-coff-sequence (length &optional offset)
  (when offset (setq *current-offset* offset))
  (let ((result (make-array (list length) :element-type '(unsigned-byte 8)
			    :displaced-to *current-source*
			    :displaced-index-offset *current-offset*)))
    (incf *current-offset* length)
    result))

;;;
;;; Header Handling and Reading
;;;

(defun check-pecoff-file-p ()
  (and
   ;; MZ Stub
   (eql (read-coff-byte 0) (char-code #\M))
   (eql (read-coff-byte) (char-code #\Z))
   (let ((pe-offset (read-coff-word #x3c)))
     (and 
      (eql (read-coff-byte pe-offset) (char-code #\P))
      (eql (read-coff-byte) (char-code #\E))
      (eql (read-coff-byte) 0)
      (eql (read-coff-byte) 0)))))

(defun read-mz-stub ()
  (let ((identifier (read-coff-sequence 2 0))
        (pecoff-offset (read-coff-word #x3c)))
    (unless (every #'eql identifier (map 'list #'char-code "MZ"))
      (pecoff-processing-error "Invalid MZ MAGIC:~{ 0x~2,'0X~}"
                               (coerce identifier 'list)))
    (make-instance 'mz-stub
                   :raw-magic identifier
                   :pecoff-offset pecoff-offset
                   :raw-data
                   (read-coff-sequence pecoff-offset 0))))

(defun read-pecoff-signature (&optional offset)
  (let ((signature (read-coff-sequence 4 offset)))
    (unless (every #'eql signature
                   (map 'list #'char-code
                        (list #\P #\E #\NULL #\NULL)))
      (pecoff-processing-error "Invalid PE Signature:~{ 0x~2,'0X~}"
                               (coerce signature 'list)))
    signature))

(defun read-header (&optional offset)
  (let ((machine-type (read-coff-half offset))
        (number-of-sections (read-coff-half))
        (time-date-stamp (read-coff-word))
        (offset-to-symbol-table (read-coff-word))
        (number-of-symbols (read-coff-word))
        (size-of-optional-header (read-coff-half))
        (characteristics (read-coff-half)))
    (make-instance 'header
                   :machine-type machine-type
                   :number-of-sections number-of-sections
                   :time-date-stamp time-date-stamp
                   :offset-to-symbol-table offset-to-symbol-table
                   :number-of-symbols number-of-symbols
                   :size-of-optional-header size-of-optional-header
                   :characteristics characteristics)))

(defun read-optional-header (size &optional offset)
  (unless (zerop size)
    (let ((magic (read-coff-half offset)))
      (cond
       ((eql magic #x010b)
        (read-optional-header-pe32))
       ((eql magic #x020b)
        (read-optional-header-pe32+))
       (t
        (pecoff-processing-error "Unsupported optional header magic 0x~4,'0X" magic))))))

(defun read-optional-header-pe32 ()
  (apply
   #'make-instance
   'optional-header-pe32
   :linker-major-version (read-coff-byte)
   :linker-minor-version (read-coff-byte)
   :size-of-code (read-coff-word)
   :size-of-initialized-data (read-coff-word)
   :size-of-uninitialized-data (read-coff-word)
   :address-of-entry-point (read-coff-word)
   :base-of-code (read-coff-word)
   :base-of-data (read-coff-word)
   (with-pecoff-reading-kind :pe32
     (read-optional-header-windows))))

(defun read-optional-header-pe32+ ()
  (apply
   #'make-instance
   'optional-header-pe32+
   :linker-major-version (read-coff-byte)
   :linker-minor-version (read-coff-byte)
   :size-of-code (read-coff-word)
   :size-of-initialized-data (read-coff-word)
   :size-of-uninitialized-data (read-coff-word)
   :address-of-entry-point (read-coff-word)
   :base-of-code (read-coff-word)
   (with-pecoff-reading-kind :pe32+
     (read-optional-header-windows))))

(defun read-optional-header-windows ()
  (list*
   :image-base (read-coff-word/xword)
   :section-alignment (read-coff-word)
   :file-alignment (read-coff-word)
   :major-os-version (read-coff-half)
   :minor-os-version (read-coff-half)
   :major-image-version (read-coff-half)
   :minor-image-version (read-coff-half)
   :major-subsystem-version (read-coff-half)
   :minor-subsystem-version (read-coff-half)
   :win32-version-value (read-coff-word)
   :size-of-image (read-coff-word)
   :size-of-headers (read-coff-word)
   :checksum (read-coff-word)
   :subsystem (read-coff-half)
   :dll-characteristics (read-coff-half)
   :size-of-stack-reserve (read-coff-word/xword)
   :size-of-stack-commit (read-coff-word/xword)
   :size-of-heap-reserve (read-coff-word/xword)
   :size-of-heap-commit (read-coff-word/xword)
   :loader-flags (read-coff-word)
   (let ((entries (read-coff-word)))
     (list :number-of-rva-and-sizes entries
           :data-directories
           (coerce
            (loop repeat entries
                  collect
                  (make-instance 'data-directory
                                 :virtual-address (read-coff-word)
                                 :size (read-coff-word)))
            'vector)))))

(defun read-section (&optional offset)
  (let ((result
         (make-instance 
          'section
          :name (read-coff-sequence 8 offset)
          :virtual-size (read-coff-word)
          :virtual-address (read-coff-word)
          :size-of-raw-data (read-coff-word)
          :pointer-to-raw-data (read-coff-word)
          :pointer-to-relocations (read-coff-word)
          :pointer-to-line-numbers (read-coff-word)
          :number-of-relocations (read-coff-half)
          :number-of-line-numbers (read-coff-half)
          :characteristics (read-coff-word))))
    (unless (zerop (section-size-of-raw-data result))
      (let ((*current-offset* *current-offset*))
        (setf (section-raw-data result)
              (read-coff-sequence (section-size-of-raw-data result)
                                  (section-pointer-to-raw-data result)))))
    result))

;;;
;;; Parsing of directories
;;;

(defun parse-export-directory (file)
  (let ((raw-data (access-data-directory-contents file 0)))
    (with-pecoff-reading raw-data
      (let* ((export-flags (read-coff-word))
             (time-date-stamp (read-coff-word))
             (major-version (read-coff-half))
             (minor-version (read-coff-half))
             (name-rva (read-coff-word))
             (ordinal-base (read-coff-word))
             (address-table-entries (read-coff-word))
             (name-table-entries (read-coff-word))
             (address-table-rva (read-coff-word))
             (name-table-rva (read-coff-word))
             (ordinal-table-rva (read-coff-word)))
        (make-instance
         'export-directory
         :export-flags export-flags
         :time-date-stamp time-date-stamp
         :major-version major-version
         :minor-version minor-version
         :name-rva name-rva
         :ordinal-base ordinal-base
         :address-table-entries address-table-entries
         :name-table-entries name-table-entries
         :address-table-rva address-table-rva
         :name-table-rva name-table-rva
         :ordinal-table-rva ordinal-table-rva
         :raw-data raw-data
         :address-table
         (loop with result = (make-array address-table-entries)
               for i from 0 below address-table-entries
               do (setf (aref result i) (read-coff-word))
               finally (return result))
         :name-table
         (loop with result = (make-array name-table-entries)
               for i from 0 below name-table-entries
               for rva = (read-coff-word)
               do (setf (aref result i) 
                        (get-string-from-rva file rva))
               finally (return result))
         :ordinal-table
         (loop with result = (make-array address-table-entries)
               for i from 0 below address-table-entries
               do (setf (aref result i) (read-coff-half))
               finally (return result)))))))

;;;
;;; Checksum Updates
;;;

(defun calculate-checksum-offset (file)
  (+ (mz-stub-pecoff-offset (file-mz-stub file))
     4
     20
     64))

(defun calculate-checksum (file)
  (loop with raw-data = (file-raw-data file)
        with checksum-offset = (calculate-checksum-offset file)
        with checksum = 0
        for offset upfrom 0 below (length raw-data) by 2
        do
        (unless (<= checksum-offset offset (+ 2 checksum-offset))
          (incf checksum
                (dpb (aref raw-data (1+ offset))
                     (byte 8 8)
                     (aref raw-data offset)))
          (setq checksum
                (+ (ldb (byte 16 16) checksum)
                   (ldb (byte 16 0) checksum))))
        finally
        (return
         (ldb (byte 32 0)
              (+
               (ldb (byte 16 0)
                    (+ (ldb (byte 16 16) checksum)
                       checksum))
               (length raw-data))))))

(defun update-checksum (file)
  (let ((raw-data (file-raw-data file))
        (checksum (calculate-checksum file))
        (offset (calculate-checksum-offset file)))
    (setf (aref raw-data offset) (ldb (byte 8  0) checksum)
          (aref raw-data (+ offset 1)) (ldb (byte 8  8) checksum)
          (aref raw-data (+ offset 2)) (ldb (byte 8 16) checksum)
          (aref raw-data (+ offset 3)) (ldb (byte 8 24) checksum))))

;;;
;;; Main Entry Points
;;;

(defun read-image-file (pathname)
  (with-open-file (in pathname :element-type '(unsigned-byte 8))
    (let ((raw-data (make-array (list (file-length in))
				:element-type '(unsigned-byte 8))))
      (read-sequence raw-data in)
      (with-pecoff-reading raw-data
        (unless (check-pecoff-file-p)
          (pecoff-processing-error "File is not in PE/COFF format!"))
        (let* ((mz-stub (read-mz-stub))
               (signature (read-pecoff-signature))
               (header (read-header))
               (optional-header (read-optional-header (header-size-of-optional-header header)))
               (sections (coerce
                          (loop repeat (header-number-of-sections header)
                                collect (read-section))
                          'vector)))
          (make-instance 'file
                         :mz-stub mz-stub
                         :signature signature
                         :header header
                         :optional-header optional-header
                         :sections sections
                         :raw-data raw-data))))))

(defun write-image-file (file pathname)
  (with-open-file (out pathname :element-type '(unsigned-byte 8)
                       :direction :output)
    (write-sequence (file-raw-data file) out)))