html/escape.go

package html

import (
	"bufio"
	"bytes"
	"errors"
	"io"
	"unicode"
)

const unicodeNBSP = 0xa0

type escapeWriter struct {
	out                     *bufio.Writer
	nonbreakSpaces          bool
	spaceStarted, lastSpace bool
	err                     error
}

func newEscapeWriter(out io.Writer, nonbreakSpaces bool) *escapeWriter {
	return &escapeWriter{
		out:            bufio.NewWriter(out),
		nonbreakSpaces: nonbreakSpaces,
	}
}

func (w *escapeWriter) write(r ...rune) {
	if w.err != nil {
		return
	}

	for _, ri := range r {
		if _, err := w.out.WriteRune(ri); err != nil {
			w.err = err
			return
		}
	}
}

func (w *escapeWriter) Write(p []byte) (int, error) {
	if w.err != nil {
		return 0, w.err
	}

	runes := bytes.NewBuffer(nil)
	if n, err := runes.Write(p); err != nil {
		w.err = err
		return n, w.err
	}

	for {
		r, _, err := runes.ReadRune()
		if errors.Is(err, io.EOF) {
			return len(p), nil
		}

		if err != nil {
			w.err = err
			return len(p), w.err
		}

		if r == unicode.ReplacementChar {
			continue
		}

		space := w.nonbreakSpaces && r == ' '
		switch {
		case space && w.spaceStarted:
			w.write([]rune("  ")...)
		case space && w.lastSpace:
			w.write([]rune(" ")...)
		case w.spaceStarted:
			w.write(' ')
		}

		w.spaceStarted = space && !w.lastSpace
		w.lastSpace = space
		if space {
			continue
		}

		switch r {
		case '<':
			w.write([]rune("&lt;")...)
		case '>':
			w.write([]rune("&gt;")...)
		case '&':
			w.write([]rune("&amp;")...)
		case unicodeNBSP:
			w.write([]rune("&nbsp;")...)
		default:
			w.write(r)
		}
	}

	return len(p), w.err
}

func (w *escapeWriter) Flush() error {
	if w.err != nil {
		return w.err
	}

	if w.spaceStarted {
		w.write(' ')
		w.spaceStarted = false
	}

	w.err = w.out.Flush()
	return w.err
}

func escape(s string, nonbreakSpaces bool) string {
	var b bytes.Buffer
	w := newEscapeWriter(&b, nonbreakSpaces)
	w.Write([]byte(s))
	w.Flush()
	return b.String()
}

func escapeAttribute(value string) string {
	var rr []rune
	r := []rune(value)
	for i := range r {
		switch r[i] {
		case '"':
			rr = append(rr, []rune("&quot;")...)
		case '&':
			rr = append(rr, []rune("&amp;")...)
		default:
			rr = append(rr, r[i])
		}
	}

	return string(rr)
}