From 6d954f2ad68b5c0589314fa40eab0c2543362887 Mon Sep 17 00:00:00 2001
From: Thomas de Zeeuw <thomas@slight.dev>
Date: Thu, 5 Jan 2023 12:11:27 +0100
Subject: [PATCH] Don't panic if driverRows.Next is called after Close

Previously this would panic with an index out of range error. Which was
misleading as the indexing itself wasn't the problem. The problem was
that fetch returns a nil error if the error channel was closed, which
led the Next function to believe there was more data ready to be
processed, but this was not the case.

The error channel is only closed if Close is called on driverRows or
driverStmt, so this panic was caused by a premature closing of the rows
or statement, but that hard to tell from the panic message.

This commit changes removes the panic and instead returns an io.EOF
error if we detect that fetch was called on a closed statement or rows.
---
 trino/integration_test.go | 32 ++++++++++++++++++++++++++++++++
 trino/trino.go            |  6 +++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/trino/integration_test.go b/trino/integration_test.go
index 13cb8d1..a56c49a 100644
--- a/trino/integration_test.go
+++ b/trino/integration_test.go
@@ -17,8 +17,10 @@ package trino
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
 	"flag"
+	"io"
 	"log"
 	"net/http"
 	"os"
@@ -527,6 +529,36 @@ func TestIntegrationQueryParametersSelect(t *testing.T) {
 	}
 }
 
+func TestIntegrationQueryNextAfterClose(t *testing.T) {
+	// NOTE: This is testing invalid behaviour. It ensures that we don't
+	// panic if we call driverRows.Next after we closed the driverStmt.
+
+	ctx := context.Background()
+	conn, err := (&Driver{}).Open(*integrationServerFlag)
+	defer conn.Close()
+
+	stmt, err := conn.(driver.ConnPrepareContext).PrepareContext(ctx, "SELECT 1")
+	if err != nil {
+		t.Fatalf("Failed preparing query: %v", err)
+	}
+
+	rows, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, []driver.NamedValue{})
+	if err != nil {
+		t.Fatalf("Failed running query: %v", err)
+	}
+	defer rows.Close()
+
+	stmt.Close() // NOTE: the important bit.
+
+	var result driver.Value
+	if err := rows.Next([]driver.Value{result}); err != nil {
+		t.Fatalf("unexpected result: %+v, no error was expected", err)
+	}
+	if err := rows.Next([]driver.Value{result}); err != io.EOF {
+		t.Fatalf("unexpected result: %+v, expected io.EOF", err)
+	}
+}
+
 func TestIntegrationExec(t *testing.T) {
 	db := integrationOpen(t)
 	defer db.Close()
diff --git a/trino/trino.go b/trino/trino.go
index ce0a966..98af17c 100644
--- a/trino/trino.go
+++ b/trino/trino.go
@@ -1144,7 +1144,11 @@ func (qr *driverRows) fetch() error {
 				return nil
 			}
 		case err = <-qr.stmt.errors:
-			if err == context.Canceled {
+			if err == nil {
+				// Channel was closed, which means the statement
+				// or rows were closed.
+				err = io.EOF
+			} else if err == context.Canceled {
 				qr.Close()
 			}
 			qr.err = err