vxcommon/lua/module.go

package lua

import (
	"errors"
	"runtime"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/sirupsen/logrus"
	"github.com/vxcontrol/luar"
	"github.com/vxcontrol/vxcommon/agent"
	"github.com/vxcontrol/vxcommon/vxproto"
)

type recvCallbacks struct {
	recvData   *luaCallback
	recvText   *luaCallback
	recvFile   *luaCallback
	recvMsg    *luaCallback
	controlMsg *luaCallback
}

// Module is struct that used for internal communication logic with Lua state
type Module struct {
	state    *State
	logger   *logrus.Entry
	socket   vxproto.IModuleSocket
	result   string
	cbs      recvCallbacks
	waitTime int64
	wgRun    sync.WaitGroup
	agents   map[string]*vxproto.AgentInfo
	args     map[string][]string
	quit     chan struct{}
	closed   bool
}

// IsClose is nonblocked function which check a state of module
func (m *Module) IsClose() bool {
	return m.closed || m.state == nil
}

// GetResult is nonblocked function which return result from module
func (m *Module) GetResult() string {
	return m.result
}

// Start is function which prepare state for module
func (m *Module) Start() {
	if m.state == nil {
		return
	}
	if m.closed && m.state.closed {
		m.quit = make(chan struct{})
		m.closed = false
	} else {
		return
	}

	m.wgRun.Add(1)
	go func() {
		defer m.wgRun.Done()
		m.recvPacket()
	}()
	m.logger.Info("the module was started")
	defer m.logger.Info("the module was stopped")

	var err error
	m.result = ""
	m.wgRun.Add(1)
	defer m.wgRun.Done()
	for m.result, err = m.state.Exec(); err != nil && !m.closed; {
		// TODO: here need using store message about problem into DB
		msg := "error executing the module code on the lua state"
		m.logger.WithError(err).WithField("result", m.result).Error(msg)
		time.Sleep(time.Second * time.Duration(5))
		m.state.L.SetTop(0)
	}
}

// Stop is function which set close state for module
func (m *Module) Stop() {
	if m.closed || m.state == nil {
		return
	}

	m.logger.Info("the module wants to stop")
	defer m.logger.Info("the module stopping has done")

	m.closed = true
	m.controlMsgCb("quit", "")
	close(m.quit)

	m.wgRun.Wait()
	m.delCbs([]interface{}{"data", "text", "file", "msg", "control"})
}

// Close is function which release lua state for module
func (m *Module) Close() {
	if m.state == nil {
		return
	}

	m.Stop()
	luar.Register(m.state.L, "__api", luar.Map{})
	luar.Register(m.state.L, "__agents", luar.Map{})
	luar.Register(m.state.L, "__routes", luar.Map{})
	luar.Register(m.state.L, "__imc", luar.Map{})
	m.state = nil
}

// SetAgents is function for storing agent list into module state
func (m *Module) SetAgents(agents map[string]*vxproto.AgentInfo) {
	m.agents = agents
}

// ControlMsg is function for send control message to module state
func (m *Module) ControlMsg(mtype, data string) bool {
	if m.closed || m.state == nil {
		return false
	}

	return m.controlMsgCb(mtype, data)
}

// Set timeout for all blocked functions
// If used timeout variable in non-zero value, it will wake up after timeout
// Timeout variable uses in milliseconds and -1 value means infinity
func (m *Module) setRecvTimeout(timeout int64) {
	m.waitTime = timeout
}

// await is blocked function which wait a close of module
// If used timeout variable in non-zero value, it will wake up after timeout
// Timeout variable uses in milliseconds and -1 value means infinity
func (m *Module) await(timeout int64) {
	if m.closed {
		return
	}

	m.state.L.Unlock()
	defer m.state.L.Lock()
	runtime.Gosched()

	if timeout >= 0 {
		select {
		case <-time.NewTimer(time.Millisecond * time.Duration(timeout)).C:
		case <-m.quit:
		}
	} else if timeout < 0 {
		<-m.quit
	}
}

func (m *Module) getName() string {
	return m.socket.GetName()
}

func (m *Module) getOS() string {
	return runtime.GOOS
}

