github.com/Cloud-Foundations/Dominator@v0.3.4/lib/srpc/serverutil/perUserMethodLimiter.go (about)

     1  package serverutil
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
     7  )
     8  
     9  func newPerUserMethodLimiter(
    10  	perUserMethodLimitsInput map[string]uint) *PerUserMethodLimiter {
    11  	perUserMethodLimits := make(map[string]uint, len(perUserMethodLimitsInput))
    12  	for method, limit := range perUserMethodLimitsInput {
    13  		perUserMethodLimits[method] = limit
    14  	}
    15  	return &PerUserMethodLimiter{
    16  		perUserMethodCounts: make(map[userMethodType]uint,
    17  			len(perUserMethodLimits)),
    18  		perUserMethodLimits: perUserMethodLimits,
    19  	}
    20  }
    21  
    22  func (limiter *PerUserMethodLimiter) blockMethod(methodName string,
    23  	authInfo *srpc.AuthInformation) (func(), error) {
    24  	if authInfo.HaveMethodAccess {
    25  		return nil, nil
    26  	}
    27  	limiter.mutex.Lock()
    28  	defer limiter.mutex.Unlock()
    29  	if limit := limiter.perUserMethodLimits[methodName]; limit < 1 {
    30  		return nil, nil
    31  	} else {
    32  		userMethod := userMethodType{
    33  			method:   methodName,
    34  			username: authInfo.Username,
    35  		}
    36  		if count := limiter.perUserMethodCounts[userMethod]; count >= limit {
    37  			return nil, fmt.Errorf("%s reached limit of %d calls for %s",
    38  				authInfo.Username, limit, methodName)
    39  		} else {
    40  			limiter.perUserMethodCounts[userMethod] = count + 1
    41  			return func() {
    42  				limiter.mutex.Lock()
    43  				defer limiter.mutex.Unlock()
    44  				if count := limiter.perUserMethodCounts[userMethod]; count < 1 {
    45  					panic(fmt.Sprintf("%s has no %s calls to release",
    46  						authInfo.Username, methodName))
    47  				} else {
    48  					if count -= 1; count < 1 {
    49  						delete(limiter.perUserMethodCounts, userMethod)
    50  					} else {
    51  						limiter.perUserMethodCounts[userMethod] = count
    52  					}
    53  				}
    54  			}, nil
    55  		}
    56  	}
    57  }