mirror of
https://github.com/shadow1ng/fscan.git
synced 2026-09-26 21:21:53 +08:00
refactor(services): 统一数据库插件的DBWrapper
4个数据库插件(MySQL、PostgreSQL、MSSQL、Oracle)都有相同的sql.DB包装代码, 合并为通用的SQLDBWrapper,减少重复。
This commit is contained in:
@@ -2,6 +2,7 @@ package services
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"database/sql"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -369,3 +370,17 @@ func matchIgnoreCase(a, b string) bool {
|
|||||||
}
|
}
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// 通用数据库连接包装
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
// SQLDBWrapper 包装 sql.DB 以实现 io.Closer
|
||||||
|
// 用于 MySQL、PostgreSQL、MSSQL、Oracle 等数据库插件的连接返回
|
||||||
|
type SQLDBWrapper struct {
|
||||||
|
*sql.DB
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *SQLDBWrapper) Close() error {
|
||||||
|
return w.DB.Close()
|
||||||
|
}
|
||||||
|
|||||||
@@ -98,21 +98,12 @@ func (p *MSSQLPlugin) doMSSQLAuth(ctx context.Context, info *common.HostInfo, cr
|
|||||||
|
|
||||||
return &AuthResult{
|
return &AuthResult{
|
||||||
Success: true,
|
Success: true,
|
||||||
Conn: &mssqlDBWrapper{db},
|
Conn: &SQLDBWrapper{db},
|
||||||
ErrorType: ErrorTypeUnknown,
|
ErrorType: ErrorTypeUnknown,
|
||||||
Error: nil,
|
Error: nil,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// mssqlDBWrapper 包装 sql.DB 以实现 io.Closer
|
|
||||||
type mssqlDBWrapper struct {
|
|
||||||
*sql.DB
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *mssqlDBWrapper) Close() error {
|
|
||||||
return w.DB.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// classifyMSSQLErrorType MSSQL错误分类
|
// classifyMSSQLErrorType MSSQL错误分类
|
||||||
func classifyMSSQLErrorType(err error) ErrorType {
|
func classifyMSSQLErrorType(err error) ErrorType {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|||||||
@@ -104,24 +104,14 @@ func (p *MySQLPlugin) doMySQLAuth(ctx context.Context, info *common.HostInfo, cr
|
|||||||
|
|
||||||
state.IncrementTCPSuccessPacketCount()
|
state.IncrementTCPSuccessPacketCount()
|
||||||
|
|
||||||
// MySQL 使用 sql.DB,包装为 io.Closer
|
|
||||||
return &AuthResult{
|
return &AuthResult{
|
||||||
Success: true,
|
Success: true,
|
||||||
Conn: &sqlDBWrapper{db},
|
Conn: &SQLDBWrapper{db},
|
||||||
ErrorType: ErrorTypeUnknown,
|
ErrorType: ErrorTypeUnknown,
|
||||||
Error: nil,
|
Error: nil,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// sqlDBWrapper 包装 sql.DB 以实现 io.Closer
|
|
||||||
type sqlDBWrapper struct {
|
|
||||||
*sql.DB
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *sqlDBWrapper) Close() error {
|
|
||||||
return w.DB.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// classifyMySQLErrorType MySQL错误分类
|
// classifyMySQLErrorType MySQL错误分类
|
||||||
func classifyMySQLErrorType(err error) ErrorType {
|
func classifyMySQLErrorType(err error) ErrorType {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|||||||
@@ -106,7 +106,7 @@ func (p *OraclePlugin) doOracleAuth(ctx context.Context, info *common.HostInfo,
|
|||||||
|
|
||||||
return &AuthResult{
|
return &AuthResult{
|
||||||
Success: true,
|
Success: true,
|
||||||
Conn: &oracleDBWrapper{db},
|
Conn: &SQLDBWrapper{db},
|
||||||
ErrorType: ErrorTypeUnknown,
|
ErrorType: ErrorTypeUnknown,
|
||||||
Error: nil,
|
Error: nil,
|
||||||
}
|
}
|
||||||
@@ -120,15 +120,6 @@ func (p *OraclePlugin) doOracleAuth(ctx context.Context, info *common.HostInfo,
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// oracleDBWrapper 包装 sql.DB 以实现 io.Closer
|
|
||||||
type oracleDBWrapper struct {
|
|
||||||
*sql.DB
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *oracleDBWrapper) Close() error {
|
|
||||||
return w.DB.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// classifyOracleErrorType Oracle错误分类
|
// classifyOracleErrorType Oracle错误分类
|
||||||
func classifyOracleErrorType(err error) ErrorType {
|
func classifyOracleErrorType(err error) ErrorType {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|||||||
@@ -104,21 +104,12 @@ func (p *PostgreSQLPlugin) doPostgreSQLAuth(ctx context.Context, info *common.Ho
|
|||||||
|
|
||||||
return &AuthResult{
|
return &AuthResult{
|
||||||
Success: true,
|
Success: true,
|
||||||
Conn: &pgDBWrapper{db},
|
Conn: &SQLDBWrapper{db},
|
||||||
ErrorType: ErrorTypeUnknown,
|
ErrorType: ErrorTypeUnknown,
|
||||||
Error: nil,
|
Error: nil,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// pgDBWrapper 包装 sql.DB 以实现 io.Closer
|
|
||||||
type pgDBWrapper struct {
|
|
||||||
*sql.DB
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *pgDBWrapper) Close() error {
|
|
||||||
return w.DB.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// classifyPostgreSQLErrorType PostgreSQL错误分类
|
// classifyPostgreSQLErrorType PostgreSQL错误分类
|
||||||
func classifyPostgreSQLErrorType(err error) ErrorType {
|
func classifyPostgreSQLErrorType(err error) ErrorType {
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
|||||||
Reference in New Issue
Block a user