github.com/crowdsecurity/crowdsec@v1.6.1/pkg/csplugin/watcher.go (about)

     1  package csplugin
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	log "github.com/sirupsen/logrus"
     8  	"gopkg.in/tomb.v2"
     9  
    10  	"github.com/crowdsecurity/crowdsec/pkg/models"
    11  )
    12  
    13  /*
    14   PluginWatcher is here to allow grouping and threshold features for notification plugins :
    15   by frequency : it will signal the plugin to deliver notifications at this frequency (watchPluginTicker)
    16   by threshold : it will signal the plugin to deliver notifications when the number of alerts for this plugin reaches this threshold (watchPluginAlertCounts)
    17  */
    18  
    19  // TODO: When we start using go 1.18, consider moving this struct in some utils pkg. Make the implementation more generic using generics :)
    20  type alertCounterByPluginName struct {
    21  	sync.Mutex
    22  	data map[string]int
    23  }
    24  
    25  func newAlertCounterByPluginName() alertCounterByPluginName {
    26  	return alertCounterByPluginName{
    27  		data: make(map[string]int),
    28  	}
    29  }
    30  
    31  func (acp *alertCounterByPluginName) Init() {
    32  	acp.data = make(map[string]int)
    33  }
    34  
    35  func (acp *alertCounterByPluginName) Get(key string) (int, bool) {
    36  	acp.Lock()
    37  	val, ok := acp.data[key]
    38  	acp.Unlock()
    39  	return val, ok
    40  }
    41  
    42  func (acp *alertCounterByPluginName) Set(key string, val int) {
    43  	acp.Lock()
    44  	acp.data[key] = val
    45  	acp.Unlock()
    46  }
    47  
    48  type PluginWatcher struct {
    49  	PluginConfigByName     map[string]PluginConfig
    50  	AlertCountByPluginName alertCounterByPluginName
    51  	PluginEvents           chan string
    52  	Inserts                chan string
    53  	tomb                   *tomb.Tomb
    54  }
    55  
    56  var DefaultEmptyTicker = time.Second * 1
    57  
    58  func (pw *PluginWatcher) Init(configs map[string]PluginConfig, alertsByPluginName map[string][]*models.Alert) {
    59  	pw.PluginConfigByName = configs
    60  	pw.PluginEvents = make(chan string)
    61  	pw.AlertCountByPluginName = newAlertCounterByPluginName()
    62  	pw.Inserts = make(chan string)
    63  	for name := range alertsByPluginName {
    64  		pw.AlertCountByPluginName.Set(name, 0)
    65  	}
    66  }
    67  
    68  func (pw *PluginWatcher) Start(tomb *tomb.Tomb) {
    69  	pw.tomb = tomb
    70  	for name := range pw.PluginConfigByName {
    71  		pname := name
    72  		pw.tomb.Go(func() error {
    73  			pw.watchPluginTicker(pname)
    74  			return nil
    75  		})
    76  	}
    77  
    78  	pw.tomb.Go(func() error {
    79  		pw.watchPluginAlertCounts()
    80  		return nil
    81  	})
    82  }
    83  
    84  func (pw *PluginWatcher) watchPluginTicker(pluginName string) {
    85  	var watchTime time.Duration
    86  	watchCount := -1
    87  	// Threshold can be set : by time, by count, or both
    88  	// if only time is set, honor it
    89  	// if only count is set, put timer to 1 second and just check size
    90  	// if both are set, set timer to 1 second, but check size && time
    91  	interval := pw.PluginConfigByName[pluginName].GroupWait
    92  	threshold := pw.PluginConfigByName[pluginName].GroupThreshold
    93  
    94  	//only size is set
    95  	if threshold > 0 && interval == 0 {
    96  		watchCount = threshold
    97  		watchTime = DefaultEmptyTicker
    98  	} else if interval != 0 && threshold == 0 {
    99  		//only time is set
   100  		watchTime = interval
   101  	} else if interval != 0 && threshold != 0 {
   102  		//both are set
   103  		watchTime = DefaultEmptyTicker
   104  		watchCount = threshold
   105  	} else {
   106  		//none are set, we sent every event we receive
   107  		watchTime = DefaultEmptyTicker
   108  		watchCount = 1
   109  	}
   110  
   111  	ticker := time.NewTicker(watchTime)
   112  	lastSend := time.Now()
   113  	for {
   114  		select {
   115  		case <-ticker.C:
   116  			send := false
   117  			//if count threshold was set, honor no matter what
   118  			if pc, _ := pw.AlertCountByPluginName.Get(pluginName); watchCount > 0 && pc >= watchCount {
   119  				log.Tracef("[%s] %d alerts received, sending\n", pluginName, pc)
   120  				send = true
   121  				pw.AlertCountByPluginName.Set(pluginName, 0)
   122  			}
   123  			//if time threshold only was set
   124  			if watchTime > 0 && watchTime == interval {
   125  				log.Tracef("sending alerts to %s, duration %s elapsed", pluginName, interval)
   126  				send = true
   127  			}
   128  
   129  			//if we hit timer because it was set low to honor count, check if we should trigger
   130  			if watchTime == DefaultEmptyTicker && watchTime != interval && interval != 0 {
   131  				if lastSend.Add(interval).Before(time.Now()) {
   132  					log.Tracef("sending alerts to %s, duration %s elapsed", pluginName, interval)
   133  					send = true
   134  					lastSend = time.Now()
   135  				}
   136  			}
   137  			if send {
   138  				log.Tracef("sending alerts to %s", pluginName)
   139  				pw.PluginEvents <- pluginName
   140  			}
   141  		case <-pw.tomb.Dying():
   142  			ticker.Stop()
   143  			// emptying
   144  			// no lock here because we have the broker still listening even in dying state before killing us
   145  			pw.PluginEvents <- pluginName
   146  			return
   147  		}
   148  	}
   149  }
   150  
   151  func (pw *PluginWatcher) watchPluginAlertCounts() {
   152  	for {
   153  		select {
   154  		case pluginName := <-pw.Inserts:
   155  			//we only "count" pending alerts, and watchPluginTicker is actually going to send it
   156  			if _, ok := pw.PluginConfigByName[pluginName]; ok {
   157  				curr, _ := pw.AlertCountByPluginName.Get(pluginName)
   158  				pw.AlertCountByPluginName.Set(pluginName, curr+1)
   159  			}
   160  		case <-pw.tomb.Dying():
   161  			return
   162  		}
   163  	}
   164  }