Add DisableFunction API

This allows built-in functions to be disabled.

Updates tailscale/corp#31396

Signed-off-by: Percy Wegmann <[email protected]>

Author: oxtoacart

cgosqlite/cgosqlite.go

@@ -203,6 +203,10 @@ func (db *DB) Prepare(query string, prepFlags sqliteh.PrepareFlags) (stmt sqlite
 	return &Stmt{db: db, stmt: cStmtFromPtr(cstmt)}, remainingQuery, nil
 }
 
+func (db *DB) DisableFunction(name string, numArgs int) error {
+	return errCode(C.ts_sqlite3_disable_function(db.db, C.CString(name), C.int(numArgs)))
+}
+
 func (stmt *Stmt) DBHandle() sqliteh.DB {
 	cdb := C.sqlite3_db_handle(stmt.stmt.ptr())
 	if cdb != nil {

cgosqlite/cgosqlite.h

@@ -146,3 +146,7 @@ static double ts_sqlite3_column_double(handle_sqlite3_stmt stmt, int iCol) {
 static sqlite3_int64 ts_sqlite3_column_int64(handle_sqlite3_stmt stmt, int iCol) {
 	return sqlite3_column_int64((sqlite3_stmt*)(stmt), iCol);
 }
+
+static int ts_sqlite3_disable_function(sqlite3 *db, const char *zFunctionName, int nArg) {
+	return sqlite3_create_function(db, zFunctionName, nArg, SQLITE_ANY, NULL, NULL, NULL, NULL);
+}

sqlite.go

@@ -1036,6 +1036,18 @@ func Checkpoint(sqlconn SQLConn, dbName string, mode sqliteh.Checkpoint) (numFra
 	return numFrames, numFramesCheckpointed, err
 }
 
+// DisableFunction disables the named function on the given connection.
+// numArgs must match the number of args of the function to be disabled.
+func DisableFunction(sqlconn SQLConn, name string, numArgs int) error {
+	return sqlconn.Raw(func(driverConn any) error {
+		c, ok := driverConn.(*conn)
+		if !ok {
+			return fmt.Errorf("sqlite.DisableFunction: sql.Conn is not the sqlite driver: %T", driverConn)
+		}
+		return c.db.DisableFunction(name, numArgs)
+	})
+}
+
 // WithPersist makes a ctx instruct the sqlite driver to persist a prepared query.
 //
 // This should be used with recurring queries to avoid constant parsing and

sqlite_test.go

@@ -1330,3 +1330,28 @@ func TestRegression(t *testing.T) {
 		t.Log("OK") // Reaching here at all means we didn't panic.
 	})
 }
+
+func TestDisableFunction(t *testing.T) {
+	db := openTestDB(t)
+
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	exec(t, conn, "CREATE TABLE t (c)")
+	ctx := context.Background()
+
+	if _, err := conn.ExecContext(ctx, "SELECT LOWER('Hi') FROM t"); err != nil {
+		t.Fatal("Attempting to use the LOWER function before disabling should have been allowed")
+	}
+
+	if err := DisableFunction(conn, "lower", 1); err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := conn.ExecContext(ctx, "SELECT LOWER('Hi') FROM t"); err == nil {
+		t.Fatal("Attempting to use the LOWER function after disabling should have failed")
+	}
+}

sqliteh/sqliteh.go

@@ -60,6 +60,10 @@ type DB interface {
 	//
 	// If hook is nil, the hook is removed.
 	SetWALHook(hook func(dbName string, pages int))
+	// DisableFunction allows disabling an existing function using
+	// sqlite3_create_function.
+	//
+	DisableFunction(name string, numArgs int) error
 }
 
 // Stmt is an sqlite3_stmt* database connection object.