lua-resty-maxminddb.lua

lua-resty-maxminddb.lua

-- Copyright (C) Anjia (anjia0532)

local json                = require('cjson') 
local json_encode         = json.encode
local json_decode         = json.decode

local ngx_log             = ngx.log
local ngx_ERR             = ngx.ERR
local ngx_CRIT            = ngx.CRIT
local ngx_INFO            = ngx.INFO


local ffi                 = require ('ffi')
local ffi_new             = ffi.new
local ffi_str             = ffi.string
local ffi_cast            = ffi.cast

local _M    ={}
_M._VERSION = '1.0.2'
local mt = { __index = _M }

ffi.cdef[[

typedef unsigned int mmdb_uint128_t __attribute__ ((__mode__(TI)));

typedef struct MMDB_entry_s {
  struct MMDB_s *mmdb;
  uint32_t offset;
} MMDB_entry_s;

typedef struct MMDB_lookup_result_s {
  bool found_entry;
  MMDB_entry_s entry;
  uint16_t netmask;
} MMDB_lookup_result_s;

typedef struct MMDB_entry_data_s {
  bool has_data;
  union {
    uint32_t pointer;
    const char *utf8_string;
    double double_value;
    const uint8_t *bytes;
    uint16_t uint16;
    uint32_t uint32;
    int32_t int32;
    uint64_t uint64;
    mmdb_uint128_t uint128;
    bool boolean;
    float float_value;
  };

  uint32_t offset;
  uint32_t offset_to_next;
  uint32_t data_size;
  uint32_t type;
} MMDB_entry_data_s;

typedef struct MMDB_entry_data_list_s {
  MMDB_entry_data_s entry_data;
  struct MMDB_entry_data_list_s *next;
} MMDB_entry_data_list_s;

typedef struct MMDB_description_s {
  const char *language;
  const char *description;
} MMDB_description_s;

typedef struct MMDB_metadata_s {
  uint32_t node_count;
  uint16_t record_size;
  uint16_t ip_version;
  const char *database_type;
  struct {
    size_t count;
    const char **names;
  } languages;
  uint16_t binary_format_major_version;
  uint16_t binary_format_minor_version;
  uint64_t build_epoch;
  struct {
    size_t count;
    MMDB_description_s **descriptions;
  } description;
} MMDB_metadata_s;

typedef struct MMDB_ipv4_start_node_s {
  uint16_t netmask;
  uint32_t node_value;
} MMDB_ipv4_start_node_s;

typedef struct MMDB_s {
  uint32_t flags;
  const char *filename;
  ssize_t file_size;
  const uint8_t *file_content;
  const uint8_t *data_section;
  uint32_t data_section_size;
  const uint8_t *metadata_section;
  uint32_t metadata_section_size;
  uint16_t full_record_byte_size;
  uint16_t depth;
  MMDB_ipv4_start_node_s ipv4_start_node;
  MMDB_metadata_s metadata;
} MMDB_s;

typedef  char * pchar;

MMDB_lookup_result_s MMDB_lookup_string(MMDB_s *const mmdb,   const char *const ipstr, int *const gai_error,int *const mmdb_error);
int MMDB_open(const char *const filename, uint32_t flags, MMDB_s *const mmdb);
int MMDB_aget_value(MMDB_entry_s *const start,  MMDB_entry_data_s *const entry_data,  const char *const *const path);
char *MMDB_strerror(int error_code);
int MMDB_get_entry_data_list(MMDB_entry_s *start, MMDB_entry_data_list_s **const entry_data_list);
void MMDB_free_entry_data_list(MMDB_entry_data_list_s *const entry_data_list);
void MMDB_close(MMDB_s *const mmdb);
]]

