go-omni/database/mysql/mysql.go

// Package mysql provides an optional MySQL adapter over database.OpenConnector.
// Each pool owns independent driver settings; no global driver registry changes
// are made. It never migrates tables or retries application transactions.
package mysql

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

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

type Config struct {
	Host      string `json:"host"`
	Port      int    `json:"port"`
	Socket    string `json:"socket"`
	User      string `json:"user"`
	Password  string `json:"password" env:",secret"`
	Database  string `json:"database"`
	ParseTime bool   `json:"parseTime"`
	Location  string `json:"location"`
	Charset   string `json:"charset"`
	Collation string `json:"collation"`
	// Session values are SQL expressions, as required by the driver, for
	// example time_zone: "'+00:00'". They are trusted deployment configuration.
	Session         map[string]string   `json:"session"`
	ConnectTimeout  time.Duration       `json:"connectTimeout"`
	ReadTimeout     time.Duration       `json:"readTimeout"`
	WriteTimeout    time.Duration       `json:"writeTimeout"`
	MultiStatements bool                `json:"multiStatements"`
	TLS             TLSConfig           `json:"tls"`
	Pool            database.PoolConfig `json:"pool"`
}
type TLSConfig struct {
	Mode       string `json:"mode"` // required (default) or disabled
	CAFile     string `json:"caFile"`
	CertFile   string `json:"certFile"`
	KeyFile    string `json:"keyFile"`
	ServerName string `json:"serverName"`
}

func DefaultConfig() Config {
	return Config{Host: "127.0.0.1", Port: 3306, ParseTime: true, Location: "UTC", Charset: "utf8mb4", ConnectTimeout: 5 * time.Second, ReadTimeout: 5 * time.Second, WriteTimeout: 5 * time.Second, TLS: TLSConfig{Mode: "required"}, Pool: database.DefaultPoolConfig()}
}