func (m *Module) getArch() string {
	return runtime.GOARCH
}

func (m *Module) getAgents() map[string]vxproto.AgentInfo {
	agents := make(map[string]vxproto.AgentInfo, 0)
	for t, a := range m.agents {
		ca := *a
		if a.Info != nil {
			ca.Info = proto.Clone(a.Info).(*agent.Information)
			if a.Info.Os != nil {
				ca.Info.Os = proto.Clone(a.Info.Os).(*agent.Information_OS)
			}
			if a.Info.User != nil {
				ca.Info.User = proto.Clone(a.Info.User).(*agent.Information_User)
			}
		}
		agents[t] = ca
	}
	return agents
}

func (m *Module) getAgentsCount() int {
	return len(m.agents)
}

func (m *Module) getAgentsByID(agentID string) map[string]vxproto.AgentInfo {
	agents := m.getAgents()
	filtered := make(map[string]vxproto.AgentInfo, 0)
	for src, info := range agents {
		if info.ID == agentID {
			filtered[src] = info
		}
	}
	return filtered
}

func (m *Module) getAgentsBySrc(srcToken string) map[string]vxproto.AgentInfo {
	agents := m.getAgents()
	filtered := make(map[string]vxproto.AgentInfo, 0)
	for src, info := range agents {
		if info.Src == srcToken {
			filtered[src] = info
		}
	}
	return filtered
}

func (m *Module) getAgentsByDst(dstToken string) map[string]vxproto.AgentInfo {
	agents := m.getAgents()
	filtered := make(map[string]vxproto.AgentInfo, 0)
	for src, info := range agents {
		if info.Dst == dstToken {
			filtered[src] = info
		}
	}
	return filtered
}

func (m *Module) getIMCToken() string {
	return m.socket.GetIMCToken()
}

func (m *Module) getIMCTokenInfo(token string) (string, string, bool) {
	ms := m.socket.GetIMCModuleSocket(token)
	if ms == nil {
		return "", "", false
	}
	return ms.GetAgentID(), ms.GetName(), true
}

func (m *Module) isIMCTokenExist(token string) bool {
	ms := m.socket.GetIMCModuleSocket(token)
	if ms == nil {
		return false
	}
	return true
}

func (m *Module) makeIMCToken(agentID, moduleName string) string {
	return m.socket.MakeIMCToken(agentID, moduleName)
}

func (m *Module) getRoutes() map[string]string {
	return m.socket.GetRoutes()
}

func (m *Module) getRoutesCount() int {
	return len(m.socket.GetRoutes())
}

func (m *Module) getRoute(dst string) string {
	return m.socket.GetRoute(dst)
}

func (m *Module) addRoute(dst, src string) bool {
	m.logger.WithFields(logrus.Fields{
		"src": src,
		"dst": dst,
	}).Debug("the module added the new route")
	return m.socket.AddRoute(dst, src) == nil
}

func (m *Module) delRoute(dst string) bool {
	m.logger.WithFields(logrus.Fields{
		"dst": dst,
	}).Debug("the module deleted the route")
	return m.socket.DelRoute(dst) == nil
}

func (m *Module) sendDataTo(dst, data string) bool {
	if len(data) == 0 {
		return false
	}

	sdata := &vxproto.Data{
		Data: []byte(data),
	}

	m.logger.WithFields(logrus.Fields{
		"len": len(data),
		"dst": dst,
	}).Debug("the module sent data")
	return m.socket.SendDataTo(dst, sdata) == nil
}

func (m *Module) sendFileTo(dst, data, name string) bool {
	if len(data) == 0 || name == "" {
		return false
	}

	sfile := &vxproto.File{
		Data: []byte(data),
		Name: name,
	}

	m.logger.WithFields(logrus.Fields{
		"len":  len(data),
		"name": name,
		"dst":  dst,
	}).Debug("the module sent file from data")
	return m.socket.SendFileTo(dst, sfile) == nil
}

func (m *Module) sendFileFromFSTo(dst, path, name string) bool {
	if path == "" || name == "" {
		return false
	}

	sfile := &vxproto.File{
		Path: name,
		Name: name,
	}

	m.logger.WithFields(logrus.Fields{
		"path": path,
		"name": name,
		"dst":  dst,
	}).Debug("the module sent file from fs")
	return m.socket.SendFileTo(dst, sfile) == nil
}

