path: root/geoip/ip.lisp

;;;; -*- mode: lisp -*-

(ql:quickload '(cffi cffi-libffi split-sequence mmap ironclad babel nibbles fiveam local-time))

(defpackage :geoip
  (:use :common-lisp :cffi :split-sequence))
(in-package :geoip)

(defun octet-to-int (seq &key (big-endian t))
  (reduce (lambda (acc n)
            (+ (ash acc 8) n))
          (if big-endian
              seq (reverse seq))
          :initial-value 0))

(5am:test t-ints
  (5am:is (= (octet-to-int '(1 5) :big-endian t) #x105))
  (5am:is (= (octet-to-int '(1 5) :big-endian nil) #x501)))

(defun parse-ipv4 (ip-address)
  (reduce (lambda (acc part)
            (let ((n (parse-integer part)))
              (assert (<= 0 n 255))
              (+ (ash acc 8) n)))
          (split-sequence #\. ip-address)
          :initial-value 0))

(defun parse-ipv6 (ip-address)
  (reduce (lambda (acc part)
            (let ((n (if (string= part "") 0
                         (parse-integer part :radix 16))))
              (assert (<= 0 n #xFFFF))
              (+ (ash acc 16) n)))
          (split-sequence #\: ip-address)
          :initial-value 0))


(defun parse-ip (ip-address)
  (if (find #\. ip-address)
      (cons 4 (parse-ipv4 ip-address))
      (cons 6 (parse-ipv6 ip-address))))

(5am:test ipv4
  (5am:is (= (parse-ipv4 "8.8.8.8") 134744072))
  (5am:is (= (parse-ipv6 "::1") 1)))

(defparameter *metadata-marker*
  (concatenate 'vector
               #(#xab #xcd #xef)
               (map 'vector #'char-code "MaxMind.com")))


(defstruct maxmind-database
  (filename nil :read-only t)
  (ptr nil :read-only t)
  (fd nil :read-only t)
  (size 0 :read-only t)
  (node-count 0 :read-only t :type fixnum)
  (record-size 0 :read-only t :type fixnum)
  (ip-version 0 :read-only t :type fixnum)
  (data-offset -1 :type fixnum)
  (metadata nil))

(defstruct db-reader
  (db-ptr)
  (head-ptr)
  (data-offset))

(defstruct maxmind-database-metadata
  (node-count 0 :read-only t :type fixnum)
  (record-size 0 :read-only t :type fixnum)
  (ip-version 0 :read-only t :type fixnum)
  (database-type nil :read-only t :type string)
  (languages nil :read-only t :type list)
  (binary-format-major-version 0 :read-only t :type fixnum)
  (binary-format-minor-version 0 :read-only t :type fixnum)
  (build-epoch 0 :read-only t :type fixnum)
  (description nil :read-only t :type list))


(defun get-val (alist key &key (test #'equal))
  (cdr (assoc key alist :test test)))

(defun make-metadata (data)
  (make-maxmind-database-metadata
   :binary-format-major-version (get-val data :binary-format-major-version)
   :binary-format-minor-version (get-val data :binary-format-minor-version)
   :build-epoch (get-val data :build-epoch)
   :database-type (get-val data :database-type)
   :description (get-val data :description)
   :ip-version (get-val data :ip-version)
   :languages (get-val data :languages)
   :node-count (get-val data :node-count)
   :record-size (get-val data :record-size)))

(defun metadata-marker-p (ptr offset)
  (loop for i below #.(length *metadata-marker*)
        always (eql (mem-ref ptr :uchar (+ i offset))
                    (svref *metadata-marker* i))))

(defun find-metadata-start (ptr size)
  (loop for offset from (- size 1 #.(length *metadata-marker*))
          downto (- size 1 (* 128 1024))
        when (metadata-marker-p ptr offset)
          return (+ offset #.(length *metadata-marker*))))

(defun read-db-char (reader)
  (with-slots (db-ptr) reader
    (prog1 (mem-ref db-ptr :uchar)
      (incf-pointer db-ptr))))

(defun bytes-from-foreign (reader data-length)
  (let ((bytes (make-array data-length :element-type '(unsigned-byte 8))))
    (with-slots (db-ptr) reader
      (foreign-array-to-lisp db-ptr
                             `(:array :uchar ,data-length)
                             :element-type '(unsigned-byte 8)
                             :displaced-to bytes)
      (incf-pointer db-ptr data-length))
    bytes))

(defun mread-uft8 (reader length)
  (with-slots (db-ptr) reader
    (prog1 (foreign-string-to-lisp db-ptr :count length :encoding :utf-8)
      (incf-pointer db-ptr length))))

(defmacro mem-uint (ptr size)
  `(let ((r 0))
     (dotimes (i ,size r)
       (setf r (+ (ash r 8) (mem-ref ,ptr :uint8 i))))))

(defun mread-unsigned (reader size)
  (with-slots (db-ptr) reader
    (prog1 (mem-uint db-ptr size)
      (incf-pointer db-ptr size))))

(defun mread-int32 (db-ptr length)
  (let ((uval (mread-unsigned db-ptr length)))
    (if (and (= length 4) (logbitp 31 uval))
        (- uval #.(expt 2 32))
        uval)))

(5am:test reader
  (with-foreign-array (a #(0 0 1 23 126 195 159 195 156) '(:array :uint8 9))
    (let ((r (make-db-reader :db-ptr a)))
      (5am:is (= 279 (mread-unsigned r 4)))
      (5am:is (equal "~ßÜ" (mread-uft8 r 5))))))

(5am:run-all-tests)


(defun mread-datafield-metadata (db-ptr)
  (let* ((control-byte (read-db-char db-ptr))
         (type (ldb (byte 3 5) control-byte))
         (length (ldb (byte 5 0) control-byte)))
    (when (zerop type)
      (setf type (+ 7 (read-db-char db-ptr))))
    (when (= type 1)
      (incf-pointer (db-reader-db-ptr db-ptr) -1))
    (list type
          (cond
            ((or (= type 1) (< length 29))
             length)
            ((= length 29)
             (+ 29 (read-db-char db-ptr)))
            ((= length 30)
             (+ 285 (mread-unsigned db-ptr 2)))
            ((= length 31)
             (+ 65821 (mread-unsigned db-ptr 3)))))))


(defun mread-map (db-ptr length)
  (loop repeat length
        collect
        (cons
         (intern (string-upcase (substitute #\- #\_ (mread-data db-ptr))) :keyword)
         (mread-data db-ptr))))

(defun mread-list (db-ptr length)
  (loop repeat length collect (mread-data db-ptr)))

(defun read-pointer (reader)
  (let* ((control-byte (read-db-char reader))
         (size-bits (ldb (byte 2 3) control-byte))
         (value-bits (ldb (byte 3 0) control-byte)))
    (ecase size-bits
      (0 (+ (ash value-bits 8) (read-db-char reader)))
      (1 (+ (ash value-bits 16) (mread-unsigned reader 2) 2048))
      (2 (+ (ash value-bits 24) (mread-unsigned reader 3) 526336))
      (3 (mread-unsigned reader 4)))))

(defun mread-pointer (reader)
  (with-slots (db-ptr head-ptr data-offset) reader
    (let* ((target (+ (read-pointer reader) data-offset))
           (curr-ptr db-ptr)
           (result))
      (setf db-ptr (inc-pointer head-ptr target))
      (setf result (mread-data reader))
      (setf db-ptr curr-ptr)
      result)))

(defun mread-data (db-ptr)
  (destructuring-bind (type length) (mread-datafield-metadata db-ptr)
    (ecase type
      (1 (mread-pointer db-ptr))
      (2 (mread-uft8 db-ptr length))
      (3 (cons 'double (mread-unsigned db-ptr 8)))
      (4 (bytes-from-foreign db-ptr length))
      ((5 6 9 10) (mread-unsigned db-ptr length))
      (7 (mread-map db-ptr length))
      (8 (mread-int32 db-ptr length))
      (11 (mread-list db-ptr length))
      (14 (< 0 length)) ; bool
      )))

(defun read-node-record (mmdb node-number bit)
  "Read left (bit=0) or right (bit=1) record from node"
  (with-slots (ptr record-size) mmdb
    (let* ((node-bytes (/ record-size 4))
           (record-bytes (floor record-size 8))
           (node-offset (* node-bytes node-number)))
      (if (evenp node-bytes)
          ;; Simple case: records align to byte boundaries
          (mem-uint (inc-pointer ptr (+ node-offset (* bit record-bytes))) record-bytes)
          ;; Complex case: records share a byte
          (let* ((base-value (mem-uint (inc-pointer ptr (+ node-offset (* bit (1+ record-bytes)))) record-bytes))
                 (shared-byte (mem-ref (inc-pointer ptr (+ node-offset record-bytes)) :uchar))
                 (high-bits (if (zerop bit) (ash shared-byte -4) (logand shared-byte #xf))))
            (+ (ash high-bits (* record-bytes 8)) base-value))))))

(defun record-value (mmdb ip-bits)
  (with-slots (node-count) mmdb
    (loop with node = 0
          for bit in ip-bits
          for next-node = (read-node-record mmdb node bit)
          do (setf node next-node)
          when (= node node-count)
            return nil
          when (> node node-count)
            return node)))

(defun lookup-ip (mmdb ip-bits)
  (alexandria:when-let ((record (record-value mmdb ip-bits)))
    (with-slots (ptr data-offset node-count) mmdb
      (mread-data
       (make-db-reader
        :db-ptr (inc-pointer ptr (+ data-offset (- record node-count 16)))
        :head-ptr ptr
        :data-offset data-offset)))))

(defun integer-to-bits (n bit-count)
  "Convert integer to list of bits (MSB first)"
  (loop for i from (1- bit-count) downto 0
        collect (if (logbitp i n) 1 0)))

(defun mmap->mmdb (file ptr fd size)
  (let* ((metadata-start (find-metadata-start ptr size))
         (db-reader (make-db-reader :db-ptr (inc-pointer ptr metadata-start)))
         (metadata (make-metadata (mread-data db-reader))))
    (with-slots (record-size node-count ip-version) metadata
      (make-maxmind-database
       :filename file
       :ptr ptr :fd fd :size size
       :metadata metadata
       :node-count node-count
       :record-size record-size
       :ip-version ip-version
       ;; :data-offset (+ 16 (* (/ (* 2 record-size) 8) node-count))
       :data-offset (+ 16 (/ (* record-size node-count) 4))))))

(defun make-mmdb (file)
  (multiple-value-bind (ptr fd size) (mmap:mmap file)
    (mmap->mmdb file ptr fd size)))

(defun close-mmdb (mmdb)
  (with-slots (ptr fd size) mmdb
    (mmap:munmap ptr fd size)))
(defvar *mmdb* (make-mmdb  "GeoLite2-Country.mmdb"))

(with-slots (metadata) *mmdb*
  (with-slots (build-epoch) metadata
    (local-time:unix-to-timestamp
     build-epoch)))