<commit_before>--[[
Copyright 2012 Rackspace

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS-IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
--]]

local os = require('os')
local timer = require('timer')
local AgentProtocol = require('./protocol')
local Emitter = require('core').Emitter
local Error = require('core').Error
local JSON = require('json')
local fmt = require('string').format
local logging = require('logging')
local msg = require ('./messages')
local table = require('table')
local utils = require('utils')

-- Response timeouts in ms
local HANDSHAKE_TIMEOUT = 30000

local STATES = {}
STATES.INITIAL = 1
STATES.HANDSHAKE = 2
STATES.RUNNING = 3

local AgentProtocolConnection = Emitter:extend()

function AgentProtocolConnection:initialize(log, myid, token, conn)
  assert(conn ~= nil)
  assert(myid ~= nil)

  self._log = log
  self._myid = myid
  self._token = token
  self._conn = conn
  self._conn:on('data', utils.bind(AgentProtocolConnection._onData, self))
  self._buf = ""
  self._msgid = 0
  self._endpoints = { }
  self._target = 'endpoint'
  self._timeoutIds = {}
  self._completions = {}
  self:setState(STATES.INITIAL)
end

function AgentProtocolConnection:_onData(data)
  local client = self._conn, obj
  newline = data:find("\n")
  if newline then
    -- TODO: use a better buffer
    self._buf = self._buf .. data:sub(1, newline - 1)
    self._log(logging.DEBUG, fmt('RECV: %s', self._buf))
    obj = JSON.parse(self._buf)
    self:_processMessage(obj)
    self._buf = data:sub(newline + 1)
  else
    self._buf = self._buf .. data
  end
end

function AgentProtocolConnection:_processMessage(msg)
  -- request
  if msg.method ~= nil then
    self:emit('message', msg)
  else
    -- response
    local key = msg.source .. ':' .. msg.id
    local callback = self._completions[key]
    if callback then
      self._completions[key] = nil
      callback(null, msg)
    end
  end
end

function AgentProtocolConnection:_send(msg, timeout, expectedCode, callback)
  msg.target = 'endpoint'
  msg.source = self._myid
  local data = JSON.stringify(msg) .. '\n'
  local key = msg.target .. ':' .. msg.id

  if not expectedCode then expectedCode = 200 end

  self._log(logging.DEBUG, fmt('SEND: %s', data))

  if timeout then
    self:_setCommandTimeoutHandler(key, timeout, callback)
  end

  if callback then
    self._completions[key] = function(err, msg)
      local result = nil

      if msg and msg.result then result = msg.result end

      if self._timeoutIds[key] ~= nil then
        timer.clearTimer(self._timeoutIds[key])
      end

      if not err and msg and result and result.code and result.code ~= expectedCode then
        err = Error:new(fmt('Unexpected status code returned: code=%s, message=%s', result.code, result.message))
      end

      callback(err, msg)
    end
  end

  self._conn:write(data)
  self._msgid = self._msgid + 1
end

--[[
Set a timeout handler for a function.

key - Command key.
timeout - Timeout in ms.
callback - Callback which is called with (err) if timeout has been reached.
]]--
function AgentProtocolConnection:_setCommandTimeoutHandler(key, timeout, callback)
  local timeoutId

  timeoutId = timer.setTimeout(timeout, function()
    callback(Error:new(fmt('Command timeout, haven\'t received response in %d ms', timeout)))
  end)
  self._timeoutIds[key] = timeoutId
end

--[[ Protocol Functions ]]--

function AgentProtocolConnection:sendHandshakeHello(agentId, token, callback)
  local m = msg.HandshakeHello:new(token, agentId)
  self:_send(m:serialize(self._msgid), HANDSHAKE_TIMEOUT, 200, callback)
end

