MAYBE: Move base64 handling into antietcd

antietcd.js

@@ -1,7 +1,10 @@
 const fs = require('fs');
+const fsp = require('fs').promises;
 const { URL } = require('url');
 const http = require('http');
 const https = require('https');
+const crypto = require('crypto');
+const zlib = require('zlib');
 
 const ws = require('ws');
 
@@ -19,27 +22,82 @@ class RequestError
 class AntiEtcd
 {
     constructor(cfg)
-    {
-        this.cfg = cfg;
-    }
-
-    run()
     {
         this.clients = {};
         this.client_id = 1;
         this.etctree = new EtcTree(true);
+        this.cfg = cfg;
+        this.stopped = false;
+        this.cfg.use_base64 = true;
+    }
+
+    async run()
+    {
+        if (this.cfg.filename)
+        {
+            // Load data from disk
+            const [ err, stat ] = await new Promise(ok => fs.stat(this.cfg.filename, (err, stat) => ok([ err, stat ])));
+            if (!err)
+            {
+                const data = await fsp.readFile(this.cfg.filename);
+                data = JSON.parse(zlib.Gunzip(data));
+                this.etctree.load(data);
+            }
+            else if (err.code != ENOENT)
+            {
+                throw err;
+            }
+            // Set exit hook
+            const on_stop_cb = () => this.on_stop();
+            process.on('SIGINT', on_stop_cb);
+            process.on('SIGTERM', on_stop_cb);
+            process.on('SIGQUIT', on_stop_cb);
+        }
         if (this.cfg.cert)
         {
-            this.tls = { key: fs.readFileSync(this.cfg.key), cert: fs.readFileSync(this.cfg.cert) };
+            this.tls = { key: await fsp.readFile(this.cfg.key), cert: await fsp.readFile(this.cfg.cert) };
             this.server = https.createServer(this.tls, (req, res) => this.handleRequest(req, res));
         }
         else
+        {
             this.server = http.createServer((req, res) => this.handleRequest(req, res));
+        }
         this.wss = new ws.WebSocketServer({ server: this.server });
         this.wss.on('connection', (conn, req) => this.startWebsocket(conn, req));
         this.server.listen(this.cfg.port || 2379);
     }
 
+    async on_stop()
+    {
+        if (this.stopped)
+        {
+            return;
+        }
+        this.stopped = true;
+        // Wait until all requests complete
+        while (this.inflight > 0)
+        {
+            await new Promise(ok => this.waitInflight.push(ok));
+        }
+        await this.persist();
+    }
+
+    async persist()
+    {
+        if (!this.cfg.filename)
+        {
+            return;
+        }
+        let dump = this.etctree.dump(true);
+        dump = JSON.stringify(dump);
+        dump = zlib.Gzip(dump);
+        const fd = await fsp.open(this.cfg.filename+'.tmp', 'w');
+        await fsp.write(fd, dump);
+        await fsp.fsync(fd);
+        await fsp.close(fd);
+        await fsp.rename(this.cfg.filename+'.tmp', this.cfg.filename);
+    }
+
     fail(res, code, text)
     {
         res.writeHead(code);
@@ -51,7 +109,7 @@ class AntiEtcd
     {
         let data = [];
         req.on('data', (chunk) => data.push(chunk));
-        req.on('end', () =>
+        req.on('end', async () =>
         {
             data = Buffer.concat(data);
             let body = '';
@@ -76,7 +134,7 @@ class AntiEtcd
                 {
                     throw new RequestError(400, 'body should be JSON object');
                 }
-                reply = this.runHandler(req, data, res);
+                reply = await this.runHandler(req, data, res);
                 reply = JSON.stringify(reply);
             }
             catch (e)
@@ -93,21 +151,28 @@ class AntiEtcd
                     reply = 'Internal error: '+e.message;
                 }
             }
-            // Access log
-            console.log(
-                new Date().toISOString()+
-                ' '+(req.headers['x-forwarded-for'] || req.socket.remoteAddress)+
-                ' '+req.method+' '+req.url+' '+code+'\n  '+body.replace(/\n/g, '\\n')+
-                '\n  '+reply.replace(/\n/g, '\\n')
-            );
-            // FIXME: Access log :req[X-Forwarded-For] [:date[clf]] pid=:pid ":method :url HTTP/:http-version" :status :res[content-length] ":referrer" ":user-agent" - :response-time ms
-            reply = Buffer.from(reply);
-            res.writeHead(200, {
-                'Content-Type': 'application/json',
-                'Content-Length': reply.length,
-            });
-            res.write(reply);
-            res.end();
+            try
+            {
+                // Access log
+                console.log(
+                    new Date().toISOString()+
+                    ' '+(req.headers['x-forwarded-for'] || req.socket.remoteAddress)+
+                    ' '+req.method+' '+req.url+' '+code+'\n  '+body.replace(/\n/g, '\\n')+
+                    '\n  '+reply.replace(/\n/g, '\\n')
+                );
+                // FIXME: Access log :req[X-Forwarded-For] [:date[clf]] pid=:pid ":method :url HTTP/:http-version" :status :res[content-length] ":referrer" ":user-agent" - :response-time ms
+                reply = Buffer.from(reply);
+                res.writeHead(200, {
+                    'Content-Type': 'application/json',
+                    'Content-Length': reply.length,
+                });
+                res.write(reply);
+                res.end();
+            }
+            catch (e)
+            {
+                console.error(e);
+            }
         });
     }
 