func (m *Module) sendTextTo(dst, data, name string) bool {
	if len(data) == 0 || name == "" {
		return false
	}

	stext := &vxproto.Text{
		Data: []byte(data),
		Name: name,
	}

	m.logger.WithFields(logrus.Fields{
		"len":  len(data),
		"name": name,
		"dst":  dst,
	}).Debug("the module sent text from data")
	return m.socket.SendTextTo(dst, stext) == nil
}

func (m *Module) sendMsgTo(dst, data string, mtype int32) bool {
	if len(data) == 0 || mtype < 0 || mtype > 3 {
		return false
	}

	msg := &vxproto.Msg{
		Data:  []byte(data),
		MType: vxproto.MsgType(mtype),
	}

	m.logger.WithFields(logrus.Fields{
		"len":  len(data),
		"type": vxproto.MsgType(mtype).String(),
		"dst":  dst,
	}).Debug("the module sent message")
	return m.socket.SendMsgTo(dst, msg) == nil
}

func (m *Module) recvDataCb(src string, data *vxproto.Data) bool {
	if m.cbs.recvData != nil {
		res := new(bool)
		m.cbs.recvData.Call(&res, src, string(data.Data[:]))
		return *res
	}

	return false
}

func (m *Module) recvFileCb(src string, file *vxproto.File) bool {
	if m.cbs.recvFile != nil {
		res := new(bool)
		m.cbs.recvFile.Call(&res, src, file.Path, file.Name)
		return *res
	}

	return false
}

func (m *Module) recvTextCb(src string, text *vxproto.Text) bool {
	if m.cbs.recvText != nil {
		res := new(bool)
		m.cbs.recvText.Call(&res, src, string(text.Data[:]), text.Name)
		return *res
	}

	return false
}

func (m *Module) recvMsgCb(src string, msg *vxproto.Msg) bool {
	if m.cbs.recvMsg != nil {
		res := new(bool)
		m.cbs.recvMsg.Call(&res, src, string(msg.Data[:]), int32(msg.MType))
		return *res
	}

	return false
}

func (m *Module) controlMsgCb(mtype, data string) bool {
	if m.cbs.controlMsg != nil {
		res := new(bool)
		m.cbs.controlMsg.Call(&res, mtype, data)
		return *res
	}

	return false
}

func (m *Module) recvData() (string, string, bool) {
	src, data, err := m.socket.RecvData(m.waitTime)
	if err != nil {
		return "", "", false
	}

	return src, string(data.Data[:]), true
}

func (m *Module) recvFile() (string, string, string, bool) {
	src, file, err := m.socket.RecvFile(m.waitTime)
	if err != nil {
		return "", "", "", false
	}

	return src, file.Path, file.Name, true
}

func (m *Module) recvText() (string, string, string, bool) {
	src, text, err := m.socket.RecvText(m.waitTime)
	if err != nil {
		return "", "", "", false
	}

	return src, string(text.Data[:]), text.Name, true
}

func (m *Module) recvMsg() (string, string, int32, bool) {
	src, msg, err := m.socket.RecvMsg(m.waitTime)
	if err != nil {
		return "", "", 0, false
	}

	return src, string(msg.Data[:]), int32(msg.MType), true
}

func (m *Module) recvDataFrom(src string) (string, bool) {
	data, err := m.socket.RecvDataFrom(src, m.waitTime)
	if err != nil {
		return "", false
	}

	return string(data.Data[:]), true
}

func (m *Module) recvFileFrom(src string) (string, string, bool) {
	file, err := m.socket.RecvFileFrom(src, m.waitTime)
	if err != nil {
		return "", "", false
	}

	return file.Path, file.Name, true
}

func (m *Module) recvTextFrom(src string) (string, string, bool) {
	text, err := m.socket.RecvTextFrom(src, m.waitTime)
	if err != nil {
		return "", "", false
	}

	return string(text.Data[:]), text.Name, true
}

