/*
 * Copyright (C) 2013 - Alejandro T. Colombini
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <string.h>
#include <mysql.h>
#include <mysql/errmsg.h>
#include <mysql/plugin_auth.h>
#include <my_sys.h>
#include <glib-unix.h>

#define CONFIG_FILE			_CONFIG_DIR"/proxy_auth.ini"
#define SQL_FILE			_SQL_DIR"/proxy_auth.sql"

typedef struct
{
	gboolean initialized;
	gchar * socket;
	gchar * host;
	gchar * user;
	gchar * pass;
	gchar * schema;
	guint port;
	guint last_error;
	GArray * fields;
	gchar * stmt;
	gsize stmt_len;
	
	gint pool_size;
	gint max_connections;
	GAsyncQueue * conn_pool;
	GMutex mutex;
}
ProxyAuth;

enum
{
	FIELD_USER,
	FIELD_PASS
};

typedef struct
{
	gboolean initialized;
	MYSQL * conn;
	MYSQL_STMT * stmt;
	gboolean connected;

	MYSQL_BIND * params;
	char user[MYSQL_USERNAME_LENGTH];
	gsize user_len;
	char pass[256];
	gsize pass_len;
	my_bool pass_is_null;

	MYSQL_BIND * result;
	char proxy_user[MYSQL_USERNAME_LENGTH];
	gsize proxy_user_len;
	my_bool proxy_user_is_null;
}
ConnData;

//++++++++++++++++++++++++++++++++++++++++++++++++++++++++ ConnData

static ConnData *
conn_data_new ()
{
	ConnData * self = g_new (ConnData, 1);
	self->initialized = FALSE;
	self->connected = FALSE;
	self->conn = mysql_init (NULL);
	return self;
}

static void
conn_data_deinit (ConnData * self)
{
	if (!self->initialized)
		return;

	if (self->stmt)
		mysql_stmt_close (self->stmt);

	g_free (self->params);
	g_free (self->result);
}

static void
conn_data_free (ConnData * self)
{
	conn_data_deinit (self);
	mysql_close (self->conn);
	g_free (self);
}

static gboolean
conn_data_reconnect (ConnData * self, ProxyAuth * pauth)
{
	if (self->connected)
		return TRUE;

	conn_data_deinit (self);
	self->initialized = TRUE;
	self->connected = FALSE;
	self->stmt = NULL;
	self->params = NULL;
	self->result = NULL;

	gboolean connected = mysql_real_connect (self->conn,
		pauth->host, pauth->user, pauth->pass, pauth->schema, pauth->port, pauth->socket, 0) != NULL;
	guint conn_error = mysql_errno (self->conn);
	
	if (conn_error && pauth->last_error != conn_error)
		g_warning ("ProxyAuth: Can't connect to database: %s",
			mysql_error (self->conn));

	pauth->last_error = conn_error;
	
	if (!connected)
		return FALSE;

	// Initializes the prepared statement

	self->stmt = mysql_stmt_init (self->conn);
	
	if (!self->stmt)
	{
		g_warning ("ProxyAuth: Error initializing statement: Out of memory");
		return FALSE;	
	}
	
	if (mysql_stmt_prepare (self->stmt, pauth->stmt, pauth->stmt_len))
	{
		g_warning ("ProxyAuth: Error preparing statement: %s",
			mysql_stmt_error (self->stmt));
		return FALSE;	
	}

	// Creates the query params

	gint i;
	MYSQL_BIND * params = self->params = g_new0 (MYSQL_BIND, pauth->fields->len);
	
	for (i = 0; i < pauth->fields->len; i++)
	{
		params[i].buffer_type = MYSQL_TYPE_STRING;
		
		if (g_array_index (pauth->fields, gint, i) == FIELD_USER)
		{
			params[i].buffer = self->user;
			params[i].length = &self->user_len;
		}
		else
		{ 
			params[i].buffer = self->pass;
			params[i].length = &self->pass_len;
			params[i].is_null = &self->pass_is_null;
		}
	}

	if (mysql_stmt_bind_param (self->stmt, params))
	{
		g_warning ("ProxyAuth: Error binding statement params: %s",
			mysql_stmt_error (self->stmt));
		return FALSE;	
	}

	// Binds the result
	
	MYSQL_BIND * result = self->result = g_new0 (MYSQL_BIND, 1);
	result->buffer_type = MYSQL_TYPE_STRING;
	result->buffer = self->proxy_user;
	result->buffer_length = MYSQL_USERNAME_LENGTH;
	result->length = &self->proxy_user_len;
	result->is_null = &self->proxy_user_is_null;

	if (mysql_stmt_bind_result (self->stmt, result))
	{
		g_warning ("ProxyAuth: Error binding statement result: %s",
			mysql_stmt_error (self->stmt));
		return FALSE;	
	}

	self->connected = TRUE;
	return TRUE;
}

//++++++++++++++++++++++++++++++++++++++++++++++++++++++++ ProxyAuth

ProxyAuth *
proxy_auth_new ()
{
	ProxyAuth * self = g_new (ProxyAuth, 1);
	self->initialized = FALSE;
	g_mutex_init (&self->mutex);
	return self;
}

void
proxy_auth_deinit (ProxyAuth * self)
{
	if (!self->initialized)
		return;

	g_free (self->socket);
	g_free (self->host);
	g_free (self->user);
	g_free (self->pass);
	g_free (self->schema);
	g_free (self->stmt);

	if (self->conn_pool)
		g_async_queue_unref (self->conn_pool);
	if (self->fields)
		g_array_free (self->fields, TRUE);
}

void
proxy_auth_free (ProxyAuth * self)
{
	proxy_auth_deinit (self);
	g_mutex_clear (&self->mutex);
	g_free (self);
}

static gboolean
proxy_auth_regex_func (const GMatchInfo * info, GString * res, gpointer data)
{
	GArray * fields = (GArray *) data;
	gchar * match = g_match_info_fetch (info, 0);
	gint field = -1;

	if (!g_strcmp0 (match, "#user"))
		field = FIELD_USER;
	else if (!g_strcmp0 (match, "#pass"))
		field = FIELD_PASS;
		
	if (field != -1)
	{
		g_string_append_c (res, '?');
		g_array_append_val (fields, field);
	}
	else
		g_string_append (res, match);

	g_free (match);
	return FALSE;
}

gboolean
proxy_auth_init (ProxyAuth * self)
{
	gboolean res = FALSE;
	GError * error = NULL;
	gchar * queryt = NULL;
	GRegex * regex = NULL;

	proxy_auth_deinit (self);
	self->initialized = TRUE;
	self->socket = NULL;
	self->host = NULL;
	self->user = NULL;
	self->pass = NULL;
	self->schema = NULL;
	self->conn_pool = NULL;
	self->fields = NULL;
	self->stmt = NULL;
	self->stmt_len = 0;

	// Reading the configuration file

	GKeyFile * key_file = g_key_file_new ();

	if (!g_key_file_load_from_file (key_file, CONFIG_FILE, G_KEY_FILE_NONE, &error))
	{
		g_warning ("ProxyAuth: Can't open configuration file: %s", error->message);
		goto end;
	}

	self->socket = NULL;
	self->host = g_key_file_get_string (key_file, "db", "host", NULL);

	if (!self->host)
		self->socket = g_key_file_get_string (key_file, "db", "socket", NULL);
	
	self->user = g_key_file_get_string (key_file, "db", "user", NULL);
	self->pass = g_key_file_get_string (key_file, "db", "pass", NULL);
	self->schema = g_key_file_get_string (key_file, "db", "schema", NULL);
	self->port = (guint) g_key_file_get_integer (key_file, "db", "port", NULL);
	self->max_connections = g_key_file_get_integer (key_file, "db", "maxConnections", NULL);

	// Reading the query template

	gsize queryt_len;

	if (!g_file_get_contents (SQL_FILE, &queryt, &queryt_len, &error))
	{
		g_warning ("ProxyAuth: Can't read the query template: %s", error->message);
		goto end;
	}

	// Creates the regular expression

	regex = g_regex_new ("#\\w+",
		G_REGEX_OPTIMIZE | G_REGEX_MULTILINE, 0, &error);

	if (error)
	{
		g_warning ("ProxyAuth: Can't create the regex: %s", error->message);
		goto end;
	}
	
	self->fields = g_array_sized_new (FALSE, FALSE, sizeof (gint), 2);
	self->stmt = g_regex_replace_eval (regex,
		queryt, queryt_len, 0, 0, proxy_auth_regex_func, self->fields, &error);
	self->stmt_len = strlen (self->stmt);

	if (error)
	{
		g_warning ("ProxyAuth: Can't evaluate regex: %s", error->message);
		goto end;
	}
	
	// Creates the connection pool

	self->conn_pool = g_async_queue_new_full ((GDestroyNotify) conn_data_free);
	self->pool_size = 0;
	self->last_error = 0;

	if (self->max_connections < 1)
		self->max_connections = 1;
	
	res = TRUE;
end:
	if (regex)
		g_regex_unref (regex);

	g_clear_error (&error);
	g_free (queryt);
	g_key_file_free (key_file);
	return res;
}

int
proxy_auth_authenticate (ProxyAuth * self, MYSQL_PLUGIN_VIO * vio, MYSQL_SERVER_AUTH_INFO * info)
{
	int res = CR_ERROR;
	gchar * query = NULL;
	unsigned char * pkt;

	//	Check for the username
	
	guint user_len = info->user_name_length;

	if (info->user_name == NULL
	&& (user_len = vio->read_packet (vio, &pkt)) < 0)
		return CR_ERROR;

	if (user_len > MYSQL_USERNAME_LENGTH)
		return CR_ERROR;

	//	Read the password and check if it's valid

	guint pass_len = vio->read_packet (vio, &pkt) - 1;

	if (pass_len < 0)
		return CR_ERROR;

	if (!pass_len || *pkt == '\0')
	{
		info->password_used = PASSWORD_USED_NO;
		return CR_ERROR;
	}

	char pass[pass_len + 1];
	memcpy (pass, pkt, pass_len);
	pass[pass_len] = '\0';

	info->password_used = PASSWORD_USED_YES;

	// Gets a connection from the pool

	ConnData * conn_data = (ConnData *) g_async_queue_try_pop (self->conn_pool);
	
	if (!conn_data)
	{
		g_mutex_lock (&self->mutex);
		
		if (self->pool_size < self->max_connections)
		{
			self->pool_size++;
			g_mutex_unlock (&self->mutex);
			conn_data = conn_data_new ();
		}
		else
		{
			g_mutex_unlock (&self->mutex);
			conn_data = (ConnData *) g_async_queue_pop (self->conn_pool);
		}
	}

	if (!conn_data_reconnect (conn_data, self))
		goto end;

	// Sends the query to the database
	
	strncpy (conn_data->user, info->user_name, user_len);
	conn_data->user_len = user_len;
	strncpy (conn_data->pass, pass, pass_len);
	conn_data->pass_len = pass_len;
	conn_data->pass_is_null = pass_len <= 0;
	
	MYSQL_STMT * stmt = conn_data->stmt;
	
	if (!mysql_stmt_execute (stmt)
	&& !mysql_stmt_fetch (stmt))
	{
		strcpy (info->external_user, info->user_name);
		strncpy (info->authenticated_as,
			conn_data->proxy_user, conn_data->proxy_user_len);
		res = CR_OK;

		mysql_stmt_free_result (stmt);
	}
	else
	{
		switch (mysql_stmt_errno (stmt))
		{
			case CR_SERVER_LOST:
			case CR_SERVER_GONE_ERROR:
				conn_data->connected = FALSE;
				break;
		}

		g_warning ("ProxyAuth: Error executing query: %s",
			mysql_stmt_error (stmt));
	}
end:
	if (conn_data)
		g_async_queue_push (self->conn_pool, conn_data);
	
	g_free (query);
	return res;
}

//++++++++++++++++++++++++++++++++++++++++++++++++++++++++ Plugin

ProxyAuth * pauth;

static int
proxy_auth_plugin_main (MYSQL_PLUGIN_VIO * vio, MYSQL_SERVER_AUTH_INFO * info)
{
	my_init ();

	if (mysql_thread_init ())
		g_warning ("ProxyAuth: Can't initialize MySQL thread");
	
	return proxy_auth_authenticate (pauth, vio, info);;
}

static int
proxy_auth_plugin_init ()
{
	if (mysql_library_init (0, NULL, NULL))
		g_warning ("ProxyAuth: Can't initialize MySQL library");

	pauth = proxy_auth_new ();
	return proxy_auth_init (pauth) ? 0 : 1;
}

static int
proxy_auth_plugin_deinit ()
{
	proxy_auth_free (pauth);
	return 0;
}

static struct st_mysql_auth proxy_auth_handler =
{
	 MYSQL_AUTHENTICATION_INTERFACE_VERSION
	,"mysql_clear_password" // Cleartext plugin required in the client
	,proxy_auth_plugin_main
};

mysql_declare_plugin(proxy_auth)
{
	 MYSQL_AUTHENTICATION_PLUGIN
	,&proxy_auth_handler
	,"proxy_auth"
	,"Alejandro T. Colombini"
	,"Proxy user authentication server-side plugin"
	,PLUGIN_LICENSE_GPL
	,proxy_auth_plugin_init
	,proxy_auth_plugin_deinit
	,0x0100 // version 1.0
	,NULL
	,NULL
	,NULL
	,0
}
mysql_declare_plugin_end;