github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/action_notify.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package waf
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
     7  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     8  	"github.com/TeaOSLab/EdgeNode/internal/events"
     9  	"github.com/TeaOSLab/EdgeNode/internal/goman"
    10  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
    11  	"github.com/TeaOSLab/EdgeNode/internal/rpc"
    12  	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
    13  	"github.com/iwind/TeaGo/types"
    14  	"net/http"
    15  	"time"
    16  )
    17  
    18  type notifyTask struct {
    19  	ServerId                int64
    20  	HttpFirewallPolicyId    int64
    21  	HttpFirewallRuleGroupId int64
    22  	HttpFirewallRuleSetId   int64
    23  	CreatedAt               int64
    24  }
    25  
    26  var notifyChan = make(chan *notifyTask, 128)
    27  
    28  func init() {
    29  	if !teaconst.IsMain {
    30  		return
    31  	}
    32  
    33  	events.On(events.EventLoaded, func() {
    34  		goman.New(func() {
    35  			rpcClient, err := rpc.SharedRPC()
    36  			if err != nil {
    37  				remotelogs.Error("WAF_NOTIFY_ACTION", "create rpc client failed: "+err.Error())
    38  				return
    39  			}
    40  
    41  			for task := range notifyChan {
    42  				_, err = rpcClient.FirewallRPC.NotifyHTTPFirewallEvent(rpcClient.Context(), &pb.NotifyHTTPFirewallEventRequest{
    43  					ServerId:                task.ServerId,
    44  					HttpFirewallPolicyId:    task.HttpFirewallPolicyId,
    45  					HttpFirewallRuleGroupId: task.HttpFirewallRuleGroupId,
    46  					HttpFirewallRuleSetId:   task.HttpFirewallRuleSetId,
    47  					CreatedAt:               task.CreatedAt,
    48  				})
    49  				if err != nil {
    50  					remotelogs.Error("WAF_NOTIFY_ACTION", "notify failed: "+err.Error())
    51  				}
    52  			}
    53  		})
    54  	})
    55  }
    56  
    57  type NotifyAction struct {
    58  	BaseAction
    59  }
    60  
    61  func (this *NotifyAction) Init(waf *WAF) error {
    62  	return nil
    63  }
    64  
    65  func (this *NotifyAction) Code() string {
    66  	return ActionNotify
    67  }
    68  
    69  func (this *NotifyAction) IsAttack() bool {
    70  	return false
    71  }
    72  
    73  // WillChange determine if the action will change the request
    74  func (this *NotifyAction) WillChange() bool {
    75  	return false
    76  }
    77  
    78  // Perform the action
    79  func (this *NotifyAction) Perform(waf *WAF, group *RuleGroup, set *RuleSet, request requests.Request, writer http.ResponseWriter) PerformResult {
    80  	select {
    81  	case notifyChan <- &notifyTask{
    82  		ServerId:                request.WAFServerId(),
    83  		HttpFirewallPolicyId:    types.Int64(waf.Id),
    84  		HttpFirewallRuleGroupId: types.Int64(group.Id),
    85  		HttpFirewallRuleSetId:   types.Int64(set.Id),
    86  		CreatedAt:               time.Now().Unix(),
    87  	}:
    88  	default:
    89  
    90  	}
    91  
    92  	return PerformResult{
    93  		ContinueRequest: true,
    94  	}
    95  }