Переглянути джерело

feature: Rate limit for actions (#76) (#275)

James Read 2 роки тому
батько
коміт
30e2aa141b

+ 7 - 0
internal/config/config.go

@@ -19,6 +19,7 @@ type Action struct {
 	ExecOnCalendarFile     string
 	ExecOnCalendarFile     string
 	Trigger                string
 	Trigger                string
 	MaxConcurrent          int
 	MaxConcurrent          int
+	MaxRate                []RateSpec
 	Arguments              []ActionArgument
 	Arguments              []ActionArgument
 	PopupOnStart           string
 	PopupOnStart           string
 }
 }
@@ -41,6 +42,12 @@ type ActionArgumentChoice struct {
 	Title string
 	Title string
 }
 }
 
 
+// RateSpec allows you to set a max frequency for an action.
+type RateSpec struct {
+	Limit    int
+	Duration string
+}
+
 // Entity represents a "thing" that can have multiple actions associated with it.
 // Entity represents a "thing" that can have multiple actions associated with it.
 // for example, a media player with a start and stop action.
 // for example, a media player with a start and stop action.
 type EntityFile struct {
 type EntityFile struct {

+ 67 - 7
internal/executor/executor.go

@@ -30,7 +30,8 @@ var (
 // Executor represents a helper class for executing commands. It's main method
 // Executor represents a helper class for executing commands. It's main method
 // is ExecRequest
 // is ExecRequest
 type Executor struct {
 type Executor struct {
-	Logs map[string]*InternalLogEntry
+	Logs           map[string]*InternalLogEntry
+	LogsByActionId map[string][]*InternalLogEntry
 
 
 	listeners []listener
 	listeners []listener
 
 
@@ -58,8 +59,8 @@ type ExecutionRequest struct {
 // state of execution (even if the command is not executed). It's designed to be
 // state of execution (even if the command is not executed). It's designed to be
 // easily serializable.
 // easily serializable.
 type InternalLogEntry struct {
 type InternalLogEntry struct {
-	DatetimeStarted     string
-	DatetimeFinished    string
+	DatetimeStarted     time.Time
+	DatetimeFinished    time.Time
 	Stdout              string
 	Stdout              string
 	Stderr              string
 	Stderr              string
 	StdoutBuffer        io.ReadCloser
 	StdoutBuffer        io.ReadCloser
@@ -89,10 +90,12 @@ type executorStepFunc func(*ExecutionRequest) bool
 func DefaultExecutor() *Executor {
 func DefaultExecutor() *Executor {
 	e := Executor{}
 	e := Executor{}
 	e.Logs = make(map[string]*InternalLogEntry)
 	e.Logs = make(map[string]*InternalLogEntry)
+	e.LogsByActionId = make(map[string][]*InternalLogEntry)
 
 
 	e.chainOfCommand = []executorStepFunc{
 	e.chainOfCommand = []executorStepFunc{
 		stepRequestAction,
 		stepRequestAction,
 		stepConcurrencyCheck,
 		stepConcurrencyCheck,
+		stepRateCheck,
 		stepACLCheck,
 		stepACLCheck,
 		stepParseArgs,
 		stepParseArgs,
 		stepLogStart,
 		stepLogStart,
@@ -123,7 +126,7 @@ func (e *Executor) ExecRequest(req *ExecutionRequest) (*sync.WaitGroup, string)
 	// duplicate UUIDs (or just random strings), but this is the only way.
 	// duplicate UUIDs (or just random strings), but this is the only way.
 
 
 	req.logEntry = &InternalLogEntry{
 	req.logEntry = &InternalLogEntry{
-		DatetimeStarted:     time.Now().Format("2006-01-02 15:04:05"),
+		DatetimeStarted:     time.Now(),
 		ExecutionTrackingID: req.TrackingID,
 		ExecutionTrackingID: req.TrackingID,
 		Stdout:              "",
 		Stdout:              "",
 		Stderr:              "",
 		Stderr:              "",
@@ -171,8 +174,8 @@ func (e *Executor) execChain(req *ExecutionRequest) {
 func getConcurrentCount(req *ExecutionRequest) int {
 func getConcurrentCount(req *ExecutionRequest) int {
 	concurrentCount := 0
 	concurrentCount := 0
 
 
-	for _, log := range req.executor.Logs {
-		if log.ActionId == req.Action.ID && !log.ExecutionFinished {
+	for _, log := range req.executor.LogsByActionId[req.Action.ID] {
+		if !log.ExecutionFinished {
 			concurrentCount += 1
 			concurrentCount += 1
 		}
 		}
 	}
 	}
@@ -199,6 +202,57 @@ func stepConcurrencyCheck(req *ExecutionRequest) bool {
 	return true
 	return true
 }
 }
 
 
+func parseDuration(rate config.RateSpec) time.Duration {
+	duration, err := time.ParseDuration(rate.Duration)
+
+	if err != nil {
+		log.Warnf("Could not parse duration: %v", rate.Duration)
+
+		return -1 * time.Minute
+	}
+
+	return duration
+}
+
+func getExecutionsCount(rate config.RateSpec, req *ExecutionRequest) int {
+	executions := -1 // Because we will find ourself when checking execution logs
+
+	duration := parseDuration(rate)
+
+	then := time.Now().Add(-duration)
+
+	for _, logEntry := range req.executor.LogsByActionId[req.Action.ID] {
+		log.Debugf("Rate check %v", logEntry)
+
+		if logEntry.DatetimeStarted.After(then) && !logEntry.Blocked {
+
+			executions += 1
+		}
+	}
+
+	return executions
+}
+
+func stepRateCheck(req *ExecutionRequest) bool {
+	for _, rate := range req.Action.MaxRate {
+		executions := getExecutionsCount(rate, req)
+
+		if executions >= rate.Limit {
+			msg := fmt.Sprintf("Blocked from executing. This action has run %d out of %d allowed times in the last %s.", executions, rate.Limit, rate.Duration)
+
+			log.WithFields(log.Fields{
+				"actionTitle": req.logEntry.ActionTitle,
+			}).Infof(msg)
+
+			req.logEntry.Stdout = msg
+			req.logEntry.Blocked = true
+			return false
+		}
+	}
+
+	return true
+}
+
 func stepACLCheck(req *ExecutionRequest) bool {
 func stepACLCheck(req *ExecutionRequest) bool {
 	return acl.IsAllowedExec(req.Cfg, req.AuthenticatedUser, req.Action)
 	return acl.IsAllowedExec(req.Cfg, req.AuthenticatedUser, req.Action)
 }
 }
@@ -245,6 +299,12 @@ func stepRequestAction(req *ExecutionRequest) bool {
 	req.logEntry.ActionIcon = req.Action.Icon
 	req.logEntry.ActionIcon = req.Action.Icon
 	req.logEntry.ActionId = req.Action.ID
 	req.logEntry.ActionId = req.Action.ID
 
 
+	if _, containsKey := req.executor.LogsByActionId[req.Action.ID]; !containsKey {
+		req.executor.LogsByActionId[req.Action.ID] = make([]*InternalLogEntry, 0)
+	}
+
+	req.executor.LogsByActionId[req.Action.ID] = append(req.executor.LogsByActionId[req.Action.ID], req.logEntry)
+
 	log.WithFields(log.Fields{
 	log.WithFields(log.Fields{
 		"actionTitle": req.logEntry.ActionTitle,
 		"actionTitle": req.logEntry.ActionTitle,
 	}).Infof("Action requested")
 	}).Infof("Action requested")
@@ -324,7 +384,7 @@ func stepExec(req *ExecutionRequest) bool {
 	}
 	}
 
 
 	req.logEntry.Tags = req.Tags
 	req.logEntry.Tags = req.Tags
-	req.logEntry.DatetimeFinished = time.Now().Format("2006-01-02 15:04:05")
+	req.logEntry.DatetimeFinished = time.Now()
 
 
 	return true
 	return true
 }
 }

+ 2 - 2
internal/grpcapi/grpcApi.go

@@ -128,8 +128,8 @@ func internalLogEntryToPb(logEntry *executor.InternalLogEntry) *pb.LogEntry {
 		ActionTitle:         logEntry.ActionTitle,
 		ActionTitle:         logEntry.ActionTitle,
 		ActionIcon:          logEntry.ActionIcon,
 		ActionIcon:          logEntry.ActionIcon,
 		ActionId:            logEntry.ActionId,
 		ActionId:            logEntry.ActionId,
-		DatetimeStarted:     logEntry.DatetimeStarted,
-		DatetimeFinished:    logEntry.DatetimeFinished,
+		DatetimeStarted:     logEntry.DatetimeStarted.Format("2006-01-02 15:04:05"),
+		DatetimeFinished:    logEntry.DatetimeFinished.Format("2006-01-02 15:04:05"),
 		Stdout:              logEntry.Stdout,
 		Stdout:              logEntry.Stdout,
 		Stderr:              logEntry.Stderr,
 		Stderr:              logEntry.Stderr,
 		TimedOut:            logEntry.TimedOut,
 		TimedOut:            logEntry.TimedOut,

+ 2 - 2
internal/websocket/websocket.go

@@ -59,8 +59,8 @@ func (WebsocketExecutionListener) OnExecutionFinished(logEntry *executor.Interna
 		ActionTitle:         logEntry.ActionTitle,
 		ActionTitle:         logEntry.ActionTitle,
 		ActionIcon:          logEntry.ActionIcon,
 		ActionIcon:          logEntry.ActionIcon,
 		ActionId:            logEntry.ActionId,
 		ActionId:            logEntry.ActionId,
-		DatetimeStarted:     logEntry.DatetimeStarted,
-		DatetimeFinished:    logEntry.DatetimeFinished,
+		DatetimeStarted:     logEntry.DatetimeStarted.Format("2006-01-02 15:04:05"),
+		DatetimeFinished:    logEntry.DatetimeFinished.Format("2006-01-02 15:04:05"),
 		Stdout:              logEntry.Stdout,
 		Stdout:              logEntry.Stdout,
 		Stderr:              logEntry.Stderr,
 		Stderr:              logEntry.Stderr,
 		TimedOut:            logEntry.TimedOut,
 		TimedOut:            logEntry.TimedOut,