Implement an auto-escaping Format native for SQL query construction

core/logic/smn_database.cpp

@@ -32,6 +32,7 @@
 #include "common_logic.h"
 #include "Database.h"
 #include "ExtensionSys.h"
+#include "sprintf.h"
 #include "stringutil.h"
 #include "ISourceMod.h"
 #include "AutoHandleRooter.h"
@@ -732,6 +733,24 @@ static cell_t SQL_QuoteString(IPluginContext *pContext, const cell_t *params)
 	return s ? 1 : 0;
 }
 
+static cell_t SQL_FormatQuery(IPluginContext *pContext, const cell_t *params)
+{
+	IDatabase *db = NULL;
+	HandleError err;
+
+	if ((err = g_DBMan.ReadHandle(params[1], DBHandle_Database, (void **)&db))
+		!= HandleError_None)
+	{
+		return pContext->ThrowNativeError("Invalid database Handle %x (error: %d)", params[1], err);
+	}
+
+	g_FormatEscapeDatabase = db;
+	cell_t result = InternalFormat(pContext, params, 1);
+	g_FormatEscapeDatabase = NULL;
+
+	return result;
+}
+
 static cell_t SQL_FastQuery(IPluginContext *pContext, const cell_t *params)
 {
 	IDatabase *db = NULL;
@@ -1814,6 +1833,7 @@ REGISTER_NATIVES(dbNatives)
 	{"Database.Driver.get",				Database_Driver_get},
 	{"Database.SetCharset",				SQL_SetCharset},
 	{"Database.Escape",					SQL_QuoteString},
+	{"Database.Format",					SQL_FormatQuery},
 	{"Database.IsSameConnection",		SQL_IsSameConnection},
 	{"Database.Execute",				SQL_ExecuteTransaction},
 
@@ -1827,6 +1847,7 @@ REGISTER_NATIVES(dbNatives)
 	{"SQL_Connect",				SQL_Connect},
 	{"SQL_ConnectEx",			SQL_ConnectEx},
 	{"SQL_EscapeString",		SQL_QuoteString},
+	{"SQL_FormatQuery",			SQL_FormatQuery},
 	{"SQL_Execute",				SQL_Execute},
 	{"SQL_FastQuery",			SQL_FastQuery},
 	{"SQL_FetchFloat",			SQL_FetchFloat},

core/logic/sprintf.cpp

@@ -30,15 +30,19 @@
 #include "sprintf.h"
 #include <am-float.h>
 #include <am-string.h>
+#include <IDBDriver.h>
 #include <ITranslator.h>
 #include <bridge/include/IScriptManager.h>
 #include <bridge/include/CoreProvider.h>
 
 using namespace SourceMod;
 
+IDatabase *g_FormatEscapeDatabase = NULL;
+
 #define LADJUST			0x00000001		/* left adjustment */
 #define ZEROPAD			0x00000002		/* zero (as opposed to blank) pad */
 #define UPPERDIGITS		0x00000004		/* make alpha digits uppercase */
+#define NOESCAPE		0x00000008		/* do not escape strings (they are only escaped if a database connection is provided) */
 #define to_digit(c)		((c) - '0')
 #define is_digit(c)		((unsigned)to_digit(c) <= 9)
 
@@ -158,6 +162,7 @@ bool AddString(char **buf_p, size_t &maxlen, const char *string, int width, int
 	{
 		string = nlstr;
 		prec = -1;
+		flags |= NOESCAPE;
 	}
 
 	if (prec >= 0)
@@ -181,12 +186,44 @@ bool AddString(char **buf_p, size_t &maxlen, const char *string, int width, int
 		size = maxlen;
 	}
 
-	maxlen -= size;
 	width -= size;
 
-	while (size--)
+	if (g_FormatEscapeDatabase && (flags & NOESCAPE) == 0)
 	{
-		*buf++ = *string++;
+		char *tempBuffer = NULL;
+		if (prec != -1)
+		{
+			// I doubt anyone will ever do this, so just allocate.
+			tempBuffer = new char[maxlen + 1];
+			memcpy(tempBuffer, string, size);
+			tempBuffer[size] = '\0';
+		}
+
+		size_t newSize;
+		bool ret = g_FormatEscapeDatabase->QuoteString(tempBuffer ? tempBuffer : string, buf, maxlen + 1, &newSize);
+
+		if (tempBuffer)
+		{
+			delete[] tempBuffer;
+		}
+
+		if (!ret)
+		{
+			return false;
+		}
+
+		maxlen -= newSize;
+		buf += newSize;
+		size = 0; // Consistency.
+	}
+	else
+	{
+		maxlen -= size;
+
+		while (size--)
+		{
+			*buf++ = *string++;
+		}
 	}
 
 	while ((width-- > 0) && maxlen)
@@ -213,7 +250,7 @@ void AddFloat(char **buf_p, size_t &maxlen, double fval, int width, int prec, in
 
 	if (ke::IsNaN(fval))
 	{
-		AddString(buf_p, maxlen, "NaN", width, prec, flags);
+		AddString(buf_p, maxlen, "NaN", width, prec, flags | NOESCAPE);
 		return;
 	}
 
@@ -1031,6 +1068,11 @@ reswitch:
 				flags |= LADJUST;
 				goto rflag;
 			}
+		case '!':
+			{
+				flags |= NOESCAPE;
+				goto rflag;
+			}
 		case '.':
 			{
 				n = 0;

core/logic/sprintf.h

@@ -30,6 +30,7 @@
 #include <sp_vm_api.h>
 
 namespace SourceMod {
+class IDatabase;
 class IPhraseCollection;
 }
 
@@ -57,4 +58,6 @@ bool gnprintf(char *buffer,
               size_t *pOutLength,
               const char **pFailPhrase);
 
+extern SourceMod::IDatabase *g_FormatEscapeDatabase;
+
 #endif // _include_sourcemod_core_logic_sprintf_h_