knowledge-governance/internal/workflows/artifactfetch/handler.go

// Package artifactfetch turns a saved policy intent into verified bytes: it
// shells out to scripts/artifact_fetch.py (which already holds every source
// adapter: AList direct download, PP ledger, DingTalk export-tree/live export)
// and records the verified content against the pending source item version.
package artifactfetch

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"knowledge-governance/internal/jobs"
)

type Request struct {
	ConnectionID string `json:"source_connection_id"`
	NodeID       string `json:"node_id"`
	VersionID    string `json:"version_id"`
}

type Result struct {
	OK              bool   `json:"ok"`
	SpecID          string `json:"spec_id"`
	NodeID          string `json:"node_id"`
	ConnectionID    string `json:"connection_id"`
	SHA256          string `json:"sha256"`
	SizeBytes       int64  `json:"size_bytes"`
	SpoolRef        string `json:"spool_ref"`
	ValidatedFormat string `json:"validated_format"`
	MatchedLocal    bool   `json:"matched_local"`
	Missing         bool   `json:"missing"`
	Retryable       bool   `json:"retryable"`
	Code            string `json:"code"`
	Error           string `json:"error"`
}

type Manifest struct {
	Request
	SourceType   string
	SourceName   string
	RelativePath string
	Name         string
	SizeBytes    *int64
	ModifiedAt   *time.Time
	LocatorHash  string
}

type Store interface {
	LoadArtifactFetchManifest(context.Context, string, string) (Manifest, error)
	MarkArtifactFetched(context.Context, string, string, int64, string, string, bool, time.Time) error
	MarkArtifactFetchMissing(context.Context, string, string, time.Time) error
	MarkArtifactFetchFailed(context.Context, string, string, string, time.Time) error
}

type Runner interface {
	Run(context.Context, Manifest) (Result, error)
}

type Config struct {
	Store  Store
	Runner Runner
	Now    func() time.Time
}

type Handler struct {
	store  Store
	runner Runner
	now    func() time.Time
}

func NewHandler(config Config) (*Handler, error) {
	if config.Store == nil || config.Runner == nil {
		return nil, errors.New("artifact fetch dependencies are required")
	}
	if config.Now == nil {
		config.Now = time.Now
	}
	return &Handler{store: config.Store, runner: config.Runner, now: config.Now}, nil
}

func (handler *Handler) Handle(ctx context.Context, job jobs.Job) error {
	var request Request
	decoder := json.NewDecoder(bytes.NewReader(job.Payload))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&request); err != nil ||
		request.ConnectionID == "" || request.NodeID == "" || request.VersionID == "" {
		return jobs.Permanent(jobs.ErrCodeArtifactInvalid, errors.New("artifact_fetch_invalid_payload"))
	}
	manifest, err := handler.store.LoadArtifactFetchManifest(ctx, request.ConnectionID, request.NodeID)
	if err != nil {
		if errors.Is(err, ErrManifestGone) {
			return jobs.Permanent(jobs.ErrCodeStateConflict, err)
		}
		return err
	}
	result, err := handler.runner.Run(ctx, manifest)
	now := handler.now()
	if err != nil {
		_ = handler.store.MarkArtifactFetchFailed(ctx, request.ConnectionID, request.NodeID, err.Error(), now)
		return err
	}
	switch {
	case result.Missing:
		if err := handler.store.MarkArtifactFetchMissing(ctx, request.ConnectionID, request.NodeID, now); err != nil {
			return err
		}
		return jobs.Permanent(jobs.ErrCodeArtifactInvalid, errors.New("artifact_fetch_missing_source: "+safeMessage(result.Error)))
	case !result.OK:
		_ = handler.store.MarkArtifactFetchFailed(ctx, request.ConnectionID, request.NodeID, safeMessage(result.Error), now)
		if result.Retryable {
			return errors.New("artifact_fetch_retryable: " + safeMessage(result.Code))
		}
		return jobs.Permanent(jobs.ErrCodeArtifactInvalid, errors.New("artifact_fetch_rejected: "+safeMessage(result.Code)))
	}
	if result.SHA256 == "" || result.SpoolRef == "" || result.SizeBytes <= 0 || result.NodeID != request.NodeID || result.ValidatedFormat == "" {
		_ = handler.store.MarkArtifactFetchFailed(ctx, request.ConnectionID, request.NodeID, "artifact fetch returned an inconsistent result", now)
		return errors.New("artifact_fetch_invalid_output")
	}
	if err := handler.store.MarkArtifactFetched(ctx, request.VersionID, strings.ToLower(result.SHA256),
		result.SizeBytes, result.SpoolRef, result.ValidatedFormat, result.MatchedLocal, now); err != nil {
		return err
	}
	return nil
}

