A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
8.0 kB · 199 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200package controller
import ( "bytes" "context" "encoding/json" "errors" "fmt" "io" "path/filepath" "slices" "strings"
"connectrpc.com/connect"
agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" "forgejo.alexma.top/alexma233/composia/internal/core/config" "forgejo.alexma.top/alexma233/composia/internal/core/repo" "forgejo.alexma.top/alexma233/composia/internal/core/task" "forgejo.alexma.top/alexma233/composia/internal/platform/rpcutil" "forgejo.alexma.top/alexma233/composia/internal/platform/store")
type bundleServer struct { db *store.DB cfg *config.ControllerConfig}
func (server *bundleServer) serviceTask(ctx context.Context, taskID, executionID string) (task.Record, serviceTaskParams, error) { var params serviceTaskParams if taskID == "" { return task.Record{}, params, connect.NewError(connect.CodeInvalidArgument, errors.New("task_id is required")) } if err := ensureTaskNodeMatch(ctx, server.db, taskID); err != nil { return task.Record{}, params, err } nodeID, _ := rpcutil.BearerSubject(ctx) record, err := server.db.ValidateTaskExecution(ctx, taskID, executionID, nodeID) if err != nil { return task.Record{}, params, connect.NewError(connect.CodeFailedPrecondition, err) } if err := json.Unmarshal([]byte(record.ParamsJSON), ¶ms); err != nil { return task.Record{}, params, connect.NewError(connect.CodeInternal, fmt.Errorf("decode service task params: %w", err)) } return record, params, nil}
func (server *bundleServer) GetServiceManifest(ctx context.Context, req *connect.Request[agentv1.GetServiceManifestRequest]) (*connect.Response[agentv1.GetServiceManifestResponse], error) { record, params, err := server.serviceTask(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId()) if err != nil { return nil, err } if record.RepoRevision == "" { return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("service manifests require a service-scoped task with a repo revision")) } primaryDir, _ := cleanBundleServiceDir(params.ServiceDir) params.ServiceDir, err = authorizedBundleServiceDir(server.cfg, record, params, req.Msg.GetServiceDir()) if err != nil { return nil, err } service, err := repo.FindServiceAtRevision(server.cfg.RepoDir, record.RepoRevision, params.ServiceDir, configuredNodeIDs(server.cfg)) if err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } if params.ServiceDir == primaryDir && service.Name != record.ServiceName { return nil, connect.NewError(connect.CodePermissionDenied, errors.New("service directory does not match task service")) } extras, err := persistentServiceExtraFiles(server.cfg, record.RepoRevision, params.ServiceDir) if err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } files, err := repo.ServiceManifest(ctx, server.cfg.RepoDir, record.RepoRevision, params.ServiceDir, extras) if err != nil { return nil, connect.NewError(connect.CodeInternal, err) } currentRevision, err := repo.CurrentRevision(server.cfg.RepoDir) if err != nil { return nil, connect.NewError(connect.CodeInternal, err) } if currentRevision != record.RepoRevision { currentExtras, err := persistentServiceExtraFiles(server.cfg, currentRevision, params.ServiceDir) if err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } currentFiles, err := repo.ServiceManifest(ctx, server.cfg.RepoDir, currentRevision, params.ServiceDir, currentExtras) if err != nil || !slices.Equal(files, currentFiles) { return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("controller service files changed since this check was queued; deploy or update and check again")) } } response := &agentv1.GetServiceManifestResponse{RepoRevision: record.RepoRevision, RelativeRoot: params.ServiceDir} for _, file := range files { response.Files = append(response.Files, &agentv1.ServiceManifestFile{Path: file.Path, Sha256: file.SHA256, Mode: file.Mode}) } return connect.NewResponse(response), nil}
func (server *bundleServer) GetServiceBundle(ctx context.Context, req *connect.Request[agentv1.GetServiceBundleRequest], stream *connect.ServerStream[agentv1.GetServiceBundleResponse]) error { record, params, err := server.serviceTask(ctx, req.Msg.GetTaskId(), req.Msg.GetExecutionId()) if err != nil { return err } if record.Type != task.TypeDeploy && record.Type != task.TypeUpdate { return connect.NewError(connect.CodeFailedPrecondition, errors.New("only deploy and update tasks can install service bundles; upgrade the agent to use local service files")) } params.ServiceDir, err = authorizedBundleServiceDir(server.cfg, record, params, req.Msg.GetServiceDir()) if err != nil { return err } extraFiles, err := persistentServiceExtraFiles(server.cfg, record.RepoRevision, params.ServiceDir) if err != nil { if errors.Is(err, errSecretsNotConfigured) { return connect.NewError(connect.CodeFailedPrecondition, err) } return connect.NewError(connect.CodeInternal, err) }
pipeReader, pipeWriter := io.Pipe() go func() { pipeWriter.CloseWithError(repo.StreamServiceBundleWithExtras(ctx, server.cfg.RepoDir, record.RepoRevision, params.ServiceDir, extraFiles, pipeWriter)) }() defer func() { _ = pipeReader.Close() }()
buffer := make([]byte, 32*1024) firstChunk := true for { count, err := pipeReader.Read(buffer) if count > 0 { response := &agentv1.GetServiceBundleResponse{Data: bytes.Clone(buffer[:count])} if firstChunk { response.ServiceName = record.ServiceName response.RepoRevision = record.RepoRevision response.RelativeRoot = params.ServiceDir firstChunk = false } if sendErr := stream.Send(response); sendErr != nil { return sendErr } } if err != nil { if errors.Is(err, io.EOF) { return nil } return connect.NewError(connect.CodeInternal, fmt.Errorf("read bundle stream: %w", err)) } }}
func authorizedBundleServiceDir(cfg *config.ControllerConfig, record task.Record, params serviceTaskParams, requested string) (string, error) { serviceDir, err := cleanBundleServiceDir(params.ServiceDir) if err != nil && requested == "" { return "", connect.NewError(connect.CodeFailedPrecondition, errors.New("task is missing a valid service_dir")) } if requested == "" { return serviceDir, nil } requested, err = cleanBundleServiceDir(requested) if err != nil { return "", connect.NewError(connect.CodePermissionDenied, err) } if requested == serviceDir { return requested, nil }
switch record.Type { case task.TypeCaddySync: for _, allowed := range params.ServiceDirs { cleaned, cleanErr := cleanBundleServiceDir(allowed) if cleanErr == nil && requested == cleaned { return requested, nil } } case task.TypeBackup, task.TypeRestore: rustic, findErr := repo.FindRusticInfraServiceAtRevision(cfg.RepoDir, record.RepoRevision, configuredNodeIDs(cfg)) if findErr != nil { return "", connect.NewError(connect.CodeInternal, findErr) } if err := validateRusticServiceTargetNode(rustic, record.NodeID); err != nil { return "", connect.NewError(connect.CodePermissionDenied, err) } rusticDir, relErr := filepath.Rel(cfg.RepoDir, rustic.Directory) if relErr != nil { return "", connect.NewError(connect.CodeInternal, relErr) } if cleaned, cleanErr := cleanBundleServiceDir(rusticDir); cleanErr == nil && requested == cleaned { return requested, nil } default: }
return "", connect.NewError(connect.CodePermissionDenied, fmt.Errorf("service_dir %q is not authorized for task %q", requested, record.TaskID))}
func cleanBundleServiceDir(value string) (string, error) { cleaned := filepath.ToSlash(filepath.Clean(strings.TrimSpace(value))) if cleaned == "" || cleaned == "." || filepath.IsAbs(cleaned) || cleaned == ".." || strings.HasPrefix(cleaned, "../") { return "", fmt.Errorf("service_dir %q must be a repo-relative directory", value) } return cleaned, nil}