From 1e09e1e9661bd67d46f62b050fd1afd622d97223 Mon Sep 17 00:00:00 2001
From: "Andrew Jackson (Ajax)" <snadrus@gmail.com>
Date: Thu, 7 Dec 2023 16:01:28 -0600
Subject: [PATCH] detect unsafe code uses

---
 lib/harmony/harmonydb/harmonydb.go |  3 ++
 lib/harmony/harmonydb/userfuncs.go | 48 ++++++++++++++++--------------
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/lib/harmony/harmonydb/harmonydb.go b/lib/harmony/harmonydb/harmonydb.go
index 0fed176d2..caa5d1daa 100644
--- a/lib/harmony/harmonydb/harmonydb.go
+++ b/lib/harmony/harmonydb/harmonydb.go
@@ -10,6 +10,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	logging "github.com/ipfs/go-log/v2"
@@ -33,6 +34,8 @@ type DB struct {
 	cfg       *pgxpool.Config
 	schema    string
 	hostnames []string
+	BTFPOnce  sync.Once
+	BTFP      uintptr
 }
 
 var logger = logging.Logger("harmonydb")
diff --git a/lib/harmony/harmonydb/userfuncs.go b/lib/harmony/harmonydb/userfuncs.go
index add50c269..8abce73d0 100644
--- a/lib/harmony/harmonydb/userfuncs.go
+++ b/lib/harmony/harmonydb/userfuncs.go
@@ -9,6 +9,7 @@ import (
 	"github.com/jackc/pgerrcode"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/samber/lo"
 )
 
 var inTxErr = errors.New("Cannot use a non-transaction func in a transaction")
@@ -25,7 +26,7 @@ type rawStringOnly string
 // Note, for CREATE & DROP please keep these permanent and express
 // them in the ./sql/ files (next number).
 func (db *DB) Exec(ctx context.Context, sql rawStringOnly, arguments ...any) (count int, err error) {
-	if usedInTransaction() {
+	if db.usedInTransaction() {
 		return 0, inTxErr
 	}
 	res, err := db.pgx.Exec(ctx, string(sql), arguments...)
@@ -61,7 +62,7 @@ type Query struct {
 //	   fmt.Println(id, name)
 //	}
 func (db *DB) Query(ctx context.Context, sql rawStringOnly, arguments ...any) (*Query, error) {
-	if usedInTransaction() {
+	if db.usedInTransaction() {
 		return &Query{}, inTxErr
 	}
 	q, err := db.pgx.Query(ctx, string(sql), arguments...)
@@ -76,7 +77,8 @@ type Row interface {
 }
 
 type rowErr struct{}
-func (rowErr) Scan(..any) error { return inTxErr }
+
+func (rowErr) Scan(_ ...any) error { return inTxErr }
 
 // QueryRow gets 1 row using column order matching.
 // This is a timesaver for the special case of wanting the first row returned only.
@@ -86,10 +88,10 @@ func (rowErr) Scan(..any) error { return inTxErr }
 //	var ID = 123
 //	err := db.QueryRow(ctx, "SELECT name, pet FROM users WHERE ID=?", ID).Scan(&name, &pet)
 func (db *DB) QueryRow(ctx context.Context, sql rawStringOnly, arguments ...any) Row {
-	if usedInTransaction() {
+	if db.usedInTransaction() {
 		return rowErr{}
 	}
-		return db.pgx.QueryRow(ctx, string(sql), arguments...)
+	return db.pgx.QueryRow(ctx, string(sql), arguments...)
 }
 
 /*
@@ -107,7 +109,7 @@ Ex:
 	err := db.Select(ctx, &users, "SELECT name, id, tel_no FROM customers WHERE pet=?", pet)
 */
 func (db *DB) Select(ctx context.Context, sliceOfStructPtr any, sql rawStringOnly, arguments ...any) error {
-	if usedInTransaction() {
+	if db.usedInTransaction() {
 		return inTxErr
 	}
 	return pgxscan.Select(ctx, db.pgx, sliceOfStructPtr, string(sql), arguments...)
@@ -118,29 +120,31 @@ type Tx struct {
 	ctx context.Context
 }
 
-// usedInTransaction is a helper to prevent nesting transactions 
-// & non-transaction calls in transactions. In the case of a stack read error,
-// it will return false, so only use true for a course of action.
-func usedInTransaction() bool {
-	ok := true
-	fn := ""
-	for v:=2; ok; v++ {
-		_,_,fn,ok = runtime.Caller(v)
-		if strings.Contains(fn, "BeginTransaction") {
-			return true
-		}
-	}
-	return false
+// usedInTransaction is a helper to prevent nesting transactions
+// & non-transaction calls in transactions. It only checks 20 frames.
+// Fast: This memory should all be in CPU Caches.
+func (db *DB) usedInTransaction() bool {
+	var framePtrs = (&[20]uintptr{})[:]
+	runtime.Callers(3, framePtrs)
+	return lo.Contains(framePtrs, db.BTFP)
 }
 
 // BeginTransaction is how you can access transactions using this library.
 // The entire transaction happens in the function passed in.
 // The return must be true or a rollback will occur.
 // Be sure to test the error for IsErrSerialization() if you want to retry
-//  when there is a DB serialization error. 
+//
+//	when there is a DB serialization error.
+//
+//go:noinline
 func (db *DB) BeginTransaction(ctx context.Context, f func(*Tx) (commit bool, err error)) (didCommit bool, retErr error) {
-	if usedInTransaction() {
-		return 0, inTxErr
+	db.BTFPOnce.Do(func() {
+		fp := make([]uintptr, 20)
+		runtime.Callers(1, fp)
+		db.BTFP = fp[0]
+	})
+	if db.usedInTransaction() {
+		return false, inTxErr
 	}
 	tx, err := db.pgx.BeginTx(ctx, pgx.TxOptions{})
 	if err != nil {