func (m *Module) recvMsgFrom(src string) (string, int32, bool) {
	msg, err := m.socket.RecvMsgFrom(src, m.waitTime)
	if err != nil {
		return "", 0, false
	}

	return string(msg.Data[:]), int32(msg.MType), true
}

func (m *Module) addCbs(callbackTable interface{}) bool {
	callbackMap, ok := callbackTable.(map[string]interface{})
	if !ok {
		return false
	}

	for name, callback := range callbackMap {
		switch name {
		case "data":
			if cb, ok := callback.(*luar.LuaObject); ok {
				cb.Push()
				m.cbs.recvData = newLuaCallback(m.state.L)
				cb.Close()
				m.logger.Debug("the module added receive data callback")
			}
		case "text":
			if cb, ok := callback.(*luar.LuaObject); ok {
				cb.Push()
				m.cbs.recvText = newLuaCallback(m.state.L)
				cb.Close()
				m.logger.Debug("the module added receive text callback")
			}
		case "file":
			if cb, ok := callback.(*luar.LuaObject); ok {
				cb.Push()
				m.cbs.recvFile = newLuaCallback(m.state.L)
				cb.Close()
				m.logger.Debug("the module added receive file callback")
			}
		case "msg":
			if cb, ok := callback.(*luar.LuaObject); ok {
				cb.Push()
				m.cbs.recvMsg = newLuaCallback(m.state.L)
				cb.Close()
				m.logger.Debug("the module added receive message callback")
			}
		case "control":
			if cb, ok := callback.(*luar.LuaObject); ok {
				cb.Push()
				m.cbs.controlMsg = newLuaCallback(m.state.L)
				cb.Close()
				m.logger.Debug("the module added receive control message callback")
			}
		default:
		}
	}

	return true
}

func (m *Module) delCbs(CallbackTable interface{}) bool {
	callbackMap, ok := CallbackTable.([]interface{})
	if !ok {
		return false
	}

	for _, name := range callbackMap {
		switch name.(string) {
		case "data":
			if m.cbs.recvData != nil {
				m.cbs.recvData.Close()
				m.cbs.recvData = nil
				m.logger.Debug("the module deleted receive data callback")
			}
		case "text":
			if m.cbs.recvText != nil {
				m.cbs.recvText.Close()
				m.cbs.recvText = nil
				m.logger.Debug("the module deleted receive text callback")
			}
		case "file":
			if m.cbs.recvFile != nil {
				m.cbs.recvFile.Close()
				m.cbs.recvFile = nil
				m.logger.Debug("the module deleted receive file callback")
			}
		case "msg":
			if m.cbs.recvMsg != nil {
				m.cbs.recvMsg.Close()
				m.cbs.recvMsg = nil
				m.logger.Debug("the module deleted receive message callback")
			}
		case "control":
			if m.cbs.controlMsg != nil {
				m.cbs.controlMsg.Close()
				m.cbs.controlMsg = nil
				m.logger.Debug("the module deleted receive control message callback")
			}
		default:
		}
	}

	return true
}

func (m *Module) recvPacket() error {
	defer m.logger.Info("packet receiver was stopped")

	m.logger.Info("packet receiver was started")
	receiver := m.socket.GetReceiver()
	if receiver == nil {
		m.logger.Error("failed to initialize packet receiver")
		return errors.New("failed to initialize packet receiver")
	}
	for !m.closed {
		var packet *vxproto.Packet
		select {
		case packet = <-receiver:
		case <-m.quit:
			m.logger.Info("got signal to quit from channel")
			return nil
		}

		if packet == nil {
			m.logger.Error("failed receive packet")
			return errors.New("failed receive packet")
		}

		logger := m.logger.WithFields(logrus.Fields{
			"type": packet.PType.String(),
			"src":  packet.Src,
			"dst":  packet.Dst,
		})
		logger.Debug("packet receiver got new packet")
		switch packet.PType {
		case vxproto.PTData:
			m.recvDataCb(packet.Src, packet.GetData())
		case vxproto.PTFile:
			m.recvFileCb(packet.Src, packet.GetFile())
		case vxproto.PTText:
			m.recvTextCb(packet.Src, packet.GetText())
		case vxproto.PTMsg:
			m.recvMsgCb(packet.Src, packet.GetMsg())
		case vxproto.PTControl:
			msg := packet.GetControlMsg()
			switch msg.MsgType {
			case vxproto.AgentConnected:
				logger.Info("agent connected to the module")
				m.agents[msg.AgentInfo.Dst] = msg.AgentInfo
				m.controlMsgCb("agent_connected", msg.AgentInfo.Dst)
			case vxproto.AgentDisconnected:
				logger.Info("agent disconnected from the module")
				m.controlMsgCb("agent_disconnected", msg.AgentInfo.Dst)
				delete(m.agents, msg.AgentInfo.Dst)
			case vxproto.StopModule:
				logger.Info("got packet with signal to stop module")
				return nil
			}
		default:
			logger.Error("got packet has unexpected packet type")
			return errors.New("unexpected packet type")
		}
	}

	return nil
}