@@ -163,7 +228,7 @@ class AntiEtcd
         });
     }
 
-    runHandler(req, data, res)
+    async runHandler(req, data, res)
     {
         // v3/kv/txn
         // v3/kv/range
@@ -187,30 +252,74 @@ class AntiEtcd
                 return cb.call(this, data);
             }
         }
+        else if (requestUrl.pathname == '/dump')
+        {
+            return this.handle_dump(data);
+        }
         throw new RequestError(404, 'Supported APIs: /v3/kv/txn, /v3/kv/range, /v3/kv/put, /v3/kv/deleterange, '+
             '/v3/lease/grant, /v3/lease/revoke, /v3/kv/lease/revoke, /v3/lease/keepalive');
     }
 
     handle_kv_txn(data)
     {
-        return this.etctree.api_txn(data);
+        if (this.cfg.use_base64)
+        {
+            for (const item of data.compare||[])
+            {
+                if (item.key != null)
+                    item.key = this.de64(item.key);
+            }
+            for (const items of [ data.success, data.failure ])
+            {
+                for (const item of items||[])
+                {
+                    const req = item.request_range || item.requestRange ||
+                        item.request_put || item.requestPut ||
+                        item.request_delete_range || item.requestDeleteRange;
+                    if (req.key != null)
+                        req.key = this.de64(req.key);
+                    if (req.range_end != null)
+                        req.range_end = this.de64(req.range_end);
+                    if (req.value != null)
+                        req.value = this.de64(req.value);
+                }
+            }
+        }
+        const result = this.etctree.api_txn(data);
+        if (this.cfg.use_base64)
+        {
+            for (const item of result.responses||[])
+            {
+                if (item.response_range)
+                {
+                    for (const kv of item.response_range.kvs)
+                    {
+                        if (kv.key != null)
+                            kv.key = this.b64(kv.key);
+                        if (kv.value != null)
+                            kv.value = this.b64(kv.value);
+                    }
+                }
+            }
+        }
+        return result;
     }
 
     handle_kv_range(data)
     {
-        const r = this.etctree.api_txn({ success: [ { request_range: data } ] });
+        const r = this.handle_kv_txn({ success: [ { request_range: data } ] });
         return { header: r.header, ...r.responses[0].response_range };
     }
 
     handle_kv_put(data)
     {
-        const r = this.etctree.api_txn({ success: [ { request_put: data } ] });
+        const r = this.handle_kv_txn({ success: [ { request_put: data } ] });
         return { header: r.header, ...r.responses[0].response_put };
     }
 
     handle_kv_deleterange(data)
     {
-        const r = this.etctree.api_txn({ success: [ { request_delete_range: data } ] });
+        const r = this.handle_kv_txn({ success: [ { request_delete_range: data } ] });
         return { header: r.header, ...r.responses[0].response_delete_range };
     }
 
@@ -234,6 +343,11 @@ class AntiEtcd
         return this.etctree.api_keepalive_lease(data);
     }
 
+    handle_dump(data)
+    {
+        return this.etctree.dump();
+    }
+
     handleMessage(client_id, msg, socket)
     {
         if (msg.create_request)
@@ -242,9 +356,15 @@ class AntiEtcd
             const create_request = msg.create_request;
             if (!create_request.watch_id || !this.clients[client_id].watches[create_request.watch_id])
             {
-                const watch = this.etctree.api_create_watch(
-                    { ...create_request, watch_id: null }, (msg) => socket.send(JSON.stringify(msg))
-                );
+                const req = { ...create_request, watch_id: null };
+                if (this.cfg.use_base64)
+                {
+                    if (req.key != null)
+                        req.key = this.de64(req.key);
+                    if (req.range_end != null)
+                        req.range_end = this.de64(req.range_end);
+                }
+                const watch = this.etctree.api_create_watch(req, (msg) => this.sendToSocket(socket, msg));
                 if (!watch.created)
                 {
                     socket.send(JSON.stringify({ result: { header: { revision: this.etctree.mod_revision }, watch_id: create_request.watch_id, ...watch } }));
@@ -273,6 +393,31 @@ class AntiEtcd
         }
     }
 
+    sendToSocket(socket, msg)
+    {
+        if (this.cfg.use_base64 && msg.result && msg.result.events)
+        {
+            msg = { ...msg, result: { ...msg.result, events: msg.result.events.map(ev =>
+            {
+                if (ev.kv)
+                {
+                    const kv = { ...ev.kv };
+                    if (kv.key != null)
+                    {
+                        kv.key = this.b64(kv.key);
+                    }
+                    if (kv.value != null)
+                    {
+                        kv.value = this.b64(kv.value);
+                    }
+                    return { ...ev, kv };
+                }
+                return ev;
+            }) } };
+        }
+        socket.send(JSON.stringify(msg));
+    }
+
     unsubscribeClient(client_id)
     {
         if (!this.clients[client_id])
@@ -283,6 +428,16 @@ class AntiEtcd
             this.etctree.api_cancel_watch({ watch_id: mapped_id });
         }
     }
