go-omni/database/mysqlConnector_test.go

package database_test

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"errors"
	"net"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/database"
	"github.com/go-sql-driver/mysql"
)

type connectorDriverLog struct{ count atomic.Int32 }

func (l *connectorDriverLog) Print(...any) { l.count.Add(1) }

func TestMySQLConnectorDriverLoggerIsolation(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	stopped := make(chan struct{})
	go func() {
		defer close(stopped)
		for {
			c, err := listener.Accept()
			if err != nil {
				return
			}
			c.Close()
		}
	}()
	defer func() { listener.Close(); <-stopped }()
	logs := [2]connectorDriverLog{}
	for i := range logs {
		cfg := mysql.NewConfig()
		cfg.Net = "tcp"
		cfg.Addr = listener.Addr().String()
		cfg.Timeout = time.Second
		cfg.ReadTimeout = time.Second
		cfg.Logger = &logs[i]
		connector, err := mysql.NewConnector(cfg)
		if err != nil {
			t.Fatal(err)
		}
		before := [2]int32{logs[0].count.Load(), logs[1].count.Load()}
		db, err := database.OpenConnector(context.Background(), connector, database.DefaultPoolConfig())
		if err == nil {
			db.Close()
			t.Fatal("handshake to closed socket succeeded")
		}
		if logs[i].count.Load() <= before[i] || logs[1-i].count.Load() != before[1-i] {
			t.Fatal("driver logger did not stay on its connector")
		}
	}
}

func TestMySQLConnectorTLSAndPoolIsolation(t *testing.T) {
	if os.Getenv("GO_OMNI_CONNECTOR_MYSQL_TEST") != "1" {
		t.Skip("use scripts/verifyConnector.py for isolated MySQL TLS verification")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	ca, err := os.ReadFile(os.Getenv("GO_OMNI_CONNECTOR_CA_FILE"))
	if err != nil {
		t.Fatal(err)
	}
	roots := x509.NewCertPool()
	if !roots.AppendCertsFromPEM(ca) {
		t.Fatal("invalid test CA")
	}
	cfg := mysql.NewConfig()
	cfg.Net = "tcp"
	cfg.Addr = os.Getenv("GO_OMNI_CONNECTOR_ADDR")
	cfg.User = os.Getenv("GO_OMNI_CONNECTOR_USER")
	cfg.Passwd = os.Getenv("GO_OMNI_CONNECTOR_PASSWORD")
	cfg.DBName = os.Getenv("GO_OMNI_CONNECTOR_DATABASE")
	cfg.Timeout = time.Second
	cfg.ReadTimeout = 2 * time.Second
	cfg.WriteTimeout = 2 * time.Second
	cfg.TLS = &tls.Config{MinVersion: tls.VersionTLS12, ServerName: "localhost", RootCAs: roots}
	pool := database.DefaultPoolConfig()
	pool.MaxOpenConns = 2
	pool.MaxIdleConns = 2
	pool.PingTimeout = 2 * time.Second
	open := func(cfg *mysql.Config) (*database.DB, error) {
		connector, err := mysql.NewConnector(cfg)
		if err != nil {
			return nil, err
		}
		return database.OpenConnector(ctx, connector, pool)
	}
	cfg.Params = map[string]string{"time_zone": "'+00:00'"}
	first, err := open(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	other := cfg.Clone()
	other.Params["time_zone"] = "'+08:00'"
	second, err := open(other)
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
	for i, db := range []*database.DB{first, second} {
		var name, cipher, zone string
		if err := db.QueryRowContext(ctx, "SHOW SESSION STATUS LIKE 'Ssl_cipher'").Scan(&name, &cipher); err != nil {
			t.Fatal(err)
		}
		if cipher == "" {
			t.Fatal("connection is not encrypted")
		}
		if err := db.QueryRowContext(ctx, "SELECT @@session.time_zone").Scan(&zone); err != nil {
			t.Fatal(err)
		}
		if zone != []string{"+00:00", "+08:00"}[i] {
			t.Fatalf("session config leaked: %s", zone)
		}
		if db.Stats().MaxOpenConnections != 2 {
			t.Fatal("pool configuration ignored")
		}
	}
	for _, bad := range []string{"untrustedCA", "wrongHostname"} {
		t.Run(bad, func(t *testing.T) {
			invalid := cfg.Clone()
			if bad == "untrustedCA" {
				invalid.TLS.RootCAs = x509.NewCertPool()
			} else {
				invalid.TLS.ServerName = "wrong.invalid"
			}
			db, err := open(invalid)
			if err == nil {
				db.Close()
				t.Fatal("invalid server certificate accepted")
			}
			var verification *tls.CertificateVerificationError
			if !errors.As(err, &verification) {
				t.Fatalf("not a TLS verification failure: %v", err)
			}
			if err := first.Health(ctx); err != nil {
				t.Fatalf("failed connector affected existing pool: %v", err)
			}
		})
	}
	if _, err := first.ExecContext(ctx, "CREATE TABLE connector_check (id INT PRIMARY KEY)"); err != nil {
		t.Fatal(err)
	}
	if err := first.InTx(ctx, nil, func(tx *sql.Tx) error {
		_, err := tx.ExecContext(ctx, "INSERT INTO connector_check VALUES (1)")
		return err
	}); err != nil {
		t.Fatal(err)
	}
	rollbackErr := errors.New("rollback requested")
	if err := first.InTx(ctx, nil, func(tx *sql.Tx) error {
		if _, err := tx.ExecContext(ctx, "INSERT INTO connector_check VALUES (2)"); err != nil {
			return err
		}
		return rollbackErr
	}); !errors.Is(err, rollbackErr) {
		t.Fatal(err)
	}
	var count int
	if err := second.QueryRowContext(ctx, "SELECT COUNT(*) FROM connector_check").Scan(&count); err != nil || count != 1 {
		t.Fatalf("transaction count=%d err=%v", count, err)
	}
	if err := first.Close(); err != nil {
		t.Fatal(err)
	}
	if err := second.Health(ctx); err != nil {
		t.Fatal("closing first pool affected second:", err)
	}
}