func safeMessage(value string) string {
	value = strings.ReplaceAll(strings.TrimSpace(value), "\n", " ")
	lower := strings.ToLower(value)
	for _, forbidden := range []string{"token", "secret", "password", "authorization", "bearer", "cookie", "/run/source-secrets", "http://", "https://"} {
		if strings.Contains(lower, forbidden) {
			return "source reported an error (details suppressed)"
		}
	}
	if len([]rune(value)) > 200 {
		value = string([]rune(value)[:200])
	}
	return value
}

// CommandRunner executes the fetch script once and parses its single JSON line.
type CommandRunner struct {
	PythonPath  string
	ScriptPath  string
	SpoolRoot   string
	SpoolPrefix string
	MaxBytes    int64
}

func NewCommandRunner(pythonPath, scriptPath, spoolRoot, spoolPrefix string, maxBytes int64) (*CommandRunner, error) {
	if !filepath.IsAbs(pythonPath) || !filepath.IsAbs(scriptPath) || !filepath.IsAbs(spoolRoot) || maxBytes <= 0 {
		return nil, errors.New("artifact fetch command paths must be absolute")
	}
	if strings.Contains(spoolPrefix, "..") || strings.Contains(spoolPrefix, "\\") {
		return nil, errors.New("artifact fetch spool prefix is invalid")
	}
	spoolPrefix = strings.Trim(spoolPrefix, "/")
	if spoolPrefix != "" {
		spoolPrefix += "/"
	}
	return &CommandRunner{PythonPath: pythonPath, ScriptPath: scriptPath, SpoolRoot: spoolRoot, SpoolPrefix: spoolPrefix, MaxBytes: maxBytes}, nil
}

func (runner *CommandRunner) Run(ctx context.Context, manifest Manifest) (Result, error) {
	manifestJSON, err := json.Marshal(map[string]any{
		"spec_id": manifest.VersionID, "node_id": manifest.NodeID, "connection_id": manifest.ConnectionID,
		"source_type": manifest.SourceType, "source_name": manifest.SourceName,
		"relative_path": manifest.RelativePath, "name": manifest.Name,
		"size_bytes": manifest.SizeBytes, "locator_hash": manifest.LocatorHash,
	})
	if err != nil {
		return Result{}, errors.New("artifact_fetch_manifest_failed")
	}
	command := exec.CommandContext(ctx, runner.PythonPath, runner.ScriptPath, string(manifestJSON))
	// The literal guard keeps secrets out of the child environment: the script
	// only ever reads pre-mounted files.
	command.Env = []string{
		"PATH=/usr/bin:/bin",
		"PYTHONUNBUFFERED=1",
		"INGESTION_SPOOL_ROOT=" + runner.SpoolRoot,
		"SPOOL_REF_PREFIX=" + runner.SpoolPrefix,
		"ARTIFACT_FETCH_MAX_BYTES=" + strconv.FormatInt(runner.MaxBytes, 10),
		"TZ=UTC",
	}
	stdout, err := command.StdoutPipe()
	if err != nil {
		return Result{}, errors.New("artifact_fetch_start_failed")
	}
	command.Stderr = &limitedWriter{buffer: bytes.NewBuffer(nil), limit: 4096}
	if err := command.Start(); err != nil {
		return Result{}, errors.New("artifact_fetch_start_failed")
	}
	var result Result
	decoded := false
	scanner := bufio.NewScanner(stdout)
	scanner.Buffer(make([]byte, 4096), 1<<20)
	for scanner.Scan() {
		line := bytes.TrimSpace(scanner.Bytes())
		if len(line) == 0 || !bytes.HasPrefix(line, []byte("{")) {
			continue
		}
		decoder := json.NewDecoder(bytes.NewReader(line))
		decoder.DisallowUnknownFields()
		if err := decoder.Decode(&result); err != nil {
			_ = command.Process.Kill()
			_ = command.Wait()
			return Result{}, errors.New("artifact_fetch_invalid_output")
		}
		decoded = true
		break
	}
	waitErr := command.Wait()
	if !decoded {
		if waitErr != nil {
			return Result{OK: false, Error: "artifact fetch exited without a result"}, errors.New("artifact_fetch_failed")
		}
		return Result{}, errors.New("artifact_fetch_no_output")
	}
	return result, nil
}

var ErrManifestGone = errors.New("artifact fetch manifest is gone")

// limitedWriter captures at most limit bytes of script stderr; full output is
// never retained (it could echo source details we do not want in failures).
type limitedWriter struct {
	buffer *bytes.Buffer
	limit  int
}

func (writer *limitedWriter) Write(payload []byte) (int, error) {
	room := writer.limit - writer.buffer.Len()
	if room > 0 {
		if len(payload) > room {
			writer.buffer.Write(payload[:room])
			return len(payload), nil
		}
		writer.buffer.Write(payload)
	}
	return len(payload), nil
}