forked from facebookarchive/iTorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
util.lua
98 lines (92 loc) · 3.18 KB
/
util.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
--[[
* Copyright (c) 2015, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree. An additional grant
* of patent rights can be found in the PATENTS file in the same directory.
*
]]--
local zmq = require 'lzmq'
local zassert = zmq.assert
local json = require 'cjson'
local uuid = require 'uuid'
local crypto = require 'crypto'
local util = {}
--------------------------------------------------------------
-- Signature Key and its setter
local session_key = ''
local function setSessionKey(key)
session_key = key
end
--------------------------------------------------------------
-- Common decoder function for all messages (except heartbeats which are just looped back)
local function ipyDecode(sock, m)
m = m or zassert(sock:recv_all())
-- print('incoming:')
-- print(m)
local o = {}
o.uuid = {}
local i = -1
for k,v in ipairs(m) do
if v == '<IDS|MSG>' then i = k+1; break; end
o.uuid[k] = v
end
assert(i ~= -1, 'Failed parsing till <IDS|MSG>')
-- json decode
for j=i+1,i+4 do if m[j] == '{}' then m[j] = nil; else m[j] = json.decode(m[j]); end; end
-- populate headers
o.header = m[i+1]
o.parent_header = m[i+2]
o.metadata = m[i+3]
o.content = m[i+4]
for j=i+5,#m do o.blob = (o.blob or '') .. m[j] end -- process blobs
return o
end
-- Common encoder function for all messages (except heartbeats which are just looped back)
-- See http://ipython.org/ipython-doc/stable/development/messaging.html
local function ipyEncodeAndSend(sock, m)
-- Message digest (for HMAC signature)
local d = crypto.hmac.new('sha256', session_key)
d:update(json.encode(m.header))
if m.parent_header then d:update(json.encode(m.parent_header)) else d:update('{}') end
if m.metadata then d:update(json.encode(m.metadata)) else d:update('{}') end
if m.content then d:update(json.encode(m.content)) else d:update('{}') end
local o = {}
for k,v in ipairs(m.uuid) do o[#o+1] = v end
o[#o+1] = '<IDS|MSG>'
o[#o+1] = d:final()
o[#o+1] = json.encode(m.header)
if m.parent_header then o[#o+1] = json.encode(m.parent_header) else o[#o+1] = '{}' end
if m.metadata then o[#o+1] = json.encode(m.metadata) else o[#o+1] = '{}' end
if m.content then o[#o+1] = json.encode(m.content) else o[#o+1] = '{}' end
if m.blob then o[#o+1] = m.blob end
-- print('outgoing:')
-- print(o)
zassert(sock:send_all(o))
end
local session_id = uuid.new()
-- function for creating a new message object
local function msg(msg_type, parent)
local m = {}
m.header = {}
if parent then
m.uuid = parent.uuid
m.parent_header = parent.header
else
m.parent_header = {}
end
m.header.msg_id = uuid.new()
m.header.msg_type = msg_type
m.header.session = session_id
m.header.date = os.date("%Y-%m-%dT%H:%M:%S")
m.header.username = 'itorch'
m.content = {}
return m
end
---------------------------------------------------------------------------
util.ipyDecode = ipyDecode
util.ipyEncodeAndSend = ipyEncodeAndSend
util.msg = msg
util.setSessionKey = setSessionKey
return util