diff --git a/external/Makefile b/external/Makefile index 8d0467dda55..edf89954e3d 100644 --- a/external/Makefile +++ b/external/Makefile @@ -10,6 +10,7 @@ SUBDIRS = # sort extention by names, less git conflict SUBDIRS += polar_feature_utils SUBDIRS += polar_io_stat +SUBDIRS += polar_login_history SUBDIRS += polar_monitor SUBDIRS += polar_monitor_preload SUBDIRS += polar_parameter_manager diff --git a/external/polar_login_history/Makefile b/external/polar_login_history/Makefile new file mode 100644 index 00000000000..f9b20f7d251 --- /dev/null +++ b/external/polar_login_history/Makefile @@ -0,0 +1,20 @@ +# external/polar_login_history/Makefile + +MODULE_big = polar_login_history +OBJS = polar_login_history.o +EXTENSION = polar_login_history +NO_INSTALLCHECK = 1 + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = external/polar_login_history +top_builddir = ../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif + +check: + $(prove_check) diff --git a/external/polar_login_history/polar_login_history.c b/external/polar_login_history/polar_login_history.c new file mode 100644 index 00000000000..bc4b2d18c23 --- /dev/null +++ b/external/polar_login_history/polar_login_history.c @@ -0,0 +1,876 @@ +/*------------------------------------------------------------------------- + * + * polar_login_history.c + * + * IDENTIFICATION + * external/polar_login_history/polar_login_history.c + * + *------------------------------------------------------------------------- + */ +#include "postgres.h" + +#include + +#include "access/xact.h" +#include "commands/user.h" +#include "libpq/libpq-be.h" +#include "miscadmin.h" +#include "replication/walsender.h" +#include "storage/fd.h" +#include "storage/ipc.h" +#include "storage/lwlock.h" +#include "storage/shmem.h" +#include "storage/spin.h" +#include "utils/acl.h" +#include "utils/guc.h" +#include "utils/memutils.h" +#include "utils/snapmgr.h" +#include "utils/syscache.h" +#include "utils/timestamp.h" + +PG_MODULE_MAGIC; + + +#define POLAR_LOGIN_HISTORY_FILE "polar_login_history" + +/* GUC variable */ +static bool polar_login_history_enable; +static int polar_login_history_maxnum; + +/**************************************************************************************************/ + +#define POLAR_LOGIN_HISTORY_STRING_LENGTH 256 + +#define POLAR_LOGIN_HISTORY_VALID 1 + +typedef struct polar_login_history_info +{ + Oid useroid; /* user oid */ + TimestampTz logintime; /* last login time */ + char ip[POLAR_LOGIN_HISTORY_STRING_LENGTH]; /* last login ip */ + char application[POLAR_LOGIN_HISTORY_STRING_LENGTH]; /* the application used + * for the last login */ + uint8 failcount; /* the number of failed login attempts since + * the last successful login */ +} polar_login_history_info; + +typedef struct polar_login_history_array_struct +{ + slock_t spinlock; + uint8 state; /* identifies whether the element is valid */ + int freenext; /* if in the free list, points to the index of + * the next free element */ + polar_login_history_info logininfo; /* login information of the user */ +} polar_login_history_array_struct; + +/* Use a fixed-length array to record the login information of all users */ +static polar_login_history_array_struct * polar_login_history_array = NULL; + +/**************************************************************************************************/ + +#define POLAR_LOGIN_HISTORY_LOCK_NUM (polar_login_history_maxnum / 64 * 2) + +typedef struct polar_login_history_lock_struct +{ + LWLock *lwlock; +} polar_login_history_lock_struct; + +/* The hash table is partitioned, and different areas are protected with different locks */ +static polar_login_history_lock_struct * polar_login_history_lock = NULL; + +typedef struct polar_login_history_hash_entry +{ + Oid useroid; + int index; /* index in the array */ +} polar_login_history_hash_entry; + +/* The mapping between user oids and array subscripts is maintained through a hash table */ +static HTAB *polar_login_history_hash = NULL; + +/**************************************************************************************************/ + +#define POLAR_LOGIN_HISTORY_FREENEXT_END_OF_LIST (-1) +#define POLAR_LOGIN_HISTORY_FREENEXT_NOT_IN_LIST (-2) + +typedef struct polar_login_history_freelist_struct +{ + slock_t spinlock; + int firstfree; /* head of the free list */ + int lastfree; /* tail of the free list */ +} polar_login_history_freelist_struct; + +/* A list of free array elements that can be used directly when generating new user login information */ +static polar_login_history_freelist_struct * polar_login_history_freelist = NULL; + +/**************************************************************************************************/ + +typedef struct polar_login_history_callback_arg +{ + TransactionId xid; /* current transaction id */ + List *useroid; /* a list of users to be dropped */ +} polar_login_history_callback_arg; + +typedef struct polar_login_history_xact_callback_item +{ + struct polar_login_history_xact_callback_item *next; + XactCallback callback; + void *arg; +} polar_login_history_xact_callback_item; + +/* + * Records the function that requires a callback when a transaction is committed, + * mainly to clean up the login information of the dropped user. + */ +static polar_login_history_xact_callback_item * polar_login_history_xact_callback = NULL; + +/**************************************************************************************************/ + +static shmem_startup_hook_type prev_shmem_startup_hook = NULL; +static shmem_request_hook_type prev_shmem_request_hook = NULL; + +void _PG_init(void); + + +/* + * Estimate shared memory space needed. + */ +static Size +polar_login_history_memsize(void) +{ + Size size; + + size = MAXALIGN(mul_size(POLAR_LOGIN_HISTORY_LOCK_NUM, sizeof(polar_login_history_lock_struct))); + size = add_size(size, hash_estimate_size(polar_login_history_maxnum, sizeof(polar_login_history_hash_entry))); + size = add_size(size, MAXALIGN(mul_size(polar_login_history_maxnum, sizeof(polar_login_history_array_struct)))); + size = add_size(size, MAXALIGN(sizeof(polar_login_history_freelist_struct))); + + return size; +} + +/* + * shmem_request hook: request additional shared resources. We'll allocate or + * attach to the shared resources in polar_login_history_shmem_startup(). + */ +static void +polar_login_history_shmem_request(void) +{ + if (prev_shmem_request_hook) + prev_shmem_request_hook(); + + if (!polar_login_history_enable) + return; + + RequestAddinShmemSpace(polar_login_history_memsize()); + RequestNamedLWLockTranche("polar_login_history", POLAR_LOGIN_HISTORY_LOCK_NUM); +} + +static void +polar_save_information(void) +{ + FILE *file = NULL; + + /* Safety check ... shouldn't get here unless shmem is set up */ + if (!polar_login_history_hash || !polar_login_history_array) + return; + + file = AllocateFile(POLAR_LOGIN_HISTORY_FILE ".tmp", PG_BINARY_W); + if (file == NULL) + goto error; + + /* Serialize to disk */ + for (int i = 0; i < polar_login_history_maxnum; i++) + { + polar_login_history_array_struct *elem = &polar_login_history_array[i]; + + if (elem->state & POLAR_LOGIN_HISTORY_VALID) + { + if (fwrite(&elem->logininfo, sizeof(polar_login_history_info), 1, file) != 1) + goto error; + } + } + + FreeFile(file); + /* Rename file into place, so we atomically replace any old one */ + (void) durable_rename(POLAR_LOGIN_HISTORY_FILE ".tmp", POLAR_LOGIN_HISTORY_FILE, LOG); + return; + +error: + ereport(WARNING, + (errcode_for_file_access(), + errmsg("could not write file \"%s\": %m", + POLAR_LOGIN_HISTORY_FILE ".tmp"))); + if (file) + FreeFile(file); + unlink(POLAR_LOGIN_HISTORY_FILE ".tmp"); +} + +static void +polar_load_information(void) +{ + FILE *file = NULL; + int num = 0; + + /* Attempt to load old statistics from the dump file */ + file = AllocateFile(POLAR_LOGIN_HISTORY_FILE, PG_BINARY_R); + if (file == NULL) + { + if (errno != ENOENT) + goto read_error; + return; + } + + /* Check whether the file is empty */ + fseek(file, 0, SEEK_END); + if (ftell(file) == 0) + { + FreeFile(file); + unlink(POLAR_LOGIN_HISTORY_FILE); + return; + } + rewind(file); + + do + { + polar_login_history_info info; + polar_login_history_hash_entry *entry; + + if (num >= polar_login_history_maxnum) + { + /* + * The login information saved in the file exceeds the allocated + * shared memory size, so the user needs to reallocate it. Setting + * NULL prevents files from being updated when the database is + * stopped. + */ + FreeFile(file); + polar_login_history_hash = NULL; + polar_login_history_array = NULL; + ereport(FATAL, (errmsg("Out of memory, please adjust polar_login_history.maxnum"))); + } + + if (fread(&info, sizeof(polar_login_history_info), 1, file) != 1) + goto read_error; + + polar_login_history_freelist->firstfree = polar_login_history_array[num].freenext; + + polar_login_history_array[num].state = POLAR_LOGIN_HISTORY_VALID; + polar_login_history_array[num].freenext = POLAR_LOGIN_HISTORY_FREENEXT_NOT_IN_LIST; + polar_login_history_array[num].logininfo = info; + + /* Create an entry with desired hash code */ + entry = (polar_login_history_hash_entry *) + hash_search(polar_login_history_hash, &info.useroid, HASH_ENTER, NULL); + entry->useroid = info.useroid; + entry->index = num; + + num++; + } while (!feof(file)); + + FreeFile(file); + + /* + * Remove the persisted stats file so it's not included in + * backups/replication standbys, etc + */ + unlink(POLAR_LOGIN_HISTORY_FILE); + return; + +read_error: + ereport(LOG, + (errcode_for_file_access(), + errmsg("could not read file \"%s\": %m", + POLAR_LOGIN_HISTORY_FILE))); + if (file) + FreeFile(file); + /* If possible, throw away the bogus file; ignore any error */ + unlink(POLAR_LOGIN_HISTORY_FILE); +} + +/* + * shmem_shutdown hook: Dump login information into file. + * + * Note: we don't bother with acquiring lock, because there should be no + * other processes running when this is called. + */ +static void +polar_login_history_shmem_shutdown(int code, Datum arg) +{ + /* Don't try to dump during a crash. */ + if (code) + return; + + polar_save_information(); +} + +/* + * shmem_startup hook: allocate or attach to shared memory, + * then load any pre-existing login information from file. + */ +static void +polar_login_history_shmem_startup(void) +{ + bool found_lock; + bool found_array; + bool found_freelist; + HASHCTL info; + + if (prev_shmem_startup_hook) + prev_shmem_startup_hook(); + + if (!polar_login_history_enable) + return; + + /* + * Create or attach to the shared memory state, including hash table. + */ + LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); + + polar_login_history_lock = (polar_login_history_lock_struct *) + ShmemInitStruct("polar_login_history_lock", + POLAR_LOGIN_HISTORY_LOCK_NUM * sizeof(polar_login_history_lock_struct), + &found_lock); + polar_login_history_array = (polar_login_history_array_struct *) + ShmemInitStruct("polar_login_history_array", + polar_login_history_maxnum * sizeof(polar_login_history_array_struct), + &found_array); + polar_login_history_freelist = (polar_login_history_freelist_struct *) + ShmemInitStruct("polar_login_history_freelist", + sizeof(polar_login_history_freelist_struct), + &found_freelist); + + if (found_lock || found_array || found_freelist) + { + /* should find all of these, or none of them */ + if (!(found_lock && found_array && found_freelist)) + { + polar_login_history_lock = NULL; + polar_login_history_array = NULL; + polar_login_history_freelist = NULL; + ereport(WARNING, (errmsg("The login history function fails to allocate shared memory"))); + return; + } + } + else + { + /* First time through ... */ + LWLockPadded *lwlock = GetNamedLWLockTranche("polar_login_history"); + + for (int i = 0; i < POLAR_LOGIN_HISTORY_LOCK_NUM; i++) + polar_login_history_lock[i].lwlock = &(lwlock[i].lock); + + for (int i = 0; i < polar_login_history_maxnum; i++) + { + SpinLockInit(&polar_login_history_array[i].spinlock); + polar_login_history_array[i].state = 0; + /* Initially link all array elements together as unused */ + polar_login_history_array[i].freenext = i + 1; + } + /* Correct last entry of linked list */ + polar_login_history_array[polar_login_history_maxnum - 1].freenext = POLAR_LOGIN_HISTORY_FREENEXT_END_OF_LIST; + + SpinLockInit(&polar_login_history_freelist->spinlock); + polar_login_history_freelist->firstfree = 0; + polar_login_history_freelist->lastfree = polar_login_history_maxnum - 1; + } + + info.keysize = sizeof(Oid); + info.entrysize = sizeof(polar_login_history_hash_entry); + info.num_partitions = POLAR_LOGIN_HISTORY_LOCK_NUM; + polar_login_history_hash = ShmemInitHash("polar_login_history_hash", + polar_login_history_maxnum, + polar_login_history_maxnum, + &info, + HASH_ELEM | HASH_BLOBS | HASH_PARTITION); + + LWLockRelease(AddinShmemInitLock); + + /* + * If we're in the postmaster (or a standalone backend...), set up a shmem + * exit hook to dump the statistics to disk. + */ + if (!IsUnderPostmaster) + on_shmem_exit(polar_login_history_shmem_shutdown, (Datum) 0); + + /* + * Done if some other process already completed our initialization. + */ + if (found_lock && found_array && found_freelist) + return; + + /* + * Note: we don't bother with locks here, because there should be no other + * processes running when this code is reached. + */ + polar_load_information(); +} + +static bool +polar_login_history_check(void) +{ + if (!polar_login_history_enable + || !polar_login_history_lock + || !polar_login_history_hash + || !polar_login_history_array) + return false; + + return true; +} + +/* + * Show the user the last login information and update the current login information. + */ +static void +polar_login_history_show_info(polar_login_history_info * info, Oid useroid, bool login_success, bool first_login) +{ + /* Print login information */ + if (login_success && !first_login) + ereport(INFO, + (errmsg("\n" \ + "Last login: %s from %s using %s\n" \ + "The number of failures since the last successful login is %d", + timestamptz_to_str(info->logintime), + info->ip, + info->application, + info->failcount))); + + /* Update login information */ + info->useroid = useroid; + info->logintime = MyStartTimestamp ? MyStartTimestamp : GetCurrentTimestamp(); + strncpy(info->ip, + MyProcPort->remote_host ? MyProcPort->remote_host : "unrecognized", + POLAR_LOGIN_HISTORY_STRING_LENGTH - 1); + info->ip[POLAR_LOGIN_HISTORY_STRING_LENGTH - 1] = '\0'; + strncpy(info->application, + MyProcPort->application_name ? MyProcPort->application_name : "unrecognized", + POLAR_LOGIN_HISTORY_STRING_LENGTH - 1); + info->application[POLAR_LOGIN_HISTORY_STRING_LENGTH - 1] = '\0'; + info->failcount = login_success ? 0 : (first_login ? 1 : info->failcount + 1); + + /* Print login information */ + if (login_success && first_login) + ereport(INFO, + (errmsg("\n" \ + "First login: %s from %s using %s", + timestamptz_to_str(info->logintime), + info->ip, + info->application))); +} + +/* + * Gets an array element from the free list. + */ +static int +polar_login_history_get_element(void) +{ + polar_login_history_array_struct *elem; + int result = -1; + + SpinLockAcquire(&polar_login_history_freelist->spinlock); + + /* There is no element in the free list */ + if (polar_login_history_freelist->firstfree < 0) + { + SpinLockRelease(&polar_login_history_freelist->spinlock); + ereport(WARNING, + (errmsg("The login information of the user cannot be recorded. " \ + "Please adjust polar_login_history.maxnum to increase the memory space."))); + return result; + } + + result = polar_login_history_freelist->firstfree; + + elem = &polar_login_history_array[polar_login_history_freelist->firstfree]; + polar_login_history_freelist->firstfree = elem->freenext; + elem->freenext = POLAR_LOGIN_HISTORY_FREENEXT_NOT_IN_LIST; + + SpinLockRelease(&polar_login_history_freelist->spinlock); + return result; +} + +/* + * An array element is invalid, add it to the free list header. + */ +static void +polar_login_history_free_element(int id) +{ + polar_login_history_array_struct *elem = &polar_login_history_array[id]; + + SpinLockAcquire(&polar_login_history_freelist->spinlock); + + /* + * It is possible that we are told to put something in the freelist that + * is already in it; don't screw up the list if so. + */ + if (elem->freenext == POLAR_LOGIN_HISTORY_FREENEXT_NOT_IN_LIST) + { + elem->freenext = polar_login_history_freelist->firstfree; + if (elem->freenext < 0) + polar_login_history_freelist->lastfree = id; + polar_login_history_freelist->firstfree = id; + } + + SpinLockRelease(&polar_login_history_freelist->spinlock); +} + +/* + * When a user logs in to the database, the login information is recorded regardless + * of whether the login succeeds or fails. However, only after a successful login, + * the database will show the user the last login information. + */ +static void +polar_update_login_history(bool login_success) +{ + bool intrans; + Oid useroid; + uint32 hashcode; + LWLock *partitionlock; + polar_login_history_hash_entry *entry; + bool found; + polar_login_history_array_struct *elem; + bool first_login = false; + + if (!polar_login_history_check()) + return; + + /* Only the postgres process calls this function during login */ + if (am_walsender || !polar_login_flag) + return; + + /* + * To avoid repeated records of login information, that is, to ensure that + * the function is called only once during login, set the flag to false + * immediately to mark the end of the login. + */ + polar_login_flag = false; + + if (!MyProcPort || !MyProcPort->user_name) + return; + + /* User information needs to be queried in a transaction */ + intrans = IsTransactionState(); + if (intrans) + useroid = get_role_oid(MyProcPort->user_name, true); + else + { + StartTransactionCommand(); + useroid = get_role_oid(MyProcPort->user_name, true); + } + /* If the user does not exist, return instead of reporting an error */ + if (!OidIsValid(useroid)) + return; + + hashcode = get_hash_value(polar_login_history_hash, &useroid); + partitionlock = polar_login_history_lock[hashcode % POLAR_LOGIN_HISTORY_LOCK_NUM].lwlock; + + /* Check whether the hash table contains information about the user */ + LWLockAcquire(partitionlock, LW_SHARED); + entry = (polar_login_history_hash_entry *) + hash_search_with_hash_value(polar_login_history_hash, &useroid, hashcode, HASH_FIND, &found); + if (!found) + { + /* New user information needs to be inserted */ + LWLockRelease(partitionlock); + LWLockAcquire(partitionlock, LW_EXCLUSIVE); + + /* Find again to prevent having been inserted by another process */ + entry = (polar_login_history_hash_entry *) + hash_search_with_hash_value(polar_login_history_hash, &useroid, hashcode, HASH_FIND, &found); + if (!found) + { + uint32 hash_value; + bool user_exist; + int new_id; + + /* + * Check whether the user exists to prevent it from being dropped + * by other processes. Invalidates cached data first. + */ + hash_value = GetSysCacheHashValue1(AUTHOID, ObjectIdGetDatum(useroid)); + SysCacheInvalidate(AUTHOID, hash_value); + user_exist = SearchSysCacheExists1(AUTHOID, ObjectIdGetDatum(useroid)); + if (!intrans) + CommitTransactionCommand(); + if (!user_exist) + { + LWLockRelease(partitionlock); + ereport(WARNING, (errmsg("role %u was concurrently dropped", useroid))); + return; + } + + /* Looks for an empty array element */ + new_id = polar_login_history_get_element(); + if (new_id < 0) + { + /* + * There is not enough memory space and the error has been + * printed + */ + LWLockRelease(partitionlock); + return; + } + + entry = (polar_login_history_hash_entry *) + hash_search_with_hash_value(polar_login_history_hash, &useroid, hashcode, HASH_ENTER, NULL); + entry->useroid = useroid; + entry->index = new_id; + + first_login = true; + } + else + { + /* Just update the user information in the array */ + if (!intrans) + CommitTransactionCommand(); + } + } + else + { + /* Just update the user information in the array */ + if (!intrans) + CommitTransactionCommand(); + } + + elem = &polar_login_history_array[entry->index]; + SpinLockAcquire(&elem->spinlock); + + LWLockRelease(partitionlock); + + polar_login_history_show_info(&elem->logininfo, useroid, login_success, first_login); + elem->state |= POLAR_LOGIN_HISTORY_VALID; + SpinLockRelease(&elem->spinlock); +} + +/* + * When dropping a user, the login information is also deleted. + */ +static void +polar_delete_login_history(Oid useroid) +{ + uint32 hashcode; + LWLock *partitionlock; + polar_login_history_hash_entry *entry; + bool found; + + if (!OidIsValid(useroid)) + return; + + hashcode = get_hash_value(polar_login_history_hash, &useroid); + partitionlock = polar_login_history_lock[hashcode % POLAR_LOGIN_HISTORY_LOCK_NUM].lwlock; + + /* Check whether the hash table contains information about the user */ + LWLockAcquire(partitionlock, LW_SHARED); + entry = (polar_login_history_hash_entry *) + hash_search_with_hash_value(polar_login_history_hash, &useroid, hashcode, HASH_FIND, &found); + if (!found) + { + /* If not found, no processing is done */ + LWLockRelease(partitionlock); + } + else + { + LWLockRelease(partitionlock); + + /* Preparing to delete login information */ + LWLockAcquire(partitionlock, LW_EXCLUSIVE); + if (hash_search_with_hash_value(polar_login_history_hash, &useroid, hashcode, HASH_REMOVE, NULL) == NULL) + { + /* It has been deleted by another process */ + LWLockRelease(partitionlock); + } + else + { + polar_login_history_array_struct *elem = &polar_login_history_array[entry->index]; + + SpinLockAcquire(&elem->spinlock); + + LWLockRelease(partitionlock); + + /* Set the array element to invalid and add it to the free list */ + elem->state &= ~POLAR_LOGIN_HISTORY_VALID; + SpinLockRelease(&elem->spinlock); + polar_login_history_free_element(entry->index); + } + } +} + +/* + * When the top-level transaction is committed, the login information of the user + * associated with the committed transaction (including sub-transactions) is deleted. + */ +static void +polar_delete_login_history_callback(XactEvent event, void *arg) +{ + polar_login_history_callback_arg *content = (polar_login_history_callback_arg *) arg; + + if (event == XACT_EVENT_COMMIT) + { + /* Submitted by the top-level transaction */ + if (GetCurrentTransactionId() == content->xid) + { + ListCell *item; + + foreach(item, content->useroid) + polar_delete_login_history(lfirst_oid(item)); + } + else + { + TransactionId *children; + int nchildren = xactGetCommittedChildren(&children); + + for (int i = 0; i < nchildren; i++) + { + /* + * Indicates that the transaction has committed and the login + * information can be removed + */ + if (children[i] == content->xid) + { + ListCell *item; + + foreach(item, content->useroid) + polar_delete_login_history(lfirst_oid(item)); + break; + } + } + } + } +} + +/* + * Register a callback function when dropping users so that the associated login information + * is not deleted until the transaction is committed. + */ +static void +polar_register_delete_login_history(List *userlist) +{ + MemoryContext oldContext; + polar_login_history_callback_arg *arg; + polar_login_history_xact_callback_item *item; + + if (!polar_login_history_check() || userlist == NIL) + return; + + oldContext = MemoryContextSwitchTo(TopMemoryContext); + + arg = (polar_login_history_callback_arg *) palloc0(sizeof(polar_login_history_callback_arg)); + arg->xid = GetCurrentTransactionId(); + arg->useroid = list_copy(userlist); + + item = (polar_login_history_xact_callback_item *) palloc0(sizeof(polar_login_history_xact_callback_item)); + item->callback = polar_delete_login_history_callback; + item->arg = arg; + item->next = polar_login_history_xact_callback; + polar_login_history_xact_callback = item; + + MemoryContextSwitchTo(oldContext); +} + +/* + * When a transaction is committed or dropped, this function is called to clean up the login information. + */ +static void +polar_callback_delete_login_history(XactEvent event) +{ + if (!polar_login_history_check() || polar_login_history_xact_callback == NULL) + return; + + while (polar_login_history_xact_callback) + { + polar_login_history_xact_callback_item *item = polar_login_history_xact_callback; + + item->callback(event, item->arg); + + polar_login_history_xact_callback = item->next; + list_free(((polar_login_history_callback_arg *) item->arg)->useroid); + pfree(item->arg); + pfree(item); + } +} + +/* + * User login information is periodically written to a file by the polar_worker process. + */ +static void +polar_flush_login_history(void) +{ + FILE *file = NULL; + + if (!polar_login_history_check()) + return; + + file = AllocateFile(POLAR_LOGIN_HISTORY_FILE ".tmp", PG_BINARY_W); + if (file == NULL) + goto error; + + for (int i = 0; i < polar_login_history_maxnum; i++) + { + polar_login_history_array_struct *elem = &polar_login_history_array[i]; + + SpinLockAcquire(&elem->spinlock); + if (elem->state & POLAR_LOGIN_HISTORY_VALID) + { + polar_login_history_info info = elem->logininfo; + + SpinLockRelease(&elem->spinlock); + if (fwrite(&info, sizeof(polar_login_history_info), 1, file) != 1) + goto error; + } + else + SpinLockRelease(&elem->spinlock); + } + + FreeFile(file); + /* Rename file into place, so we atomically replace any old one */ + (void) durable_rename(POLAR_LOGIN_HISTORY_FILE ".tmp", POLAR_LOGIN_HISTORY_FILE, LOG); + return; + +error: + ereport(LOG, + (errcode_for_file_access(), + errmsg("could not write file \"%s\": %m", + POLAR_LOGIN_HISTORY_FILE ".tmp"))); + if (file) + FreeFile(file); + unlink(POLAR_LOGIN_HISTORY_FILE ".tmp"); +} + +void +_PG_init(void) +{ + /* Define custom GUC variables */ + DefineCustomBoolVariable("polar_login_history.enable", + "Whether to enable the login history function", + "After the function is enabled, historical login information is displayed to users", + &polar_login_history_enable, + false, + PGC_POSTMASTER, + POLAR_GUC_IS_VISIBLE | POLAR_GUC_IS_CHANGABLE, + NULL, + NULL, + NULL); + + DefineCustomIntVariable("polar_login_history.maxnum", + "The maximum number of users that this function can record", + NULL, + &polar_login_history_maxnum, + 512, + 64, + INT_MAX / 2, + PGC_POSTMASTER, + POLAR_GUC_IS_VISIBLE | POLAR_GUC_IS_CHANGABLE, + NULL, + NULL, + NULL); + + /* Install hooks */ + prev_shmem_request_hook = shmem_request_hook; + shmem_request_hook = polar_login_history_shmem_request; + + prev_shmem_startup_hook = shmem_startup_hook; + shmem_startup_hook = polar_login_history_shmem_startup; + + polar_update_login_history_hook = polar_update_login_history; + polar_register_delete_login_history_hook = polar_register_delete_login_history; + polar_callback_delete_login_history_hook = polar_callback_delete_login_history; + polar_flush_login_history_hook = polar_flush_login_history; +} diff --git a/external/polar_login_history/polar_login_history.control b/external/polar_login_history/polar_login_history.control new file mode 100644 index 00000000000..8e0796c6b89 --- /dev/null +++ b/external/polar_login_history/polar_login_history.control @@ -0,0 +1,4 @@ +comment = 'record user login information' +default_version = '1.0' +module_pathname = '$libdir/polar_login_history' +relocatable = true diff --git a/external/polar_login_history/t/polar_login_history.pl b/external/polar_login_history/t/polar_login_history.pl new file mode 100644 index 00000000000..d903dc4172e --- /dev/null +++ b/external/polar_login_history/t/polar_login_history.pl @@ -0,0 +1,101 @@ +#!/usr/bin/perl + +# polar_login_history.pl +# +# IDENTIFICATION +# external/polar_login_history/t/polar_login_history.pl + +use strict; +use warnings; +use PostgreSQL::Test::Cluster; +use PostgreSQL::Test::Utils; +use Test::More; + +############### init primary/replica/standby ############## +my $node_primary = PostgreSQL::Test::Cluster->new('primary'); +$node_primary->polar_init_primary; +$node_primary->append_conf('postgresql.conf', + "shared_preload_libraries = 'polar_login_history'"); + +my $node_replica = PostgreSQL::Test::Cluster->new('replica'); +$node_replica->polar_init_replica($node_primary); + +my $node_standby = PostgreSQL::Test::Cluster->new('standby'); +$node_standby->polar_init_standby($node_primary); + +$node_primary->start; + +$node_primary->polar_create_slot($node_replica->name); +$node_primary->polar_create_slot($node_standby->name); + +$node_replica->start; +$node_standby->start; + +$node_standby->polar_drop_all_slots; + + +############### enable login history ############## +$node_primary->append_conf('postgresql.conf', + "polar_login_history.enable = on"); +$node_primary->restart; +is( $node_primary->safe_psql('postgres', 'show polar_login_history.enable;'), + 'on', + 'login history function is enabled'); + +$node_replica->append_conf('postgresql.conf', + "polar_login_history.enable = on"); +$node_replica->restart; +is( $node_replica->safe_psql('postgres', 'show polar_login_history.enable;'), + 'on', + 'login history function is enabled'); + +$node_standby->append_conf('postgresql.conf', + "polar_login_history.enable = on"); +$node_standby->restart; +is( $node_standby->safe_psql('postgres', 'show polar_login_history.enable;'), + 'on', + 'login history function is enabled'); + + +############### log in to the database with the new user ############## +$node_primary->safe_psql('postgres', 'create user zhangsan;'); +$node_replica->restart; + +is( $node_primary->psql( + 'postgres', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); +isnt( + $node_primary->psql( + 'postgres1', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); + +is( $node_replica->psql( + 'postgres', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); +isnt( + $node_replica->psql( + 'postgres1', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); + +is( $node_standby->psql( + 'postgres', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); +isnt( + $node_standby->psql( + 'postgres1', undef, extra_params => [ '-U', 'zhangsan' ]), + 0, + 'login success'); + +$node_primary->safe_psql('postgres', 'drop user zhangsan;'); + + +############### stop ############## +$node_primary->stop; +$node_replica->stop; +$node_standby->stop; +done_testing(); diff --git a/external/polar_worker/polar_worker.c b/external/polar_worker/polar_worker.c index ae7efb697dc..0b6f4a88a81 100644 --- a/external/polar_worker/polar_worker.c +++ b/external/polar_worker/polar_worker.c @@ -346,6 +346,10 @@ polar_worker_handler_main(Datum main_arg) ProcessConfigFile(PGC_SIGHUP); } + /* Write the user's login information to the file */ + if (polar_flush_login_history_hook) + polar_flush_login_history_hook(); + if (!enable_polar_worker) { /* We don't want to prealloc wal file now, so just wait forever. */ diff --git a/src/backend/access/transam/xact.c b/src/backend/access/transam/xact.c index 1d3be6c0d1a..83fd36be2ea 100644 --- a/src/backend/access/transam/xact.c +++ b/src/backend/access/transam/xact.c @@ -337,6 +337,13 @@ static SubXactCallbackItem *SubXact_callbacks = NULL; polar_unsplittable_reason_t polar_unsplittable_reason; XLogRecPtr polar_xact_split_wait_lsn = InvalidXLogRecPtr; +/* + * POLAR: login history + */ +polar_callback_delete_login_history_hook_type polar_callback_delete_login_history_hook = NULL; + +/* POLAR end */ + /* local function prototypes */ static void AssignTransactionId(TransactionState s); static void AbortTransaction(void); @@ -2351,6 +2358,13 @@ CommitTransaction(void) CallXactCallbacks(is_parallel_worker ? XACT_EVENT_PARALLEL_COMMIT : XACT_EVENT_COMMIT); + /* + * POLAR: login history Delete the login information of the user. + */ + if (polar_callback_delete_login_history_hook) + polar_callback_delete_login_history_hook(is_parallel_worker ? XACT_EVENT_PARALLEL_COMMIT : XACT_EVENT_COMMIT); + /* POLAR end */ + ResourceOwnerRelease(TopTransactionResourceOwner, RESOURCE_RELEASE_BEFORE_LOCKS, true, true); @@ -2913,6 +2927,14 @@ AbortTransaction(void) else CallXactCallbacks(XACT_EVENT_ABORT); + /* + * POLAR: login history Cancel the callback function for deleting user + * login information. + */ + if (polar_callback_delete_login_history_hook) + polar_callback_delete_login_history_hook(is_parallel_worker ? XACT_EVENT_PARALLEL_ABORT : XACT_EVENT_ABORT); + /* POLAR end */ + ResourceOwnerRelease(TopTransactionResourceOwner, RESOURCE_RELEASE_BEFORE_LOCKS, false, true); diff --git a/src/backend/commands/user.c b/src/backend/commands/user.c index 2ce21c1b045..1afe7b13ef7 100644 --- a/src/backend/commands/user.c +++ b/src/backend/commands/user.c @@ -50,6 +50,13 @@ int Password_encryption = PASSWORD_TYPE_SCRAM_SHA_256; /* Hook to check passwords in CreateRole() and AlterRole() */ check_password_hook_type check_password_hook = NULL; +/* + * POLAR: login history + */ +polar_register_delete_login_history_hook_type polar_register_delete_login_history_hook = NULL; + +/* POLAR end */ + static void AddRoleMems(const char *rolename, Oid roleid, List *memberSpecs, List *memberIds, Oid grantorId, bool admin_opt); @@ -913,6 +920,13 @@ DropRole(DropRoleStmt *stmt) pg_auth_members_rel; ListCell *item; + /* + * POLAR: login history + */ + List *roleidlist = NIL; + + /* POLAR end */ + if (!have_createrole_privilege()) ereport(ERROR, (errcode(ERRCODE_INSUFFICIENT_PRIVILEGE), @@ -1072,8 +1086,22 @@ DropRole(DropRoleStmt *stmt) * itself.) */ CommandCounterIncrement(); + + /* + * POLAR: login history + */ + roleidlist = lappend_oid(roleidlist, roleid); + /* POLAR end */ } + /* + * POLAR: login history Register a callback function to delete the login + * information for these users when the transaction is committed. + */ + if (polar_register_delete_login_history_hook) + polar_register_delete_login_history_hook(roleidlist); + /* POLAR end */ + /* * Now we can clean up; but keep locks until commit. */ diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c index 81c027c0621..a097bee3d66 100644 --- a/src/backend/tcop/postgres.c +++ b/src/backend/tcop/postgres.c @@ -4510,6 +4510,12 @@ PostgresMain(const char *dbname, const char *username) long secs = 0; int microsecs = 0; + /* + * POLAR: login history Indicate that the login process starts. + */ + polar_login_flag = true; + /* POLAR end */ + AssertArg(dbname != NULL); AssertArg(username != NULL); @@ -4678,6 +4684,14 @@ PostgresMain(const char *dbname, const char *username) initStringInfo(&row_description_buf); MemoryContextSwitchTo(TopMemoryContext); + /* + * POLAR: login history Record the successful login information of the + * user. + */ + if (polar_update_login_history_hook) + polar_update_login_history_hook(true); + /* POLAR end */ + /* * POSTGRES main processing loop begins here * diff --git a/src/backend/utils/error/elog.c b/src/backend/utils/error/elog.c index 874e12a363f..c3fbab4dd03 100644 --- a/src/backend/utils/error/elog.c +++ b/src/backend/utils/error/elog.c @@ -136,6 +136,15 @@ int polar_audit_log_flush_timeout = 0; /* POLAR end */ +/* + * POLAR: login history + */ +bool polar_login_flag = false; +polar_update_login_history_hook_type polar_update_login_history_hook = NULL; +polar_flush_login_history_hook_type polar_flush_login_history_hook = NULL; + +/* POLAR end */ + #ifdef HAVE_SYSLOG /* @@ -541,6 +550,14 @@ errfinish(const char *filename, int lineno, const char *funcname) MemoryContext oldcontext; ErrorContextCallback *econtext; + /* + * POLAR: login history Record the failed login information of the user. + * Most of the errors caused by users during login are fatal. + */ + if (edata->elevel == FATAL && polar_update_login_history_hook) + polar_update_login_history_hook(false); + /* POLAR end */ + recursion_depth++; CHECK_STACK_DEPTH(); diff --git a/src/include/access/xact.h b/src/include/access/xact.h index ce0e5f66a2d..895d9f9182a 100644 --- a/src/include/access/xact.h +++ b/src/include/access/xact.h @@ -558,4 +558,11 @@ extern void polar_stat_update_xact_split_info(void); /* POLAR end */ +/* + * POLAR: login history + */ +typedef void (*polar_callback_delete_login_history_hook_type) (XactEvent); +extern PGDLLIMPORT polar_callback_delete_login_history_hook_type polar_callback_delete_login_history_hook; +/* POLAR end */ + #endif /* XACT_H */ diff --git a/src/include/commands/user.h b/src/include/commands/user.h index d3dd8303d28..cdf44714013 100644 --- a/src/include/commands/user.h +++ b/src/include/commands/user.h @@ -34,4 +34,11 @@ extern void DropOwnedObjects(DropOwnedStmt *stmt); extern void ReassignOwnedObjects(ReassignOwnedStmt *stmt); extern List *roleSpecsToIds(List *memberNames); +/* + * POLAR: login history + */ +typedef void (*polar_register_delete_login_history_hook_type) (List *); +extern PGDLLIMPORT polar_register_delete_login_history_hook_type polar_register_delete_login_history_hook; +/* POLAR end */ + #endif /* USER_H */ diff --git a/src/include/utils/elog.h b/src/include/utils/elog.h index 5995448cd48..420924b7e05 100644 --- a/src/include/utils/elog.h +++ b/src/include/utils/elog.h @@ -536,4 +536,14 @@ extern void polar_reset_program_error_handler(void); /* POLAR end */ +/* + * POLAR: login history + */ +extern PGDLLIMPORT bool polar_login_flag; +typedef void (*polar_update_login_history_hook_type) (bool); +extern PGDLLIMPORT polar_update_login_history_hook_type polar_update_login_history_hook; +typedef void (*polar_flush_login_history_hook_type) (void); +extern PGDLLIMPORT polar_flush_login_history_hook_type polar_flush_login_history_hook; +/* POLAR end */ + #endif /* ELOG_H */