gitea/routers/api/actions/runner/interceptor.go

// Copyright 2022 The Gitea Authors. All rights reserved.
// SPDX-License-Identifier: MIT

package runner

import (
	"context"
	"crypto/sha256"
	"errors"
	"strings"
	"time"

	"gitea.dev/actionslib/pkg/protocol"
	actions_model "gitea.dev/models/actions"
	auth_model "gitea.dev/models/auth"
	"gitea.dev/modules/log"
	"gitea.dev/modules/timeutil"
	"gitea.dev/modules/util"

	"connectrpc.com/connect"
	lru "github.com/hashicorp/golang-lru/v2"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	uuidHeaderKey  = protocol.UUIDHeader
	tokenHeaderKey = protocol.TokenHeader
)

type runnerTokenHash struct{ salt, hash string }

var verifiedRunnerTokens, _ = lru.New[runnerTokenHash, [sha256.Size]byte](10000) // not auth_model.TokenCache, whose small default size thrashes once more runners poll in turn

var withRunner = connect.WithInterceptors(connect.UnaryInterceptorFunc(func(unaryFunc connect.UnaryFunc) connect.UnaryFunc {
	return func(ctx context.Context, request connect.AnyRequest) (connect.AnyResponse, error) {
		methodName := getMethodName(request)
		if methodName == "Register" {
			return unaryFunc(ctx, request)
		}
		uuid := request.Header().Get(uuidHeaderKey)
		token := request.Header().Get(tokenHeaderKey)

		runner, err := actions_model.GetRunnerByUUID(ctx, uuid)
		if err != nil {
			if errors.Is(err, util.ErrNotExist) {
				return nil, status.Error(codes.Unauthenticated, "unregistered runner")
			}
			return nil, status.Error(codes.Internal, err.Error())
		}
		if !verifyRunnerToken(runner, token) {
			return nil, status.Error(codes.Unauthenticated, "unregistered runner")
		}

		now := time.Now()
		cols := make([]string, 0, 2)
		// Debounce last_active too: while a runner streams logs, UpdateLog fires
		// many times per second and writing on each is a major source of DB load.
		// Persist only when stale enough to affect the active/idle status.
		if (methodName == "UpdateTask" || methodName == "UpdateLog") &&
			actions_model.ShouldPersistLastActive(runner.LastActive, now) {
			runner.LastActive = timeutil.TimeStamp(now.Unix())
			cols = append(cols, "last_active")
		}
		// Debounce last_online: writing on every poll is a major source of DB load
		// with many runners. Persist only when stale enough to affect offline status.
		if actions_model.ShouldPersistLastOnline(runner.LastOnline, now) {
			runner.LastOnline = timeutil.TimeStamp(now.Unix())
			cols = append(cols, "last_online")
		}
		if len(cols) > 0 {
			if err := actions_model.UpdateRunner(ctx, runner, cols...); err != nil {
				log.Error("can't update runner status: %v", err)
			}
		}

		ctx = context.WithValue(ctx, runnerCtxKey{}, runner)
		return unaryFunc(ctx, request)
	}
}))

func verifyRunnerToken(runner *actions_model.ActionRunner, token string) bool {
	key, digest := runnerTokenHash{runner.TokenSalt, runner.TokenHash}, sha256.Sum256([]byte(token))
	if cached, ok := verifiedRunnerTokens.Get(key); ok && util.CryptoConstTimeEqual(cached[:], digest[:]) {
		return true
	}
	if !util.CryptoConstTimeEqual(runner.TokenHash, auth_model.HashToken(token, runner.TokenSalt)) {
		return false
	}
	verifiedRunnerTokens.Add(key, digest)
	return true
}

func getMethodName(req connect.AnyRequest) string {
	splits := strings.Split(req.Spec().Procedure, "/")
	if len(splits) > 0 {
		return splits[len(splits)-1]
	}
	return ""
}

type runnerCtxKey struct{}

func GetRunner(ctx context.Context) *actions_model.ActionRunner {
	if v := ctx.Value(runnerCtxKey{}); v != nil {
		if r, ok := v.(*actions_model.ActionRunner); ok {
			return r
		}
	}
	return nil
}