Add timeouts scaling, add queue size metrics

internal/exploit/metrics.go

@@ -9,12 +9,14 @@ import (
 type Metrics struct {
 	FlagsSubmitted *prometheus.CounterVec
 	Teams          prometheus.Gauge
+	Queue          *prometheus.GaugeVec
 }
 
 func NewMetrics(namespace string) *Metrics {
 	const subsystem = "exploit_runner"
 	targetLabels := []string{"target_id", "target_ip"}
 	exploitLabels := []string{"exploit_id", "exploit_version", "exploit_type"}
+	queueLabels := []string{"type"}
 
 	return &Metrics{
 		FlagsSubmitted: promauto.NewCounterVec(
@@ -33,5 +35,15 @@ func NewMetrics(namespace string) *Metrics {
 			Name:      "teams",
 			Help:      "Number of teams scheduled for the current runner",
 		}),
+
+		Queue: promauto.NewGaugeVec(
+			prometheus.GaugeOpts{
+				Namespace: namespace,
+				Subsystem: subsystem,
+				Name:      "queue",
+				Help:      "Number of exploits in the queue",
+			},
+			queueLabels,
+		),
 	}
 }

internal/exploit/models.go

@@ -42,15 +42,16 @@ func (r *FullResult) MetricLabels() prometheus.Labels {
 }
 
 type State struct {
-	ID       string
-	Version  int64
-	Dir      string
-	Path     string
-	Disabled bool
-	Endless  bool
-	RunEvery time.Duration
-	LastRun  time.Time
-	Timeout  time.Duration
+	ID              string
+	Version         int64
+	Dir             string
+	Path            string
+	Disabled        bool
+	Endless         bool
+	RunEvery        time.Duration
+	LastRun         time.Time
+	Timeout         time.Duration
+	OriginalTimeout time.Duration
 }
 
 func (s *State) ExploitType() models.ExploitType {

internal/exploit/runner.go

@@ -371,6 +371,7 @@ func (r *Runner) onServerStateUpdate(ctx context.Context, state *epb.ServerState
 	}
 
 	if r.storage.UpdateExploits(ctx, state.Exploits) {
+		r.storage.ScaleTimeouts(r.maxJobs, len(r.teams))
 		r.logger.Info("Exploits changed, scheduling loops restart")
 		r.restartLoops()
 	}

internal/exploit/storage.go

@@ -61,6 +61,30 @@ func (s *Storage) UpdateExploits(ctx context.Context, exs []*epb.ExploitState) b
 	return true
 }
 
+func (s *Storage) ScaleTimeouts(workers, teams int) {
+	alpha := 0.0
+	for _, ex := range s.cache.Exploits() {
+		if ex.Endless {
+			continue
+		}
+		alpha += ex.OriginalTimeout.Seconds() / ex.RunEvery.Seconds()
+	}
+	alpha = alpha * float64(teams) / float64(workers)
+	logrus.Infof("Scaling timeouts: alpha = %.2f", alpha)
+	for _, ex := range s.cache.Exploits() {
+		if ex.Endless {
+			continue
+		}
+		newTimeout := time.Duration(float64(ex.OriginalTimeout) / alpha)
+
+		// Round down to nearest second.
+		newTimeout -= newTimeout % time.Second
+
+		logrus.Infof("Scaling timeout for exploit %s: %s -> %s", ex.ID, ex.Timeout, newTimeout)
+		ex.Timeout = newTimeout
+	}
+}
+
 func (s *Storage) updateExploit(ctx context.Context, exploitID string) (*State, error) {
 	// Download the current exploit state.
 	resp, err := s.client.Exploit(ctx, exploitID)
@@ -115,14 +139,15 @@ func (s *Storage) updateExploit(ctx context.Context, exploitID string) (*State,
 	}
 
 	res := &State{
-		ID:       state.ExploitId,
-		Version:  state.Version,
-		Dir:      "",
-		Path:     entryPath,
-		Disabled: state.Config.Disabled,
-		Endless:  state.Config.Endless,
-		RunEvery: state.Config.RunEvery.AsDuration(),
-		Timeout:  state.Config.Timeout.AsDuration(),
+		ID:              state.ExploitId,
+		Version:         state.Version,
+		Dir:             "",
+		Path:            entryPath,
+		Disabled:        state.Config.Disabled,
+		Endless:         state.Config.Endless,
+		RunEvery:        state.Config.RunEvery.AsDuration(),
+		Timeout:         state.Config.Timeout.AsDuration(),
+		OriginalTimeout: state.Config.Timeout.AsDuration(),
 	}
 	if state.Config.IsArchive {
 		res.Dir = oPath

internal/exploit/storage_test.go

@@ -7,6 +7,7 @@ import (
 	"path"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
@@ -168,3 +169,55 @@ func Test_prepareEntry(t *testing.T) {
 	// Check that file is executable.
 	require.NotZero(t, fi.Mode()&0111)
 }
+
+func TestStorage_Scale(t *testing.T) {
+	st, cleanup := mockStorage()
+	defer func() {
+		require.NoError(t, cleanup())
+	}()
+
+	// This exploit's timeout should be halved, as teams = 2 * workers.
+	st.cache.Update([]*State{
+		{
+			ID:              "1",
+			Version:         1,
+			RunEvery:        time.Minute,
+			Timeout:         time.Minute,
+			OriginalTimeout: time.Minute,
+		},
+	})
+	st.ScaleTimeouts(10, 20)
+	res, ok := st.Exploit("1")
+	require.True(t, ok)
+	require.EqualValues(t, 1, res.Version)
+	require.EqualValues(t, time.Minute, res.RunEvery)
+	require.EqualValues(t, 30*time.Second, res.Timeout)
+
+	// Now it should be doubled, as workers = 2 * teams.
+	st.ScaleTimeouts(20, 10)
+	res, ok = st.Exploit("1")
+	require.True(t, ok)
+	require.EqualValues(t, time.Minute, res.RunEvery)
+	require.EqualValues(t, 2*time.Minute, res.Timeout)
+
+	// Add another exploit, expect scale to work proportionally to original timeouts.
+	st.cache.Update([]*State{
+		{
+			ID:              "2",
+			Version:         1,
+			RunEvery:        time.Minute,
+			Timeout:         time.Minute,
+			OriginalTimeout: time.Minute,
+		},
+	})
+	st.ScaleTimeouts(20, 10)
+	res, ok = st.Exploit("1")
+	require.True(t, ok)
+	require.EqualValues(t, time.Minute, res.RunEvery)
+	require.EqualValues(t, time.Minute, res.Timeout)
+
+	res, ok = st.Exploit("2")
+	require.True(t, ok)
+	require.EqualValues(t, time.Minute, res.RunEvery)
+	require.EqualValues(t, time.Minute, res.Timeout)
+}

internal/exploit/submit_loop.go

@@ -126,6 +126,7 @@ func (l *submitLoop) Start(ctx context.Context) {
 			}
 		case <-t.C:
 			flush()
+			l.metrics.Queue.WithLabelValues(string(l.q.Type())).Set(float64(l.q.Size()))
 		case <-ctx.Done():
 			return
 		}

internal/exploit/submit_loop_test.go

@@ -330,6 +330,10 @@ func (m *mockQueue) Type() queue.Type {
 	return "mock"
 }
 
+func (m *mockQueue) Size() int {
+	return len(m.in)
+}
+
 func (m *mockQueue) Start(ctx context.Context) {
 	<-ctx.Done()
 }

internal/queue/endless.go

@@ -50,6 +50,10 @@ func (q *endlessQueue) Type() Type {
 	return TypeEndless
 }
 
+func (q *endlessQueue) Size() int {
+	return len(q.c)
+}
+
 // Start is synchronous.
 // Cancel the start's context to stop the queue.
 func (q *endlessQueue) Start(ctx context.Context) {

internal/queue/queue.go

@@ -31,6 +31,7 @@ type Queue interface {
 	Add(*Job) error
 	Results() <-chan *Output
 	Type() Type
+	Size() int
 
 	fmt.Stringer
 }

internal/queue/simple.go

@@ -49,6 +49,10 @@ func (q *simpleQueue) Type() Type {
 	return TypeSimple
 }
 
+func (q *simpleQueue) Size() int {
+	return len(q.c)
+}
+
 // Start is synchronous.
 // Cancel the start's context to stop the queue.
 func (q *simpleQueue) Start(ctx context.Context) {