/*
 * Copyright (C) 2017 Oracle.
 *
 * 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 2
 * 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/copyleft/gpl.txt
 */

#include "smatch.h"
#include "smatch_extra.h"

static int my_id;

struct allocator {
	const char *func;
	int param;
	int param2;
};

static struct allocator generic_allocator_table[] = {
	{"malloc", 0},
	{"memdup", 1},
	{"realloc", 1},
};

static struct allocator kernel_allocator_table[] = {
	{"kmalloc", 0},
	{"kzalloc", 0},
	{"vmalloc", 0},
	{"__vmalloc", 0},
	{"vzalloc", 0},
	{"sock_kmalloc", 1},
	{"kmemdup", 1},
	{"kmemdup_user", 1},
	{"dma_alloc_attrs", 1},
	{"pci_alloc_consistent", 1},
	{"pci_alloc_coherent", 1},
	{"devm_kmalloc", 1},
	{"devm_kzalloc", 1},
	{"krealloc", 1},
};

static struct allocator calloc_table[] = {
	{"calloc", 0, 1},
	{"kcalloc", 0, 1},
	{"kmalloc_array", 0, 1},
	{"devm_kcalloc", 1, 2},
};

static int bytes_per_element(struct expression *expr)
{
	struct symbol *type;

	type = get_type(expr);
	if (!type)
		return 0;

	if (type->type != SYM_PTR && type->type != SYM_ARRAY)
		return 0;

	type = get_base_type(type);
	return type_bytes(type);
}

static void save_constraint_required(struct expression *pointer, int op, struct expression *constraint)
{
	char *data, *limit;

	data = get_constraint_str(pointer);
	if (!data)
		return;

	limit = get_constraint_str(constraint);
	if (!limit) {
		// FIXME deal with <= also
		if (op == '<')
			set_state_expr(my_id, constraint, alloc_state_expr(pointer));
		goto free_data;
	}

	sql_save_constraint_required(data, op, limit);

	free_string(limit);
free_data:
	free_string(data);
}

static int handle_zero_size_arrays(struct expression *pointer, struct expression *size)
{
	struct expression *left, *right;
	struct symbol *type, *array, *array_type;
	sval_t struct_size;
	char *limit;
	char data[128];

	if (size->type != EXPR_BINOP || size->op != '+')
		return 0;

	type = get_type(pointer);
	if (!type || type->type != SYM_PTR)
		return 0;
	type = get_real_base_type(type);
	if (!type || !type->ident || type->type != SYM_STRUCT)
		return 0;
	if (!last_member_is_resizable(type))
		return 0;
	array = last_ptr_list((struct ptr_list *)type->symbol_list);
	if (!array || !array->ident)
		return 0;
	array_type = get_real_base_type(array);
	if (!array_type || array_type->type != SYM_ARRAY)
		return 0;
	array_type = get_real_base_type(array_type);

	left = strip_expr(size->left);
	right = strip_expr(size->right);

	if (!get_implied_value(left, &struct_size))
		return 0;
	if (struct_size.value != type_bytes(type))
		return 0;

	if (right->type == EXPR_BINOP && right->op == '*') {
		struct expression *mult_left, *mult_right;
		sval_t sval;

		mult_left = strip_expr(right->left);
		mult_right = strip_expr(right->right);

		if (get_implied_value(mult_left, &sval) &&
		    sval.value == type_bytes(array_type))
			size = mult_right;
		else if (get_implied_value(mult_right, &sval) &&
		    sval.value == type_bytes(array_type))
			size = mult_left;
		else
			return 0;
	}

	snprintf(data, sizeof(data), "(struct %s)->%s", type->ident->name, array->ident->name);
	limit = get_constraint_str(size);
	if (!limit) {
		set_state_expr(my_id, size, alloc_state_expr(
			       member_expression(deref_expression(pointer), '*', array->ident)));
		return 1;
	}

	sql_save_constraint_required(data, '<', limit);

	free_string(limit);
	return 1;
}

