-- mod_http_roster_admin

-- Description: Allow user rosters to be sourced from a remote HTTP API

--

-- Version: 1.0

-- Date: 2015-03-06

-- Author: Matthew Wild <matthew@prosody.im>

-- License: MPLv2

--

-- Requirements:

--   Prosody config:

--     storage = { roster = "memory" }

--     modules_disabled = { "roster" }

--   Dependencies:

--     Prosody 0.9

--     lua-cjson (Debian/Ubuntu/LuaRocks: lua-cjson)

local http = require "net.http";

local json = require "cjson";

local it = require "util.iterators";

local set = require "util.set";

local rm = require "core.rostermanager";

local st = require "util.stanza";

local array = require "util.array";

local new_id = require "util.id".short;

local host = module.host;

local sessions = hosts[host].sessions;

local roster_url = module:get_option_string("http_roster_url", "http://localhost/%s");

-- Send a roster push to the named user, with the given roster, for the specified

-- contact's roster entry. Used to notify clients of changes/removals.

local function roster_push(username, roster, contact_jid)

	local stanza = st.iq({type="set", id=new_id()})

		:tag("query", {xmlns = "jabber:iq:roster" });

	local item = roster[contact_jid];

	if item then

		stanza:tag("item", {jid = contact_jid, subscription = item.subscription, name = item.name, ask = item.ask});

		for group in pairs(item.groups) do

			stanza:tag("group"):text(group):up();

		end

	else

		stanza:tag("item", {jid = contact_jid, subscription = "remove"});

	end

	stanza:up():up(); -- move out from item

	for _, session in pairs(hosts[host].sessions[username].sessions) do

		if session.interested then

			session.send(stanza);

		end

	end

end

-- Send latest presence from the named local user to a contact.

local function send_presence(username, contact_jid, available)

	module:log("debug", "Sending %savailable presence from %s to contact %s", (available and "" or "un"), username, contact_jid);

	for resource, session in pairs(sessions[username].sessions) do

		local pres;

		if available then

			pres = st.clone(session.presence);

			pres.attr.to = contact_jid;

		else

			pres = st.presence({ to = contact_jid, from = session.full_jid, type = "unavailable" });

		end

		module:send(pres);

	end

end

-- Converts a 'friend' object from the API to a Prosody roster item object

local function friend_to_roster_item(friend)

	return {

		name = friend.name;

		subscription = "both";

		groups = friend.groups or {};

	};

end

-- Returns a handler function to consume the data returned from

-- the API, compare it to the user's current roster, and perform

-- any actions necessary (roster pushes, presence probes) to

-- synchronize them.

local function updated_friends_handler(username, cb)

	return (function (ok, code, friends)

		if not ok then

			cb(false, code);

		end

		local user = sessions[username];

		local roster = user.roster;

		local old_contacts = set.new(array.collect(it.keys(roster)));

		local new_contacts = set.new(array.collect(it.keys(friends)));

		-- These two entries are not real contacts, ignore them

		old_contacts:remove(false);

		old_contacts:remove("pending");

		module:log("debug", "New friends list of %s: %s", username, json.encode(friends));

		-- Calculate which contacts have been added/removed since

		-- the last time we fetched the roster

		local added_contacts = new_contacts - old_contacts;

		local removed_contacts = old_contacts - new_contacts;

		local added, removed = 0, 0;

		-- Add new contacts and notify connected clients

		for contact_jid in added_contacts do

			module:log("debug", "Processing new friend of %s: %s", username, contact_jid);

			roster[contact_jid] = friend_to_roster_item(friends[contact_jid]);

			roster_push(username, roster, contact_jid);

			send_presence(username, contact_jid, true);

			added = added + 1;

		end

		-- Remove contacts and notify connected clients

		for contact_jid in removed_contacts do

			module:log("debug", "Processing removed friend of %s: %s", username, contact_jid);

			roster[contact_jid] = nil;

			roster_push(username, roster, contact_jid);

			send_presence(username, contact_jid, false);

			removed = removed + 1;

		end

		module:log("debug", "User %s: added %d new contacts, removed %d contacts", username, added, removed);

		if cb ~= nil then

			cb(true);

		end

	end);

end

-- Fetch the named user's roster from the API, call callback (cb)

-- with status and result (friends list) when received.

