12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364 |
- #define LUA_LIB
- #include <lua.h>
- #include <lauxlib.h>
- #include <time.h>
- #include <unistd.h>
- #include <stdint.h>
- #include <stdlib.h>
- #include <string.h>
- #include <stdbool.h>
- #include "atomic.h"
- #define DEFAULT_CAP 64
- #define MAX_NUMBER 1024
- // avoid circular reference while encodeing
- #define MAX_DEPTH 128
- #define BSON_REAL 1
- #define BSON_STRING 2
- #define BSON_DOCUMENT 3
- #define BSON_ARRAY 4
- #define BSON_BINARY 5
- #define BSON_UNDEFINED 6
- #define BSON_OBJECTID 7
- #define BSON_BOOLEAN 8
- #define BSON_DATE 9
- #define BSON_NULL 10
- #define BSON_REGEX 11
- #define BSON_DBPOINTER 12
- #define BSON_JSCODE 13
- #define BSON_SYMBOL 14
- #define BSON_CODEWS 15
- #define BSON_INT32 16
- #define BSON_TIMESTAMP 17
- #define BSON_INT64 18
- #define BSON_MINKEY 255
- #define BSON_MAXKEY 127
- #define BSON_TYPE_SHIFT 5
- static char bson_numstrs[MAX_NUMBER][4];
- static int bson_numstr_len[MAX_NUMBER];
- struct bson {
- int size;
- int cap;
- uint8_t *ptr;
- uint8_t buffer[DEFAULT_CAP];
- };
- struct bson_reader {
- const uint8_t * ptr;
- int size;
- };
- static inline int32_t
- get_length(const uint8_t * data) {
- const uint8_t * b = (const uint8_t *)data;
- int32_t len = b[0] | b[1]<<8 | b[2]<<16 | b[3]<<24;
- return len;
- }
- static inline void
- bson_destroy(struct bson *b) {
- if (b->ptr != b->buffer) {
- free(b->ptr);
- }
- }
- static inline void
- bson_create(struct bson *b) {
- b->size = 0;
- b->cap = DEFAULT_CAP;
- b->ptr = b->buffer;
- }
- static inline void
- bson_reserve(struct bson *b, int sz) {
- if (b->size + sz <= b->cap)
- return;
- do {
- b->cap *= 2;
- } while (b->cap <= b->size + sz);
- if (b->ptr == b->buffer) {
- b->ptr = (uint8_t*)malloc(b->cap);
- memcpy(b->ptr, b->buffer, b->size);
- } else {
- b->ptr = (uint8_t*)realloc(b->ptr, b->cap);
- }
- }
- static inline void
- check_reader(lua_State *L, struct bson_reader *br, int sz) {
- if (br->size < sz) {
- luaL_error(L, "Invalid bson block (%d:%d)", br->size, sz);
- }
- }
- static inline int
- read_byte(lua_State *L, struct bson_reader *br) {
- check_reader(L, br, 1);
- const uint8_t * b = br->ptr;
- int r = b[0];
- ++br->ptr;
- --br->size;
- return r;
- }
- static inline int32_t
- read_int32(lua_State *L, struct bson_reader *br) {
- check_reader(L, br, 4);
- const uint8_t * b = br->ptr;
- uint32_t v = b[0] | b[1]<<8 | b[2]<<16 | b[3]<<24;
- br->ptr+=4;
- br->size-=4;
- return (int32_t)v;
- }
- static inline int64_t
- read_int64(lua_State *L, struct bson_reader *br) {
- check_reader(L, br, 8);
- const uint8_t * b = br->ptr;
- uint32_t lo = b[0] | b[1]<<8 | b[2]<<16 | b[3]<<24;
- uint32_t hi = b[4] | b[5]<<8 | b[6]<<16 | b[7]<<24;
- uint64_t v = (uint64_t)lo | (uint64_t)hi<<32;
- br->ptr+=8;
- br->size-=8;
- return (int64_t)v;
- }
- static inline lua_Number
- read_double(lua_State *L, struct bson_reader *br) {
- check_reader(L, br, 8);
- union {
- uint64_t i;
- double d;
- } v;
- const uint8_t * b = br->ptr;
- uint32_t lo = b[0] | b[1]<<8 | b[2]<<16 | b[3]<<24;
- uint32_t hi = b[4] | b[5]<<8 | b[6]<<16 | b[7]<<24;
- v.i = (uint64_t)lo | (uint64_t)hi<<32;
- br->ptr+=8;
- br->size-=8;
- return v.d;
- }
- static inline const void *
- read_bytes(lua_State *L, struct bson_reader *br, int sz) {
- const void * r = br->ptr;
- check_reader(L, br, sz);
- br->ptr+=sz;
- br->size-=sz;
- return r;
- }
- static inline const char *
- read_cstring(lua_State *L, struct bson_reader *br, size_t *sz) {
- int i;
- for (i=0;;i++) {
- if (i==br->size) {
- luaL_error(L, "Invalid bson block : cstring");
- }
- if (br->ptr[i] == '\0') {
- break;
- }
- }
- *sz = i;
- const char * r = (const char *)br->ptr;
- br->ptr += i+1;
- br->size -= i+1;
- return r;
- }
- static inline void
- write_byte(struct bson *b, uint8_t v) {
- bson_reserve(b,1);
- b->ptr[b->size++] = v;
- }
- static inline void
- write_int32(struct bson *b, int32_t v) {
- uint32_t uv = (uint32_t)v;
- bson_reserve(b,4);
- b->ptr[b->size++] = uv & 0xff;
- b->ptr[b->size++] = (uv >> 8)&0xff;
- b->ptr[b->size++] = (uv >> 16)&0xff;
- b->ptr[b->size++] = (uv >> 24)&0xff;
- }
- static inline void
- write_length(struct bson *b, int32_t v, int off) {
- uint32_t uv = (uint32_t)v;
- b->ptr[off++] = uv & 0xff;
- b->ptr[off++] = (uv >> 8)&0xff;
- b->ptr[off++] = (uv >> 16)&0xff;
- b->ptr[off++] = (uv >> 24)&0xff;
- }
- #define MAXUNICODE 0x10FFFF
- static int
- utf8_copy(const char *s, char *d, size_t limit) {
- static const unsigned int limits[] = {0xFF, 0x7F, 0x7FF, 0xFFFF};
- unsigned int c = s[0];
- unsigned int res = 0;
- if (limit < 1)
- return 0;
- d[0] = s[0];
- if (c < 0x80) {
- return 1;
- } else {
- int count = 0;
- while (c & 0x40) {
- int cc = s[++count];
- if (limit <= count || (cc & 0xC0) != 0x80)
- return 0;
- d[count] = s[count];
- res = (res << 6) | (cc & 0x3F);
- c <<= 1;
- }
- res |= ((c & 0x7F) << (count * 5));
- if (count > 3 || res > MAXUNICODE || res <= limits[count])
- return 0;
- return count+1;
- }
- }
- static void
- write_string(struct bson *b, lua_State *L, const char *key, size_t sz) {
- bson_reserve(b,sz+1);
- char *dst = (char *)(b->ptr + b->size);
- const char *src = key;
- size_t n = sz;
- while(n > 0) {
- int c = utf8_copy(src, dst, n);
- if (c == 0) {
- luaL_error(L, "Invalid utf8 string");
- }
- src += c;
- dst += c;
- n -= c;
- }
- b->ptr[b->size+sz] = '\0';
- b->size+=sz+1;
- }
- static inline int
- reserve_length(struct bson *b) {
- int sz = b->size;
- bson_reserve(b,4);
- b->size +=4;
- return sz;
- }
- static inline void
- write_int64(struct bson *b, int64_t v) {
- uint64_t uv = (uint64_t)v;
- int i;
- bson_reserve(b,8);
- for (i=0;i<64;i+=8) {
- b->ptr[b->size++] = (uv>>i) & 0xff;
- }
- }
- static inline void
- write_double(struct bson *b, lua_Number d) {
- union {
- double d;
- uint64_t i;
- } v;
- v.d = d;
- int i;
- bson_reserve(b,8);
- for (i=0;i<64;i+=8) {
- b->ptr[b->size++] = (v.i>>i) & 0xff;
- }
- }
- static inline void
- append_key(struct bson *bs, lua_State *L, int type, const char *key, size_t sz) {
- write_byte(bs, type);
- write_string(bs, L, key, sz);
- }
- static inline int
- is_32bit(int64_t v) {
- return v >= INT32_MIN && v <= INT32_MAX;
- }
- static void
- append_number(struct bson *bs, lua_State *L, const char *key, size_t sz) {
- if (lua_isinteger(L, -1)) {
- int64_t i = lua_tointeger(L, -1);
- if (is_32bit(i)) {
- append_key(bs, L, BSON_INT32, key, sz);
- write_int32(bs, i);
- } else {
- append_key(bs, L, BSON_INT64, key, sz);
- write_int64(bs, i);
- }
- } else {
- lua_Number d = lua_tonumber(L,-1);
- append_key(bs, L, BSON_REAL, key, sz);
- write_double(bs, d);
- }
- }
- static void append_table(struct bson *bs, lua_State *L, const char *key, size_t sz, int depth);
- static void
- write_binary(struct bson *b, const void * buffer, size_t sz) {
- int length = reserve_length(b);
- bson_reserve(b,sz);
- memcpy(b->ptr + b->size, buffer, sz); // include sub type
- b->size+=sz;
- write_length(b, sz-1, length); // not include sub type
- }
- static void
- append_one(struct bson *bs, lua_State *L, const char *key, size_t sz, int depth) {
- int vt = lua_type(L,-1);
- switch(vt) {
- case LUA_TNUMBER:
- append_number(bs, L, key, sz);
- break;
- case LUA_TUSERDATA: {
- append_key(bs, L, BSON_DOCUMENT, key, sz);
- int32_t * doc = (int32_t*)lua_touserdata(L,-1);
- int32_t sz = *doc;
- bson_reserve(bs,sz);
- memcpy(bs->ptr + bs->size, doc, sz);
- bs->size += sz;
- break;
- }
- case LUA_TSTRING: {
- size_t len;
- const char * str = lua_tolstring(L,-1,&len);
- if (len > 1 && str[0]==0) {
- int subt = (uint8_t)str[1];
- append_key(bs, L, subt, key, sz);
- switch(subt) {
- case BSON_BINARY:
- write_binary(bs, str+2, len-2);
- break;
- case BSON_OBJECTID:
- if (len != 2+12) {
- luaL_error(L, "Invalid object id %s", str+2);
- }
- // go though
- case BSON_JSCODE:
- case BSON_DBPOINTER:
- case BSON_SYMBOL:
- case BSON_CODEWS:
- bson_reserve(bs,len-2);
- memcpy(bs->ptr + bs->size, str+2, len-2);
- bs->size += len-2;
- break;
- case BSON_DATE: {
- if (len != 2+4) {
- luaL_error(L, "Invalid date");
- }
- const uint32_t * ts = (const uint32_t *)(str + 2);
- int64_t v = (int64_t)*ts * 1000;
- write_int64(bs, v);
- break;
- }
- case BSON_TIMESTAMP: {
- if (len != 2+8) {
- luaL_error(L, "Invalid timestamp");
- }
- const uint32_t * inc = (const uint32_t *)(str + 2);
- const uint32_t * ts = (const uint32_t *)(str + 6);
- write_int32(bs, *inc);
- write_int32(bs, *ts);
- break;
- }
- case BSON_REGEX: {
- str+=2;
- len-=3;
- size_t i;
- for (i=0;i<len;i++) {
- if (str[len-i-1]==0) {
- break;
- }
- }
- write_string(bs, L, str, len-i-1);
- write_string(bs, L, str + len-i, i);
- break;
- }
- case BSON_MINKEY:
- case BSON_MAXKEY:
- case BSON_NULL:
- break;
- case BSON_INT64: {
- if (len != 2 + 8) {
- luaL_error(L, "Invalid int64");
- }
- const int64_t * v = (const int64_t *)(str + 2);
- write_int64(bs, *v);
- break;
- }
- default:
- luaL_error(L,"Invalid subtype %d", subt);
- }
- } else {
- size_t len;
- const char * str = lua_tolstring(L,-1,&len);
- append_key(bs, L, BSON_STRING, key, sz);
- int off = reserve_length(bs);
- write_string(bs, L, str, len);
- write_length(bs, len+1, off);
- }
- break;
- }
- case LUA_TTABLE:
- append_table(bs, L, key, sz, depth+1);
- break;
- case LUA_TBOOLEAN:
- append_key(bs, L, BSON_BOOLEAN, key, sz);
- write_byte(bs, lua_toboolean(L,-1));
- break;
- case LUA_TNIL:
- luaL_error(L, "Bson array has a hole (nil), Use bson.null instead");
- default:
- luaL_error(L, "Invalid value type : %s", lua_typename(L,vt));
- }
- }
- static inline int
- bson_numstr( char *str, unsigned int i ) {
- if ( i < MAX_NUMBER) {
- memcpy( str, bson_numstrs[i], 4 );
- return bson_numstr_len[i];
- } else {
- return sprintf( str,"%u", i );
- }
- }
- static void
- pack_array(lua_State *L, struct bson *b, int depth, size_t len) {
- int length = reserve_length(b);
- size_t i;
- for (i=1;i<=len;i++) {
- char numberkey[32];
- size_t sz = bson_numstr(numberkey, i - 1);
- const char * key = numberkey;
- lua_geti(L, -1, i);
- append_one(b, L, key, sz, depth);
- lua_pop(L, 1);
- }
- write_byte(b,0);
- write_length(b, b->size - length, length);
- }
- static void
- pack_dict_data(lua_State *L, struct bson *b, int depth, int kt) {
- const char * key = NULL;
- size_t sz;
- switch(kt) {
- case LUA_TNUMBER:
- luaL_error(L, "Bson dictionary's key can't be number");
- break;
- case LUA_TSTRING:
- key = lua_tolstring(L,-2,&sz);
- append_one(b, L, key, sz, depth);
- lua_pop(L,1);
- break;
- default:
- luaL_error(L, "Invalid key type : %s", lua_typename(L, kt));
- return;
- }
- }
- static void
- pack_simple_dict(lua_State *L, struct bson *b, int depth) {
- int length = reserve_length(b);
- lua_pushnil(L);
- while(lua_next(L,-2) != 0) {
- int kt = lua_type(L, -2);
- pack_dict_data(L, b, depth, kt);
- }
- write_byte(b,0);
- write_length(b, b->size - length, length);
- }
- static void
- pack_meta_dict(lua_State *L, struct bson *b, int depth) {
- int length = reserve_length(b);
- lua_pushvalue(L, -2); // push meta_obj
- lua_call(L, 1, 3); // call __pairs_func => next_func, t_data, first_k
- for(;;) {
- lua_pushvalue(L, -2); // copy data
- lua_pushvalue(L, -2); // copy k
- lua_copy(L, -5, -3); // copy next_func replace old_k
- lua_call(L, 2, 2); // call next_func
- int kt = lua_type(L, -2);
- if (kt == LUA_TNIL) {
- lua_pop(L, 4); // pop all k, v, next_func, obj
- break;
- }
- pack_dict_data(L, b, depth, kt);
- }
- write_byte(b,0);
- write_length(b, b->size - length, length);
- }
- static bool
- is_rawarray(lua_State *L) {
- lua_pushnil(L);
- if (lua_next(L, -2) == 0) {
- // empty table
- return false;
- }
- lua_Integer firstkey = lua_isinteger(L, -2) ? lua_tointeger(L, -2) : 0;
- lua_pop(L, 2);
- if (firstkey <= 1) {
- return firstkey > 0;
- }
- return firstkey <= lua_rawlen(L, -1);
- }
- static void
- append_table(struct bson *bs, lua_State *L, const char *key, size_t sz, int depth) {
- if (depth > MAX_DEPTH) {
- luaL_error(L, "Too depth while encoding bson");
- }
- luaL_checkstack(L, 16, NULL); // reserve enough stack space to pack table
- if (luaL_getmetafield(L, -1, "__len") != LUA_TNIL) {
- lua_pushvalue(L, -2);
- lua_call(L, 1, 1);
- if (!lua_isinteger(L, -1)) {
- luaL_error(L, "__len should return integer");
- }
- size_t len = lua_tointeger(L, -1);
- lua_pop(L, 1);
- append_key(bs, L, BSON_ARRAY, key, sz);
- pack_array(L, bs, depth, len);
- } else if (luaL_getmetafield(L, -1, "__pairs") != LUA_TNIL) {
- append_key(bs, L, BSON_DOCUMENT, key, sz);
- pack_meta_dict(L, bs, depth);
- } else if (is_rawarray(L)) {
- append_key(bs, L, BSON_ARRAY, key, sz);
- pack_array(L, bs, depth, lua_rawlen(L, -1));
- } else {
- append_key(bs, L, BSON_DOCUMENT, key, sz);
- pack_simple_dict(L, bs, depth);
- }
- }
- static void
- pack_ordered_dict(lua_State *L, struct bson *b, int n, int depth) {
- int length = reserve_length(b);
- int i;
- size_t sz;
- // the first key is at index n
- const char * key = lua_tolstring(L, n, &sz);
- for (i=0;i<n;i+=2) {
- if (key == NULL) {
- luaL_error(L, "Argument %d need a string", i+1);
- }
- lua_pushvalue(L, i+1);
- append_one(b, L, key, sz, depth);
- lua_pop(L,1);
- key = lua_tolstring(L, i+2, &sz); // next key
- }
- write_byte(b,0);
- write_length(b, b->size - length, length);
- }
-
- static int
- ltostring(lua_State *L) {
- size_t sz = lua_rawlen(L, 1);
- void * ud = lua_touserdata(L,1);
- lua_pushlstring(L, (const char*)ud, sz);
- return 1;
- }
- static int
- llen(lua_State *L) {
- size_t sz = lua_rawlen(L, 1);
- lua_pushinteger(L, sz);
- return 1;
- }
- static void
- make_object(lua_State *L, int type, const void * ptr, size_t len) {
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, type);
- luaL_addlstring(&b, (const char*)ptr, len);
- luaL_pushresult(&b);
- }
- static void
- unpack_dict(lua_State *L, struct bson_reader *br, bool array) {
- luaL_checkstack(L, 16, NULL); // reserve enough stack space to unpack table
- int sz = read_int32(L, br);
- const void * bytes = read_bytes(L, br, sz-5);
- struct bson_reader t = { (const uint8_t*)bytes, sz-5 };
- int end = read_byte(L, br);
- if (end != '\0') {
- luaL_error(L, "Invalid document end");
- }
- lua_newtable(L);
- for (;;) {
- if (t.size == 0)
- break;
- int bt = read_byte(L, &t);
- size_t klen = 0;
- const char * key = read_cstring(L, &t, &klen);
- if (array) {
- int id = strtol(key, NULL, 10) + 1;
- lua_pushinteger(L,id);
- } else {
- lua_pushlstring(L, key, klen);
- }
- switch (bt) {
- case BSON_REAL:
- lua_pushnumber(L, read_double(L, &t));
- break;
- case BSON_BOOLEAN:
- lua_pushboolean(L, read_byte(L, &t));
- break;
- case BSON_STRING: {
- int sz = read_int32(L, &t);
- if (sz <= 0) {
- luaL_error(L, "Invalid bson string , length = %d", sz);
- }
- lua_pushlstring(L, (const char*)read_bytes(L, &t, sz), sz-1);
- break;
- }
- case BSON_DOCUMENT:
- unpack_dict(L, &t, false);
- break;
- case BSON_ARRAY:
- unpack_dict(L, &t, true);
- break;
- case BSON_BINARY: {
- int sz = read_int32(L, &t);
- int subtype = read_byte(L, &t);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_BINARY);
- luaL_addchar(&b, subtype);
- luaL_addlstring(&b, (const char*)read_bytes(L, &t, sz), sz);
- luaL_pushresult(&b);
- break;
- }
- case BSON_OBJECTID:
- make_object(L, BSON_OBJECTID, read_bytes(L, &t, 12), 12);
- break;
- case BSON_DATE: {
- int64_t date = read_int64(L, &t);
- uint32_t v = date / 1000;
- make_object(L, BSON_DATE, &v, 4);
- break;
- }
- case BSON_MINKEY:
- case BSON_MAXKEY:
- case BSON_NULL: {
- char key[] = { 0, (char)bt };
- lua_pushlstring(L, key, sizeof(key));
- break;
- }
- case BSON_REGEX: {
- size_t rlen1=0;
- size_t rlen2=0;
- const char * r1 = read_cstring(L, &t, &rlen1);
- const char * r2 = read_cstring(L, &t, &rlen2);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_REGEX);
- luaL_addlstring(&b, r1, rlen1);
- luaL_addchar(&b,0);
- luaL_addlstring(&b, r2, rlen2);
- luaL_addchar(&b,0);
- luaL_pushresult(&b);
- break;
- }
- case BSON_INT32:
- lua_pushinteger(L, read_int32(L, &t));
- break;
- case BSON_TIMESTAMP: {
- int32_t inc = read_int32(L, &t);
- int32_t ts = read_int32(L, &t);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_TIMESTAMP);
- luaL_addlstring(&b, (const char *)&inc, 4);
- luaL_addlstring(&b, (const char *)&ts, 4);
- luaL_pushresult(&b);
- break;
- }
- case BSON_INT64:
- lua_pushinteger(L, read_int64(L, &t));
- break;
- case BSON_DBPOINTER: {
- const void * ptr = t.ptr;
- int sz = read_int32(L, &t);
- read_bytes(L, &t, sz+12);
- make_object(L, BSON_DBPOINTER, ptr, sz + 16);
- break;
- }
- case BSON_JSCODE:
- case BSON_SYMBOL: {
- const void * ptr = t.ptr;
- int sz = read_int32(L, &t);
- read_bytes(L, &t, sz);
- make_object(L, bt, ptr, sz + 4);
- break;
- }
- case BSON_CODEWS: {
- const void * ptr = t.ptr;
- int sz = read_int32(L, &t);
- read_bytes(L, &t, sz-4);
- make_object(L, bt, ptr, sz);
- break;
- }
- default:
- // unsupported
- luaL_error(L, "Invalid bson type : %d", bt);
- lua_pop(L,1);
- continue;
- }
- lua_rawset(L,-3);
- }
- }
- static int
- lmakeindex(lua_State *L) {
- int32_t *bson = (int32_t*)luaL_checkudata(L,1,"bson");
- const uint8_t * start = (const uint8_t *)bson;
- struct bson_reader br = { start+4, get_length(start) - 5 };
- lua_newtable(L);
- for (;;) {
- if (br.size == 0)
- break;
- int bt = read_byte(L, &br);
- size_t klen = 0;
- const char * key = read_cstring(L, &br, &klen);
- int field_size = 0;
- switch (bt) {
- case BSON_INT64:
- case BSON_TIMESTAMP:
- case BSON_DATE:
- case BSON_REAL:
- field_size = 8;
- break;
- case BSON_BOOLEAN:
- field_size = 1;
- break;
- case BSON_JSCODE:
- case BSON_SYMBOL:
- case BSON_STRING: {
- int sz = read_int32(L, &br);
- read_bytes(L, &br, sz);
- break;
- }
- case BSON_CODEWS:
- case BSON_ARRAY:
- case BSON_DOCUMENT: {
- int sz = read_int32(L, &br);
- read_bytes(L, &br, sz-4);
- break;
- }
- case BSON_BINARY: {
- int sz = read_int32(L, &br);
- read_bytes(L, &br, sz+1);
- break;
- }
- case BSON_OBJECTID:
- field_size = 12;
- break;
- case BSON_MINKEY:
- case BSON_MAXKEY:
- case BSON_NULL:
- break;
- case BSON_REGEX: {
- size_t rlen1=0;
- size_t rlen2=0;
- read_cstring(L, &br, &rlen1);
- read_cstring(L, &br, &rlen2);
- break;
- }
- case BSON_INT32:
- field_size = 4;
- break;
- case BSON_DBPOINTER: {
- int sz = read_int32(L, &br);
- read_bytes(L, &br, sz+12);
- break;
- }
- default:
- // unsupported
- luaL_error(L, "Invalid bson type : %d", bt);
- lua_pop(L,1);
- continue;
- }
- if (field_size > 0) {
- int id = bt | (int)(br.ptr - start) << BSON_TYPE_SHIFT;
- read_bytes(L, &br, field_size);
- lua_pushlstring(L, key, klen);
- lua_pushinteger(L,id);
- lua_rawset(L,-3);
- }
- }
- lua_setiuservalue(L,1,1);
- lua_settop(L,1);
- return 1;
- }
- static void
- replace_object(lua_State *L, int type, struct bson * bs) {
- size_t len = 0;
- const char * data = luaL_checklstring(L,3, &len);
- if (len < 6 || data[0] != 0 || data[1] != type) {
- luaL_error(L, "Type mismatch, need bson type %d", type);
- }
- switch (type) {
- case BSON_OBJECTID:
- if (len != 2+12) {
- luaL_error(L, "Invalid object id");
- }
- memcpy(bs->ptr, data+2, 12);
- break;
- case BSON_DATE: {
- if (len != 2+4) {
- luaL_error(L, "Invalid date");
- }
- const uint32_t * ts = (const uint32_t *)(data + 2);
- int64_t v = (int64_t)*ts * 1000;
- write_int64(bs, v);
- break;
- }
- case BSON_TIMESTAMP: {
- if (len != 2+8) {
- luaL_error(L, "Invalid timestamp");
- }
- const uint32_t * inc = (const uint32_t *)(data + 2);
- const uint32_t * ts = (const uint32_t *)(data + 6);
- write_int32(bs, *inc);
- write_int32(bs, *ts);
- break;
- }
- }
- }
- static int
- lreplace(lua_State *L) {
- lua_getiuservalue(L,1,1);
- if (!lua_istable(L,-1)) {
- return luaL_error(L, "call makeindex first");
- }
- lua_pushvalue(L,2);
- if (lua_rawget(L, -2) != LUA_TNUMBER) {
- return luaL_error(L, "Can't replace key : %s", lua_tostring(L,2));
- }
- int id = lua_tointeger(L, -1);
- int type = id & ((1<<(BSON_TYPE_SHIFT)) - 1);
- int offset = id >> BSON_TYPE_SHIFT;
- uint8_t * start = (uint8_t*)lua_touserdata(L,1);
- struct bson b = { 0,16, start + offset };
- switch (type) {
- case BSON_REAL:
- write_double(&b, luaL_checknumber(L, 3));
- break;
- case BSON_BOOLEAN:
- write_byte(&b, lua_toboolean(L,3));
- break;
- case BSON_OBJECTID:
- case BSON_DATE:
- case BSON_TIMESTAMP:
- replace_object(L, type, &b);
- break;
- case BSON_INT32: {
- if (!lua_isinteger(L, 3)) {
- luaL_error(L, "%f must be a 32bit integer ", lua_tonumber(L, 3));
- }
- int32_t i = lua_tointeger(L,3);
- write_int32(&b, i);
- break;
- }
- case BSON_INT64: {
- if (!lua_isinteger(L, 3)) {
- luaL_error(L, "%f must be a 64bit integer ", lua_tonumber(L, 3));
- }
- int64_t i = lua_tointeger(L,3);
- write_int64(&b, i);
- break;
- }
- default:
- luaL_error(L, "Can't replace type %d", type);
- break;
- }
- return 0;
- }
- static int
- ldecode(lua_State *L) {
- const int32_t * data = (const int32_t*)lua_touserdata(L,1);
- if (data == NULL) {
- return 0;
- }
- const uint8_t * b = (const uint8_t *)data;
- int32_t len = get_length(b);
- struct bson_reader br = { b , len };
- unpack_dict(L, &br, false);
- return 1;
- }
- static void
- bson_meta(lua_State *L) {
- if (luaL_newmetatable(L, "bson")) {
- luaL_Reg l[] = {
- { "decode", ldecode },
- { "makeindex", lmakeindex },
- { NULL, NULL },
- };
- luaL_newlib(L,l);
- lua_setfield(L, -2, "__index");
- lua_pushcfunction(L, ltostring);
- lua_setfield(L, -2, "__tostring");
- lua_pushcfunction(L, llen);
- lua_setfield(L, -2, "__len");
- lua_pushcfunction(L, lreplace);
- lua_setfield(L, -2, "__newindex");
- }
- lua_setmetatable(L, -2);
- }
- static int
- encode_bson(lua_State *L) {
- struct bson *b = (struct bson*)lua_touserdata(L, 2);
- lua_settop(L, 1);
- if (luaL_getmetafield(L, -1, "__pairs") != LUA_TNIL) {
- pack_meta_dict(L, b, 0);
- } else {
- pack_simple_dict(L, b, 0);
- }
- void * ud = lua_newuserdatauv(L, b->size, 1);
- memcpy(ud, b->ptr, b->size);
- return 1;
- }
- static int
- lencode(lua_State *L) {
- struct bson b;
- lua_settop(L,1);
- luaL_checktype(L, 1, LUA_TTABLE);
- bson_create(&b);
- lua_pushcfunction(L, encode_bson);
- lua_pushvalue(L, 1);
- lua_pushlightuserdata(L, &b);
- if (lua_pcall(L, 2, 1, 0) != LUA_OK) {
- bson_destroy(&b);
- return lua_error(L);
- }
- bson_destroy(&b);
- bson_meta(L);
- return 1;
- }
- static int
- encode_bson_byorder(lua_State *L) {
- int n = lua_gettop(L);
- struct bson *b = (struct bson*)lua_touserdata(L, n);
- lua_settop(L, --n);
- pack_ordered_dict(L, b, n, 0);
- lua_settop(L,0);
- void * ud = lua_newuserdatauv(L, b->size, 1);
- memcpy(ud, b->ptr, b->size);
- return 1;
- }
- static int
- lencode_order(lua_State *L) {
- struct bson b;
- int n = lua_gettop(L);
- if (n%2 != 0) {
- return luaL_error(L, "Invalid ordered dict");
- }
- bson_create(&b);
- lua_pushvalue(L, 1); // copy the first arg to n
- lua_pushcfunction(L, encode_bson_byorder);
- lua_replace(L, 1);
- lua_pushlightuserdata(L, &b);
- if (lua_pcall(L, n+1, 1, 0) != LUA_OK) {
- bson_destroy(&b);
- return lua_error(L);
- }
- bson_destroy(&b);
- bson_meta(L);
- return 1;
- }
- static int
- ldate(lua_State *L) {
- int d = luaL_checkinteger(L,1);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_DATE);
- luaL_addlstring(&b, (const char *)&d, sizeof(d));
- luaL_pushresult(&b);
- return 1;
- }
- static int
- lint64(lua_State *L) {
- int64_t d = luaL_checkinteger(L, 1);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_INT64);
- luaL_addlstring(&b, (const char *)&d, sizeof(d));
- luaL_pushresult(&b);
- return 1;
- }
- static int
- ltimestamp(lua_State *L) {
- int d = luaL_checkinteger(L,1);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_TIMESTAMP);
- if (lua_isnoneornil(L,2)) {
- static uint32_t inc = 0;
- luaL_addlstring(&b, (const char *)&inc, sizeof(inc));
- ++inc;
- } else {
- uint32_t i = (uint32_t)lua_tointeger(L,2);
- luaL_addlstring(&b, (const char *)&i, sizeof(i));
- }
- luaL_addlstring(&b, (const char *)&d, sizeof(d));
- luaL_pushresult(&b);
- return 1;
- }
- static int
- lregex(lua_State *L) {
- luaL_checkstring(L,1);
- if (lua_gettop(L) < 2) {
- lua_pushliteral(L,"");
- }
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_REGEX);
- lua_pushvalue(L,1);
- luaL_addvalue(&b);
- luaL_addchar(&b,0);
- lua_pushvalue(L,2);
- luaL_addvalue(&b);
- luaL_addchar(&b,0);
- luaL_pushresult(&b);
- return 1;
- }
- static int
- lbinary(lua_State *L) {
- lua_settop(L,1);
- luaL_Buffer b;
- luaL_buffinit(L, &b);
- luaL_addchar(&b, 0);
- luaL_addchar(&b, BSON_BINARY);
- luaL_addchar(&b, 0); // sub type
- lua_pushvalue(L,1);
- luaL_addvalue(&b);
- luaL_pushresult(&b);
- return 1;
- }
- static int
- lsubtype(lua_State *L, int subtype, const uint8_t * buf, size_t sz) {
- switch(subtype) {
- case BSON_BINARY:
- lua_pushvalue(L, lua_upvalueindex(6));
- lua_pushlstring(L, (const char *)buf+1, sz-1);
- lua_pushinteger(L, buf[0]);
- return 3;
- case BSON_OBJECTID: {
- if (sz != 12) {
- return luaL_error(L, "Invalid object id");
- }
- char oid[24];
- int i;
- const uint8_t * id = buf;
- static const char *hex = "0123456789abcdef";
- for (i=0;i<12;i++) {
- oid[i*2] = hex[id[i] >> 4];
- oid[i*2+1] = hex[id[i] & 0xf];
- }
- lua_pushvalue(L, lua_upvalueindex(7));
- lua_pushlstring(L, oid, 24);
- return 2;
- }
- case BSON_DATE: {
- if (sz != 4) {
- return luaL_error(L, "Invalid date");
- }
- int d = *(const int *)buf;
- lua_pushvalue(L, lua_upvalueindex(9));
- lua_pushinteger(L, d);
- return 2;
- }
- case BSON_TIMESTAMP: {
- if (sz != 8) {
- return luaL_error(L, "Invalid timestamp");
- }
- const uint32_t * ts = (const uint32_t *)buf;
- lua_pushvalue(L, lua_upvalueindex(8));
- lua_pushinteger(L, (lua_Integer)ts[1]);
- lua_pushinteger(L, (lua_Integer)ts[0]);
- return 3;
- }
- case BSON_REGEX: {
- --sz;
- size_t i;
- const uint8_t *str = buf;
- for (i=0;i<sz;i++) {
- if (str[sz-i-1]==0) {
- break;
- }
- }
- lua_pushvalue(L, lua_upvalueindex(10));
- if (i==sz) {
- return luaL_error(L, "Invalid regex");
- }
- lua_pushlstring(L, (const char *)str, sz - i - 1);
- lua_pushlstring(L, (const char *)str+sz-i, i);
- return 3;
- }
- case BSON_MINKEY:
- lua_pushvalue(L, lua_upvalueindex(11));
- return 1;
- case BSON_MAXKEY:
- lua_pushvalue(L, lua_upvalueindex(12));
- return 1;
- case BSON_NULL:
- lua_pushvalue(L, lua_upvalueindex(4));
- return 1;
- case BSON_JSCODE:
- case BSON_DBPOINTER:
- case BSON_SYMBOL:
- case BSON_CODEWS:
- lua_pushvalue(L, lua_upvalueindex(14));
- lua_pushlstring(L, (const char *)buf, sz);
- return 2;
- case BSON_INT64: {
- if (sz != 8) {
- return luaL_error(L, "Invalid int64");
- }
- int64_t d = *(const int64_t *)buf;
- lua_pushvalue(L, lua_upvalueindex(13));
- lua_pushinteger(L, d);
- return 2;
- }
- default:
- return luaL_error(L, "Invalid subtype %d", subtype);
- }
- }
- static int
- ltype(lua_State *L) {
- int t = lua_type(L,1);
- int type = 0;
- switch (t) {
- case LUA_TNUMBER:
- type = 1;
- break;
- case LUA_TBOOLEAN:
- type = 2;
- break;
- case LUA_TTABLE:
- type = 3;
- break;
- case LUA_TNIL:
- lua_pushvalue(L, lua_upvalueindex(4));
- return 1;
- case LUA_TSTRING: {
- size_t len = 0;
- const char * str = lua_tolstring(L,1,&len);
- if (str[0] == 0 && len >= 2) {
- return lsubtype(L, (uint8_t)str[1], (const uint8_t *)str+2, len-2);
- } else {
- type = 5;
- break;
- }
- }
- default:
- return luaL_error(L, "Invalid type %s",lua_typename(L,t));
- }
- lua_pushvalue(L, lua_upvalueindex(type));
- lua_pushvalue(L,1);
- return 2;
- }
- static void
- typeclosure(lua_State *L) {
- static const char * typename_[] = {
- "number", // 1
- "boolean", // 2
- "table", // 3
- "nil", // 4
- "string", // 5
- "binary", // 6
- "objectid", // 7
- "timestamp", // 8
- "date", // 9
- "regex", // 10
- "minkey", // 11
- "maxkey", // 12
- "int64", // 13
- "unsupported", // 14
- };
- int i;
- int n = sizeof(typename_)/sizeof(typename_[0]);
- for (i=0;i<n;i++) {
- lua_pushstring(L, typename_[i]);
- }
- lua_pushcclosure(L, ltype, n);
- }
- static uint8_t oid_header[5];
- static ATOM_ULONG oid_counter;
- static void
- init_oid_header() {
- if (ATOM_LOAD(&oid_counter)) {
- // already init
- return;
- }
- pid_t pid = getpid();
- uint32_t h = 0;
- char hostname[256];
- if (gethostname(hostname, sizeof(hostname))==0) {
- int i;
- for (i=0;i<sizeof(hostname) && hostname[i];i++) {
- h = h ^ ((h<<5)+(h>>2)+hostname[i]);
- }
- h ^= i;
- }
- oid_header[0] = h & 0xff;
- oid_header[1] = (h>>8) & 0xff;
- oid_header[2] = (h>>16) & 0xff;
- oid_header[3] = pid & 0xff;
- oid_header[4] = (pid >> 8) & 0xff;
-
- unsigned long c = h ^ time(NULL) ^ (uintptr_t)&h;
- if (c == 0) {
- c = 1;
- }
- ATOM_STORE(&oid_counter, c);
- }
- static inline int
- hextoint(char c) {
- if (c>='0' && c<='9')
- return c-'0';
- if (c>='a' && c<='z')
- return c-'a'+10;
- if (c>='A' && c<='Z')
- return c-'A'+10;
- return 0;
- }
- static int
- lobjectid(lua_State *L) {
- uint8_t oid[14] = { 0, BSON_OBJECTID };
- if (lua_isstring(L,1)) {
- size_t len;
- const char * str = lua_tolstring(L,1,&len);
- if (len != 24) {
- return luaL_error(L, "Invalid objectid %s", str);
- }
- int i;
- for (i=0;i<12;i++) {
- oid[i+2] = hextoint(str[i*2]) << 4 | hextoint(str[i*2+1]);
- }
- } else {
- time_t ti = time(NULL);
- // old_counter is a static var, use atom inc.
- uint32_t id = ATOM_FINC(&oid_counter);
- oid[2] = (ti>>24) & 0xff;
- oid[3] = (ti>>16) & 0xff;
- oid[4] = (ti>>8) & 0xff;
- oid[5] = ti & 0xff;
- memcpy(oid+6 , oid_header, 5);
- oid[11] = (id>>16) & 0xff;
- oid[12] = (id>>8) & 0xff;
- oid[13] = id & 0xff;
- }
- lua_pushlstring( L, (const char *)oid, 14);
- return 1;
- }
- LUAMOD_API int
- luaopen_bson(lua_State *L) {
- luaL_checkversion(L);
- int i;
- for (i=0;i<MAX_NUMBER;i++) {
- char tmp[8];
- bson_numstr_len[i] = sprintf(tmp,"%d",i);
- memcpy(bson_numstrs[i], tmp, bson_numstr_len[i]);
- }
- luaL_Reg l[] = {
- { "encode", lencode },
- { "encode_order", lencode_order },
- { "date", ldate },
- { "timestamp", ltimestamp },
- { "regex", lregex },
- { "binary", lbinary },
- { "objectid", lobjectid },
- { "int64", lint64 },
- { "decode", ldecode },
- { NULL, NULL },
- };
- luaL_newlib(L,l);
- typeclosure(L);
- lua_setfield(L,-2,"type");
- char null[] = { 0, BSON_NULL };
- lua_pushlstring(L, null, sizeof(null));
- lua_setfield(L,-2,"null");
- char minkey[] = { 0, (char)BSON_MINKEY };
- lua_pushlstring(L, minkey, sizeof(minkey));
- lua_setfield(L,-2,"minkey");
- char maxkey[] = { 0, BSON_MAXKEY };
- lua_pushlstring(L, maxkey, sizeof(maxkey));
- lua_setfield(L,-2,"maxkey");
- init_oid_header();
- return 1;
- }
|