database/sql: ensure releaseConn is defined before a possible close

When running a Query on Stmt a dependency is added to the stmt and
rows. To do that it needs a reference to Rows, so the releaseConn
function is defined after the definition. However the
rows.initContextClose was set to run before the releaseConn was
set on rows, setting up a situation where the connection could
be canceled before the releaseConn was set and resulting in
a segfault.

Fixes #20160

Change-Id: I5592e7db2cf653dfc48d42cbc2b03ca20501b1a0
Reviewed-on: https://go-review.googlesource.com/42139
Run-TryBot: Daniel Theophanes <kardianos@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 09d61f12875ddf49707fd7941b38f9e3f5e07fef..03f66c6cb75720e07475b6434fdf71d6a49575b9 100644 (file)
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -2183,12 +2183,17 @@ func (s *Stmt) QueryContext(ctx context.Context, args ...interface{}) (*Rows, er
                                rowsi: rowsi,
                                // releaseConn set below
                        }
-                       rows.initContextClose(ctx)
+                       // addDep must be added before initContextClose or it could attempt
+                       // to removeDep before it has been added.
                        s.db.addDep(s, rows)
+
+                       // releaseConn must be set before initContextClose or it could
+                       // release the connection before it is set.
                        rows.releaseConn = func(err error) {
                                releaseConn(err)
                                s.db.removeDep(s, rows)
                        }
+                       rows.initContextClose(ctx)
                        return rows, nil
                }
 

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 5ea965fb28ec994e3398fb4b7c50e5b61baf6df3..2fd81f29a541b0d68412738e74c7b81f08ad9881 100644 (file)
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -3044,6 +3044,46 @@ func TestIssue18429(t *testing.T) {
        wg.Wait()
 }
 
+// TestIssue20160 attempts to test a short context life on a stmt Query.
+func TestIssue20160(t *testing.T) {
+       db := newTestDB(t, "people")
+       defer closeDB(t, db)
+
+       ctx := context.Background()
+       sem := make(chan bool, 20)
+       var wg sync.WaitGroup
+
+       const milliWait = 30
+
+       stmt, err := db.PrepareContext(ctx, "SELECT|people|name|")
+       if err != nil {
+               t.Fatal(err)
+       }
+       defer stmt.Close()
+
+       for i := 0; i < 100; i++ {
+               sem <- true
+               wg.Add(1)
+               go func() {
+                       defer func() {
+                               <-sem
+                               wg.Done()
+                       }()
+                       ctx, cancel := context.WithTimeout(ctx, time.Duration(rand.Intn(milliWait))*time.Millisecond)
+                       defer cancel()
+
+                       // This is expected to give a cancel error most, but not all the time.
+                       // Test failure will happen with a panic or other race condition being
+                       // reported.
+                       rows, _ := stmt.QueryContext(ctx)
+                       if rows != nil {
+                               rows.Close()
+                       }
+               }()
+       }
+       wg.Wait()
+}
+
 // TestIssue18719 closes the context right before use. The sql.driverConn
 // will nil out the ci on close in a lock, but if another process uses it right after
 // it will panic with on the nil ref.