commit

internal/service/types.go

@@ -183,8 +183,16 @@ type RuleRequest struct {
 	Status     int      `json:"status"`
 }
 
+type WafMode string
+
+const (
+	WafModeProtect     WafMode = "protect"
+	WafModeAlert       WafMode = "alert"
+	WafModePassthrough WafMode = "passthrough"
+)
+
 type UpdateModeReq struct {
-	Mode        string `json:"mode"`
+	Mode        WafMode `json:"mode"`
 	GatewayName string  `json:"gateway_name"`
 	Namespace   string  `json:"namespace"`
 	RegionCode  string  `json:"region_code"`

internal/service/waf.go

@@ -69,8 +69,19 @@ func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*Waf, er
 	// Get existing WAF service config if any
 	wafService := &model.WafService{}
 	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
-	if err != nil && err != gorm.ErrRecordNotFound {
-		return nil, fmt.Errorf("failed to get WAF service: %v", err)
+	if err != nil {
+		if err == gorm.ErrRecordNotFound {
+			// Create new WAF service record if not found
+			wafService = &model.WafService{
+				GatewayName: req.GatewayName,
+				Mode:        string(WafModeAlert),
+			}
+			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
@@ -112,6 +123,28 @@ func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*Waf, er
 }
 
 func (s *wafService) UpdateMode(ctx context.Context, req *UpdateModeReq) (*Waf, error) {
+	// Check if WAF service exists
+	wafService := &model.WafService{}
+	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
+	if err != nil {
+		if err == gorm.ErrRecordNotFound {
+			// Create new WAF service record if not found
+			wafService = &model.WafService{
+				GatewayName: req.GatewayName,
+				Mode:        string(req.Mode),
+			}
+			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)
+		}
+	} else {
+		// Update mode if service exists
+		if err := s.db.Model(wafService).Update("mode", string(req.Mode)).Error; err != nil {
+			return nil, fmt.Errorf("failed to update WAF service mode: %v", err)
+		}
+	}
 	return nil, nil
 }