/*-------------------------------------------------------------------------
 *
 * test_extstate.c
 *		Reproducer: planner flat copies of RelOptInfo share the
 *		extension_state array with the RelOptInfo they were copied from.
 *
 * build_grouped_rel(), create_unique_paths() and create_partial_unique_paths()
 * memcpy() a RelOptInfo and reset its path fields, but not extension_state /
 * extension_state_allocated.  The copy therefore aliases the original's
 * private-state array:
 *
 *	mode 1: state the extension stored on a joinrel in joinrel_setup_hook is
 *			returned for the grouped joinrel copied from it.
 *	mode 2: storing into a higher-numbered slot on the grouped joinrel makes
 *			SetRelOptInfoExtensionState() repalloc the shared array, which frees
 *			the original joinrel's array; the next read on the original joinrel
 *			is a use-after-free.
 *
 * Build against a PG 19 / master install with PGXS:
 *		make PG_CONFIG=/path/to/pg19/bin/pg_config install
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"

#include "fmgr.h"
#include "nodes/bitmapset.h"
#include "optimizer/extendplan.h"
#include "optimizer/pathnode.h"
#include "optimizer/paths.h"
#include "utils/guc.h"

PG_MODULE_MAGIC;

/* 0 = off, 1 = report sharing, 2 = also force the shared array to grow */
static int	test_mode = 0;

/* GetPlannerExtensionId() keeps the pointer, so the names must be constant */
static const char *const ext_names[] = {
	"test_extstate_0", "test_extstate_1", "test_extstate_2",
	"test_extstate_3", "test_extstate_4",
};
#define NUM_IDS lengthof(ext_names)
static int	ext_ids[NUM_IDS];

static joinrel_setup_hook_type prev_joinrel_setup_hook = NULL;
static join_path_setup_hook_type prev_join_path_setup_hook = NULL;

typedef struct OwnerTag
{
	RelOptInfo *owner;			/* the RelOptInfo this tag was stored in */
} OwnerTag;

static void
tes_joinrel_setup(PlannerInfo *root, RelOptInfo *joinrel,
				  RelOptInfo *outer_rel, RelOptInfo *inner_rel,
				  SpecialJoinInfo *sjinfo, List *restrictlist)
{
	if (test_mode > 0)
	{
		OwnerTag   *tag = palloc_object(OwnerTag);

		tag->owner = joinrel;
		SetRelOptInfoExtensionState(joinrel, ext_ids[0], tag);
		SetRelOptInfoExtensionState(joinrel, ext_ids[1], tag);
	}

	if (prev_joinrel_setup_hook)
		prev_joinrel_setup_hook(root, joinrel, outer_rel, inner_rel,
								sjinfo, restrictlist);
}

static void
tes_join_path_setup(PlannerInfo *root, RelOptInfo *joinrel,
					RelOptInfo *outerrel, RelOptInfo *innerrel,
					JoinType jointype, JoinPathExtraData *extra)
{
	if (test_mode > 0)
	{
		OwnerTag   *tag = GetRelOptInfoExtensionState(joinrel, ext_ids[1]);

		if (tag != NULL && tag->owner != joinrel)
			elog(WARNING, "extension state on %s joinrel %s belongs to a different RelOptInfo",
				 IS_GROUPED_REL(joinrel) ? "grouped" : "plain",
				 bmsToString(joinrel->relids));

		/*
		 * Store into a slot past the end of the array.  On the grouped rel
		 * this repallocs the array it shares with the plain joinrel, which
		 * frees the plain joinrel's array out from under it.
		 */
		if (test_mode > 1 && IS_GROUPED_REL(joinrel) &&
			tag != NULL && tag->owner != joinrel)
		{
			RelOptInfo *orig = tag->owner;
			void	  **before = joinrel->extension_state;

			SetRelOptInfoExtensionState(joinrel, ext_ids[NUM_IDS - 1], tag);

			if (orig->extension_state == before &&
				joinrel->extension_state != before)
				elog(WARNING, "plain joinrel %s still points at array %p, which repalloc just freed; slot %d now reads %p",
					 bmsToString(orig->relids), (void *) before, ext_ids[1],
					 orig->extension_state[ext_ids[1]]);
		}
	}

	if (prev_join_path_setup_hook)
		prev_join_path_setup_hook(root, joinrel, outerrel, innerrel,
								  jointype, extra);
}

void
_PG_init(void)
{
	DefineCustomIntVariable("test_extstate.mode",
							"0 = off, 1 = report shared state, 2 = force array growth.",
							NULL,
							&test_mode,
							0, 0, 2,
							PGC_USERSET,
							0,
							NULL, NULL, NULL);
	MarkGUCPrefixReserved("test_extstate");

	for (int i = 0; i < NUM_IDS; i++)
		ext_ids[i] = GetPlannerExtensionId(ext_names[i]);

	prev_joinrel_setup_hook = joinrel_setup_hook;
	joinrel_setup_hook = tes_joinrel_setup;
	prev_join_path_setup_hook = join_path_setup_hook;
	join_path_setup_hook = tes_join_path_setup;
}