orbit/pulse/worker/service.go

package worker

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"sort"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	"git.yuntuo.hk/k2safe/orbit/pulse/service"
)

var StaleSession = fault.New("STALE_SESSION", "connection is expired or superseded")
var LeaseLost = fault.New("LEASE_LOST", "attempt no longer grants execution authority")
var GateClosed = fault.New("EXECUTION_PAUSED", "queue or execution gate is paused")

type Service struct{ store Store }

func New(store Store) *Service { return &Service{store: store} }
func (s *Service) Authorize(ctx context.Context, t auth.Token) error {
	return s.store.Authorize(ctx, t)
}
func NewSecret() (string, core.Digest, error) {
	var raw [32]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", core.Digest{}, err
	}
	return base64.RawURLEncoding.EncodeToString(raw[:]), core.Digest(sha256.Sum256(raw[:])), nil
}
func SecretDigest(secret string) (core.Digest, error) {
	if len(secret) != 43 {
		return core.Digest{}, fault.Unauthorized
	}
	raw, err := base64.RawURLEncoding.Strict().DecodeString(secret)
	if err != nil || len(raw) != 32 {
		return core.Digest{}, fault.Unauthorized
	}
	return core.Digest(sha256.Sum256(raw)), nil
}
func NormalizeConnect(in ConnectInput) (ConnectInput, error) {
	if in.InstanceID.Validate() != nil || in.QueueID.Validate() != nil || in.Capacity < 1 || in.Capacity > 256 || len(in.Capabilities) < 1 || len(in.Capabilities) > 32 {
		return in, fault.Invalid("invalid worker identity, queue, capacity or capability count")
	}
	if _, err := SecretDigest(in.InstanceSecret); err != nil {
		return in, fault.Invalid("instance_secret must encode 32 random bytes")
	}
	in.Capabilities = append([]Capability(nil), in.Capabilities...)
	for _, c := range in.Capabilities {
		if !service.ValidName(c.Handler, 128) || !service.ValidName(c.HandlerVersion, 64) || !service.ValidName(c.PayloadVersion, 32) {
			return in, fault.Invalid("invalid worker capability")
		}
	}
	sort.Slice(in.Capabilities, func(i, j int) bool {
		a, b := in.Capabilities[i], in.Capabilities[j]
		if a.Handler != b.Handler {
			return a.Handler < b.Handler
		}
		if a.HandlerVersion != b.HandlerVersion {
			return a.HandlerVersion < b.HandlerVersion
		}
		return a.PayloadVersion < b.PayloadVersion
	})
	for i := 1; i < len(in.Capabilities); i++ {
		if in.Capabilities[i] == in.Capabilities[i-1] {
			return in, fault.Invalid("duplicate worker capability")
		}
	}
	return in, nil
}
func (s *Service) Connect(ctx context.Context, t auth.Token, in ConnectInput) (Binding, Welcome, error) {
	in, err := NormalizeConnect(in)
	if err != nil {
		return Binding{}, Welcome{}, err
	}
	return s.store.Connect(ctx, t, in)
}
func (s *Service) Capacity(ctx context.Context, t auth.Token, b Binding, in CapacityInput) (Welcome, error) {
	if in.Sequence == 0 || in.Capacity > 256 {
		return Welcome{}, fault.Invalid("invalid capacity sequence or value")
	}
	return s.store.Capacity(ctx, t, b, in)
}
func (s *Service) Claim(ctx context.Context, t auth.Token, b Binding) (*Assignment, error) {
	return s.store.Claim(ctx, t, b)
}
func validRef(ref AttemptRef) error {
	if ref.TaskID.Validate() != nil || ref.AttemptID.Validate() != nil || ref.Epoch == 0 {
		return fault.Invalid("invalid attempt identity")
	}
	_, err := SecretDigest(ref.LeaseSecret)
	return err
}
func validCommand(message core.ID, ref AttemptRef) error {
	if message.Validate() != nil {
		return fault.Invalid("invalid operation message id")
	}
	return validRef(ref)
}
func (s *Service) Start(ctx context.Context, t auth.Token, b Binding, message core.ID, ref AttemptRef) (CommandResult, error) {
	if err := validCommand(message, ref); err != nil {
		return CommandResult{}, err
	}
	return s.store.Start(ctx, t, b, message, ref)
}
func (s *Service) Renew(ctx context.Context, t auth.Token, b Binding, message core.ID, ref AttemptRef) (CommandResult, error) {
	if err := validCommand(message, ref); err != nil {
		return CommandResult{}, err
	}
	return s.store.Renew(ctx, t, b, message, ref)
}
func (s *Service) Finish(ctx context.Context, t auth.Token, b Binding, message core.ID, in FinishInput) (CommandResult, error) {
	if err := validCommand(message, in.Ref); err != nil {
		return CommandResult{}, err
	}
	if in.Outcome != "SUCCEEDED" && in.Outcome != "FAILED" && in.Outcome != "UNKNOWN" {
		return CommandResult{}, fault.Invalid("invalid result outcome")
	}
	if len(in.Result) == 0 || len(in.Result) > 256<<10 || bytes.Equal(bytes.TrimSpace(in.Result), []byte("null")) {
		return CommandResult{}, fault.Invalid("result must be non-null JSON up to 256 KiB")
	}
	if _, err := jsonvalue.Canonical(in.Result); err != nil {
		return CommandResult{}, fault.Invalid("invalid result JSON")
	}
	if (in.Outcome == "SUCCEEDED" && in.ErrorCode != "") || (in.Outcome != "SUCCEEDED" && !service.ValidName(in.ErrorCode, 64)) {
		return CommandResult{}, fault.Invalid("error_code must match the outcome")
	}
	return s.store.Finish(ctx, t, b, message, in)
}
func (s *Service) Resume(ctx context.Context, t auth.Token, b Binding, ref AttemptRef) (LeaseStatus, error) {
	if err := validRef(ref); err != nil {
		return LeaseStatus{}, err
	}
	return s.store.Resume(ctx, t, b, ref)
}