function AgentProtocolConnection:sendPing(timestamp, callback)
  local m = msg.Ping:new(timestamp)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendSystemInfo(request, callback)
  local m = msg.SystemInfoResponse:new(request)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendManifest(callback)
  local m = msg.Manifest:new()
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendMetrics(check, checkResults, callback)
  local m = msg.MetricsRequest:new(check, checkResults)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

--[[ Public Functions ]] --

function AgentProtocolConnection:setState(state)
  self._state = state
end

function AgentProtocolConnection:startHandshake(callback)
  self:setState(STATES.HANDSHAKE)
  self:sendHandshakeHello(self._myid, self._token, function(err, msg)
    if err then
      self._log(logging.ERR, fmt('handshake failed (message=%s)', err.message))
      callback(err, msg)
      return
    end

    if msg.result.code and msg.result.code ~= 200 then
      err = Error:new(fmt('handshake failed [message=%s,code=%s]', msg.result.message, msg.result.code))
      self._log(logging.ERR, err.message)
      callback(err, msg)
      return
    end

    self:setState(STATES.RUNNING)
    self._log(logging.INFO, fmt('handshake successful (ping_interval=%dms)', msg.result.ping_interval))
    callback(nil, msg)
  end)
end

function AgentProtocolConnection:getManifest(callback)
  self:sendManifest(function(err, response)
    if err then
      callback(err)
    else
      callback(nil, response.result)
    end
  end)
end

--[[
Process an async message

msg - The Incoming Message
]]--
function AgentProtocolConnection:execute(msg)
  if msg.method == 'system.info' then
    self:sendSystemInfo(msg)
  else
    local err = Error:new(fmt('invalid method [method=%s]', msg.method))
    self:emit('error', err)
  end
end


return AgentProtocolConnection
<commit_msg>Don't explode on corrupted JSON.<commit_after>--[[
Copyright 2012 Rackspace

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS-IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
--]]

local os = require('os')
local timer = require('timer')
local AgentProtocol = require('./protocol')
local Emitter = require('core').Emitter
local Error = require('core').Error
local JSON = require('json')
local fmt = require('string').format
local logging = require('logging')
local msg = require ('./messages')
local table = require('table')
local utils = require('utils')

-- Response timeouts in ms
local HANDSHAKE_TIMEOUT = 30000

local STATES = {}
STATES.INITIAL = 1
STATES.HANDSHAKE = 2
STATES.RUNNING = 3

local AgentProtocolConnection = Emitter:extend()

function AgentProtocolConnection:initialize(log, myid, token, conn)
  assert(conn ~= nil)
  assert(myid ~= nil)

  self._log = log
  self._myid = myid
  self._token = token
  self._conn = conn
  self._conn:on('data', utils.bind(AgentProtocolConnection._onData, self))
  self._buf = ""
  self._msgid = 0
  self._endpoints = { }
  self._target = 'endpoint'
  self._timeoutIds = {}
  self._completions = {}
  self:setState(STATES.INITIAL)
end

function AgentProtocolConnection:_onData(data)
  local client = self._conn, obj, status
  newline = data:find("\n")
  if newline then
    -- TODO: use a better buffer
    self._buf = self._buf .. data:sub(1, newline - 1)
    self._log(logging.DEBUG, fmt('RECV: %s', self._buf))
    status, obj = pcall(JSON.parse, self._buf)

    self._buf = data:sub(newline + 1)

    if not status then
      self._log(logging.ERROR, fmt('Failed to parse incoming line: line=%s,err=%s', self._buf, obj))
      return
    end

    self:_processMessage(obj)
  else
    self._buf = self._buf .. data
  end
end

function AgentProtocolConnection:_processMessage(msg)
  -- request
  if msg.method ~= nil then
    self:emit('message', msg)
  else
    -- response
    local key = msg.source .. ':' .. msg.id
    local callback = self._completions[key]
    if callback then
      self._completions[key] = nil
      callback(null, msg)
    end
  end
end