function fetch_roster(username, cb)

	local x = {headers = {}};

	x["headers"]["ACCEPT"] = "application/json, text/plain, */*";

	module:log("debug", "Fetching roster at URL: %s", roster_url:format(username));

	local ok, err = http.request(

		roster_url:format(username),

		x,

		function (roster_data, code)

			if code ~= 200 then

				module:log("error", "Error fetching roster from %s (code %d): %s", roster_url:format(username), code, tostring(roster_data):sub(1, 40):match("^[^\r\n]+"));

				if code ~= 0 then

					cb(nil, code, roster_data);

				end

				return;

			end

			module:log("debug", "Successfully fetched roster for %s", username);

			module:log("debug", "The roster data is %s", roster_data);

			cb(true, code, json.decode(roster_data));

		end

	);

	if not ok then

		module:log("error", "Failed to connect to roster API at %s: %s", roster_url:format(username), err);

		cb(false, 0, err);

	end

end

-- Fetch the named user's roster from the API, synchronize it with

-- the user's current roster. Notify callback (cb) with true/false

-- depending on success or failure.

function refresh_roster(username, cb)

	local user = sessions[username];

	if not (user and user.roster) then

		module:log("debug", "User's (%q) roster updated, but they are not online - ignoring", username);

		cb(true);

		return;

	end

	fetch_roster(username, updated_friends_handler(username, cb));

end

--- Roster protocol handling ---

-- Build a reply to a "roster get" request

local function build_roster_reply(stanza, roster_data)

	local roster = st.reply(stanza)

		:tag("query", { xmlns = "jabber:iq:roster" });

	for jid, item in pairs(roster_data) do

		if jid and jid ~= "pending" then

			roster:tag("item", {

				jid = jid,

				subscription = item.subscription,

				ask = item.ask,

				name = item.name,

			});

			for group in pairs(item.groups) do

				roster:tag("group"):text(group):up();

			end

			roster:up(); -- move out from item

		end

	end

	return roster;

end

-- Handle clients requesting their roster (generally at login)

-- This will not work if mod_roster is loaded (in 0.9).

module:hook("iq-get/self/jabber:iq:roster:query", function(event)

	local session, stanza = event.origin, event.stanza;

	session.interested = true; -- resource is interested in roster updates

	local roster = session.roster;

	if roster[false].downloaded then

		return session.send(build_roster_reply(stanza, roster));

	end

	-- It's possible that we can call this more than once for a new roster

	-- Should happen rarely (multiple clients of the same user request the

	-- roster in the time it takes the API to respond). Currently we just

	-- issue multiple requests, as it's harmless apart from the wasted

	-- requests.

	fetch_roster(session.username, function (ok, code, friends)

		if not ok then

			session.send(st.error_reply(stanza, "cancel", "internal-server-error"));

			session:close("internal-server-error");

			return;

		end

		-- Are we the first callback to handle the downloaded roster?

		local first = roster[false].downloaded == nil;

		if first then

			-- Fill out new roster

			for jid, friend in pairs(friends) do

				roster[jid] = friend_to_roster_item(friend);

			end

		end

		roster[false].downloaded = true;

		-- Send full roster to client

		session.send(build_roster_reply(stanza, roster));

		if not first then

			-- We already had a roster, make sure to handle any changes...

			updated_friends_handler(session.username, nil)(ok, code, friends);

		end

	end);

	return true;

end);

-- Prevent client from making changes to the roster. This will not

-- work if mod_roster is loaded (in 0.9).

module:hook("iq-set/self/jabber:iq:roster:query", function(event)

	local session, stanza = event.origin, event.stanza;

	return session.send(st.error_reply(stanza, "cancel", "service-unavailable"));

end);

--- HTTP endpoint to trigger roster refresh ---

-- Handles updating for a single user: GET /roster_admin/refresh/USERNAME

function handle_refresh_single(event, username)

	refresh_roster(username, function (ok, code, err)

		event.response.headers["Content-Type"] = "application/json";

		event.response:send(json.encode({

			status = ok and "ok" or "error";

			message = err or "roster update complete";

		}));

	end);

	return true;

end

-- Handles updating for multiple users: POST /roster_admin/refresh

-- Payload should be a JSON array of usernames, e.g. ["user1", "user2", "user3"]

function handle_refresh_multi(event)

	local users = json.decode(event.request.body);

	if not users then

		module:log("warn", "Multi-user refresh attempted with missing/invalid payload");

		event.response:send(400);

		return true;

	end

	local count, count_err = 0, 0;

	local function cb(ok)

		count = count + 1;

		if not ok then

			count_err = count_err + 1;

		end

		if count == #users then

			event.response.headers["Content-Type"] = "application/json";

			event.response:send(json.encode({

				status = "ok";

				message = "roster update complete";

				updated = count - count_err;

				errors = count_err;

			}));

		end

	end

	for _, username in ipairs(users) do

		refresh_roster(username, cb);

	end

	return true;

end

module:depends("http");

module:provides("http", {

	route = {

		["POST /refresh"] = handle_refresh_multi;

		["GET /refresh/*"] = handle_refresh_single;

	};

});