Add hosts column to gateway_listeners and update related WAF service methods

- Modified SQL migration to add 'hosts' column to gateway_listeners table
- Updated model, service, and controller to support hosts in gateway listener operations
- Added new GetWafGatewayInfo method to retrieve WAF gateway information
- Refactored WAF service methods to handle hosts and improve flexibility

api/waf.go

@@ -12,6 +12,7 @@ func SetWafRouter(e *gin.Engine, clusterClientManager *utils.ClusterClientManage
 
 	wafController := controller.NewWafController(clusterClientManager, db)
 	v1.GET("/:region_code/:namespace/:gateway_name", wafController.Waf)
+	// v1.POST("/gateway", wafController.GetWafGatewayInfo)
 	v1.POST("/", wafController.CreateWaf)
 	v1.PUT("mode", wafController.UpdateMode)
 	v1.PUT("rules", wafController.UpdateRule)

db/migrations/create_waf_tables.sql

@@ -42,7 +42,8 @@ CREATE TABLE gateway_listeners (
     namespace VARCHAR(255) NOT NULL,
     region_code VARCHAR(50) NOT NULL,
     port INTEGER NOT NULL,
-    enable BOOLEAN NOT NULL
+    enable BOOLEAN NOT NULL,
+    hosts TEXT
 );
 
 -- Add indexes for better query performance

internal/controller/waf.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"context"
+	"errors"
 	"strconv"
 	"time"
 
@@ -40,6 +41,30 @@ func (c *WafController) Waf(ctx *gin.Context) {
 	utils.AssembleResponse(ctx, resp, nil)
 }
 
+func (c *WafController) GetWafGatewayInfo(ctx *gin.Context) {
+	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	cookie := ctx.Request.Header.Get("Cookie")
+	if cookie == "" {
+		utils.AssembleResponse(ctx, nil, errors.New("cookie is required"))
+		return
+	}
+
+	var req service.GetWafGatewayInfoReq
+	if err := ctx.BindJSON(&req); err != nil {
+		utils.AssembleResponse(ctx, nil, err)
+		return
+	}
+
+	info, err := c.service.GetWafGatewayInfo(ctx1, &req)
+	if err != nil {
+		utils.AssembleResponse(ctx, nil, err)
+		return
+	}
+	utils.AssembleResponse(ctx, info, nil)
+}
+
 func (c *WafController) CreateWaf(ctx *gin.Context) {
 	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
@@ -148,6 +173,7 @@ func (c *WafController) EnableGatewayWaf(ctx *gin.Context) {
 		utils.AssembleResponse(ctx, nil, err)
 		return
 	}
+	req.Cookie = ctx.Request.Header.Get("Cookie")
 
 	err := c.service.EnableGatewayWaf(ctx1, &req)
 	if err != nil {

internal/model/waf.go

@@ -56,7 +56,7 @@ type WafService struct {
 	RuleNum            int                 `gorm:"column:rule_num"`
 	AttackNum          int                 `gorm:"column:attack_num"`
 	RuleCategoryStatus *RuleCategoryStatus `gorm:"column:rule_category_status;type:json"`
-	Host               HostList            `gorm:"column:host"`
+	// Host               HostList            `gorm:"column:host"`
 }
 
 func (WafService) TableName() string {
@@ -123,6 +123,7 @@ type GatewayListener struct {
 	RegionCode  string   `gorm:"column:region_code"`
 	Port        int      `gorm:"column:port"`
 	Enable      bool     `gorm:"column:enable"`
+	Hosts       HostList `gorm:"column:hosts"`
 }
 
 func (GatewayListener) TableName() string {

internal/service/service.go

@@ -5,6 +5,7 @@ import "context"
 type Service interface {
 	// QueryIP(ip string) (*model.IPInfo, error)
 	GetWaf(ctx context.Context, regionCode, namespace, gatewayName string) (*WafService, error)
+	GetWafGatewayInfo(ctx context.Context, req *GetWafGatewayInfoReq) (*WafService, error)
 	CreateWaf(ctx context.Context, req *CreateWafReq) (*WafService, error)
 	UpdateMode(ctx context.Context, req *UpdateModeReq) (*WafService, error)
 	UpdateRule(ctx context.Context, req *RuleRequest) error

internal/service/types.go

@@ -175,7 +175,12 @@ type GatewateInfo struct {
 	Namespace     string `json:"namespace"`
 	RegionCode    string `json:"region_code"`
 	ApiGatewayCrn string `json:"gateway_crn"`
-	Hosts         []string `json:"hosts"`
+	// Hosts         []string `json:"hosts"`
+}
+
+type GetWafGatewayInfoReq struct {
+	GatewateInfo
+	Cookie string `json:"cookie"`
 }
 
 type CreateWafReq struct {
@@ -213,12 +218,14 @@ type UpdateModeReq struct {
 type EnableListenerWafReq struct {
 	GatewateInfo
 	Enable bool     `json:"enable"`
+	Hosts  []string `json:"hosts"`
 	Port   int      `json:"port"`
 }
 
 type EnableGatewayWafReq struct {
 	GatewateInfo
 	Enable bool   `json:"enable"`
+	Cookie string `json:"cookie"`
 }
 
 type WafRule struct {
@@ -266,3 +273,29 @@ type DeleteListenerReq struct {
 	GatewateInfo
 	Port int `json:"port"`
 }
+
+type GatewayRespListenerData struct {
+	GatewayName       string   `json:"gateway_name"`
+	Namespace         string   `json:"namespace"`
+	ListenerName      string   `json:"listener_name"`
+	ApiGatewayCrn     string   `json:"apigateway_crn"`
+	CreateAccountName string   `json:"create_account_name"`
+	CreateAccountID   string   `json:"create_account_id"`
+	Hosts             []string `json:"hosts"`
+	Port              int      `json:"port"`
+}
+
+type GatewayResponseBase struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+}
+
+type GatewayListenerResponse struct {
+	GatewayResponseBase
+	Data GatewayRespListenerData `json:"data"`
+}
+
+type GatewayListenerResponseList struct {
+	GatewayResponseBase
+	Data []GatewayRespListenerData `json:"data"`
+}

internal/service/waf.go

@@ -5,7 +5,9 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
+	"net/url"
 	"os"
 	"strconv"
 	"strings"
@@ -39,7 +41,6 @@ func (s *wafService) GetWaf(ctx context.Context, regionCode, namespace, gatewayN
 				Namespace:   namespace,
 				GatewayName: gatewayName,
 				Mode:        string(WafModeAlert),
-				Host:        model.HostList([]string{"*"}),
 			}
 			if err := s.db.Create(wafService).Error; err != nil {
 				return nil, fmt.Errorf("failed to create WAF service: %v", err)
@@ -56,6 +57,55 @@ func (s *wafService) GetWaf(ctx context.Context, regionCode, namespace, gatewayN
 	}, nil
 }
 
+func (s *wafService) GetWafGatewayInfo(ctx context.Context, req *GetWafGatewayInfoReq) (*WafService, error) {
+	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 {
+			httpRequst := http.Request{
+				Method: http.MethodPost,
+				URL:    &url.URL{Scheme: "https", Host: "console.tensorsecurity.com", Path: "/api/v1/waf/gateway"},
+				Header: http.Header{
+					"Cookie": []string{req.Cookie},
+				},
+				Body: io.NopCloser(strings.NewReader(fmt.Sprintf(`{"gateway_name": "%s", "namespace": "%s", "region_code": "%s"}`, req.GatewayName, req.Namespace, req.RegionCode))),
+			}
+			resp, err := http.DefaultClient.Do(&httpRequst)
+			if err != nil {
+				return nil, fmt.Errorf("failed to get WAF service: %v", err)
+			}
+			defer resp.Body.Close()
+			body, err := io.ReadAll(resp.Body)
+			if err != nil {
+				return nil, fmt.Errorf("failed to read WAF service: %v", err)
+			}
+			var wafService model.WafService
+			err = json.Unmarshal(body, &wafService)
+			if err != nil {
+				return nil, fmt.Errorf("failed to unmarshal WAF service: %v", err)
+			}
+			wafService.ID = 0
+			wafService.RuleCategoryStatus = nil
+			wafService.RuleNum = 0
+			wafService.AttackNum = 0
+			// wafService.Host = model.HostList([]string{"*"})
+			wafService.Mode = string(WafModeAlert)
+			err = s.db.Create(wafService).Error
+			if 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)
+		}
+	}
+	return &WafService{
+		GatewayName: wafService.GatewayName,
+		Mode:        wafService.Mode,
+		RuleNum:     wafService.RuleNum,
+		AttackNum:   wafService.AttackNum,
+	}, nil
+}
+
 func (s *wafService) getRulesForService(req *CreateWafReq) ([]v1alpha1.Rule, error) {
 	rules := []v1alpha1.Rule{}
 	ruleCategories := []model.WafRuleCategory{}
@@ -74,7 +124,6 @@ func (s *wafService) getRulesForService(req *CreateWafReq) ([]v1alpha1.Rule, err
 				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)
@@ -151,64 +200,70 @@ 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
+	// 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 {
-		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,
-			})
-		}
+		return nil, fmt.Errorf("failed to get rules for service: %v", err)
 	}
+	service.Spec.Rules = rules
 
 	// Create the WAF service in Kubernetes
 	client := s.clusterClientManager.GetClient(req.RegionCode)
@@ -389,6 +444,7 @@ func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerW
 				RegionCode:  req.RegionCode,
 				Port:        int(req.Port),
 				Enable:      req.Enable,
+				Hosts:       req.Hosts,
 			}
 			err = s.db.Model(&model.GatewayListener{}).Create(listener).Error
 			if err != nil {
@@ -405,24 +461,23 @@ 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),
-				Host:        model.HostList([]string{"*"}),
-			}
-			if err := s.db.Create(wafService).Error; err != nil {
-				return err
-			}
-		} else {
-			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)
@@ -433,7 +488,7 @@ func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerW
 				RegionCode:  req.RegionCode,
 			},
 			Port: uint32(req.Port),
