Consolidate and improve checking of key-column-attnum arguments for

dblink_build_sql_insert() and related functions.  In particular, be sure to
reject references to dropped and out-of-range column numbers.  The numbers
are still interpreted as physical column numbers, though, for backward
compatibility.

This patch replaces Joe's patch of 2010-02-03, which handled only some aspects
of the problem.

contrib/dblink/dblink.c

@@ -76,18 +76,20 @@ static HTAB *createConnHash(void);
 static void createNewConnection(const char *name, remoteConn * rconn);
 static void deleteConnection(const char *name);
 static char **get_pkey_attnames(Relation rel, int16 *numatts);
-static char *get_sql_insert(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals, char **tgt_pkattvals);
+static char *get_sql_insert(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals, char **tgt_pkattvals);
-static char *get_sql_delete(Relation rel, int2vector *pkattnums, int16 pknumatts, char **tgt_pkattvals);
+static char *get_sql_delete(Relation rel, int *pkattnums, int pknumatts, char **tgt_pkattvals);
-static char *get_sql_update(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals, char **tgt_pkattvals);
+static char *get_sql_update(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals, char **tgt_pkattvals);
 static char *quote_literal_cstr(char *rawstr);
 static char *quote_ident_cstr(char *rawstr);
-static int16 get_attnum_pk_pos(int2vector *pkattnums, int16 pknumatts, int16 key);
+static int	get_attnum_pk_pos(int *pkattnums, int pknumatts, int key);
-static HeapTuple get_tuple_of_interest(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals);
+static HeapTuple get_tuple_of_interest(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals);
 static Relation get_rel_from_relname(text *relname_text, LOCKMODE lockmode, AclMode aclmode);
 static char *generate_relation_name(Relation rel);
 static char *connstr_strip_password(const char *connstr);
 static void dblink_security_check(PGconn *conn, remoteConn *rconn, const char *connstr);
-static int get_nondropped_natts(Relation rel);
+static void validate_pkattnums(Relation rel,
+				   int2vector *pkattnums_arg, int32 pknumatts_arg,
+				   int **pkattnums, int *pknumatts);
 
 /* Global */
 static remoteConn *pconn = NULL;
@@ -1150,9 +1152,10 @@ dblink_build_sql_insert(PG_FUNCTION_ARGS)
 {
 	Relation	rel;
 	text	   *relname_text;
-	int2vector *pkattnums;
+	int2vector *pkattnums_arg;
-	int			pknumatts_tmp;
+	int32		pknumatts_arg;
-	int16		pknumatts = 0;
+	int		   *pkattnums;
+	int			pknumatts;
 	char	  **src_pkattvals;
 	char	  **tgt_pkattvals;
 	ArrayType  *src_pkattvals_arry;
@@ -1169,7 +1172,6 @@ dblink_build_sql_insert(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
-	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1178,36 +1180,17 @@ dblink_build_sql_insert(PG_FUNCTION_ARGS)
 	 */
 	rel = get_rel_from_relname(relname_text, AccessShareLock, ACL_SELECT);
 
-	pkattnums = (int2vector *) PG_GETARG_POINTER(1);
-	pknumatts_tmp = PG_GETARG_INT32(2);
-	if (pknumatts_tmp <= SHRT_MAX)
-		pknumatts = pknumatts_tmp;
-	else
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("input for number of primary key " \
-						"attributes too large")));
-
 	/*
-	 * There should be at least one key attribute
+	 * Process pkattnums argument.
 	 */
-	if (pknumatts == 0)
+	pkattnums_arg = (int2vector *) PG_GETARG_POINTER(1);
-		ereport(ERROR,
+	pknumatts_arg = PG_GETARG_INT32(2);
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+	validate_pkattnums(rel, pkattnums_arg, pknumatts_arg,
-				 errmsg("number of key attributes must be > 0")));
+					   &pkattnums, &pknumatts);
 
 	src_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(4);
 
-	/*
-	 * ensure we don't ask for more pk attributes than we have
-	 * non-dropped columns
-	 */
-	nondropped_natts = get_nondropped_natts(rel);
-	if (pknumatts > nondropped_natts)
-		ereport(ERROR, (errcode(ERRCODE_SYNTAX_ERROR),
-				errmsg("number of primary key fields exceeds number of specified relation attributes")));
-
 	/*
 	 * Source array is made up of key values that will be used to locate the
 	 * tuple of interest from the local system.
@@ -1312,9 +1295,10 @@ dblink_build_sql_delete(PG_FUNCTION_ARGS)
 {
 	Relation	rel;
 	text	   *relname_text;
-	int2vector *pkattnums;
+	int2vector *pkattnums_arg;
-	int			pknumatts_tmp;
+	int32		pknumatts_arg;
-	int16		pknumatts = 0;
+	int		   *pkattnums;
+	int			pknumatts;
 	char	  **tgt_pkattvals;
 	ArrayType  *tgt_pkattvals_arry;
 	int			tgt_ndim;
@@ -1326,7 +1310,6 @@ dblink_build_sql_delete(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
-	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1335,35 +1318,16 @@ dblink_build_sql_delete(PG_FUNCTION_ARGS)
 	 */
 	rel = get_rel_from_relname(relname_text, AccessShareLock, ACL_SELECT);
 
