Jelajahi Sumber

feat(pkg): 添加 MongoDB 驱动和 WAF 格式化服务

- 新增 mongo 包,实现 MongoDB 数据库连接和管理
- 添加 WafFormatterService接口和实现,用于处理 WAF 相关的格式化操作
- 两个新功能完全独立,互不影响
fusu 2 bulan lalu
induk
melakukan
cf015831e3
2 mengubah file dengan 131 tambahan dan 0 penghapusan
  1. 52 0
      internal/service/wafformatter.go
  2. 79 0
      pkg/mongo/mongo.go

+ 52 - 0
internal/service/wafformatter.go

@@ -0,0 +1,52 @@
+package service
+
+import (
+	"context"
+	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
+	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
+
+	"strconv"
+)
+
+type WafFormatterService interface {
+	require(ctx context.Context, req v1.GlobalRequire, category string) (v1.GlobalRequire, error)
+}
+func NewWafFormatterService(
+    service *Service,
+	globalRep repository.GlobalLimitRepository,
+	hostRep repository.HostRepository,
+) WafFormatterService {
+	return &wafFormatterService{
+		Service:        service,
+		globalRep: globalRep,
+		hostRep: hostRep,
+	}
+}
+
+type wafFormatterService struct {
+	*Service
+	globalRep repository.GlobalLimitRepository
+	hostRep repository.HostRepository
+}
+
+func (s *wafFormatterService) require(ctx context.Context,req v1.GlobalRequire,category string) (v1.GlobalRequire, error) {
+	RuleIds, err := s.globalRep.GetGlobalLimitByHostId(ctx, int64(req.HostId))
+	if err != nil {
+		return v1.GlobalRequire{}, err
+	}
+	req.WafGatewayGroupId = RuleIds.GatewayGroupId
+	switch category {
+	case "tcp":
+		req.LimitRuleId = RuleIds.TcpLimitRuleId
+	case "udp":
+		req.LimitRuleId = RuleIds.UdpLimitRuleId
+	case "web":
+		req.LimitRuleId = RuleIds.WebLimitRuleId
+	}
+	domain, err := s.hostRep.GetDomainById(ctx, req.HostId)
+	if err != nil {
+		return v1.GlobalRequire{}, err
+	}
+	req.Tag = strconv.Itoa(req.Uid) + "_" + strconv.Itoa(req.HostId) + "_" + domain + "_" + req.Comment
+	return req, nil
+}

+ 79 - 0
pkg/mongo/mongo.go

@@ -0,0 +1,79 @@
+package mongo
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/readpref"
+)
+
+// Config MongoDB配置
+type Config struct {
+	URI         string        `mapstructure:"uri"`
+	Database    string        `mapstructure:"database"`
+	Timeout     time.Duration `mapstructure:"timeout"`
+	MaxPoolSize uint64        `mapstructure:"max_pool_size"`
+}
+
+// MongoDB连接管理器
+type MongoDB struct {
+	config   *Config
+	client   *mongo.Client
+	database *mongo.Database
+}
+
+// New 创建新的MongoDB客户端
+func New(config *Config) (*MongoDB, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), config.Timeout)
+	defer cancel()
+
+	// 创建MongoDB客户端选项
+	clientOptions := options.Client().
+		ApplyURI(config.URI).
+		SetMaxPoolSize(config.MaxPoolSize)
+
+	// 连接到MongoDB
+	client, err := mongo.Connect(ctx, clientOptions)
+	if err != nil {
+		return nil, fmt.Errorf("连接MongoDB失败: %w", err)
+	}
+
+	// 验证连接
+	if err := client.Ping(ctx, readpref.Primary()); err != nil {
+		return nil, fmt.Errorf("MongoDB连接测试失败: %w", err)
+	}
+
+	// 获取数据库
+	database := client.Database(config.Database)
+
+	return &MongoDB{
+		config:   config,
+		client:   client,
+		database: database,
+	}, nil
+}
+
+// Close 关闭MongoDB连接
+func (m *MongoDB) Close() error {
+	ctx, cancel := context.WithTimeout(context.Background(), m.config.Timeout)
+	defer cancel()
+	return m.client.Disconnect(ctx)
+}
+
+// GetDatabase 获取数据库实例
+func (m *MongoDB) GetDatabase() *mongo.Database {
+	return m.database
+}
+
+// GetCollection 获取集合实例
+func (m *MongoDB) GetCollection(name string) *mongo.Collection {
+	return m.database.Collection(name)
+}
+
+// Client 获取原始MongoDB客户端
+func (m *MongoDB) Client() *mongo.Client {
+	return m.client
+}