static void match_alloc_helper(struct expression *pointer, struct expression *size, int recurse)
{
	struct expression *size_orig, *tmp;
	sval_t sval;
	int cnt = 0;

	pointer = strip_expr(pointer);
	size = strip_expr(size);
	if (!size || !pointer)
		return;

	size_orig = size;
	if (recurse) {
		while ((tmp = get_assigned_expr(size))) {
			size = strip_expr(tmp);
			if (cnt++ > 5)
				break;
		}
		if (size != size_orig) {
			match_alloc_helper(pointer, size, 0);
			size = size_orig;
		}
	}

	if (handle_zero_size_arrays(pointer, size))
		return;

	if (size->type == EXPR_BINOP && size->op == '*') {
		struct expression *mult_left, *mult_right;

		mult_left = strip_expr(size->left);
		mult_right = strip_expr(size->right);

		if (get_implied_value(mult_left, &sval) &&
		    sval.value == bytes_per_element(pointer))
			size = mult_right;
		else if (get_implied_value(mult_right, &sval) &&
		    sval.value == bytes_per_element(pointer))
			size = mult_left;
		else
			return;
	}

	if (size->type == EXPR_BINOP && size->op == '+' &&
	    get_implied_value(size->right, &sval) &&
	    sval.value == 1)
		save_constraint_required(pointer, SPECIAL_LTE, size->left);
	else
		save_constraint_required(pointer, '<', size);
}

static void match_alloc(const char *fn, struct expression *expr, void *_size_arg)
{
	int size_arg = PTR_INT(_size_arg);
	struct expression *call, *arg;

	call = strip_expr(expr->right);
	arg = get_argument_from_call_expr(call->args, size_arg);

	match_alloc_helper(expr->left, arg, 1);
}

static void match_calloc(const char *fn, struct expression *expr, void *_start_arg)
{
	struct expression *pointer, *call, *size;
	struct expression *count = NULL;
	int start_arg = PTR_INT(_start_arg);
	sval_t sval;

	pointer = strip_expr(expr->left);
	call = strip_expr(expr->right);

	size = get_argument_from_call_expr(call->args, start_arg);
	if (get_implied_value(size, &sval) &&
	    sval.value == bytes_per_element(pointer))
		count = get_argument_from_call_expr(call->args, start_arg + 1);
	else {
		size = get_argument_from_call_expr(call->args, start_arg + 1);
		if (get_implied_value(size, &sval) &&
		    sval.value == bytes_per_element(pointer))
			count = get_argument_from_call_expr(call->args, start_arg);
	}

	if (!count)
		return;

	save_constraint_required(pointer, '<', count);
}

static void add_allocation_function(const char *func, void *call_back, int param)
{
	add_function_assign_hook(func, call_back, INT_PTR(param));
}

static void match_assign_size(struct expression *expr)
{
	struct smatch_state *state;
	char *data, *limit;

	state = get_state_expr(my_id, expr->right);
	if (!state || !state->data)
		return;

	data = get_constraint_str(state->data);
	if (!data)
		return;

	limit = get_constraint_str(expr->left);
	if (!limit)
		goto free_data;

	sql_save_constraint_required(data, '<', limit);

	free_string(limit);
free_data:
	free_string(data);
}

static void match_assign_has_buf_comparison(struct expression *expr)
{
	struct expression *size;
	int limit_type;

	if (expr->op != '=')
		return;
	if (expr->right->type == EXPR_CALL)
		return;
	size = get_size_variable(expr->right, &limit_type);
	if (!size)
		return;
	if (limit_type != ELEM_COUNT)
		return;
	match_alloc_helper(expr->left, size, 1);
}

static void match_assign_data(struct expression *expr)
{
	struct expression *right, *arg, *tmp;
	int i;
	int size_arg;
	int size_arg2 = -1;

	if (expr->op != '=')
		return;

	/* Direct calls are handled else where (for now at least) */
	tmp = get_assigned_expr(expr->right);
	if (!tmp)
		return;

	right = strip_expr(tmp);
	if (right->type != EXPR_CALL)
		return;

	if (right->fn->type != EXPR_SYMBOL ||
	    !right->fn->symbol ||
	    !right->fn->symbol->ident)
		return;

	for (i = 0; i < ARRAY_SIZE(generic_allocator_table); i++) {
		if (strcmp(right->fn->symbol->ident->name,
			   generic_allocator_table[i].func) == 0) {
			size_arg = generic_allocator_table[i].param;
			goto found;
		}
	}

	if (option_project != PROJ_KERNEL)
		return;

	for (i = 0; i < ARRAY_SIZE(kernel_allocator_table); i++) {
		if (strcmp(right->fn->symbol->ident->name,
			   kernel_allocator_table[i].func) == 0) {
			size_arg = kernel_allocator_table[i].param;
			goto found;
		}
	}

	for (i = 0; i < ARRAY_SIZE(calloc_table); i++) {
		if (strcmp(right->fn->symbol->ident->name,
			   calloc_table[i].func) == 0) {
			size_arg = calloc_table[i].param;
			size_arg2 = calloc_table[i].param2;
			goto found;
		}
	}

	return;

found:
	arg = get_argument_from_call_expr(right->args, size_arg);
	match_alloc_helper(expr->left, arg, 1);
	if (size_arg2 == -1)
		return;
	arg = get_argument_from_call_expr(right->args, size_arg2);
	match_alloc_helper(expr->left, arg, 1);
}

static void match_assign_ARRAY_SIZE(struct expression *expr)
{
	struct expression *array;
	char *data, *limit;
	const char *macro;

	macro = get_macro_name(expr->right->pos);
	if (!macro || strcmp(macro, "ARRAY_SIZE") != 0)
		return;
	array = strip_expr(expr->right);
	if (array->type != EXPR_BINOP || array->op != '+')
		return;
	array = strip_expr(array->left);
	if (array->type != EXPR_BINOP || array->op != '/')
		return;
	array = strip_expr(array->left);
	if (array->type != EXPR_SIZEOF)
		return;
	array = strip_expr(array->cast_expression);
	if (array->type != EXPR_PREOP || array->op != '*')
		return;
	array = strip_expr(array->unop);

	data = get_constraint_str(array);
	limit = get_constraint_str(expr->left);
	if (!data || !limit)
		goto free;

	sql_save_constraint_required(data, '<', limit);

free:
	free_string(data);
	free_string(limit);
}

static void match_assign_buf_comparison(struct expression *expr)
{
	struct expression *pointer;

	if (expr->op != '=')
		return;
	pointer = get_array_variable(expr->right);
	if (!pointer)
		return;

	match_alloc_helper(pointer, expr->right, 1);
}

static int constraint_found(void *_found, int argc, char **argv, char **azColName)
{
	int *found = _found;

	*found = 1;
	return 0;
}

static int has_constraint(struct expression *expr, const char *constraint)
{
	int found = 0;

	if (get_state_expr(my_id, expr))
		return 1;

	run_sql(constraint_found, &found,
		"select data from constraints_required where bound = '%q' limit 1",
		escape_newlines(constraint));

	return found;
}

static void match_assign_constraint(struct expression *expr)
{
	struct symbol *type;
	char *left, *right;

	if (expr->op != '=')
		return;

	type = get_type(expr->left);
	if (!type || type->type != SYM_BASETYPE)
		return;

	left = get_constraint_str(expr->left);
	if (!left)
		return;
	right = get_constraint_str(expr->right);
	if (!right)
		goto free;
	if (!has_constraint(expr->right, right))
		return;
	sql_copy_constraint_required(left, right);
free:
	free_string(right);
	free_string(left);
}

void register_constraints_required(int id)
{
	my_id = id;

	set_dynamic_states(my_id);
	add_hook(&match_assign_size, ASSIGNMENT_HOOK);
	add_hook(&match_assign_data, ASSIGNMENT_HOOK);
	add_hook(&match_assign_has_buf_comparison, ASSIGNMENT_HOOK);

	add_hook(&match_assign_ARRAY_SIZE, ASSIGNMENT_HOOK);
	add_hook(&match_assign_ARRAY_SIZE, GLOBAL_ASSIGNMENT_HOOK);
	add_hook(&match_assign_buf_comparison, ASSIGNMENT_HOOK);
	add_hook(&match_assign_constraint, ASSIGNMENT_HOOK);

	add_allocation_function("malloc", &match_alloc, 0);
	add_allocation_function("memdup", &match_alloc, 1);
	add_allocation_function("realloc", &match_alloc, 1);
	add_allocation_function("realloc", &match_calloc, 0);
	if (option_project == PROJ_KERNEL) {
		add_allocation_function("kmalloc", &match_alloc, 0);
		add_allocation_function("kzalloc", &match_alloc, 0);
		add_allocation_function("vmalloc", &match_alloc, 0);
		add_allocation_function("__vmalloc", &match_alloc, 0);
		add_allocation_function("vzalloc", &match_alloc, 0);
		add_allocation_function("sock_kmalloc", &match_alloc, 1);
		add_allocation_function("kmemdup", &match_alloc, 1);
		add_allocation_function("kmemdup_user", &match_alloc, 1);
		add_allocation_function("dma_alloc_attrs", &match_alloc, 1);
		add_allocation_function("pci_alloc_consistent", &match_alloc, 1);
		add_allocation_function("pci_alloc_coherent", &match_alloc, 1);
		add_allocation_function("devm_kmalloc", &match_alloc, 1);
		add_allocation_function("devm_kzalloc", &match_alloc, 1);
		add_allocation_function("kcalloc", &match_calloc, 0);
		add_allocation_function("kmalloc_array", &match_calloc, 0);
		add_allocation_function("devm_kcalloc", &match_calloc, 1);
		add_allocation_function("krealloc", &match_alloc, 1);
	}
}