// NewModule is function which constructed Module object
func NewModule(args map[string][]string, state *State, socket vxproto.IModuleSocket) (*Module, error) {
	if socket == nil {
		logrus.Error("failed to make new module because socket object unset")
		return nil, errors.New("socket object not initialized")
	}

	m := &Module{
		socket: socket,
		state:  state,
		args:   args,
		agents: make(map[string]*vxproto.AgentInfo),
		closed: true,
		logger: logrus.WithFields(logrus.Fields{
			"component": "module",
			"module":    socket.GetName(),
			"agent":     socket.GetAgentID(),
		}),
	}

	luar.Register(state.L, "__api", luar.Map{
		// Functions
		"await":    m.await,
		"is_close": m.IsClose,
		"get_name": m.getName,
		"get_os":   m.getOS,
		"get_arch": m.getArch,
		"unsafe": luar.Map{
			"lock":   func() { m.state.L.Lock() },
			"unlock": func() { m.state.L.Unlock() },
		},

		"add_cbs":          m.addCbs,
		"del_cbs":          m.delCbs,
		"set_recv_timeout": m.setRecvTimeout,

		"send_data_to":         m.sendDataTo,
		"send_file_to":         m.sendFileTo,
		"send_text_to":         m.sendTextTo,
		"send_msg_to":          m.sendMsgTo,
		"send_file_from_fs_to": m.sendFileFromFSTo,

		"recv_data": m.recvData,
		"recv_file": m.recvFile,
		"recv_text": m.recvText,
		"recv_msg":  m.recvMsg,

		"recv_data_from": m.recvDataFrom,
		"recv_file_from": m.recvFileFrom,
		"recv_text_from": m.recvTextFrom,
		"recv_msg_from":  m.recvMsgFrom,
	})

	luar.Register(state.L, "__agents", luar.Map{
		// Functions
		"dump":       m.getAgents,
		"count":      m.getAgentsCount,
		"get_by_id":  m.getAgentsByID,
		"get_by_src": m.getAgentsBySrc,
		"get_by_dst": m.getAgentsByDst,
	})

	luar.Register(state.L, "__routes", luar.Map{
		// Functions
		"dump":  m.getRoutes,
		"count": m.getRoutesCount,
		"get":   m.getRoute,
		"add":   m.addRoute,
		"del":   m.delRoute,
	})

	luar.Register(state.L, "__imc", luar.Map{
		// Functions
		"get_token":  m.getIMCToken,
		"get_info":   m.getIMCTokenInfo,
		"is_exist":   m.isIMCTokenExist,
		"make_token": m.makeIMCToken,
	})

	luar.GoToLua(state.L, args)
	state.L.SetGlobal("__args")

	// TODO: change it to native load function
	state.L.DoString(`
	io.stdout:setvbuf('no')

	function __api.async(f, ...)
		local glue = require("glue")
		__api.unsafe.unlock()
		t = glue.pack(f(...))
		__api.unsafe.lock()
		return glue.unpack(t)
	end

	function __api.sync(f, ...)
		local glue = require("glue")
		__api.unsafe.lock()
		t = glue.pack(f(...))
		__api.unsafe.unlock()
		return glue.unpack(t)
	end
	`)

	m.logger.Info("the module was created")
	return m, nil
}