-			Host: wafService.Host,
+			Host: req.Hosts,
 		})
 		if err != nil {
 			return err
@@ -452,36 +507,49 @@ func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerW
 			return err
 		}
 	}
-
 	return nil
 }
 
-func (s *wafService) EnableGatewayWaf(ctx context.Context, req *EnableGatewayWafReq) error {
-	if req.Enable {
-		// Get listener list from API
+func getGatewayNameFromCrn(crn string) string {
+	// crn:ucs::apigateway:lf-tst7:214613666997:instance/testaaa
+	parts := strings.Split(crn, "/")
+	return parts[len(parts)-1]
+}
+
+func (s *wafService) ListListenerFromApiGateway(ctx context.Context, apiGatewayCrn string, regionCode string, cookie string) ([]GatewayRespListenerData, error) {
 	body, err := json.Marshal(map[string]string{
-			"apigateway_crn": req.ApiGatewayCrn,
-			"region_code":    req.RegionCode,
+		"apigateway_crn": apiGatewayCrn,
+		"region_code":    regionCode,
 	})
 	if err != nil {
-			return fmt.Errorf("failed to marshal request body: %v", err)
+		return nil, fmt.Errorf("failed to marshal request body: %v", err)
 	}
-		resp, err := http.Post("https://csm.console.test.tg.unicom.local/apigatewaymng/listener/lf-tst7/list_listeners", "application/json", bytes.NewBuffer(body))
+	request, err := http.NewRequest("POST", "https://csm.console.test.tg.unicom.local/apigatewaymng/listener/lf-tst7/list_listeners", bytes.NewBuffer(body))
 	if err != nil {
-			return fmt.Errorf("failed to get listener list: %v", err)
+		return nil, fmt.Errorf("failed to create request: %v", err)
+	}
+	request.Header.Set("Cookie", cookie)
+	resp, err := http.DefaultClient.Do(request)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get listener list: %v", err)
 	}
 	defer resp.Body.Close()
 
 	// Parse response
-		var listeners []struct {
-			ApigatewayCrn string   `json:"apigateway_crn"`
-			Port          uint32   `json:"port"`
-			Hosts         []string `json:"hosts"`
-		}
-		if err := json.NewDecoder(resp.Body).Decode(&listeners); err != nil {
-			return fmt.Errorf("failed to parse listener list: %v", err)
+	var response GatewayListenerResponseList
+
+	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+		return nil, fmt.Errorf("failed to parse listener list: %v", err)
 	}
+	return response.Data, nil
+}
 
+func (s *wafService) EnableGatewayWaf(ctx context.Context, req *EnableGatewayWafReq) error {
+	if req.Enable {
+		listeners, err := s.ListListenerFromApiGateway(ctx, req.ApiGatewayCrn, req.RegionCode, req.Cookie)
+		if err != nil {
+			return fmt.Errorf("failed to get listener list: %v", err)
+		}
 		// Create WAF for each listener
 		for _, listener := range listeners {
 			if _, err := s.CreateWaf(ctx, &CreateWafReq{