github.com/crowdsecurity/crowdsec@v1.6.1/pkg/leakybucket/uniq.go (about)

     1  package leakybucket
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/antonmedv/expr"
     7  	"github.com/antonmedv/expr/vm"
     8  
     9  	"github.com/crowdsecurity/crowdsec/pkg/exprhelpers"
    10  	"github.com/crowdsecurity/crowdsec/pkg/types"
    11  )
    12  
    13  // Uniq creates three new functions that share the same initialisation and the same scope.
    14  // They are triggered respectively:
    15  // on pour
    16  // on overflow
    17  // on leak
    18  
    19  var uniqExprCache map[string]vm.Program
    20  var uniqExprCacheLock sync.Mutex
    21  
    22  type Uniq struct {
    23  	DistinctCompiled *vm.Program
    24  	KeyCache         map[string]bool
    25  	CacheMutex       sync.Mutex
    26  }
    27  
    28  func (u *Uniq) OnBucketPour(bucketFactory *BucketFactory) func(types.Event, *Leaky) *types.Event {
    29  	return func(msg types.Event, leaky *Leaky) *types.Event {
    30  		element, err := getElement(msg, u.DistinctCompiled)
    31  		if err != nil {
    32  			leaky.logger.Errorf("Uniq filter exec failed : %v", err)
    33  			return &msg
    34  		}
    35  		leaky.logger.Tracef("Uniq '%s' -> '%s'", bucketFactory.Distinct, element)
    36  		u.CacheMutex.Lock()
    37  		defer u.CacheMutex.Unlock()
    38  		if _, ok := u.KeyCache[element]; !ok {
    39  			leaky.logger.Debugf("Uniq(%s) : ok", element)
    40  			u.KeyCache[element] = true
    41  			return &msg
    42  
    43  		} else {
    44  			leaky.logger.Debugf("Uniq(%s) : ko, discard event", element)
    45  			return nil
    46  		}
    47  	}
    48  }
    49  
    50  func (u *Uniq) OnBucketOverflow(bucketFactory *BucketFactory) func(*Leaky, types.RuntimeAlert, *types.Queue) (types.RuntimeAlert, *types.Queue) {
    51  	return func(leaky *Leaky, alert types.RuntimeAlert, queue *types.Queue) (types.RuntimeAlert, *types.Queue) {
    52  		return alert, queue
    53  	}
    54  }
    55  
    56  func (u *Uniq) AfterBucketPour(bucketFactory *BucketFactory) func(types.Event, *Leaky) *types.Event {
    57  	return func(msg types.Event, leaky *Leaky) *types.Event {
    58  		return &msg
    59  	}
    60  }
    61  
    62  func (u *Uniq) OnBucketInit(bucketFactory *BucketFactory) error {
    63  	var err error
    64  	var compiledExpr *vm.Program
    65  
    66  	if uniqExprCache == nil {
    67  		uniqExprCache = make(map[string]vm.Program)
    68  	}
    69  
    70  	uniqExprCacheLock.Lock()
    71  	if compiled, ok := uniqExprCache[bucketFactory.Distinct]; ok {
    72  		uniqExprCacheLock.Unlock()
    73  		u.DistinctCompiled = &compiled
    74  	} else {
    75  		uniqExprCacheLock.Unlock()
    76  		//release the lock during compile
    77  		compiledExpr, err = expr.Compile(bucketFactory.Distinct, exprhelpers.GetExprOptions(map[string]interface{}{"evt": &types.Event{}})...)
    78  		u.DistinctCompiled = compiledExpr
    79  		uniqExprCacheLock.Lock()
    80  		uniqExprCache[bucketFactory.Distinct] = *compiledExpr
    81  		uniqExprCacheLock.Unlock()
    82  	}
    83  	u.KeyCache = make(map[string]bool)
    84  	return err
    85  }
    86  
    87  // getElement computes a string from an event and a filter
    88  func getElement(msg types.Event, cFilter *vm.Program) (string, error) {
    89  	el, err := expr.Run(cFilter, map[string]interface{}{"evt": &msg})
    90  	if err != nil {
    91  		return "", err
    92  	}
    93  	element, ok := el.(string)
    94  	if !ok {
    95  		return "", err
    96  	}
    97  	return element, nil
    98  }