-	pkattnums = (int2vector *) PG_GETARG_POINTER(1);
-	pknumatts_tmp = PG_GETARG_INT32(2);
-	if (pknumatts_tmp <= SHRT_MAX)
-		pknumatts = pknumatts_tmp;
-	else
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("input for number of primary key " \
-						"attributes too large")));
-
 	/*
-	 * There should be at least one key attribute
+	 * Process pkattnums argument.
 	 */
-	if (pknumatts == 0)
+	pkattnums_arg = (int2vector *) PG_GETARG_POINTER(1);
-		ereport(ERROR,
+	pknumatts_arg = PG_GETARG_INT32(2);
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+	validate_pkattnums(rel, pkattnums_arg, pknumatts_arg,
-				 errmsg("number of key attributes must be > 0")));
+					   &pkattnums, &pknumatts);
 
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 
-	/*
-	 * ensure we don't ask for more pk attributes than we have
-	 * non-dropped columns
-	 */
-	nondropped_natts = get_nondropped_natts(rel);
-	if (pknumatts > nondropped_natts)
-		ereport(ERROR, (errcode(ERRCODE_SYNTAX_ERROR),
-				errmsg("number of primary key fields exceeds number of specified relation attributes")));
-
 	/*
 	 * Target array is made up of key values that will be used to build the
 	 * SQL string for use on the remote system.
@@ -1439,9 +1403,10 @@ dblink_build_sql_update(PG_FUNCTION_ARGS)
 {
 	Relation	rel;
 	text	   *relname_text;
-	int2vector *pkattnums;
+	int2vector *pkattnums_arg;
-	int			pknumatts_tmp;
+	int32		pknumatts_arg;
-	int16		pknumatts = 0;
+	int		   *pkattnums;
+	int			pknumatts;
 	char	  **src_pkattvals;
 	char	  **tgt_pkattvals;
 	ArrayType  *src_pkattvals_arry;
@@ -1458,7 +1423,6 @@ dblink_build_sql_update(PG_FUNCTION_ARGS)
 	int16		typlen;
 	bool		typbyval;
 	char		typalign;
-	int			nondropped_natts;
 
 	relname_text = PG_GETARG_TEXT_P(0);
 
@@ -1467,36 +1431,17 @@ dblink_build_sql_update(PG_FUNCTION_ARGS)
 	 */
 	rel = get_rel_from_relname(relname_text, AccessShareLock, ACL_SELECT);
 
-	pkattnums = (int2vector *) PG_GETARG_POINTER(1);
-	pknumatts_tmp = PG_GETARG_INT32(2);
-	if (pknumatts_tmp <= SHRT_MAX)
-		pknumatts = pknumatts_tmp;
-	else
-		ereport(ERROR,
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-				 errmsg("input for number of primary key " \
-						"attributes too large")));
-
 	/*
-	 * There should be one source array key values for each key attnum
+	 * Process pkattnums argument.
 	 */
-	if (pknumatts == 0)
+	pkattnums_arg = (int2vector *) PG_GETARG_POINTER(1);
-		ereport(ERROR,
+	pknumatts_arg = PG_GETARG_INT32(2);
-				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+	validate_pkattnums(rel, pkattnums_arg, pknumatts_arg,
-				 errmsg("number of key attributes must be > 0")));
+					   &pkattnums, &pknumatts);
 
 	src_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(3);
 	tgt_pkattvals_arry = PG_GETARG_ARRAYTYPE_P(4);
 
