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

package json

import (
	"context"
	stdjson "encoding/json"
	"fmt"
	"mime"
	"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 Options struct{}

type Encoder struct {
	options Options
}

func New() *Encoder {
	options, _ := DecodeOptions(nil)
	return &Encoder{options: options}
}

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.RegisterBuilderWithSpec(ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.OutputEncoder, error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return &Encoder{options: options}, nil
	})
}

func validateOptions(options map[string]any) error {
	_, err := DecodeOptions(options)
	return err
}

func DecodeOptions(options map[string]any) (Options, error) {
	if err := pipeline.RejectUnknownOptions(options); err != nil {
		return Options{}, encoderErrorf("%w", err)
	}
	return Options{}, 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")
	}
	mediaType := strings.TrimSpace(payload.MediaType)
	if mediaType == "" {
		mediaType = "application/octet-stream"
	}
	if !isJSONMediaType(mediaType) {
		return contracts.OutputFile{}, encoderErrorf("normalized output %q has unsupported media type %q", name, mediaType)
	}
	var decoded any
	if err := stdjson.Unmarshal(content, &decoded); err != nil {
		return contracts.OutputFile{}, encoderErrorf("normalized output %q contains invalid JSON: %w", name, err)
	}
	pretty, err := marshalPretty(decoded)
	if err != nil {
		return contracts.OutputFile{}, err
	}
	return contracts.OutputFile{
		Name:        name,
		ContentType: mediaType,
		Bytes:       pretty,
	}, nil
}

func isJSONMediaType(mediaType string) bool {
	base, _, err := mime.ParseMediaType(strings.TrimSpace(mediaType))
	if err != nil {
		base = strings.TrimSpace(mediaType)
	}
	return strings.EqualFold(base, contentTypeJSON)
}

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 "lanes/" + 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...)
}