Commit 7ba7c191 authored by qiuqunfeng's avatar qiuqunfeng
Browse files

Refactor WAF service creation and error handling

- Improve CreateWaf method by adding validation for rule count
- Update error messages in cluster client retrieval
- Return WafService struct with basic metadata on successful creation
- Remove commented-out legacy code
- Enhance error handling in DeleteListenerWaf and CreateWaf methods
parent 22159493
......@@ -201,88 +201,37 @@ func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*WafServ
},
}
// Get enabled rule categories from DB
// var ruleCategories []model.WafRuleCategory
// if err := s.db.Model(&model.WafRuleCategory{}).Where("status = ?", 0).Find(&ruleCategories).Error; err != nil {
// return nil, fmt.Errorf("failed to get rule categories: %v", err)
// }
// // Get existing WAF service config if any
// wafService := &model.WafService{}
// err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
// if err != nil {
// if err == gorm.ErrRecordNotFound {
// // Create new WAF service record if not found
// wafService = &model.WafService{
// RegionCode: req.RegionCode,
// Namespace: req.Namespace,
// GatewayName: req.GatewayName,
// Mode: string(WafModeAlert),
// // Host: model.HostList(req.Host),
// }
// if err := s.db.Create(wafService).Error; err != nil {
// return nil, fmt.Errorf("failed to create WAF service: %v", err)
// }
// } else {
// return nil, fmt.Errorf("failed to query WAF service: %v", err)
// }
// }
// // Determine which rule categories to enable
// var enabledCategories []model.WafRuleCategory
// if wafService.RuleCategoryStatus != nil && len(wafService.RuleCategoryStatus.CategoryID) > 0 {
// // Only include categories not already enabled
// for _, category := range ruleCategories {
// for _, id := range wafService.RuleCategoryStatus.CategoryID {
// if id == category.CategoryID {
// enabledCategories = append(enabledCategories, category)
// continue
// }
// }
// }
// } else {
// // Enable all categories if none specified
// enabledCategories = ruleCategories
// }
// // Add rules from enabled categories
// for _, category := range enabledCategories {
// for _, rule := range category.Rules {
// service.Spec.Rules = append(service.Spec.Rules, v1alpha1.Rule{
// ID: rule.ID,
// Level: rule.Level,
// Name: rule.Name,
// Type: rule.Type,
// Description: rule.Description,
// Expr: rule.Expr,
// Mode: rule.Mode,
// })
// }
// }
rules, err := s.getRulesForService(req)
if err != nil {
return nil, fmt.Errorf("failed to get rules for service: %v", err)
}
service.Spec.Rules = rules
if len(service.Spec.Rules) == 0 {
return nil, fmt.Errorf("cannot create WAF service with no rules")
}
// Create the WAF service in Kubernetes
client := s.clusterClientManager.GetClient(req.RegionCode)
if client == nil {
return nil, fmt.Errorf("failed to get cluster client: %v", err)
return nil, fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
}
if _, err := client.WafV1alpha1().Services(req.Namespace).Create(ctx, service, metav1.CreateOptions{}); err != nil {
return nil, fmt.Errorf("failed to create WAF service: %v", err)
}
return nil, nil
return &WafService{
GatewayName: req.GatewayName,
Mode: service.Spec.Mode,
RuleNum: len(service.Spec.Rules),
AttackNum: 0,
}, nil
}
func (s *wafService) DeleteListenerWaf(ctx context.Context, req *DeleteListenerReq) error {
client := s.clusterClientManager.GetClient(req.RegionCode)
if client == nil {
return fmt.Errorf("failed to get cluster client")
return fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
}
name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)
if err := client.WafV1alpha1().Services(req.Namespace).Delete(ctx, name, metav1.DeleteOptions{}); err != nil {
......@@ -434,24 +383,6 @@ func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerW
return err
}
// wafService := &model.WafService{}
// err = s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
// if err != nil {
// if err == gorm.ErrRecordNotFound {
// wafService = &model.WafService{
// GatewayName: req.GatewayName,
// Namespace: req.Namespace,
// RegionCode: req.RegionCode,
// Mode: string(WafModeAlert),
// }
// if err := s.db.Create(wafService).Error; err != nil {
// return err
// }
// } else {
// return err
// }
// }
if listener.Enable {
log.Info().Msgf("Create WAF for listener %s", listener.GatewayName)
_, err := s.CreateWaf(ctx, &CreateWafReq{
......
// Package id provides unique ID generation using the Sonyflake algorithm
package id
import (
"fmt"
"log"
"math/rand"
"strconv"
......@@ -11,32 +13,44 @@ import (
var sf *sonyflake.Sonyflake
const max = 1<<16 - 1
const maxMachineID = 1<<16 - 1
func init() {
// 有一定的几率会产生相同的machineID
// 比如两个引用这个包的实例,在同一个微秒被创建...
rand.Seed(time.Now().UnixMicro())
// Initialize random source for machine ID generation
// Note: There is a small risk of machineID collisions if multiple
// instances are started within the same microsecond
r := rand.New(rand.NewSource(time.Now().UnixMicro()))
st := sonyflake.Settings{
StartTime: time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC),
MachineID: func() (uint16, error) {
return uint16(rand.Intn(max)), nil
return uint16(r.Intn(maxMachineID)), nil
},
CheckMachineID: nil,
}
sf = sonyflake.NewSonyflake(st)
if sf == nil {
log.Panicf("sonyflake init fails")
log.Panicf("failed to initialize sonyflake")
}
}
// Str returns a string representation of a unique Sonyflake ID.
// It panics if ID generation fails.
func Str() string {
ui64id, _ := sf.NextID()
return strconv.FormatUint(ui64id, 10)
id, err := sf.NextID()
if err != nil {
panic(fmt.Sprintf("failed to generate sonyflake ID: %v", err))
}
return strconv.FormatUint(id, 10)
}
// UInt64 returns a uint64 unique Sonyflake ID.
// It panics if ID generation fails.
func UInt64() uint64 {
ui64id, _ := sf.NextID()
return ui64id
id, err := sf.NextID()
if err != nil {
panic(fmt.Sprintf("failed to generate sonyflake ID: %v", err))
}
return id
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment