From caf1beae3721e5e90a5e3dc9c2bb6fb4bbcc2e91 Mon Sep 17 00:00:00 2001
From: ElecTwix <root@electwix.dev>
Date: Sun, 28 Apr 2024 17:15:07 +0300
Subject: [PATCH] Panic fix on gorilla ws when connection lost. fix linter
 errors with shadow checking upgrade linter

---
 .github/workflows/lint.yml  | 12 +++++------
 .golangci.yml               | 11 +++++-----
 db.go                       |  6 ++++++
 db_test.go                  |  9 ++++++++
 internal/mock/mock.go       |  7 ++++--
 pkg/conn/conn.go            |  1 +
 pkg/conn/gorilla/gorilla.go | 43 +++++++++++++++++++++----------------
 7 files changed, 58 insertions(+), 31 deletions(-)

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 1302d11..d377bb1 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -12,16 +12,16 @@ permissions:
 jobs:
   golangci:
     permissions:
-      contents: read  # for actions/checkout to fetch code
-      pull-requests: read  # for golangci/golangci-lint-action to fetch pull requests
+      contents: read # for actions/checkout to fetch code
+      pull-requests: read # for golangci/golangci-lint-action to fetch pull requests
     name: lint
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/setup-go@v3
+      - uses: actions/setup-go@v5
         with:
-          go-version: '>=1.18'
-      - uses: actions/checkout@v3
+          go-version: ">=1.18"
+      - uses: actions/checkout@v4
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v3
+        uses: golangci/golangci-lint-action@v5
         with:
           args: --timeout=5m
diff --git a/.golangci.yml b/.golangci.yml
index cbad97e..3b72088 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -34,15 +34,16 @@ linters-settings:
       - condition
       - return
     ignored-numbers:
-      - '0'
-      - '1'
-      - '2'
-      - '3'
+      - "0"
+      - "1"
+      - "2"
+      - "3"
     ignored-functions:
       - strings.SplitN
 
   govet:
-    check-shadowing: true
+    enable:
+      - shadow
   lll:
     line-length: 140
   misspell:
diff --git a/db.go b/db.go
index f73f2da..02c82e6 100644
--- a/db.go
+++ b/db.go
@@ -32,6 +32,12 @@ func New(url string, connection conn.Connection) (*DB, error) {
 	return &DB{connection}, nil
 }
 
+// Initialize initializes the connection to the database.
+// Can return an error if the connection fails on runtime.
+func (db *DB) Initialize() error {
+	return db.conn.Initialize()
+}
+
 // --------------------------------------------------
 // Public methods
 // --------------------------------------------------
diff --git a/db_test.go b/db_test.go
index 874c567..66cc118 100644
--- a/db_test.go
+++ b/db_test.go
@@ -121,6 +121,15 @@ func (s *SurrealDBTestSuite) openConnection() *surrealdb.DB {
 	impl := s.connImplementations[s.name]
 	require.NotNil(s.T(), impl)
 	db, err := surrealdb.New(url, impl)
+
+	go func(s *SurrealDBTestSuite) {
+		connErr := db.Initialize()
+		if connErr != nil {
+			fmt.Println(connErr)
+			require.NoError(s.T(), connErr)
+		}
+	}(s)
+
 	s.Require().NoError(err)
 	return db
 }
diff --git a/internal/mock/mock.go b/internal/mock/mock.go
index a3b7e3e..8a7593a 100644
--- a/internal/mock/mock.go
+++ b/internal/mock/mock.go
@@ -7,13 +7,16 @@ import (
 	"github.com/surrealdb/surrealdb.go/pkg/model"
 )
 
-type ws struct {
-}
+type ws struct{}
 
 func (w *ws) Connect(url string) (conn.Connection, error) {
 	return w, nil
 }
 
+func (w *ws) Initialize() error {
+	return nil
+}
+
 func (w *ws) Send(method string, params []interface{}) (interface{}, error) {
 	return nil, nil
 }
diff --git a/pkg/conn/conn.go b/pkg/conn/conn.go
index dc619ce..00f685b 100644
--- a/pkg/conn/conn.go
+++ b/pkg/conn/conn.go
@@ -4,6 +4,7 @@ import "github.com/surrealdb/surrealdb.go/pkg/model"
 
 type Connection interface {
 	Connect(url string) (Connection, error)
+	Initialize() error
 	Send(method string, params []interface{}) (interface{}, error)
 	Close() error
 	LiveNotifications(id string) (chan model.Notification, error)
diff --git a/pkg/conn/gorilla/gorilla.go b/pkg/conn/gorilla/gorilla.go
index dcfbb37..c5e5f40 100644
--- a/pkg/conn/gorilla/gorilla.go
+++ b/pkg/conn/gorilla/gorilla.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"reflect"
 	"strconv"
@@ -73,10 +74,13 @@ func (ws *WebSocket) Connect(url string) (conn.Connection, error) {
 		}
 	}
 
-	ws.initialize()
 	return ws, nil
 }
 
+func (ws *WebSocket) Initialize() error {
+	return ws.initialize()
+}
+
 func (ws *WebSocket) SetTimeOut(timeout time.Duration) *WebSocket {
 	ws.Option = append(ws.Option, func(ws *WebSocket) error {
 		ws.Timeout = timeout
@@ -234,26 +238,29 @@ func (ws *WebSocket) write(v interface{}) error {
 	return ws.Conn.WriteMessage(gorilla.TextMessage, data)
 }
 
-func (ws *WebSocket) initialize() {
-	go func() {
-		for {
-			select {
-			case <-ws.close:
-				return
-			default:
-				var res rpc.RPCResponse
-				err := ws.read(&res)
-				if err != nil {
-					if errors.Is(err, net.ErrClosed) {
-						break
-					}
-					ws.logger.Error(err.Error())
-					continue
+func (ws *WebSocket) initialize() error {
+	for {
+		select {
+		case <-ws.close:
+			return nil
+		default:
+			var res rpc.RPCResponse
+			err := ws.read(&res)
+			if err != nil {
+				// this needed because gorilla not shudown gracefully
+				if errors.Is(err, net.ErrClosed) {
+					return nil
+				}
+				// returns error if connection drop in fly
+				if gorilla.IsUnexpectedCloseError(err) {
+					return io.ErrClosedPipe
 				}
-				go ws.handleResponse(res)
+				ws.logger.Error(err.Error())
+				continue
 			}
+			go ws.handleResponse(res)
 		}
-	}()
+	}
 }
 
 func (ws *WebSocket) handleResponse(res rpc.RPCResponse) {