Fix db engine (#32351)

Fix #32349

models/db/context.go

@@ -6,6 +6,12 @@ package db
 import (
 	"context"
 	"database/sql"
+	"errors"
+	"runtime"
+	"slices"
+	"sync"
+
+	"code.gitea.io/gitea/modules/setting"
 
 	"xorm.io/builder"
 	"xorm.io/xorm"
@@ -15,45 +21,23 @@ import (
 // will be overwritten by Init with HammerContext
 var DefaultContext context.Context
 
-// contextKey is a value for use with context.WithValue.
-type contextKey struct {
-	name string
-}
+type engineContextKeyType struct{}
 
-// enginedContextKey is a context key. It is used with context.Value() to get the current Engined for the context
-var (
-	enginedContextKey         = &contextKey{"engined"}
-	_                 Engined = &Context{}
-)
+var engineContextKey = engineContextKeyType{}
 
 // Context represents a db context
 type Context struct {
 	context.Context
-	e           Engine
-	transaction bool
+	engine Engine
 }
 
-func newContext(ctx context.Context, e Engine, transaction bool) *Context {
-	return &Context{
-		Context:     ctx,
-		e:           e,
-		transaction: transaction,
-	}
-}
-
-// InTransaction if context is in a transaction
-func (ctx *Context) InTransaction() bool {
-	return ctx.transaction
-}
-
-// Engine returns db engine
-func (ctx *Context) Engine() Engine {
-	return ctx.e
+func newContext(ctx context.Context, e Engine) *Context {
+	return &Context{Context: ctx, engine: e}
 }
 
 // Value shadows Value for context.Context but allows us to get ourselves and an Engined object
 func (ctx *Context) Value(key any) any {
-	if key == enginedContextKey {
+	if key == engineContextKey {
 		return ctx
 	}
 	return ctx.Context.Value(key)
@@ -61,30 +45,66 @@ func (ctx *Context) Value(key any) any {
 
 // WithContext returns this engine tied to this context
 func (ctx *Context) WithContext(other context.Context) *Context {
-	return newContext(ctx, ctx.e.Context(other), ctx.transaction)
+	return newContext(ctx, ctx.engine.Context(other))
 }
 
-// Engined structs provide an Engine
-type Engined interface {
-	Engine() Engine
+var (
+	contextSafetyOnce          sync.Once
+	contextSafetyDeniedFuncPCs []uintptr
+)
+
+func contextSafetyCheck(e Engine) {
+	if setting.IsProd && !setting.IsInTesting {
+		return
+	}
+	if e == nil {
+		return
+	}
+	// Only do this check for non-end-users. If the problem could be fixed in the future, this code could be removed.
+	contextSafetyOnce.Do(func() {
+		// try to figure out the bad functions to deny
+		type m struct{}
+		_ = e.SQL("SELECT 1").Iterate(&m{}, func(int, any) error {
+			callers := make([]uintptr, 32)
+			callerNum := runtime.Callers(1, callers)
+			for i := 0; i < callerNum; i++ {
+				if funcName := runtime.FuncForPC(callers[i]).Name(); funcName == "xorm.io/xorm.(*Session).Iterate" {
+					contextSafetyDeniedFuncPCs = append(contextSafetyDeniedFuncPCs, callers[i])
+				}
+			}
+			return nil
+		})
+		if len(contextSafetyDeniedFuncPCs) != 1 {
+			panic(errors.New("unable to determine the functions to deny"))
+		}
+	})
+
+	// it should be very fast: xxxx ns/op
+	callers := make([]uintptr, 32)
+	callerNum := runtime.Callers(3, callers) // skip 3: runtime.Callers, contextSafetyCheck, GetEngine
+	for i := 0; i < callerNum; i++ {
+		if slices.Contains(contextSafetyDeniedFuncPCs, callers[i]) {
+			panic(errors.New("using database context in an iterator would cause corrupted results"))
+		}
+	}
 }
 
-// GetEngine will get a db Engine from this context or return an Engine restricted to this context
+// GetEngine gets an existing db Engine/Statement or creates a new Session
 func GetEngine(ctx context.Context) Engine {
-	if e := getEngine(ctx); e != nil {
+	if e := getExistingEngine(ctx); e != nil {
 		return e
 	}
 	return x.Context(ctx)
 }
 
-// getEngine will get a db Engine from this context or return nil
-func getEngine(ctx context.Context) Engine {
-	if engined, ok := ctx.(Engined); ok {
-		return engined.Engine()
+// getExistingEngine gets an existing db Engine/Statement from this context or returns nil
+func getExistingEngine(ctx context.Context) (e Engine) {
+	defer func() { contextSafetyCheck(e) }()
+	if engined, ok := ctx.(*Context); ok {
+		return engined.engine
 	}
-	enginedInterface := ctx.Value(enginedContextKey)
-	if enginedInterface != nil {
-		return enginedInterface.(Engined).Engine()
+	if engined, ok := ctx.Value(engineContextKey).(*Context); ok {
+		return engined.engine
 	}
 	return nil
 }
@@ -132,23 +152,23 @@ func (c *halfCommitter) Close() error {
 //	  d. It doesn't mean rollback is forbidden, but always do it only when there is an error, and you do want to rollback.
 func TxContext(parentCtx context.Context) (*Context, Committer, error) {
 	if sess, ok := inTransaction(parentCtx); ok {
-		return newContext(parentCtx, sess, true), &halfCommitter{committer: sess}, nil
+		return newContext(parentCtx, sess), &halfCommitter{committer: sess}, nil
 	}
 
 	sess := x.NewSession()
 	if err := sess.Begin(); err != nil {
-		sess.Close()
+		_ = sess.Close()
 		return nil, nil, err
 	}
 
-	return newContext(DefaultContext, sess, true), sess, nil
+	return newContext(DefaultContext, sess), sess, nil
 }
 
 // WithTx represents executing database operations on a transaction, if the transaction exist,
 // this function will reuse it otherwise will create a new one and close it when finished.
 func WithTx(parentCtx context.Context, f func(ctx context.Context) error) error {
 	if sess, ok := inTransaction(parentCtx); ok {
-		err := f(newContext(parentCtx, sess, true))
+		err := f(newContext(parentCtx, sess))
 		if err != nil {
 			// rollback immediately, in case the caller ignores returned error and tries to commit the transaction.
 			_ = sess.Close()
@@ -165,7 +185,7 @@ func txWithNoCheck(parentCtx context.Context, f func(ctx context.Context) error)
 		return err
 	}
 
-	if err := f(newContext(parentCtx, sess, true)); err != nil {
+	if err := f(newContext(parentCtx, sess)); err != nil {
 		return err
 	}
 
@@ -312,7 +332,7 @@ func InTransaction(ctx context.Context) bool {
 }
 
 func inTransaction(ctx context.Context) (*xorm.Session, bool) {
-	e := getEngine(ctx)
+	e := getExistingEngine(ctx)
 	if e == nil {
 		return nil, false
 	}

models/db/context_test.go

@@ -84,3 +84,47 @@ func TestTxContext(t *testing.T) {
 		}))
 	}
 }
+
+func TestContextSafety(t *testing.T) {
+	type TestModel1 struct {
+		ID int64
+	}
+	type TestModel2 struct {
+		ID int64
+	}
+	assert.NoError(t, unittest.GetXORMEngine().Sync(&TestModel1{}, &TestModel2{}))
+	assert.NoError(t, db.TruncateBeans(db.DefaultContext, &TestModel1{}, &TestModel2{}))
+	testCount := 10
+	for i := 1; i <= testCount; i++ {
+		assert.NoError(t, db.Insert(db.DefaultContext, &TestModel1{ID: int64(i)}))
+		assert.NoError(t, db.Insert(db.DefaultContext, &TestModel2{ID: int64(-i)}))
+	}
+
+	actualCount := 0
+	// here: db.GetEngine(db.DefaultContext) is a new *Session created from *Engine
+	_ = db.WithTx(db.DefaultContext, func(ctx context.Context) error {
+		_ = db.GetEngine(ctx).Iterate(&TestModel1{}, func(i int, bean any) error {
+			// here: db.GetEngine(ctx) is always the unclosed "Iterate" *Session with autoResetStatement=false,
+			// and the internal states (including "cond" and others) are always there and not be reset in this callback.
+			m1 := bean.(*TestModel1)
+			assert.EqualValues(t, i+1, m1.ID)
+
+			// here: XORM bug, it fails because the SQL becomes "WHERE id=-1", "WHERE id=-1 AND id=-2", "WHERE id=-1 AND id=-2 AND id=-3" ...
+			// and it conflicts with the "Iterate"'s internal states.
+			// has, err := db.GetEngine(ctx).Get(&TestModel2{ID: -m1.ID})
+
+			actualCount++
+			return nil
+		})
+		return nil
+	})
+	assert.EqualValues(t, testCount, actualCount)
+
+	// deny the bad usages
+	assert.PanicsWithError(t, "using database context in an iterator would cause corrupted results", func() {
+		_ = unittest.GetXORMEngine().Iterate(&TestModel1{}, func(i int, bean any) error {
+			_ = db.GetEngine(db.DefaultContext)
+			return nil
+		})
+	})
+}

models/db/engine.go

@@ -161,10 +161,7 @@ func InitEngine(ctx context.Context) error {
 // SetDefaultEngine sets the default engine for db
 func SetDefaultEngine(ctx context.Context, eng *xorm.Engine) {
 	x = eng
-	DefaultContext = &Context{
-		Context: ctx,
-		e:       x,
-	}
+	DefaultContext = &Context{Context: ctx, engine: x}
 }
 
 // UnsetDefaultEngine closes and unsets the default engine

models/db/install/db.go

@@ -11,7 +11,7 @@ import (
 )
 
 func getXORMEngine() *xorm.Engine {
-	return db.DefaultContext.(*db.Context).Engine().(*xorm.Engine)
+	return db.GetEngine(db.DefaultContext).(*xorm.Engine)
 }
 
 // CheckDatabaseConnection checks the database connection

models/db/iterate.go

@@ -11,7 +11,7 @@ import (
 	"xorm.io/builder"
 )
 
-// Iterate iterate all the Bean object
+// Iterate iterates all the Bean object
 func Iterate[Bean any](ctx context.Context, cond builder.Cond, f func(ctx context.Context, bean *Bean) error) error {
 	var start int
 	batchSize := setting.Database.IterateBufferSize