github.com/juju/juju@v0.0.0-20240327075706-a90865de2538/core/quota/kv.go (about)

     1  // Copyright 2020 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package quota
     5  
     6  import (
     7  	"reflect"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/mgo/v3/bson"
    11  )
    12  
    13  var _ Checker = (*MapKeyValueSizeChecker)(nil)
    14  
    15  // A MapKeyValueSizeChecker can be used to verify that none of the keys and
    16  // values in a map exceed a particular limit.
    17  type MapKeyValueSizeChecker struct {
    18  	maxKeySize   int
    19  	maxValueSize int
    20  	lastErr      error
    21  }
    22  
    23  // NewMapKeyValueSizeChecker returns a new MapKeyValueSizeChecker instance that
    24  // limits map keys to maxKeySize and map values to maxValueSize. Any of the
    25  // max values may also set to 0 to disable quota checks.
    26  func NewMapKeyValueSizeChecker(maxKeySize, maxValueSize int) *MapKeyValueSizeChecker {
    27  	return &MapKeyValueSizeChecker{
    28  		maxKeySize:   maxKeySize,
    29  		maxValueSize: maxValueSize,
    30  	}
    31  }
    32  
    33  // Check applies the configured size checks to v and updates the checker's
    34  // internal state. Check expects a map as an argument where both the keys and
    35  // the values can be serialized to BSON; any other value will cause an error
    36  // to be returned when Outcome is called.
    37  func (c *MapKeyValueSizeChecker) Check(v interface{}) {
    38  	if v == nil || c.lastErr != nil {
    39  		return
    40  	}
    41  
    42  	reflMap := reflect.ValueOf(v)
    43  	if reflMap.Kind() != reflect.Map {
    44  		c.lastErr = errors.NotImplementedf("key/value size check for non map-values")
    45  		return
    46  	}
    47  
    48  	for _, mapKey := range reflMap.MapKeys() {
    49  		mapVal := reflMap.MapIndex(mapKey)
    50  		if err := CheckTupleSize(mapKey.Interface(), mapVal.Interface(), c.maxKeySize, c.maxValueSize); err != nil {
    51  			c.lastErr = err
    52  			return
    53  		}
    54  	}
    55  }
    56  
    57  // Outcome returns the check outcome or whether an error occurred within a call
    58  // to the Check method.
    59  func (c *MapKeyValueSizeChecker) Outcome() error {
    60  	return c.lastErr
    61  }
    62  
    63  // CheckTupleSize checks whether the length of the provided key-value pair is
    64  // within the provided limits. If the key or value is a string, then its length
    65  // will be used for comparison purposes. Otherwise, the effective length is
    66  // calculated by serializing to BSON and counting the length of the serialized
    67  // data.
    68  //
    69  // Any of the max values can be set to zero to bypass the size check.
    70  func CheckTupleSize(key, value interface{}, maxKeyLen, maxValueLen int) error {
    71  	size, err := effectiveSize(key)
    72  	if err != nil {
    73  		return err
    74  	} else if maxKeyLen > 0 && size > maxKeyLen {
    75  		return errors.QuotaLimitExceededf("max allowed key length (%d) exceeded", maxKeyLen)
    76  	}
    77  
    78  	size, err = effectiveSize(value)
    79  	if err != nil {
    80  		return err
    81  	} else if maxValueLen > 0 && size > maxValueLen {
    82  		return errors.QuotaLimitExceededf("max allowed value length (%d) exceeded", maxValueLen)
    83  	}
    84  
    85  	return nil
    86  }
    87  
    88  func effectiveSize(v interface{}) (int, error) {
    89  	switch rawValue := v.(type) {
    90  	case string:
    91  		return len(rawValue), nil
    92  	default: // marshal non-string values to bson and return the serialized length
    93  		d, err := bson.Marshal(rawValue)
    94  		if err != nil {
    95  			return -1, errors.Annotatef(err, "marshaling value to BSON")
    96  		}
    97  		return len(d), nil
    98  	}
    99  }