From 93b66af9fb90e4c57a4091c6aa7fccb817d4193d Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <45962551+tarunKoyalwar@users.noreply.github.com>
Date: Wed, 31 Jan 2024 02:32:23 +0530
Subject: [PATCH] mysql: ignore warning log + misc updates (#4702)

---
 go.mod                                      |  2 +-
 go.sum                                      |  4 +--
 pkg/js/libs/mysql/mysql.go                  | 32 +++++++++++++++++++--
 pkg/protocols/common/protocolstate/state.go |  8 ++++++
 4 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index a6a300c81c..90cca8ba45 100644
--- a/go.mod
+++ b/go.mod
@@ -68,7 +68,7 @@ require (
 	github.com/go-git/go-git/v5 v5.11.0
 	github.com/go-ldap/ldap/v3 v3.4.5
 	github.com/go-pg/pg v8.0.7+incompatible
-	github.com/go-sql-driver/mysql v1.6.0
+	github.com/go-sql-driver/mysql v1.7.1
 	github.com/h2non/filetype v1.1.3
 	github.com/labstack/echo/v4 v4.10.2
 	github.com/lib/pq v1.10.1
diff --git a/go.sum b/go.sum
index 256ffb5740..00128e54e3 100644
--- a/go.sum
+++ b/go.sum
@@ -379,8 +379,8 @@ github.com/go-rod/rod v0.114.0/go.mod h1:aiedSEFg5DwG/fnNbUOTPMTTWX3MRj6vIs/a684
 github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU=
 github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg=
 github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
-github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE=
-github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
+github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI=
+github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
 github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
diff --git a/pkg/js/libs/mysql/mysql.go b/pkg/js/libs/mysql/mysql.go
index af5de5883b..9dd49c3bd7 100644
--- a/pkg/js/libs/mysql/mysql.go
+++ b/pkg/js/libs/mysql/mysql.go
@@ -4,11 +4,13 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"io"
+	"log"
 	"net"
 	"net/url"
 	"time"
 
-	_ "github.com/go-sql-driver/mysql"
+	"github.com/go-sql-driver/mysql"
 	"github.com/praetorian-inc/fingerprintx/pkg/plugins"
 	mysqlplugin "github.com/praetorian-inc/fingerprintx/pkg/plugins/services/mysql"
 	utils "github.com/projectdiscovery/nuclei/v3/pkg/js/utils"
@@ -66,6 +68,24 @@ func (c *MySQLClient) ConnectWithDB(host string, port int, username, password, d
 	return connect(host, port, username, password, dbName)
 }
 
+// ConnectWithDSN connects to MySQL database using given DSN.
+// we override mysql dialer with fastdialer so it respects network policy
+func (c *MySQLClient) ConnectWithDSN(dsn string) (bool, error) {
+	db, err := sql.Open("mysql", dsn)
+	if err != nil {
+		return false, err
+	}
+	defer db.Close()
+	db.SetMaxOpenConns(1)
+	db.SetMaxIdleConns(0)
+
+	_, err = db.Exec("select 1")
+	if err != nil {
+		return false, err
+	}
+	return true, nil
+}
+
 func connect(host string, port int, username, password, dbName string) (bool, error) {
 	if host == "" || port <= 0 {
 		return false, fmt.Errorf("invalid host or port")
@@ -78,7 +98,7 @@ func connect(host string, port int, username, password, dbName string) (bool, er
 
 	target := net.JoinHostPort(host, fmt.Sprintf("%d", port))
 
-	db, err := sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%s",
+	db, err := sql.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v)/%s?allowOldPasswords=1",
 		url.PathEscape(username),
 		url.PathEscape(password),
 		target,
@@ -87,6 +107,8 @@ func connect(host string, port int, username, password, dbName string) (bool, er
 		return false, err
 	}
 	defer db.Close()
+	db.SetMaxOpenConns(1)
+	db.SetMaxIdleConns(0)
 
 	_, err = db.Exec("select 1")
 	if err != nil {
@@ -115,6 +137,8 @@ func (c *MySQLClient) ExecuteQuery(host string, port int, username, password, db
 		return "", err
 	}
 	defer db.Close()
+	db.SetMaxOpenConns(1)
+	db.SetMaxIdleConns(0)
 
 	rows, err := db.Query(query)
 	if err != nil {
@@ -126,3 +150,7 @@ func (c *MySQLClient) ExecuteQuery(host string, port int, username, password, db
 	}
 	return string(resp), nil
 }
+
+func init() {
+	_ = mysql.SetLogger(log.New(io.Discard, "", 0))
+}
diff --git a/pkg/protocols/common/protocolstate/state.go b/pkg/protocols/common/protocolstate/state.go
index 6d4bd7672a..cc5845862b 100644
--- a/pkg/protocols/common/protocolstate/state.go
+++ b/pkg/protocols/common/protocolstate/state.go
@@ -1,10 +1,12 @@
 package protocolstate
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"net/url"
 
+	"github.com/go-sql-driver/mysql"
 	"github.com/pkg/errors"
 	"golang.org/x/net/proxy"
 
@@ -133,6 +135,12 @@ func Init(options *types.Options) error {
 		return errors.Wrap(err, "could not create dialer")
 	}
 	Dialer = dialer
+
+	// override dialer in mysql
+	mysql.RegisterDialContext("tcp", func(ctx context.Context, addr string) (net.Conn, error) {
+		return Dialer.Dial(ctx, "tcp", addr)
+	})
+
 	return nil
 }