actions/runners: enforce drain and hard revoke (S41j-4)

Status	File	+	-
M	`cmd/shithubd-runner/run.go`	7	0
M	`internal/runner/api/client.go`	2	0
M	`internal/runner/api/client_test.go`	8	2
M	`internal/runner/runner.go`	12	1
M	`internal/runner/runner_test.go`	21	2
M	`internal/web/handlers/api/runners.go`	89	19
M	`internal/web/handlers/api/runners_test.go`	103	1

cmd/shithubd-runner/run.gomodified

  	runnerconfig "github.com/tenseleyFlow/shithub/internal/runner/config"
  	"github.com/tenseleyFlow/shithub/internal/runner/engine"
  	"github.com/tenseleyFlow/shithub/internal/runner/workspace"
++	"github.com/tenseleyFlow/shithub/internal/version"
  )
  var runConfigPath string
  		if err != nil {
  			return err
  		}
++		hostName, err := os.Hostname()
++		if err != nil {
++			hostName = ""
++		}
  		execEngine := engine.NewDocker(engine.DockerConfig{
  			Binary:         cfg.Engine.Kind,
  			DefaultImage:   cfg.Engine.DefaultImage,
  			Logger:       logger,
  			Labels:       cfg.Runner.Labels,
  			Capacity:     cfg.Runner.Capacity,
++			HostName:     hostName,
++			Version:      version.String(),
  			PollInterval: cfg.Runner.PollInterval,
  			DefaultImage: cfg.Engine.DefaultImage,
  			Clock:        func() time.Time { return time.Now().UTC() },

internal/runner/api/client.gomodified


 type HeartbeatRequest struct {
 	Labels   []string `json:"labels"`
 	Capacity int      `json:"capacity"`
+	HostName string   `json:"host_name,omitempty"`
+	Version  string   `json:"version,omitempty"`
 }
 
 type Claim struct {

internal/runner/api/client_test.gomodified

  		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  			t.Fatalf("Decode: %v", err)
+ 		}
--		if req.Capacity != 2 || strings.Join(req.Labels, ",") != "self-hosted,linux" {
++		if req.Capacity != 2 || strings.Join(req.Labels, ",") != "self-hosted,linux" ||
++			req.HostName != "runner-host" || req.Version != "dev-test" {
  			t.Fatalf("request: %#v", req)
+ 		}
  		w.Header().Set("Content-Type", "application/json")
  	if err != nil {
  		t.Fatalf("New: %v", err)
+ 	}
--	claim, err := client.Heartbeat(t.Context(), HeartbeatRequest{Labels: []string{"self-hosted", "linux"}, Capacity: 2})
++	claim, err := client.Heartbeat(t.Context(), HeartbeatRequest{
++		Labels:   []string{"self-hosted", "linux"},
++		Capacity: 2,
++		HostName: "runner-host",
++		Version:  "dev-test",
++	})
  	if err != nil {
  		t.Fatalf("Heartbeat: %v", err)
+ 	}

internal/runner/runner.gomodified

  	Logger             *slog.Logger
  	Labels             []string
  	Capacity           int
++	HostName           string
++	Version            string
  	PollInterval       time.Duration
  	CancelPollInterval time.Duration
  	DefaultImage       string
  	logger             *slog.Logger
  	labels             []string
  	capacity           int
++	hostName           string
++	version            string
  	pollInterval       time.Duration
  	cancelPollInterval time.Duration
  	defaultImage       string
  		logger:             logger,
  		labels:             append([]string{}, opts.Labels...),
  		capacity:           capacity,
++		hostName:           opts.HostName,
++		version:            opts.Version,
  		pollInterval:       poll,
  		cancelPollInterval: cancelPoll,
  		defaultImage:       opts.DefaultImage,
  }
  func (r *Runner) RunOnce(ctx context.Context) (bool, error) {
--	claim, err := r.api.Heartbeat(ctx, api.HeartbeatRequest{Labels: r.labels, Capacity: r.capacity})
++	claim, err := r.api.Heartbeat(ctx, api.HeartbeatRequest{
++		Labels:   r.labels,
++		Capacity: r.capacity,
++		HostName: r.hostName,
++		Version:  r.version,
++	})
  	if err != nil {
  		return false, err
  	}

internal/runner/runner_test.gomodified

  type fakeAPI struct {
  	claim        *api.Claim
++	heartbeats   []api.HeartbeatRequest
  	statuses     []api.StatusRequest
  	stepStatuses []api.StatusRequest
  	logs         []api.LogRequest
  	next         int
+ }
--func (f *fakeAPI) Heartbeat(_ context.Context, _ api.HeartbeatRequest) (*api.Claim, error) {
++func (f *fakeAPI) Heartbeat(_ context.Context, req api.HeartbeatRequest) (*api.Claim, error) {
++	f.heartbeats = append(f.heartbeats, req)
  	return f.claim, nil
+ }
  func TestRunOnce_NoClaim(t *testing.T) {
  	t.Parallel()
--	r := New(Options{API: &fakeAPI{}, Engine: &fakeEngine{}, Workspaces: &fakeWorkspaces{}})
++	fapi := &fakeAPI{}
++	r := New(Options{
++		API:        fapi,
++		Engine:     &fakeEngine{},
++		Workspaces: &fakeWorkspaces{},
++		Labels:     []string{"self-hosted", "linux"},
++		Capacity:   2,
++		HostName:   "runner-host",
++		Version:    "dev-test",
++	})
  	claimed, err := r.RunOnce(t.Context())
  	if err != nil {
  		t.Fatalf("RunOnce: %v", err)
  	if claimed {
  		t.Fatal("claimed = true")
+ 	}
++	if len(fapi.heartbeats) != 1 {
++		t.Fatalf("heartbeats: %#v", fapi.heartbeats)
++	}
++	got := fapi.heartbeats[0]
++	if got.Capacity != 2 || got.HostName != "runner-host" || got.Version != "dev-test" ||
++		len(got.Labels) != 2 || got.Labels[0] != "self-hosted" || got.Labels[1] != "linux" {
++		t.Fatalf("heartbeat: %#v", got)
++	}
+ }
  func TestRunOnce_ExecutesAndCompletesSuccess(t *testing.T) {

internal/web/handlers/api/runners.gomodified

  	"strconv"
  	"strings"
  	"time"
++	"unicode/utf8"
  	"github.com/go-chi/chi/v5"
  	"github.com/jackc/pgx/v5"
  type runnerHeartbeatRequest struct {
  	Labels   []string `json:"labels"`
  	Capacity int      `json:"capacity"`
++	HostName string   `json:"host_name"`
++	Version  string   `json:"version"`
+ }
++const runnerMetadataMaxBytes = 255
++
++var errRunnerRevoked = errors.New("runner is revoked")
++
  func (h *Handlers) runnerHeartbeat(w http.ResponseWriter, r *http.Request) {
  	if h.d.RunnerJWT == nil {
  		writeAPIError(w, http.StatusServiceUnavailable, "runner API is not configured")
  		writeAPIError(w, http.StatusBadRequest, "capacity must be between 1 and 64")
  		return
+ 	}
++	hostName := cleanRunnerMetadata(body.HostName)
++	if hostName == "" {
++		hostName = runner.HostName
++	}
++	version := cleanRunnerMetadata(body.Version)
++	if version == "" {
++		version = runner.Version
++	}
--	job, steps, resolvedSecrets, claimed, err := h.claimRunnerJob(r.Context(), runner.ID, labels, int32(capacity))
++	job, steps, resolvedSecrets, claimed, err := h.claimRunnerJob(r.Context(), runner.ID, labels, int32(capacity), hostName, version)
  	if err != nil {
++		if errors.Is(err, errRunnerRevoked) {
++			metrics.ActionsRunnerHeartbeatsTotal.WithLabelValues("rejected").Inc()
++			writeAPIError(w, http.StatusUnauthorized, "runner revoked")
++			return
++		}
  		h.d.Logger.ErrorContext(r.Context(), "runner heartbeat claim failed", "runner_id", runner.ID, "error", err)
  		writeAPIError(w, http.StatusInternalServerError, "runner heartbeat failed")
  		return
  	writeJSON(w, http.StatusOK, h.presentRunnerClaim(job, steps, resolvedSecrets, token, checkoutToken, time.Unix(claims.Exp, 0)))
+ }
++func cleanRunnerMetadata(value string) string {
++	value = strings.TrimSpace(value)
++	if len(value) <= runnerMetadataMaxBytes {
++		return value
++	}
++	var b strings.Builder
++	for _, r := range value {
++		runeLen := utf8.RuneLen(r)
++		if runeLen < 0 || b.Len()+runeLen > runnerMetadataMaxBytes {
++			break
++		}
++		b.WriteRune(r)
++	}
++	return strings.TrimSpace(b.String())
++}
++
  func (h *Handlers) authenticateRunner(w http.ResponseWriter, r *http.Request) (actionsdb.GetRunnerByTokenHashRow, bool) {
  	const prefix = "Bearer "
  	authz := r.Header.Get("Authorization")
  	runnerID int64,
  	labels []string,
  	capacity int32,
++	hostName string,
++	version string,
  ) (actionsdb.ClaimQueuedWorkflowJobRow, []actionsdb.ListRunnerStepsForJobRow, map[string]string, bool, error) {
  	q := actionsdb.New()
  	tx, err := h.d.Pool.Begin(ctx)
+ 		}
  	}()
--	if _, err := q.LockRunnerByID(ctx, tx, runnerID); err != nil {
++	lockedRunner, err := q.LockRunnerByID(ctx, tx, runnerID)
++	if err != nil {
  		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 	}
++	if lockedRunner.RevokedAt.Valid {
++		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, errRunnerRevoked
++	}
  	running, err := q.CountRunningJobsForRunner(ctx, tx, runnerID)
  	if err != nil {
  		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 	}
--	if running >= capacity {
++	heartbeat := func(status actionsdb.WorkflowRunnerStatus) error {
--		if _, err := q.HeartbeatRunner(ctx, tx, actionsdb.HeartbeatRunnerParams{
++		_, err := q.HeartbeatRunner(ctx, tx, actionsdb.HeartbeatRunnerParams{
  			ID:       runnerID,
  			Labels:   labels,
  			Capacity: capacity,
--			Status:   actionsdb.WorkflowRunnerStatusBusy,
++			Status:   status,
--		}); err != nil {
++			HostName: hostName,
++			Version:  version,
++		})
++		return err
++	}
++	if lockedRunner.DrainingAt.Valid {
++		status := actionsdb.WorkflowRunnerStatusIdle
++		if running > 0 {
++			status = actionsdb.WorkflowRunnerStatusBusy
++		}
++		if err := heartbeat(status); err != nil {
++			return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
++		}
++		if err := tx.Commit(ctx); err != nil {
++			return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
++		}
++		committed = true
++		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, nil
++	}
++	if running >= capacity {
++		if err := heartbeat(actionsdb.WorkflowRunnerStatusBusy); err != nil {
  			return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 		}
  		if err := tx.Commit(ctx); err != nil {
  		if !errors.Is(err, pgx.ErrNoRows) {
  			return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 		}
--		if _, err := q.HeartbeatRunner(ctx, tx, actionsdb.HeartbeatRunnerParams{
++		if err := heartbeat(actionsdb.WorkflowRunnerStatusIdle); err != nil {
--			ID:       runnerID,
--			Labels:   labels,
--			Capacity: capacity,
--			Status:   actionsdb.WorkflowRunnerStatusIdle,
--		}); err != nil {
  			return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 		}
  		if err := tx.Commit(ctx); err != nil {
  	if running+1 >= capacity {
  		status = actionsdb.WorkflowRunnerStatusBusy
+ 	}
--	if _, err := q.HeartbeatRunner(ctx, tx, actionsdb.HeartbeatRunnerParams{
++	if err := heartbeat(status); err != nil {
--		ID:       runnerID,
--		Labels:   labels,
--		Capacity: capacity,
--		Status:   status,
--	}); err != nil {
  		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 	}
++	var claimLatencySeconds float64
++	observeClaimLatency := false
++	if job.CreatedAt.Valid {
++		if latency := time.Since(job.CreatedAt.Time); latency >= 0 {
++			claimLatencySeconds = latency.Seconds()
++			observeClaimLatency = true
++		}
++	}
  	if err := tx.Commit(ctx); err != nil {
  		return actionsdb.ClaimQueuedWorkflowJobRow{}, nil, nil, false, err
+ 	}
  	committed = true
++	if observeClaimLatency {
++		metrics.ActionsJobClaimLatencySeconds.Observe(claimLatencySeconds)
++	}
  	return job, steps, resolvedSecrets, true, nil
+ }
  		writeAPIError(w, http.StatusUnauthorized, "job token invalid")
  		return runnerJobAuth{}, false
+ 	}
--	job, err := actionsdb.New().GetWorkflowJobByID(r.Context(), h.d.Pool, pathJobID)
++	q := actionsdb.New()
++	runner, err := q.GetRunnerByID(r.Context(), h.d.Pool, runnerID)
++	if err != nil || runner.RevokedAt.Valid {
++		metrics.ActionsRunnerJWTTotal.WithLabelValues("rejected").Inc()
++		writeAPIError(w, http.StatusUnauthorized, "job token invalid")
++		return runnerJobAuth{}, false
++	}
++	job, err := q.GetWorkflowJobByID(r.Context(), h.d.Pool, pathJobID)
  	if err != nil {
  		if errors.Is(err, pgx.ErrNoRows) {
  			writeAPIError(w, http.StatusNotFound, "job not found")

internal/web/handlers/api/runners_test.gomodified

  	router := newRunnerAPIRouter(t, pool, logger, signer)
  	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/heartbeat",
--		strings.NewReader(`{"labels":["ubuntu-latest","linux"],"capacity":1}`))
++		strings.NewReader(`{"labels":["ubuntu-latest","linux"],"capacity":1,"host_name":"runner-host","version":"dev-test"}`))
  	req.Header.Set("Authorization", "Bearer "+token)
  	rr := httptest.NewRecorder()
  	router.ServeHTTP(rr, req)
  		checkoutClaims.Purpose != runnerjwt.PurposeCheckout {
  		t.Fatalf("checkout claims/job mismatch: claims=%+v job=%+v", checkoutClaims, resp.Job)
+ 	}
++	runnerRow, err := actionsdb.New().GetRunnerByID(ctx, pool, runnerID)
++	if err != nil {
++		t.Fatalf("GetRunnerByID: %v", err)
++	}
++	if runnerRow.HostName != "runner-host" || runnerRow.Version != "dev-test" {
++		t.Fatalf("runner metadata: host=%q version=%q", runnerRow.HostName, runnerRow.Version)
++	}
  	var logResp struct {
  		Accepted  bool   `json:"accepted"`
+ 	}
+ }
++func TestRunnerHeartbeatDoesNotClaimWhenDraining(t *testing.T) {
++	ctx := context.Background()
++	pool := dbtest.NewTestDB(t)
++	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
++	repoID, userID := setupRunnerAPIRepo(t, pool)
++	runID := enqueueRunnerAPIRun(t, pool, logger, repoID, userID)
++	token, runnerID := registerRunnerForTest(t, pool, []string{"ubuntu-latest", "linux"}, 1)
++	q := actionsdb.New()
++	if _, err := q.SetRunnerDraining(ctx, pool, actionsdb.SetRunnerDrainingParams{
++		ID:          runnerID,
++		DrainReason: "maintenance",
++	}); err != nil {
++		t.Fatalf("SetRunnerDraining: %v", err)
++	}
++	router := newRunnerAPIRouter(t, pool, logger, runnerAPISigner(t, time.Now()))
++
++	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/heartbeat",
++		strings.NewReader(`{"labels":["ubuntu-latest","linux"],"capacity":1,"host_name":"draining-host","version":"dev-test"}`))
++	req.Header.Set("Authorization", "Bearer "+token)
++	rr := httptest.NewRecorder()
++	router.ServeHTTP(rr, req)
++
++	if rr.Code != http.StatusNoContent {
++		t.Fatalf("status: got %d, want 204; body=%s", rr.Code, rr.Body.String())
++	}
++	jobs, err := q.ListJobsForRun(ctx, pool, runID)
++	if err != nil {
++		t.Fatalf("ListJobsForRun: %v", err)
++	}
++	if len(jobs) != 1 || jobs[0].Status != actionsdb.WorkflowJobStatusQueued {
++		t.Fatalf("job was claimed while runner drained: %#v", jobs)
++	}
++	job, err := q.GetWorkflowJobByID(ctx, pool, jobs[0].ID)
++	if err != nil {
++		t.Fatalf("GetWorkflowJobByID: %v", err)
++	}
++	if job.RunnerID.Valid {
++		t.Fatalf("job was assigned to runner while drained: %+v", job)
++	}
++	runnerRow, err := q.GetRunnerByID(ctx, pool, runnerID)
++	if err != nil {
++		t.Fatalf("GetRunnerByID: %v", err)
++	}
++	if !runnerRow.DrainingAt.Valid || runnerRow.HostName != "draining-host" {
++		t.Fatalf("runner drain/metadata not preserved: %+v", runnerRow)
++	}
++}
++
++func TestRunnerJobTokenRejectedAfterRunnerRevoked(t *testing.T) {
++	ctx := context.Background()
++	pool := dbtest.NewTestDB(t)
++	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
++	repoID, userID := setupRunnerAPIRepo(t, pool)
++	enqueueRunnerAPIRun(t, pool, logger, repoID, userID)
++	token, runnerID := registerRunnerForTest(t, pool, []string{"ubuntu-latest", "linux"}, 1)
++	router := newRunnerAPIRouter(t, pool, logger, runnerAPISigner(t, time.Now()))
++
++	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/heartbeat",
++		strings.NewReader(`{"labels":["ubuntu-latest","linux"],"capacity":1}`))
++	req.Header.Set("Authorization", "Bearer "+token)
++	rr := httptest.NewRecorder()
++	router.ServeHTTP(rr, req)
++	if rr.Code != http.StatusOK {
++		t.Fatalf("claim status: got %d, want 200; body=%s", rr.Code, rr.Body.String())
++	}
++	var claim struct {
++		Token string `json:"token"`
++		Job   struct {
++			ID int64 `json:"id"`
++		} `json:"job"`
++	}
++	if err := json.Unmarshal(rr.Body.Bytes(), &claim); err != nil {
++		t.Fatalf("decode claim: %v", err)
++	}
++	q := actionsdb.New()
++	if _, err := q.RevokeRunner(ctx, pool, actionsdb.RevokeRunnerParams{
++		ID:            runnerID,
++		RevokedReason: "compromised",
++	}); err != nil {
++		t.Fatalf("RevokeRunner: %v", err)
++	}
++	if err := q.RevokeAllTokensForRunner(ctx, pool, runnerID); err != nil {
++		t.Fatalf("RevokeAllTokensForRunner: %v", err)
++	}
++
++	statusReq := httptest.NewRequest(http.MethodPost, fmt.Sprintf("/api/v1/jobs/%d/status", claim.Job.ID),
++		strings.NewReader(`{"status":"running"}`))
++	statusReq.Header.Set("Authorization", "Bearer "+claim.Token)
++	statusRR := httptest.NewRecorder()
++	router.ServeHTTP(statusRR, statusReq)
++	if statusRR.Code != http.StatusUnauthorized {
++		t.Fatalf("status: got %d, want 401; body=%s", statusRR.Code, statusRR.Body.String())
++	}
++}
++
  func TestRunnerHeartbeatBypassesGlobalAnonAPILimit(t *testing.T) {
  	pool := dbtest.NewTestDB(t)
  	logger := slog.New(slog.NewTextHandler(io.Discard, nil))

`@@ -47,6 +47,8 @@` func New(cfg Config) (*Client, error) {
47	type HeartbeatRequest struct {	47	type HeartbeatRequest struct {
48	Labels []string `json:"labels"`	48	Labels []string `json:"labels"`
49	Capacity int `json:"capacity"`	49	Capacity int `json:"capacity"`
		50	+ HostName string `json:"host_name,omitempty"`
		51	+ Version string `json:"version,omitempty"`
50	}	52	}
51		53
52	type Claim struct {	54	type Claim struct {

tenseleyflow/shithub / `6931e9a`

7 changed files