Fix possible segfault in pg_stat_wait_get_history

ildus · Alexander Korotkov · commit 3e2c22c941c2 · 2016-02-18T14:52:39.000+03:00
diff --git a/contrib/pg_stat_wait/Makefile b/contrib/pg_stat_wait/Makefile
@@ -19,7 +19,7 @@ include $(top_builddir)/src/Makefile.global
 include $(top_srcdir)/contrib/contrib-global.mk
 endif
 
-check: regresscheck
+check: isolationcheck
 
 submake-regress:
 	$(MAKE) -C $(top_builddir)/src/test/regress all
diff --git a/contrib/pg_stat_wait/collector.c b/contrib/pg_stat_wait/collector.c
@@ -20,7 +20,7 @@
 
 #include "pg_stat_wait.h"
 
-CollectorShmqHeader *hdr;
+CollectorShmqHeader *hdr = NULL;
 
 static void *pgsw;
 shm_toc *toc;
@@ -52,13 +52,18 @@ CollectorShmemSize(void)
 	return size;
 }
 
-void
-AllocateCollectorMem(void)
+CollectorShmqHeader *
+GetCollectorMem(bool init)
 {
 	bool found;
 	Size segsize = CollectorShmemSize();
 
 	pgsw = ShmemInitStruct("pg_stat_wait", segsize, &found);
+	if (!init && !found)
+	{
+		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR),
+						errmsg("A collector memory wasn't initialized yet")));
+	}
 
 	if (!found)
 	{
@@ -76,6 +81,7 @@ AllocateCollectorMem(void)
 		toc = shm_toc_attach(PG_STAT_WAIT_MAGIC, pgsw);
 		hdr = shm_toc_lookup(toc, 0);
 	}
+	return hdr;
 }
 
 void
diff --git a/contrib/pg_stat_wait/pg_stat_wait.c b/contrib/pg_stat_wait/pg_stat_wait.c
@@ -18,7 +18,6 @@ PG_MODULE_MAGIC;
 void		_PG_init(void);
 void		_PG_fini(void);
 
-extern CollectorShmqHeader *hdr;
 extern shm_toc			   *toc;
 
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -35,7 +34,7 @@ pgsw_shmem_startup(void)
 		if (prev_shmem_startup_hook)
 			prev_shmem_startup_hook();
 
-		AllocateCollectorMem();
+		GetCollectorMem(true);
 	}
 }
 
@@ -582,13 +581,16 @@ PG_FUNCTION_INFO_V1(pg_stat_wait_get_history);
 Datum
 pg_stat_wait_get_history(PG_FUNCTION_ARGS)
 {
-	History *observations;
-	FuncCallContext *funcctx;
+	History				*observations;
+	FuncCallContext		*funcctx;
+	CollectorShmqHeader	*hdr;
 
 	if (!WaitsHistoryOn)
 		ereport(ERROR, (errcode(ERRCODE_CONFIG_FILE_ERROR),
 						errmsg("Waits history turned off")));
 
+	hdr = GetCollectorMem(false);
+
 	if (SRF_IS_FIRSTCALL())
 	{
 		shm_mq	   *mq;
diff --git a/contrib/pg_stat_wait/pg_stat_wait.h b/contrib/pg_stat_wait/pg_stat_wait.h
@@ -75,7 +75,7 @@ extern PGDLLIMPORT bool        historySkipLatch;
 
 
 Size CollectorShmemSize(void);
-void AllocateCollectorMem(void);
+CollectorShmqHeader *GetCollectorMem(bool init);
 void RegisterWaitsCollector(void);
 void AllocHistory(History *, int);
 int GetCurrentWaitsState(PGPROC *, HistoryItem *, int);

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@`
`20`	`20`
`21`	`21`	`#include "pg_stat_wait.h"`
`22`	`22`
`23`		`-CollectorShmqHeader *hdr;`
	`23`	`+CollectorShmqHeader *hdr = NULL;`
`24`	`24`
`25`	`25`	`static void *pgsw;`
`26`	`26`	`shm_toc *toc;`
`@@ -52,13 +52,18 @@ CollectorShmemSize(void)`
`52`	`52`	`return size;`
`53`	`53`	`}`
`54`	`54`
`55`		`-void`
`56`		`-AllocateCollectorMem(void)`
	`55`	`+CollectorShmqHeader *`
	`56`	`+GetCollectorMem(bool init)`
`57`	`57`	`{`
`58`	`58`	`bool found;`
`59`	`59`	`Size segsize = CollectorShmemSize();`
`60`	`60`
`61`	`61`	`pgsw = ShmemInitStruct("pg_stat_wait", segsize, &found);`
	`62`	`+ if (!init && !found)`
	`63`	`+ {`
	`64`	`+ ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR),`
	`65`	`+ errmsg("A collector memory wasn't initialized yet")));`
	`66`	`+ }`
`62`	`67`
`63`	`68`	`if (!found)`
`64`	`69`	`{`
`@@ -76,6 +81,7 @@ AllocateCollectorMem(void)`
`76`	`81`	`toc = shm_toc_attach(PG_STAT_WAIT_MAGIC, pgsw);`
`77`	`82`	`hdr = shm_toc_lookup(toc, 0);`
`78`	`83`	`}`
	`84`	`+ return hdr;`
`79`	`85`	`}`
`80`	`86`
`81`	`87`	`void`