-- error codes 
-- https://github.com/maxmind/libmaxminddb/blob/master/include/maxminddb.h#L66
local MMDB_SUCCESS                                  =   0
local MMDB_FILE_OPEN_ERROR                          =   1
local MMDB_CORRUPT_SEARCH_TREE_ERROR                =   2
local MMDB_INVALID_METADATA_ERROR                   =   3
local MMDB_IO_ERROR                                 =   4
local MMDB_OUT_OF_MEMORY_ERROR                      =   5
local MMDB_UNKNOWN_DATABASE_FORMAT_ERROR            =   6
local MMDB_INVALID_DATA_ERROR                       =   7
local MMDB_INVALID_LOOKUP_PATH_ERROR                =   8
local MMDB_LOOKUP_PATH_DOES_NOT_MATCH_DATA_ERROR    =   9
local MMDB_INVALID_NODE_NUMBER_ERROR                =   10
local MMDB_IPV6_LOOKUP_IN_IPV4_DATABASE_ERROR       =   11

-- data type
-- https://github.com/maxmind/libmaxminddb/blob/master/include/maxminddb.h#L40
local MMDB_DATA_TYPE_EXTENDED                       =   0
local MMDB_DATA_TYPE_POINTER                        =   1
local MMDB_DATA_TYPE_UTF8_STRING                    =   2
local MMDB_DATA_TYPE_DOUBLE                         =   3
local MMDB_DATA_TYPE_BYTES                          =   4
local MMDB_DATA_TYPE_UINT16                         =   5
local MMDB_DATA_TYPE_UINT32                         =   6
local MMDB_DATA_TYPE_MAP                            =   7
local MMDB_DATA_TYPE_INT32                          =   8
local MMDB_DATA_TYPE_UINT64                         =   9
local MMDB_DATA_TYPE_UINT128                        =   10
local MMDB_DATA_TYPE_ARRAY                          =   11
local MMDB_DATA_TYPE_CONTAINER                      =   12
local MMDB_DATA_TYPE_END_MARKER                     =   13
local MMDB_DATA_TYPE_BOOLEAN                        =   14
local MMDB_DATA_TYPE_FLOAT                          =   15

-- you should install the libmaxminddb to your system
local maxm                                          = ffi.load('libmaxminddb')
local mmdb                                          = ffi_new('MMDB_s')
initted                                             = false
db                                                  = nil
--https://github.com/maxmind/libmaxminddb

function _M.init(dbfile)
  initted = true
  db = dbfile
  return initted
end

function _M.initted()
    return initted
end