function AgentProtocolConnection:_send(msg, timeout, expectedCode, callback)
  msg.target = 'endpoint'
  msg.source = self._myid
  local data = JSON.stringify(msg) .. '\n'
  local key = msg.target .. ':' .. msg.id

  if not expectedCode then expectedCode = 200 end

  self._log(logging.DEBUG, fmt('SEND: %s', data))

  if timeout then
    self:_setCommandTimeoutHandler(key, timeout, callback)
  end

  if callback then
    self._completions[key] = function(err, msg)
      local result = nil

      if msg and msg.result then result = msg.result end

      if self._timeoutIds[key] ~= nil then
        timer.clearTimer(self._timeoutIds[key])
      end

      if not err and msg and result and result.code and result.code ~= expectedCode then
        err = Error:new(fmt('Unexpected status code returned: code=%s, message=%s', result.code, result.message))
      end

      callback(err, msg)
    end
  end

  self._conn:write(data)
  self._msgid = self._msgid + 1
end

--[[
Set a timeout handler for a function.

key - Command key.
timeout - Timeout in ms.
callback - Callback which is called with (err) if timeout has been reached.
]]--
function AgentProtocolConnection:_setCommandTimeoutHandler(key, timeout, callback)
  local timeoutId

  timeoutId = timer.setTimeout(timeout, function()
    callback(Error:new(fmt('Command timeout, haven\'t received response in %d ms', timeout)))
  end)
  self._timeoutIds[key] = timeoutId
end

--[[ Protocol Functions ]]--

function AgentProtocolConnection:sendHandshakeHello(agentId, token, callback)
  local m = msg.HandshakeHello:new(token, agentId)
  self:_send(m:serialize(self._msgid), HANDSHAKE_TIMEOUT, 200, callback)
end

function AgentProtocolConnection:sendPing(timestamp, callback)
  local m = msg.Ping:new(timestamp)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendSystemInfo(request, callback)
  local m = msg.SystemInfoResponse:new(request)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendManifest(callback)
  local m = msg.Manifest:new()
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

function AgentProtocolConnection:sendMetrics(check, checkResults, callback)
  local m = msg.MetricsRequest:new(check, checkResults)
  self:_send(m:serialize(self._msgid), nil, 200, callback)
end

--[[ Public Functions ]] --

function AgentProtocolConnection:setState(state)
  self._state = state
end

function AgentProtocolConnection:startHandshake(callback)
  self:setState(STATES.HANDSHAKE)
  self:sendHandshakeHello(self._myid, self._token, function(err, msg)
    if err then
      self._log(logging.ERR, fmt('handshake failed (message=%s)', err.message))
      callback(err, msg)
      return
    end

    if msg.result.code and msg.result.code ~= 200 then
      err = Error:new(fmt('handshake failed [message=%s,code=%s]', msg.result.message, msg.result.code))
      self._log(logging.ERR, err.message)
      callback(err, msg)
      return
    end

    self:setState(STATES.RUNNING)
    self._log(logging.INFO, fmt('handshake successful (ping_interval=%dms)', msg.result.ping_interval))
    callback(nil, msg)
  end)
end

function AgentProtocolConnection:getManifest(callback)
  self:sendManifest(function(err, response)
    if err then
      callback(err)
    else
      callback(nil, response.result)
    end
  end)
end

--[[
Process an async message

msg - The Incoming Message
]]--
function AgentProtocolConnection:execute(msg)
  if msg.method == 'system.info' then
    self:sendSystemInfo(msg)
  else
    local err = Error:new(fmt('invalid method [method=%s]', msg.method))
    self:emit('error', err)
  end
end


return AgentProtocolConnection

{'fraction_non_alphanumeric': 0.11376905611417451, 'fraction_numerical': 0.006811547194291274, 'mean_word_length': 4.06696795398521, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6123762', 'n_tokens_mistral': 3816, 'n_tokens_neox': 3702, 'n_words': 1066}