+
+    de64(k)
+    {
+        return Buffer.from(k, 'base64').toString();
+    }
+
+    b64(k)
+    {
+        return Buffer.from(k).toString('base64');
+    }
 }
 
-new AntiEtcd({ port: 12379 }).run();
+new AntiEtcd({ port: 12379 }).run().catch(console.error);

etctree.js

@@ -14,35 +14,20 @@ const stableStringify = require('./stable-stringify.js');
 
 class EtcTree
 {
-    constructor(use_base64)
+    constructor()
     {
         this.state = {};
         this.leases = {};
         this.watchers = {};
         this.watcher_id = 0;
         this.mod_revision = 0;
-        this.use_base64 = use_base64;
         this.paused = false;
         this.on_expire_lease = null;
     }
 
-    de64(k)
-    {
-        if (k == null) // null or undefined
-            return k;
-        return this.use_base64 ? Buffer.from(k, 'base64').toString() : k;
-    }
-
-    b64(k)
-    {
-        if (k == null) // null or undefined
-            return k;
-        return this.use_base64 ? Buffer.from(k).toString('base64') : k;
-    }
-
     check(chk)
     {
-        const parts = this.key_parts(this.de64(chk.key));
+        const parts = this.key_parts(chk.key);
         const { cur } = this.get_subtree(parts, false, false);
         let check_value, ref_value;
         if (chk.target === 'MOD')
@@ -88,8 +73,8 @@ class EtcTree
 
     get_range(req)
     {
-        const key = this.de64(req.key);
-        const end = this.de64(req.range_end);
+        const key = req.key;
+        const end = req.range_end;
         if (end != null && (key[key.length-1] != '/' || end[end.length-1] != '0' ||
             end.substr(0, end.length-1) !== key.substr(0, key.length-1)))
         {
@@ -439,9 +424,9 @@ class EtcTree
         {
             events.push({
                 type: cur.value == null ? 'DELETE' : 'PUT',
-                kv: cur.value == null ? { key: this.b64(prefix === null ? '' : prefix) } : {
-                    key: this.b64(prefix),
-                    value: this.b64(cur.value),
+                kv: cur.value == null ? { key: (prefix === null ? '' : prefix) } : {
+                    key: prefix,
+                    value: cur.value,
                 },
             });
         }
@@ -545,9 +530,9 @@ class EtcTree
         {
             const request_put = req.request_put || req.requestPut;
             // FIXME: prev_kv, ignore_value(?), ignore_lease(?)
-            const parts = this.key_parts(this.de64(request_put.key));
+            const parts = this.key_parts(request_put.key);
             const key = parts.join('/');
-            const value = this.de64(request_put.value);
+            const value = request_put.value;
             const { cur, watchers } = this.get_subtree(parts, true, true);
             if (cur.key_watchers)
             {
@@ -580,7 +565,7 @@ class EtcTree
                     cur.create_revision = cur_revision;
                 }
                 cur.value = value;
-                const notify = { watchers, key: this.b64(key), value: this.b64(value), mod_revision: cur.mod_revision };
+                const notify = { watchers, key, value, mod_revision: cur.mod_revision };
                 if (cur.lease)
                 {
                     notify.lease = cur.lease;
@@ -613,10 +598,10 @@ class EtcTree
         }
         if (cur.value != null)
         {
-            const item = { key: this.b64(prefix === null ? '' : prefix) };
+            const item = { key: (prefix === null ? '' : prefix) };
             if (!req.keys_only)
             {
-                item.value = this.b64(cur.value);
+                item.value = cur.value;
                 item.mod_revision = cur.mod_revision;
                 //item.create_revision = cur.create_revision;
                 //item.version = cur.version;
@@ -653,7 +638,7 @@ class EtcTree
             this.mod_revision = cur_revision;
             notifications.push({
                 watchers: cur.key_watchers ? [ ...watchers, ...cur.key_watchers ] : watchers,
-                key: this.b64(prefix === null ? '' : prefix),
+                key: (prefix === null ? '' : prefix),
                 mod_revision: cur_revision,
             });
         }