bloom.c 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  1. /*
  2. * Module for bloom filters
  3. *
  4. * Philip Gladstone, N1DQ
  5. */
  6. #include "module.h"
  7. #include "lauxlib.h"
  8. #include "c_types.h"
  9. #include "../crypto/sha2.h"
  10. #if defined(LUA_USE_MODULES_BLOOM) && !defined(SHA2_ENABLE)
  11. #error Must have SHA2_ENABLE set for BLOOM module
  12. #endif
  13. typedef struct {
  14. uint8 fns;
  15. uint16 size;
  16. uint32 occupancy;
  17. uint32 buf[];
  18. } bloom_t;
  19. static bool add_or_check(const uint8 *buf, size_t len, bloom_t *filter, bool add) {
  20. SHA256_CTX ctx;
  21. SHA256_Init(&ctx);
  22. SHA256_Update(&ctx, buf, len);
  23. char hash[32];
  24. SHA256_Final(hash, &ctx);
  25. int i;
  26. uint32 bits = filter->size << 5;
  27. uint8 *h = hash;
  28. bool prev = true;
  29. int hstep = filter->fns > 10 ? 2 : 3;
  30. for (i = 0; i < filter->fns; i++) {
  31. uint32 val = (((h[0] << 8) + h[1]) << 8) + h[2];
  32. h += hstep;
  33. val = val % bits;
  34. uint32 offset = val >> 5;
  35. uint32 bit = 1 << (val & 31);
  36. if (!(filter->buf[offset] & bit)) {
  37. prev = false;
  38. if (add) {
  39. filter->buf[offset] |= bit;
  40. filter->occupancy++;
  41. } else {
  42. break;
  43. }
  44. }
  45. }
  46. return prev;
  47. }
  48. static int bloom_filter_check(lua_State *L) {
  49. bloom_t *filter = (bloom_t *)luaL_checkudata(L, 1, "bloom.filter");
  50. size_t length;
  51. const uint8 *buffer = (uint8 *) luaL_checklstring(L, 2, &length);
  52. bool rc = add_or_check(buffer, length, filter, false);
  53. lua_pushboolean(L, rc);
  54. return 1;
  55. }
  56. static int bloom_filter_add(lua_State *L) {
  57. bloom_t *filter = (bloom_t *)luaL_checkudata(L, 1, "bloom.filter");
  58. size_t length;
  59. const uint8 *buffer = (uint8 *) luaL_checklstring(L, 2, &length);
  60. bool rc = add_or_check(buffer, length, filter, true);
  61. lua_pushboolean(L, rc);
  62. return 1;
  63. }
  64. static int bloom_filter_reset(lua_State *L) {
  65. bloom_t *filter = (bloom_t *)luaL_checkudata(L, 1, "bloom.filter");
  66. memset(filter->buf, 0, filter->size << 2);
  67. filter->occupancy = 0;
  68. return 0;
  69. }
  70. static int bloom_filter_info(lua_State *L) {
  71. bloom_t *filter = (bloom_t *)luaL_checkudata(L, 1, "bloom.filter");
  72. lua_pushinteger(L, filter->size << 5);
  73. lua_pushinteger(L, filter->fns);
  74. lua_pushinteger(L, filter->occupancy);
  75. // Now calculate the chance that a FP will be returned
  76. uint64 prob = 1000000;
  77. if (filter->occupancy > 0) {
  78. unsigned int ratio = (filter->size << 5) / filter->occupancy;
  79. int i;
  80. prob = ratio;
  81. for (i = 1; i < filter->fns && prob < 1000000; i++) {
  82. prob = prob * ratio;
  83. }
  84. if (prob < 1000000) {
  85. // try again with some scaling
  86. unsigned int ratio256 = (filter->size << 13) / filter->occupancy;
  87. uint64 prob256 = ratio256;
  88. for (i = 1; i < filter->fns && prob256 < 256000000; i++) {
  89. prob256 = (prob256 * ratio256) >> 8;
  90. }
  91. prob = prob256 >> 8;
  92. }
  93. }
  94. lua_pushinteger(L, prob > 1000000 ? 1000000 : (int) prob);
  95. return 4;
  96. }
  97. static int bloom_create(lua_State *L) {
  98. int items = luaL_checkinteger(L, 1);
  99. int error = luaL_checkinteger(L, 2);
  100. int n = error;
  101. int logp = 0;
  102. while (n > 0) {
  103. n = n >> 1;
  104. logp--;
  105. }
  106. int bits = -items * logp;
  107. bits += bits >> 1;
  108. bits = (bits + 31) & ~31;
  109. if (bits < 256) {
  110. bits = 256;
  111. }
  112. int size = bits >> 3;
  113. int fns = bits / items;
  114. fns = (fns >> 1) + fns / 6;
  115. if (fns < 2) {
  116. fns = 2;
  117. }
  118. if (fns > 15) {
  119. fns = 15;
  120. }
  121. bloom_t *filter = (bloom_t *) lua_newuserdata(L, sizeof(bloom_t) + size);
  122. //
  123. // Associate its metatable
  124. luaL_getmetatable(L, "bloom.filter");
  125. lua_setmetatable(L, -2);
  126. memset(filter, 0, sizeof(bloom_t) + size);
  127. filter->size = size >> 2;
  128. filter->fns = fns;
  129. return 1;
  130. }
  131. static const LUA_REG_TYPE bloom_filter_map[] = {
  132. { LSTRKEY( "add" ), LFUNCVAL( bloom_filter_add ) },
  133. { LSTRKEY( "check" ), LFUNCVAL( bloom_filter_check ) },
  134. { LSTRKEY( "reset" ), LFUNCVAL( bloom_filter_reset ) },
  135. { LSTRKEY( "info" ), LFUNCVAL( bloom_filter_info ) },
  136. { LSTRKEY( "__index" ), LROVAL( bloom_filter_map ) },
  137. { LNILKEY, LNILVAL }
  138. };
  139. // Module function map
  140. static const LUA_REG_TYPE bloom_map[] = {
  141. { LSTRKEY( "create" ), LFUNCVAL( bloom_create ) },
  142. { LNILKEY, LNILVAL }
  143. };
  144. LUALIB_API int bloom_open(lua_State *L) {
  145. luaL_rometatable(L, "bloom.filter", (void *)bloom_filter_map);
  146. return 1;
  147. }
  148. NODEMCU_MODULE(BLOOM, "bloom", bloom_map, bloom_open);