diagram.go

package main

import (
	"bytes"
	"context"
	"sync"

	graphviz "github.com/goccy/go-graphviz"
	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/parser"
	"github.com/yuin/goldmark/renderer"
	"github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"
)

// Diagrams turn a fenced code block tagged as a graph language into inline SVG,
// rendered on the server so the page needs no JavaScript:
//
//	```dot
//	digraph { a -> b -> c }
//	```
//
// The rendering engine is go-graphviz, which runs Graphviz as an embedded
// WebAssembly module under wazero: no cgo, no external dot binary, no browser.
//
// A diagram whose source does not compile falls back to echoing that source in
// a failure-tagged code block rather than panicking or failing the whole page,
// matching vitrine's fail-stable-at-request-time stance.

// diagramTags are the fenced-code languages rendered as diagrams.  dot is the
// canonical Graphviz language; graphviz is an accepted alias, mirroring how
// callouts fold warn onto warning.  Both feed the same DOT engine.
var diagramTags = map[string]bool{
	"dot":      true,
	"graphviz": true,
}

// maxDiagramSource caps the DOT handed to the engine.  A render holds a
// process-wide lock (see diagramEngine), so an oversized diagram would stall
// every other diagram render while Graphviz lays it out; this bounds the obvious
// axis of that exposure.  It does not bound layout cost in general -- a small but
// adversarial graph can still be expensive -- so it is a coarse guard, not a
// complete one.  Legitimate diagrams are far smaller than this.
const maxDiagramSource = 64 << 10

// diagramNode is the block node produced by diagramTransformer in place of a
// fenced code block, and consumed by diagramRenderer.  It carries the DOT source
// captured at transform time; it has no children.
type diagramNode struct {
	ast.BaseBlock
	Source []byte
}

var kindDiagram = ast.NewNodeKind("Diagram")

func (n *diagramNode) Kind() ast.NodeKind { return kindDiagram }

func (n *diagramNode) Dump(src []byte, level int) { ast.DumpHelper(n, src, level, nil, nil) }

// diagramEngine renders DOT to SVG through one long-lived go-graphviz instance,
// serialized by a mutex.
//
// go-graphviz drives Graphviz as a single process-global WebAssembly module
// (one instance, built in the library's init, with shared mutable state and no
// internal locking).  Concurrent renders through it corrupt that state and abort
// the process with "fatal error: concurrent map writes".  vitrine serves
// requests concurrently, so holding engineMu across every render is load-bearing,
// not a nicety.  Renders are a couple of milliseconds and diagrams are rare, so
// serializing them costs nothing that matters here.
//
// The *graphviz.Graphviz is created once and reused: reuse is safe under the
// lock and deterministic across renders, and it avoids re-running plugin setup
// against the wasm module for every diagram.  If creation ever fails the error
// is remembered rather than retried, since that failure is deterministic; every
// diagram then takes the fail-visible source-echo path.
type diagramEngine struct {
	mu    sync.Mutex
	gv    *graphviz.Graphviz
	err   error
	ready bool
}

var engine diagramEngine

func (e *diagramEngine) renderDOT(src []byte) ([]byte, error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if !e.ready {
		e.gv, e.err = graphviz.New(context.Background())
		e.ready = true
	}
	if e.err != nil {
		return nil, e.err
	}
	graph, err := graphviz.ParseBytes(src)
	if err != nil {
		return nil, err
	}
	defer graph.Close()
	var buf bytes.Buffer
	if err := e.gv.Render(context.Background(), graph, graphviz.SVG, &buf); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// diagramTransformer swaps every diagram-tagged fenced code block for a
// diagramNode.  Matching blocks are collected first and replaced after the walk,
// so the tree is not mutated while it is being traversed.
type diagramTransformer struct{}

func (diagramTransformer) Transform(doc *ast.Document, reader text.Reader, pc parser.Context) {
	source := reader.Source()
	var targets []*ast.FencedCodeBlock
	_ = ast.Walk(doc, func(n ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		if fc, ok := n.(*ast.FencedCodeBlock); ok {
			if lang := fc.Language(source); lang != nil && diagramTags[string(lang)] {
				targets = append(targets, fc)
			}
		}
		return ast.WalkContinue, nil
	})
	for _, fc := range targets {
		d := &diagramNode{Source: fencedText(fc, source)}
		fc.Parent().ReplaceChild(fc.Parent(), fc, d)
	}
}

// fencedText concatenates the raw lines of a fenced code block back into one
// buffer, the DOT source to render.
func fencedText(fc *ast.FencedCodeBlock, source []byte) []byte {
	var buf bytes.Buffer
	lines := fc.Lines()
	for i := 0; i < lines.Len(); i++ {
		seg := lines.At(i)
		buf.Write(seg.Value(source))
	}
	return buf.Bytes()
}

// diagramRenderer writes a diagramNode as inline SVG in a "diagram" container,
// so style.css can inset and scale it consistently with other article blocks.
type diagramRenderer struct{}

func (diagramRenderer) RegisterFuncs(reg renderer.NodeRendererFuncRegisterer) {
	reg.Register(kindDiagram, renderDiagram)
}

func renderDiagram(w util.BufWriter, source []byte, node ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}
	n := node.(*diagramNode)
	if len(n.Source) > maxDiagramSource {
		writeDiagramFallback(w, n.Source, "diagram source exceeds the size limit")
		return ast.WalkContinue, nil
	}
	svg, err := engine.renderDOT(n.Source)
	if err != nil {
		writeDiagramFallback(w, n.Source, err.Error())
		return ast.WalkContinue, nil
	}
	_, _ = w.WriteString(`<div class="diagram">` + "\n")
	_, _ = w.Write(stripSVGProlog(svg))
	_, _ = w.WriteString("</div>\n")
	return ast.WalkContinue, nil
}

// writeDiagramFallback renders a diagram that could not be produced as its raw
// source in a code block, with a visible note carrying the reason.  Echoing the
// source keeps a bad diagram from taking the page down and shows the author
// exactly what did not render.  The source and reason are HTML-escaped, so a
// hostile diagram cannot inject structure through the failure path.
func writeDiagramFallback(w util.BufWriter, src []byte, reason string) {
	_, _ = w.WriteString(`<div class="diagram diagram-error">` + "\n")
	_, _ = w.WriteString(`<p class="diagram-error-note">diagram failed to render: `)
	_, _ = w.Write(util.EscapeHTML([]byte(reason)))
	_, _ = w.WriteString("</p>\n<pre><code>")
	_, _ = w.Write(util.EscapeHTML(src))
	_, _ = w.WriteString("</code></pre>\n</div>\n")
}

// stripSVGProlog trims graphviz's XML declaration, DOCTYPE, and generator
// comments, returning the document from the opening <svg tag onward.  Inline SVG
// in an HTML page needs no XML prolog.  If no <svg tag is found the input is
// returned unchanged, so an unexpected output still reaches the page.
func stripSVGProlog(svg []byte) []byte {
	if i := bytes.Index(svg, []byte("<svg")); i >= 0 {
		return svg[i:]
	}
	return svg
}

// diagrams is the goldmark extender that wires the diagram transformer and
// renderer into the pipeline, at the same priority callouts use.
type diagrams struct{}

func (diagrams) Extend(m goldmark.Markdown) {
	m.Parser().AddOptions(parser.WithASTTransformers(
		util.Prioritized(diagramTransformer{}, 100),
	))
	m.Renderer().AddOptions(renderer.WithNodeRenderers(
		util.Prioritized(diagramRenderer{}, 100),
	))
}