-	/*
-	 * ensure we don't ask for more pk attributes than we have
-	 * non-dropped columns
-	 */
-	nondropped_natts = get_nondropped_natts(rel);
-	if (pknumatts > nondropped_natts)
-		ereport(ERROR, (errcode(ERRCODE_SYNTAX_ERROR),
-				errmsg("number of primary key fields exceeds number of specified relation attributes")));
-
 	/*
 	 * Source array is made up of key values that will be used to locate the
 	 * tuple of interest from the local system.
@@ -1653,7 +1598,7 @@ get_pkey_attnames(Relation rel, int16 *numatts)
 }
 
 static char *
-get_sql_insert(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals, char **tgt_pkattvals)
+get_sql_insert(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals, char **tgt_pkattvals)
 {
 	char	   *relname;
 	HeapTuple	tuple;
@@ -1662,7 +1607,7 @@ get_sql_insert(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 	StringInfo	str = makeStringInfo();
 	char	   *sql;
 	char	   *val;
-	int16		key;
+	int			key;
 	int			i;
 	bool		needComma;
 
@@ -1709,7 +1654,7 @@ get_sql_insert(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 			appendStringInfo(str, ",");
 
 		if (tgt_pkattvals != NULL)
-			key = get_attnum_pk_pos(pkattnums, pknumatts, i + 1);
+			key = get_attnum_pk_pos(pkattnums, pknumatts, i);
 		else
 			key = -1;
 
@@ -1737,11 +1682,10 @@ get_sql_insert(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 }
 
 static char *
-get_sql_delete(Relation rel, int2vector *pkattnums, int16 pknumatts, char **tgt_pkattvals)
+get_sql_delete(Relation rel, int *pkattnums, int pknumatts, char **tgt_pkattvals)
 {
 	char	   *relname;
 	TupleDesc	tupdesc;
-	int			natts;
 	StringInfo	str = makeStringInfo();
 	char	   *sql;
 	char	   *val = NULL;
@@ -1751,18 +1695,17 @@ get_sql_delete(Relation rel, int2vector *pkattnums, int16 pknumatts, char **tgt_
 	relname = generate_relation_name(rel);
 
 	tupdesc = rel->rd_att;
-	natts = tupdesc->natts;
 
 	appendStringInfo(str, "DELETE FROM %s WHERE ", relname);
 	for (i = 0; i < pknumatts; i++)
 	{
-		int16		pkattnum = pkattnums->values[i];
+		int			pkattnum = pkattnums[i];
 
 		if (i > 0)
 			appendStringInfo(str, " AND ");
 
 		appendStringInfo(str, "%s",
-		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum - 1]->attname)));
+		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum]->attname)));
 
 		if (tgt_pkattvals != NULL)
 			val = pstrdup(tgt_pkattvals[i]);
@@ -1787,7 +1730,7 @@ get_sql_delete(Relation rel, int2vector *pkattnums, int16 pknumatts, char **tgt_
 }
 
 static char *
-get_sql_update(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals, char **tgt_pkattvals)
+get_sql_update(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals, char **tgt_pkattvals)
 {
 	char	   *relname;
 	HeapTuple	tuple;
@@ -1796,7 +1739,7 @@ get_sql_update(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 	StringInfo	str = makeStringInfo();
 	char	   *sql;
 	char	   *val;
-	int16		key;
+	int			key;
 	int			i;
 	bool		needComma;
 
@@ -1827,7 +1770,7 @@ get_sql_update(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 					  quote_ident_cstr(NameStr(tupdesc->attrs[i]->attname)));
 
 		if (tgt_pkattvals != NULL)
-			key = get_attnum_pk_pos(pkattnums, pknumatts, i + 1);
+			key = get_attnum_pk_pos(pkattnums, pknumatts, i);
 		else
 			key = -1;
 
@@ -1850,18 +1793,18 @@ get_sql_update(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_
 
 	for (i = 0; i < pknumatts; i++)
 	{
-		int16		pkattnum = pkattnums->values[i];
+		int			pkattnum = pkattnums[i];
 
 		if (i > 0)
 			appendStringInfo(str, " AND ");
 
 		appendStringInfo(str, "%s",
-		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum - 1]->attname)));
+		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum]->attname)));
 
 		if (tgt_pkattvals != NULL)
 			val = pstrdup(tgt_pkattvals[i]);
 		else
-			val = SPI_getvalue(tuple, tupdesc, pkattnum);
+			val = SPI_getvalue(tuple, tupdesc, pkattnum + 1);
 
 		if (val != NULL)
 		{
@@ -1915,8 +1858,8 @@ quote_ident_cstr(char *rawstr)
 	return result;
 }
 
-static int16
+static int
-get_attnum_pk_pos(int2vector *pkattnums, int16 pknumatts, int16 key)
+get_attnum_pk_pos(int *pkattnums, int pknumatts, int key)
 {
 	int			i;
 
@@ -1924,14 +1867,14 @@ get_attnum_pk_pos(int2vector *pkattnums, int16 pknumatts, int16 key)
 	 * Not likely a long list anyway, so just scan for the value
 	 */
 	for (i = 0; i < pknumatts; i++)
-		if (key == pkattnums->values[i])
+		if (key == pkattnums[i])
 			return i;
 
 	return -1;
 }
 
 static HeapTuple