-- https://github.com/maxmind/libmaxminddb/blob/master/src/maxminddb.c#L1938 
-- LOCAL MMDB_entry_data_list_s *dump_entry_data_list( FILE *stream, MMDB_entry_data_list_s *entry_data_list, int indent, int *status)
local function _dump_entry_data_list(entry_data_list,status,resultTab)

  if not entry_data_list then
    return nil,nil,resultTab
  end
  
  if not resultTab then
    resultTab = {}
  end
  
  
  local entry_data_item = entry_data_list[0].entry_data
  local data_type = entry_data_item.type
  local data_size = entry_data_item.data_size
  
  if data_type == MMDB_DATA_TYPE_MAP then
    table.insert(resultTab,"{")
    
    local flag = false
    local size = entry_data_item.data_size
    
    entry_data_list = entry_data_list[0].next
    
    while(size > 0 and entry_data_list)
    do 
      entry_data_item = entry_data_list[0].entry_data
      data_type = entry_data_item.type
      data_size = entry_data_item.data_size
      
      if MMDB_DATA_TYPE_UTF8_STRING  ~= data_type then
        return nil,MMDB_INVALID_DATA_ERROR,resultTab
      end
      
      local val = ffi_str(entry_data_item.utf8_string,data_size)

      if not val then
        return nil,MMDB_OUT_OF_MEMORY_ERROR,resultTab
      end
      
      -- if latest str is { ,dont append `,`
      table.insert(resultTab,(flag and ',"%s":' or '"%s":'):format(val))
      
      flag = true
      
      entry_data_list = entry_data_list[0].next
      entry_data_list,status,resultTab = _dump_entry_data_list(entry_data_list,status,resultTab)
      
      if status ~= MMDB_SUCCESS then
        return nil,nil,resultTab
      end
      
      size = size -1 
    end
    table.insert(resultTab,"}")
    

  elseif entry_data_list[0].entry_data.type == MMDB_DATA_TYPE_ARRAY then
    local size = entry_data_list[0].entry_data.data_size

    table.insert(resultTab,"[")
    
    entry_data_list = entry_data_list[0].next
    
    while(size >0 and entry_data_list)
    do
      entry_data_list,status,resultTab = _dump_entry_data_list(entry_data_list,status,resultTab)
      
      if status ~= MMDB_SUCCESS then
        return nil,nil,resultTab
      end
      
      size = size -1 
      -- thanks for https://github.com/anjia0532/lua-resty-maxminddb/issues/7
      if size > 0 then
        table.insert(resultTab,",")
      end
    end
    table.insert(resultTab,"]")

    
  else
    entry_data_item = entry_data_list[0].entry_data
    data_type = entry_data_item.type
    data_size = entry_data_item.data_size
    
    local val
    -- string type "key":"val"
    -- other type "key":val
    -- default other type
    local fmt="%s"
    if data_type == MMDB_DATA_TYPE_UTF8_STRING then
      val = ffi_str(entry_data_item.utf8_string,data_size)
      fmt='"%s"'
      if not val then
        status = MMDB_OUT_OF_MEMORY_ERROR
        return nil,status,resultTab
      end
    elseif data_type == MMDB_DATA_TYPE_BYTES then
      val = ffi_str(ffi_cast('char * ',entry_data_item.bytes),data_size)
      fmt='"%s"'
      if not val then
        status = MMDB_OUT_OF_MEMORY_ERROR
        return nil,status,resultTab
      end
    elseif data_type == MMDB_DATA_TYPE_DOUBLE then
      val = entry_data_item.double_value
    elseif data_type == MMDB_DATA_TYPE_FLOAT then
      val = entry_data_item.float_value
    elseif data_type == MMDB_DATA_TYPE_UINT16 then
      val = entry_data_item.uint16
    elseif data_type == MMDB_DATA_TYPE_UINT32 then
      val = entry_data_item.uint32
    elseif data_type == MMDB_DATA_TYPE_BOOLEAN then
      val = entry_data_item.boolean == 1
    elseif data_type == MMDB_DATA_TYPE_UINT64 then
      val = entry_data_item.uint64
    elseif data_type == MMDB_DATA_TYPE_INT32 then
      val = entry_data_item.int32
    else
      return nil,MMDB_INVALID_DATA_ERROR,resultTab
    end
    
    table.insert(resultTab,(fmt):format(val))
    entry_data_list = entry_data_list[0].next
  end
  status = MMDB_SUCCESS
  return entry_data_list,status,resultTab
end

function _M.lookup(ip)

  local file_name_ip2   = ffi_new('char[?]',#db,db)
  local maxmind_ready   = maxm.MMDB_open(file_name_ip2,0,mmdb)
  
  local ip_str = ffi_cast('const char *',ffi_new('char[?]',#ip+1,ip))
  local gai_error = ffi_new('int[1]')
  local mmdb_error = ffi_new('int[1]')

  local result = maxm.MMDB_lookup_string(mmdb,ip_str,gai_error,mmdb_error)

  if mmdb_error[0] ~= MMDB_SUCCESS then
    return nil,'fail when lookup'
  end

  if gai_error[0] ~= MMDB_SUCCESS then
    return nil,'ga error'
  end

  if true ~= result.found_entry then
    ngx_log(ngx_ERR, "stream lua mmdb lookup: entry not found")
    return nil,'not found'
  end

  local entry_data_list = ffi_cast('MMDB_entry_data_list_s **const',ffi_new("MMDB_entry_data_list_s"))

  local status = maxm.MMDB_get_entry_data_list(result.entry,entry_data_list)

  if status ~= MMDB_SUCCESS then
    return nil,'MMDB_get_entry_data_list failed while populating description.'
  end  
  
  local _,status,resultTap = _dump_entry_data_list(entry_data_list)
  maxm.MMDB_free_entry_data_list(entry_data_list[0])
  
  if status ~= MMDB_SUCCESS then
    return nil,'no data'
  end
  
  maxm.MMDB_close(mmdb)
  return json_decode(table.concat(resultTap or {'{','}'})),nil -- fix https://github.com/anjia0532/lua-resty-maxminddb/issues/5#issuecomment-390845118
end
 
-- https://www.maxmind.com/en/geoip2-databases  you should download  the mmdb file from maxmind
 
return _M;