-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at http://mozilla.org/MPL/2.0/.

local class = require("class")
local check = require("check")

local ipv4 = require("protocol/ipv4")
local udp = require("protocol/udp")

local module = {}
local log = haka.log_section("udp")

local udp_connection_dissector = haka.dissector.new{
	type = haka.helper.PacketDissector,
	name = 'udp_connection'
}

udp_connection_dissector.cnx_table = ipv4.cnx_table()

udp_connection_dissector:register_event('new_connection')
udp_connection_dissector:register_event('receive_packet')
udp_connection_dissector:register_event('receive_data')
udp_connection_dissector:register_event('end_connection')

local function udp_get_cnx_key(pkt)
	return pkt.src, pkt.dst, pkt.srcport, pkt.dstport
end

function udp_connection_dissector:receive(pkt)
	local connection, direction = udp_connection_dissector.cnx_table:get(udp_get_cnx_key(pkt))
	if not connection then
		local data = haka.context.newscope()
		local self = udp_connection_dissector:new(pkt)

		haka.context:exec(data, function ()
			self:trigger('new_connection', pkt)
			class.classof(self).policies.next_dissector:apply{
				values = self:install_criterion(),
				ctx = self,
			}
			self:activate_next_dissector()
		end)

		pkt:continue()

		connection = udp_connection_dissector.cnx_table:create(udp_get_cnx_key(pkt))
		connection.data = data
		self:init(connection)
		data:createnamespace('udp_connection', self)
	end

	local dissector = connection.data:namespace('udp_connection')

	local ret, err = xpcall(function ()
		haka.context:exec(connection.data, function ()
			return dissector:emit(direction, pkt)
		end)
	end, debug.format_error)

	if not ret then
		if err then
			log.error("%s", err)
			dissector:error()
		end
	end
end

function udp_connection_dissector.method:install_criterion()
	return { port = self.dstport }
end

udp_connection_dissector.state_machine = haka.state_machine.new("udp", function ()
	state_type{
		events = { 'receive', 'drop' },
		update = function (self, state_machine, direction, pkt)
			state_machine:trigger('receive', pkt, direction)
		end,
	}

	drop = state()
	established = state()

	any:on{
		event = events.fail,
		jump = drop,
	}

	any:on{
		event = events.drop,
		jump = drop,
	}

	any:on{
		event = events.finish,
		execute = function (self)
			self:trigger('end_connection')
			self._parent:close()
		end,
	}

	drop:on{
		event = events.enter,
		execute = function (self)
			self:trigger('end_connection')
			self._dropped = true
			self._parent:drop()
		end,
	}

	drop:on{
		event = events.timeout(60),
		jump = finish,
	}

	drop:on{
		event = events.receive,
		execute = function (self, pkt, direction)
			pkt:drop()
		end,
	}

	drop:on{
		event = events.finish,
		execute = function (self)
			self._parent:close()
		end
	}

	established:on{
		event = events.receive,
		execute = function (self, pkt, direction)
			self:trigger('receive_data', pkt, direction)

			local next_dissector = self:activate_next_dissector()
			if next_dissector then
				local payload
				pkt._restore, payload = pkt.payload:select()
				return next_dissector:receive(pkt, payload, direction)
			else
				pkt:send()
			end
		end,
		jump = established,
	}

	established:on{
		event = events.timeout(60),
		jump = fail,
	}

	initial(established)
end)

udp_connection_dissector.auto_state_machine = false

function udp_connection_dissector.method:__init(pkt)
	class.super(udp_connection_dissector).__init(self)

	self._dropped = false
	self.srcip = pkt.src
	self.dstip = pkt.dst
	self.srcport = pkt.srcport
	self.dstport = pkt.dstport
end

function udp_connection_dissector.method:init(connection)
	self._parent = connection
	self._state = udp_connection_dissector.state_machine:instanciate(self)
end

function udp_connection_dissector.method:emit(direction, pkt)
	self._parent:update_stat(direction, pkt.len)
	self:trigger('receive_packet', pkt, direction)

	self._state:update(direction, pkt)
end

function udp_connection_dissector.method:send(pkt, payload, clone)
	pkt._restore:restore(payload, clone or false)
	pkt:send()
end

function udp_connection_dissector.method:drop(pkt)
	if pkt then
		return pkt:drop()
	else
		return self._state:trigger('drop')
	end
end

function udp_connection_dissector.method:can_continue()
	return not self._dropped
end

haka.policy {
	name = "udp connection",
	on = haka.dissectors.udp.policies.next_dissector,
	action = haka.dissectors.udp_connection.install
}

--
-- Helpers
--

module.helper = {}

module.helper.UdpFlowDissector =  class.class('UdpFlowDissector', haka.helper.FlowDissector)

function module.helper.UdpFlowDissector.method:__init(flow)
	check.assert(class.isa(flow, udp_connection_dissector), "invalid flow parameter")

	class.super(module.helper.UdpFlowDissector).__init(self, flow)
end

function module.helper.UdpFlowDissector.method:can_continue()
	return self._parent:can_continue()
end

function module.helper.UdpFlowDissector.method:receive(pkt, payload, direction)
	assert(self._state, "no state machine defined")

	self._state:update(payload:pos('begin'), direction, pkt, payload)
end

--
-- Console utilities
--

module.console = {}

function module.console.list_connections(show_dropped)
	local ret = {}
	for _, cnx in ipairs(udp_connection_dissector.cnx_table:all(show_dropped)) do
		if cnx.data then
			local udp_data = cnx.data:namespace('udp_connection')

			table.insert(ret, {
				_thread = haka.current_thread(),
				_id = udp_data.id,
				id = string.format("%d-%d", haka.current_thread(), udp_data.id),
				srcip = udp_data.srcip,
				srcport = udp_data.srcport,
				dstip = udp_data.dstip,
				dstport = udp_data.dstport,
				state = udp_data.state,
				in_pkts = udp_data.in_pkts,
				in_bytes = udp_data.in_bytes,
				out_pkts = udp_data.out_pkts,
				out_bytes = udp_data.out_bytes
			})
		end
	end
	return ret
end

function module.console.drop_connection(id)
	local udp_conn = udp_connection_dissector.cnx_table:get_byid(id)
	if not udp_conn then
		error("unknown udp connection")
	end

	if udp_conn.data then
		local udp_data = udp_conn.data:namespace('udp_connection')
		udp_data:drop()
	end
end

return module
