fscan/Plugins/services/mysql/connector.go

package mysql

import (
	"context"
	"database/sql"
	"fmt"
	"net"
	"strconv"
	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
)

// MySQLConnector 实现MySQL数据库服务连接器
// 遵循 base.ServiceConnector 接口规范，提供标准化的MySQL连接和认证功能

// MySQLConnector MySQL数据库连接器
type MySQLConnector struct {
	timeout time.Duration // 连接超时时间
	host    string        // 目标主机地址
	port    int           // 目标端口号
}

// NewMySQLConnector 创建新的MySQL连接器实例
// 自动注册SOCKS代理支持，配置适当的超时时间
func NewMySQLConnector() *MySQLConnector {
	connector := &MySQLConnector{
		timeout: time.Duration(common.Timeout) * time.Second,
	}
	
	// 注册SOCKS代理支持的dialer（如果配置了代理）
	connector.registerProxyDialer()
	
	return connector
}

// Connect 建立到MySQL服务的基础连接
// 实现 base.ServiceConnector 接口的 Connect 方法
func (c *MySQLConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 解析目标端口号
	port, err := strconv.Atoi(info.Ports)
	if err != nil {
		return nil, fmt.Errorf("无效的端口号: %s", info.Ports)
	}
	
	// 缓存目标信息，供认证阶段使用
	c.host = info.Host
	c.port = port
	
	// 构建基础连接字符串（无认证信息）
	connStr := c.buildConnectionString(info.Host, port, "", "")
	
	// 创建数据库连接实例
	db, err := sql.Open("mysql", connStr)
	if err != nil {
		return nil, fmt.Errorf("创建连接失败: %v", err)
	}

	// 配置连接池参数
	db.SetConnMaxLifetime(c.timeout)
	db.SetConnMaxIdleTime(c.timeout)
	db.SetMaxIdleConns(0)
	
	return db, nil
}

// Authenticate 使用凭据对MySQL服务进行身份认证
// 实现 base.ServiceConnector 接口的 Authenticate 方法
// 关键优化：使用独立的Context避免上游超时问题，并优化内存使用
func (c *MySQLConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	// 创建独立的超时上下文，避免上游Context超时问题
	// 这是解决新架构Context传递问题的关键修复
	timeout := time.Duration(common.Timeout) * time.Second
	authCtx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	
	// 内存优化：预构建连接字符串，避免重复分配
	connStr := c.buildConnectionString(c.host, c.port, cred.Username, cred.Password)
	common.LogDebug(fmt.Sprintf("MySQL尝试认证: %s@%s:%d", cred.Username, c.host, c.port))
	
	// 内存优化：直接建立连接而不创建连接池
	// 避免为单次认证创建不必要的连接池开销
	rawConn, err := c.connectDirect(authCtx, connStr)
	if err != nil {
		common.LogDebug(fmt.Sprintf("MySQL直连失败: %s@%s:%d - %v", cred.Username, c.host, c.port, err))
		return fmt.Errorf("连接失败: %v", err)
	}
	defer rawConn.Close()
	
	// 执行简单的认证验证
	err = c.validateConnection(rawConn)
	if err != nil {
		common.LogDebug(fmt.Sprintf("MySQL认证失败: %s@%s:%d - %v", cred.Username, c.host, c.port, err))
		return fmt.Errorf("认证失败: %v", err)
	}
	
	common.LogDebug(fmt.Sprintf("MySQL认证成功: %s@%s:%d", cred.Username, c.host, c.port))
	return nil
}

// Close 关闭MySQL连接
// 实现 base.ServiceConnector 接口的 Close 方法
func (c *MySQLConnector) Close(conn interface{}) error {
	if db, ok := conn.(*sql.DB); ok {
		return db.Close()
	}
	return nil
}

// connectWithCredentials 使用凭据创建新连接
func (c *MySQLConnector) connectWithCredentials(ctx context.Context, originalDB *sql.DB, cred *base.Credential) (*sql.DB, error) {
	// 从原始连接中提取主机和端口信息
	// 这里简化处理，实际应该从原始连接字符串中解析
	// 为了示例，我们假设可以从某种方式获取主机端口信息
	
	// 临时解决方案：重新构建连接字符串
	connStr := c.buildConnectionStringWithCredentials(cred)
	
	db, err := sql.Open("mysql", connStr)
	if err != nil {
		return nil, fmt.Errorf("创建认证连接失败: %v", err)
	}
	
	return db, nil
}

// buildConnectionString 构建MySQL连接字符串
// 根据是否配置SOCKS代理选择合适的连接方式
// 使用与老版本兼容的连接字符串格式，确保稳定性
func (c *MySQLConnector) buildConnectionString(host string, port int, username, password string) string {
	// 根据代理配置选择网络类型
	if common.Socks5Proxy != "" {
		// SOCKS代理连接模式
		return fmt.Sprintf("%v:%v@tcp-proxy(%v:%v)/mysql?charset=utf8&timeout=%v",
			username, password, host, port, c.timeout)
	} else {
		// 标准TCP直连模式
		return fmt.Sprintf("%v:%v@tcp(%v:%v)/mysql?charset=utf8&timeout=%v",
			username, password, host, port, c.timeout)
	}
}

// buildConnectionStringWithCredentials 构建带凭据的连接字符串
func (c *MySQLConnector) buildConnectionStringWithCredentials(cred *base.Credential) string {
	// 使用保存的主机和端口信息
	return c.buildConnectionString(c.host, c.port, cred.Username, cred.Password)
}

// connectDirect 内存优化：直接建立MySQL连接，避免连接池开销
// 用于单次认证场景，减少内存分配和资源浪费
func (c *MySQLConnector) connectDirect(ctx context.Context, connStr string) (*sql.Conn, error) {
	// 创建最小化配置的临时数据库实例
	db, err := sql.Open("mysql", connStr)
	if err != nil {
		return nil, fmt.Errorf("创建连接实例失败: %v", err)
	}
	defer db.Close() // 确保临时db实例被清理
	
	// 禁用连接池以减少内存开销
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(0)
	db.SetConnMaxLifetime(0)
	
	// 获取原始连接
	conn, err := db.Conn(ctx)
	if err != nil {
		return nil, fmt.Errorf("获取连接失败: %v", err)
	}
	
	return conn, nil
}

// validateConnection 内存优化：轻量级连接验证
// 使用最小开销的方式验证MySQL连接有效性
func (c *MySQLConnector) validateConnection(conn *sql.Conn) error {
	// 使用带超时的Ping进行快速验证
	timeout := time.Duration(common.Timeout) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	
	return conn.PingContext(ctx)
}

// registerProxyDialer 注册SOCKS代理支持的网络拨号器
// 仅在配置了SOCKS代理时才注册，避免不必要的开销
func (c *MySQLConnector) registerProxyDialer() {
	if common.Socks5Proxy == "" {
		return // 未配置代理，跳过注册
	}
	
	// 向MySQL驱动注册自定义的代理拨号器
	mysql.RegisterDialContext("tcp-proxy", func(ctx context.Context, addr string) (net.Conn, error) {
		return common.WrapperTcpWithContext(ctx, "tcp", addr)
	})
}