bitbucket.org/Aishee/synsec@v0.0.0-20210414005726-236fc01a153d/pkg/leakybucket/uniq.go (about)

     1  package leakybucket
     2  
     3  import (
     4  	"github.com/antonmedv/expr"
     5  	"github.com/antonmedv/expr/vm"
     6  
     7  	"bitbucket.org/Aishee/synsec/pkg/exprhelpers"
     8  	"bitbucket.org/Aishee/synsec/pkg/types"
     9  )
    10  
    11  // Uniq creates three new functions that share the same initialisation and the same scope.
    12  // They are triggered respectively:
    13  // on pour
    14  // on overflow
    15  // on leak
    16  
    17  type Uniq struct {
    18  	DistinctCompiled *vm.Program
    19  }
    20  
    21  func (u *Uniq) OnBucketPour(bucketFactory *BucketFactory) func(types.Event, *Leaky) *types.Event {
    22  	return func(msg types.Event, leaky *Leaky) *types.Event {
    23  		element, err := getElement(msg, u.DistinctCompiled)
    24  		if err != nil {
    25  			leaky.logger.Errorf("Uniq filter exec failed : %v", err)
    26  			return &msg
    27  		}
    28  		leaky.logger.Tracef("Uniq '%s' -> '%s'", bucketFactory.Distinct, element)
    29  		for _, evt := range leaky.Queue.GetQueue() {
    30  			if val, err := getElement(evt, u.DistinctCompiled); err == nil && val == element {
    31  				leaky.logger.Debugf("Uniq(%s) : ko, discard event", element)
    32  				return nil
    33  			}
    34  			if err != nil {
    35  				leaky.logger.Errorf("Uniq filter exec failed : %v", err)
    36  			}
    37  		}
    38  		leaky.logger.Debugf("Uniq(%s) : ok", element)
    39  		return &msg
    40  	}
    41  }
    42  
    43  func (u *Uniq) OnBucketOverflow(bucketFactory *BucketFactory) func(*Leaky, types.RuntimeAlert, *Queue) (types.RuntimeAlert, *Queue) {
    44  	return func(leaky *Leaky, alert types.RuntimeAlert, queue *Queue) (types.RuntimeAlert, *Queue) {
    45  		return alert, queue
    46  	}
    47  }
    48  
    49  func (u *Uniq) OnBucketInit(bucketFactory *BucketFactory) error {
    50  	var err error
    51  
    52  	u.DistinctCompiled, err = expr.Compile(bucketFactory.Distinct, expr.Env(exprhelpers.GetExprEnv(map[string]interface{}{"evt": &types.Event{}})))
    53  	return err
    54  }
    55  
    56  // getElement computes a string from an event and a filter
    57  func getElement(msg types.Event, cFilter *vm.Program) (string, error) {
    58  	el, err := expr.Run(cFilter, exprhelpers.GetExprEnv(map[string]interface{}{"evt": &msg}))
    59  	if err != nil {
    60  		return "", err
    61  	}
    62  	element, ok := el.(string)
    63  	if !ok {
    64  		return "", err
    65  	}
    66  	return element, nil
    67  }