bloom.c 4.1 KB

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