var identifier = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`)

// Validate does not read certificate files or open connections. All defaults
// come from DefaultConfig; explicit zero I/O timeouts mean no driver deadline.
func (c Config) Validate() error {
	if c.Socket == "" && (c.Host == "" || strings.ContainsAny(c.Host, "/\r\n\t ") || c.Port < 1 || c.Port > 65535) {
		return errors.New("mysql: invalid host or port")
	}
	if strings.ContainsRune(c.Socket, 0) || strings.ContainsRune(c.Database, 0) || strings.ContainsRune(c.User, 0) {
		return errors.New("mysql: invalid connection field")
	}
	if c.ConnectTimeout < 0 || c.ReadTimeout < 0 || c.WriteTimeout < 0 {
		return errors.New("mysql: timeouts cannot be negative")
	}
	if err := c.Pool.Validate(); err != nil {
		return err
	}
	if c.Location != "" {
		if _, err := time.LoadLocation(c.Location); err != nil {
			return errors.New("mysql: invalid time location")
		}
	}
	if c.Charset != "" && !identifier.MatchString(c.Charset) {
		return errors.New("mysql: invalid charset")
	}
	if c.Collation != "" && !identifier.MatchString(c.Collation) {
		return errors.New("mysql: invalid collation")
	}
	for key := range c.Session {
		if !identifier.MatchString(key) {
			return errors.New("mysql: invalid session variable name")
		}
	}
	if c.TLS.Mode != "" && c.TLS.Mode != "required" && c.TLS.Mode != "disabled" {
		return errors.New("mysql: TLS mode must be required or disabled")
	}
	if (c.TLS.CertFile == "") != (c.TLS.KeyFile == "") {
		return errors.New("mysql: TLS client certificate and key must be configured together")
	}
	if c.TLS.Mode == "disabled" && (c.TLS.CAFile != "" || c.TLS.CertFile != "" || c.TLS.ServerName != "") {
		return errors.New("mysql: TLS fields cannot be used when disabled")
	}
	if c.Socket != "" && c.TLS.Mode != "disabled" && c.TLS.ServerName == "" {
		return errors.New("mysql: TLS over a unix socket requires serverName")
	}
	return nil
}

// Error has a safe summary; errors.Is/As retain driver/verification causes.
// Do not log raw configuration or a driver's manually extracted error fields.
type Error struct {
	Operation string
	cause     error
}

func (e *Error) Error() string { return "mysql: " + e.Operation + " failed" }
func (e *Error) Unwrap() error { return e.cause }

type Check func(context.Context, *database.DB) error
type options struct {
	logger    *logx.Logger
	checks    []Check
	configure func(*mysqldriver.Config) error
}
type Option func(*options)

func WithLogger(l *logx.Logger) Option { return func(o *options) { o.logger = l } }

// WithStartupCheck runs after Ping, before ownership reaches the caller. A
// failure closes the pool. Checks must honor ctx and must not close the pool.
func WithStartupCheck(check Check) Option {
	return func(o *options) {
		if check != nil {
			o.checks = append(o.checks, check)
		}
	}
}

// WithDriverConfig provides an advanced, per-pool escape hatch. Configuration
// is cloned after this callback; keep credentials out of callback errors/logs.
func WithDriverConfig(fn func(*mysqldriver.Config) error) Option {
	return func(o *options) { o.configure = fn }
}

func settings(opts []Option) options {
	var o options
	for _, opt := range opts {
		if opt != nil {
			opt(&o)
		}
	}
	return o
}

// Open returns the existing Omni pool; .DB is the same standard sql.DB for
// repositories or ORMs. Ping and startup checks share one bounded deadline.
func Open(ctx context.Context, cfg Config, opts ...Option) (db *database.DB, err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err = ctx.Err(); err != nil {
		return nil, err
	}
	o := settings(opts)
	if o.logger == nil {
		o.logger = logx.FromContext(ctx)
	}
	connector, err := connector(cfg, o)
	if err != nil {
		return nil, err
	}
	return openPool(ctx, connector, cfg.Pool, o.checks)
}

func openPool(ctx context.Context, connector driver.Connector, pool database.PoolConfig, checks []Check) (db *database.DB, err error) {
	checkCtx, cancel := context.WithTimeout(ctx, pool.PingTimeout)
	defer cancel()
	db, err = database.OpenConnector(checkCtx, connector, pool)
	if err != nil {
		return nil, &Error{Operation: "connect", cause: err}
	}
	defer func() {
		if p := recover(); p != nil {
			_ = db.Close()
			panic(p)
		}
		if err != nil {
			err = errors.Join(err, db.Close())
			db = nil
		}
	}()
	for _, check := range checks {
		if err = check(checkCtx, db); err != nil {
			return db, &Error{Operation: "startup check", cause: err}
		}
		if err = checkCtx.Err(); err != nil {
			return db, &Error{Operation: "startup check", cause: err}
		}
	}
	return db, nil
}

// NewConnector builds a standard connector for callers that own pool creation.
// Startup checks are only executed by Open.
func NewConnector(cfg Config, opts ...Option) (driver.Connector, error) {
	o := settings(opts)
	if len(o.checks) > 0 {
		return nil, errors.New("mysql: startup checks require Open, not NewConnector")
	}
	return connector(cfg, o)
}
func connector(cfg Config, o options) (driver.Connector, error) {
	d, err := driverConfig(cfg, o)
	if err != nil {
		return nil, err
	}
	c, err := mysqldriver.NewConnector(d)
	if err != nil {
		return nil, &Error{Operation: "connector configuration", cause: err}
	}
	return c, nil
}

func driverConfig(cfg Config, o options) (*mysqldriver.Config, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	d := mysqldriver.NewConfig()
	d.User, d.Passwd, d.DBName = cfg.User, cfg.Password, cfg.Database
	d.Net, d.Addr = "tcp", net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port))
	if cfg.Socket != "" {
		d.Net, d.Addr = "unix", cfg.Socket
	}
	d.ParseTime = cfg.ParseTime
	if cfg.Location != "" {
		d.Loc, _ = time.LoadLocation(cfg.Location)
	}
	if cfg.Charset != "" {
		if err := d.Apply(mysqldriver.Charset(cfg.Charset, cfg.Collation)); err != nil {
			return nil, &Error{Operation: "charset", cause: err}
		}
	} else {
		d.Collation = cfg.Collation
	}
	d.Params = make(map[string]string, len(cfg.Session))
	for k, v := range cfg.Session {
		d.Params[k] = v
	}
	d.Timeout, d.ReadTimeout, d.WriteTimeout = cfg.ConnectTimeout, cfg.ReadTimeout, cfg.WriteTimeout
	d.MultiStatements = cfg.MultiStatements
	logger := o.logger
	if logger == nil {
		logger = logx.FromContext(nil)
	}
	d.Logger = driverLog{logger: logger, secret: cfg.Password}
	if cfg.TLS.Mode != "disabled" {
		name := cfg.TLS.ServerName
		if name == "" {
			name = cfg.Host
		}
		t := &tls.Config{MinVersion: tls.VersionTLS12, ServerName: name}
		if cfg.TLS.CAFile != "" {
			pem, err := readCertificate(cfg.TLS.CAFile)
			if err != nil {
				return nil, &Error{Operation: "read CA", cause: err}
			}
			t.RootCAs = x509.NewCertPool()
			if !t.RootCAs.AppendCertsFromPEM(pem) {
				return nil, errors.New("mysql: invalid CA certificate")
			}
		}
		if cfg.TLS.CertFile != "" {
			cert, err := readCertificate(cfg.TLS.CertFile)
			if err != nil {
				return nil, &Error{Operation: "read client certificate", cause: err}
			}
			key, err := readCertificate(cfg.TLS.KeyFile)
			if err != nil {
				return nil, &Error{Operation: "read client key", cause: err}
			}
			pair, err := tls.X509KeyPair(cert, key)
			if err != nil {
				return nil, &Error{Operation: "client key pair", cause: err}
			}
			t.Certificates = []tls.Certificate{pair}
		}
		d.TLS = t
	}
	if o.configure != nil {
		if err := o.configure(d); err != nil {
			return nil, &Error{Operation: "custom configuration", cause: err}
		}
	}
	return d.Clone(), nil
}

func readCertificate(path string) ([]byte, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	data, err := io.ReadAll(io.LimitReader(f, (4<<20)+1))
	if err != nil {
		return nil, err
	}
	if len(data) > 4<<20 {
		return nil, errors.New("certificate file exceeds 4 MiB")
	}
	return data, nil
}

type driverLog struct {
	logger *logx.Logger
	secret string
}

func (l driverLog) Print(values ...any) {
	message := fmt.Sprint(values...)
	if l.secret != "" {
		message = strings.ReplaceAll(message, l.secret, "[REDACTED]")
	}
	l.logger.Error("database.driver", "driver", "mysql", "error", message)
}