notarius/internal/modules/output/json/encoder.go

package json

import (
	"context"
	stdjson "encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const Key = "json"

const contentTypeJSON = "application/json"

var safeOutputFileChar = regexp.MustCompile(`[^A-Za-z0-9._-]`)

var _ contracts.OutputEncoder = (*Encoder)(nil)

type Encoder struct{}

func New() *Encoder {
	return &Encoder{}
}

func (e *Encoder) Key() string {
	return Key
}

func (e *Encoder) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	if e == nil {
		return contracts.OutputResult{}, encoderErrorf("encoder must not be nil")
	}
	if ctx == nil {
		return contracts.OutputResult{}, encoderErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.OutputResult{}, encoderErrorf("context error before encoding: %w", err)
	}

	files, err := logicalFiles(req)
	if err != nil {
		return contracts.OutputResult{}, err
	}
	return contracts.OutputResult{Files: files}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:      Key,
		Stage:    pipeline.StageOutput,
		Requires: []string{"normalized"},
		Provides: []string{"encoded"},
	}
}

func Register(registry *pipeline.OutputEncoderRegistry) error {
	return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.OutputEncoder, error) {
		return New(), nil
	})
}

type indexFile struct {
	ManifestFile string            `json:"manifest_file"`
	OutputFiles  []outputFileIndex `json:"output_files"`
	RejectedFile string            `json:"rejected_file"`
	WarningsFile string            `json:"warnings_file"`
}

type outputFileIndex struct {
	LaneID     string `json:"lane_id"`
	MediaType  string `json:"media_type,omitempty"`
	File       string `json:"file"`
	ModuleKey  string `json:"module_key,omitempty"`
	SchemaID   string `json:"schema_id,omitempty"`
	SchemaName string `json:"schema_name,omitempty"`
	SchemaVer  string `json:"schema_version,omitempty"`
}

type rejectedFile struct {
	Rejected []contracts.RejectedOutput `json:"rejected"`
}

type warningsFile struct {
	Warnings []contracts.Warning `json:"warnings"`
}

func logicalFiles(req contracts.OutputRequest) ([]contracts.OutputFile, error) {
	outputs := cloneNormalizeOutputs(req.NormalizeOutputs)
	sort.SliceStable(outputs, func(i, j int) bool {
		return outputs[i].LaneID < outputs[j].LaneID
	})

	outputIndexes := make([]outputFileIndex, 0, len(outputs))
	files := make([]contracts.OutputFile, 0, len(outputs)+4)
	manifestFile, err := jsonFile("manifest.json", req.Manifest)
	if err != nil {
		return nil, err
	}
	files = append(files, manifestFile)

	usedOutputFiles := make(map[string]string, len(outputs))
	for _, output := range outputs {
		name, err := outputFileName(output.LaneID)
		if err != nil {
			return nil, err
		}
		if existingLane, ok := usedOutputFiles[name]; ok {
			return nil, encoderErrorf("lanes %q and %q produce duplicate output file %q", existingLane, output.LaneID, name)
		}
		usedOutputFiles[name] = output.LaneID
		outputIndexes = append(outputIndexes, outputFileIndex{
			LaneID:     output.LaneID,
			MediaType:  output.Payload.MediaType,
			File:       name,
			ModuleKey:  output.NormalizerKey,
			SchemaID:   output.Schema.ID,
			SchemaName: output.Schema.Name,
			SchemaVer:  output.Schema.Version,
		})
		file, err := rawOutputFile(name, output.Payload)
		if err != nil {
			return nil, err
		}
		files = append(files, file)
	}

	index := indexFile{
		ManifestFile: "manifest.json",
		OutputFiles:  outputIndexes,
		RejectedFile: "rejected.json",
		WarningsFile: "warnings.json",
	}
	indexOutput, err := jsonFile("index.json", index)
	if err != nil {
		return nil, err
	}
	rejectedOutput, err := jsonFile("rejected.json", rejectedFile{Rejected: cloneRejected(req.Rejected)})
	if err != nil {
		return nil, err
	}
	warningsOutput, err := jsonFile("warnings.json", warningsFile{Warnings: cloneWarnings(req.Warnings)})
	if err != nil {
		return nil, err
	}
	files = append(files, indexOutput, rejectedOutput, warningsOutput)
	sort.Slice(files, func(i, j int) bool {
		return files[i].Name < files[j].Name
	})
	return files, nil
}

func rawOutputFile(name string, payload contracts.RawPayload) (contracts.OutputFile, error) {
	content := append([]byte(nil), payload.Content...)
	if len(content) == 0 {
		content = []byte("null")
	}
	if payload.MediaType == contentTypeJSON && stdjson.Valid(content) {
		var decoded any
		if err := stdjson.Unmarshal(content, &decoded); err == nil {
			pretty, err := marshalPretty(decoded)
			if err != nil {
				return contracts.OutputFile{}, err
			}
			content = pretty
		}
	}
	mediaType := strings.TrimSpace(payload.MediaType)
	if mediaType == "" {
		mediaType = "application/octet-stream"
	}
	return contracts.OutputFile{
		Name:        name,
		ContentType: mediaType,
		Bytes:       append([]byte(nil), content...),
	}, nil
}

func jsonFile(name string, value any) (contracts.OutputFile, error) {
	data, err := marshalPretty(value)
	if err != nil {
		return contracts.OutputFile{}, encoderErrorf("encode %s: %w", name, err)
	}
	return contracts.OutputFile{
		Name:        name,
		ContentType: contentTypeJSON,
		Bytes:       data,
	}, nil
}

func marshalPretty(value any) ([]byte, error) {
	data, err := stdjson.MarshalIndent(value, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(data, '\n'), nil
}

func outputFileName(laneID string) (string, error) {
	sanitized := safeOutputFileChar.ReplaceAllString(strings.TrimSpace(laneID), "_")
	for strings.Contains(sanitized, "..") {
		sanitized = strings.ReplaceAll(sanitized, "..", "__")
	}
	sanitized = strings.Trim(sanitized, "._")
	if sanitized == "" {
		return "", encoderErrorf("lane id %q cannot produce a safe file name", laneID)
	}
	return "outputs/" + sanitized + ".json", nil
}

func cloneNormalizeOutputs(outputs []contracts.NormalizeOutput) []contracts.NormalizeOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]contracts.NormalizeOutput, 0, len(outputs))
	for _, output := range outputs {
		output.Payload = cloneRawPayload(output.Payload)
		out = append(out, output)
	}
	return out
}

func cloneRawPayload(payload contracts.RawPayload) contracts.RawPayload {
	return contracts.RawPayload{
		Content:   append([]byte(nil), payload.Content...),
		MediaType: payload.MediaType,
		Metadata:  cloneMetadata(payload.Metadata),
		Warnings:  append([]contracts.Warning(nil), payload.Warnings...),
	}
}

func cloneRejected(rejected []contracts.RejectedOutput) []contracts.RejectedOutput {
	if len(rejected) == 0 {
		return []contracts.RejectedOutput{}
	}
	return append([]contracts.RejectedOutput(nil), rejected...)
}

func cloneWarnings(warnings []contracts.Warning) []contracts.Warning {
	if len(warnings) == 0 {
		return []contracts.Warning{}
	}
	return append([]contracts.Warning(nil), warnings...)
}

func cloneMetadata(metadata map[string]any) map[string]any {
	if len(metadata) == 0 {
		return nil
	}
	out := make(map[string]any, len(metadata))
	for key, value := range metadata {
		out[key] = value
	}
	return out
}

func encoderErrorf(format string, args ...any) error {
	return fmt.Errorf("json output encoder: "+format, args...)
}