-get_tuple_of_interest(Relation rel, int2vector *pkattnums, int16 pknumatts, char **src_pkattvals)
+get_tuple_of_interest(Relation rel, int *pkattnums, int pknumatts, char **src_pkattvals)
 {
 	char	   *relname;
 	TupleDesc	tupdesc;
@@ -1962,13 +1905,13 @@ get_tuple_of_interest(Relation rel, int2vector *pkattnums, int16 pknumatts, char
 
 	for (i = 0; i < pknumatts; i++)
 	{
-		int16		pkattnum = pkattnums->values[i];
+		int			pkattnum = pkattnums[i];
 
 		if (i > 0)
 			appendStringInfo(str, " AND ");
 
 		appendStringInfo(str, "%s",
-		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum - 1]->attname)));
+		   quote_ident_cstr(NameStr(tupdesc->attrs[pkattnum]->attname)));
 
 		val = pstrdup(src_pkattvals[i]);
 		if (val != NULL)
@@ -2308,24 +2251,52 @@ dblink_security_check(PGconn *conn, remoteConn *rconn, const char *connstr)
 	}
 }
 
-static int
+/*
-get_nondropped_natts(Relation rel)
+ * Validate the PK-attnums argument for dblink_build_sql_insert() and related
+ * functions, and translate to the internal representation.
+ *
+ * The user supplies an int2vector of 1-based physical attnums, plus a count
+ * argument (the need for the separate count argument is historical, but we
+ * still check it).  We check that each attnum corresponds to a valid,
+ * non-dropped attribute of the rel.  We do *not* prevent attnums from being
+ * listed twice, though the actual use-case for such things is dubious.
+ *
+ * The internal representation is a palloc'd int array of 0-based physical
+ * attnums.
+ */
+static void
+validate_pkattnums(Relation rel,
+				   int2vector *pkattnums_arg, int32 pknumatts_arg,
+				   int **pkattnums, int *pknumatts)
 {
-	int			nondropped_natts = 0;
+	TupleDesc	tupdesc = rel->rd_att;
-	TupleDesc	tupdesc;
+	int			natts = tupdesc->natts;
-	int			natts;
 	int			i;
 
-	tupdesc = rel->rd_att;
+	/* Don't take more array elements than there are */
-	natts = tupdesc->natts;
+	pknumatts_arg = Min(pknumatts_arg, pkattnums_arg->dim1);
 
-	for (i = 0; i < natts; i++)
+	/* Must have at least one pk attnum selected */
+	if (pknumatts_arg <= 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("number of key attributes must be > 0")));
+
+	/* Allocate output array */
+	*pkattnums = (int *) palloc(pknumatts_arg * sizeof(int));
+	*pknumatts = pknumatts_arg;
+
+	/* Validate attnums and convert to internal form */
+	for (i = 0; i < pknumatts_arg; i++)
 	{
-		if (tupdesc->attrs[i]->attisdropped)
+		int		pkattnum = pkattnums_arg->values[i];
-			continue;
-		nondropped_natts++;
-	}
 
-	return nondropped_natts;
+		if (pkattnum <= 0 || pkattnum > natts ||
+			tupdesc->attrs[pkattnum - 1]->attisdropped)
+			ereport(ERROR,
+					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+					 errmsg("invalid attribute number %d", pkattnum)));
+		(*pkattnums)[i] = pkattnum - 1;
+	}
 }
 

contrib/dblink/expected/dblink.out

@@ -46,7 +46,7 @@ SELECT dblink_build_sql_insert('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
 
 -- too many pk fields, should fail
 SELECT dblink_build_sql_insert('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
-ERROR:  number of primary key fields exceeds number of specified relation attributes
+ERROR:  invalid attribute number 4
 -- build an update statement based on a local tuple,
 -- replacing the primary key values with new ones
 SELECT dblink_build_sql_update('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
@@ -57,7 +57,7 @@ SELECT dblink_build_sql_update('foo','1 2',2,'{"0", "a"}','{"99", "xyz"}');
 
 -- too many pk fields, should fail
 SELECT dblink_build_sql_update('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}','{"99", "xyz", "{za0,zb0,zc0}"}');
-ERROR:  number of primary key fields exceeds number of specified relation attributes
+ERROR:  invalid attribute number 4
 -- build a delete statement based on a local tuple,
 SELECT dblink_build_sql_delete('foo','1 2',2,'{"0", "a"}');
            dblink_build_sql_delete           
@@ -67,7 +67,7 @@ SELECT dblink_build_sql_delete('foo','1 2',2,'{"0", "a"}');
 
 -- too many pk fields, should fail
 SELECT dblink_build_sql_delete('foo','1 2 3 4',4,'{"0", "a", "{a0,b0,c0}"}');
-ERROR:  number of primary key fields exceeds number of specified relation attributes
+ERROR:  invalid attribute number 4
 -- retest using a quoted and schema qualified table
 CREATE SCHEMA "MySchema";
 CREATE TABLE "MySchema"."Foo"(f1 int, f2 text, f3 text[], primary key (f1,f2));