sanitizer.go

// Copyright 2021-2022 Intrinsec. All rights reserved.

package main

import (
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/template"

	"github.com/intrinsec/protoc-gen-sanitize/sanitize"
	pgs "github.com/lyft/protoc-gen-star"
	pgsgo "github.com/lyft/protoc-gen-star/lang/go"
)

// SanitizeModule adds Sanitize methods on PB
type SanitizeModule struct {
	*pgs.ModuleBase
	ctx       pgsgo.Context
	tpl       *template.Template
	strict    bool
	hasErrors bool
}

// Sanitize returns an initialized SanitizePlugin
func Sanitize() *SanitizeModule {
	return &SanitizeModule{
		ModuleBase: &pgs.ModuleBase{},
		hasErrors:  false,
	}
}

// InitContext populates the module with needed context and fields
func (p *SanitizeModule) InitContext(c pgs.BuildContext) {
	c.Debug("InitContext")
	p.ModuleBase.InitContext(c)
	p.ctx = pgsgo.InitContext(c.Parameters())

	tpl := template.New("Sanitize").Funcs(map[string]interface{}{
		"package":           p.ctx.PackageName,
		"name":              p.ctx.Name,
		"sanitizer":         p.sanitizer,
		"initializer":       p.initializer,
		"leadingCommenter":  p.leadingCommenter,
		"isDisabledMessage": p.isDisabledMessage,
		"checkNoSanitize":   p.checkNoSanitize,
	})

	p.tpl = template.Must(tpl.Parse(sanitizeTpl))
}

// Name satisfies the generator.Plugin interface.
func (p *SanitizeModule) Name() string { return "Sanitize" }

// Execute generates sanitization code for files
func (p *SanitizeModule) Execute(targets map[string]pgs.File, pkgs map[string]pgs.Package) []pgs.Artifact {
	p.Debug("Execute")

	if ok, _ := p.Parameters().Bool("strict"); ok {
		p.strict = true
	}

	for _, t := range targets {
		if !p.doSanitize(t) {
			continue
		}
		p.generateFile(t)
	}

	return p.Artifacts()
}

func (p *SanitizeModule) ExitCheck() {
	if p.hasErrors && p.strict {
		p.Log("Sanitization strict mode enabled. Cannot generate with missing sanitize option fields or with sanitize option defined on fields of disabled message, check above errors")
		os.Exit(1)
	}
}

func (p *SanitizeModule) doSanitize(f pgs.File) bool {
	var disableFile bool

	if ok, err := f.Extension(sanitize.E_DisableFile, &disableFile); ok && err == nil && disableFile {
		p.Debug("Skipping: ", f.InputPath())
		return false
	}

	for _, m := range f.AllMessages() {
		var disableMessage bool
		if ok, err := m.Extension(sanitize.E_DisableMessage, &disableMessage); ok && err == nil && disableMessage {
			p.Debug("Skipping: ", m.Name())
			continue
		}

		for _, field := range m.Fields() {
			var rules sanitize.FieldRules
			if ok, err := field.Extension(sanitize.E_Rules, &rules); ok && err == nil {
				return true
			}
		}
	}

	p.Debug("No sanitization options encountered for:", f.InputPath())
	// Nonetheless we generate sanitization function to enable calling sanitization function of nested messages
	return true
}

func (p *SanitizeModule) generateFile(f pgs.File) {
	if len(f.Messages()) == 0 {
		return
	}
	p.Push(f.Name().String())
	defer p.Pop()
	p.Debug("File:", f.InputPath())

	name := f.InputPath().BaseName() + ".pb.sanitize.go"

	p.Debug("generate:", name)

	p.AddGeneratorTemplateFile(name, p.tpl, f)
}

func (p *SanitizeModule) leadingCommenter(f pgs.File) string {
	p.Debug("Comments:", f.SourceCodeInfo().LeadingDetachedComments())

	var comments []string
	re := regexp.MustCompile("(\r?\n)+")

	// Add default comment
	comments = append(comments, " Code generated by protoc-gen-sanitize. DO NOT EDIT.")

	for _, comment := range f.SourceCodeInfo().LeadingDetachedComments() {
		tmpCmt := re.Split(comment, -1)
		comments = append(comments, tmpCmt[:len(tmpCmt)-1]...)
	}

	return "//" + strings.Join(comments, "\n//")
}

func (p *SanitizeModule) initializer(m pgs.Message) string {
	html := false
	text := false

	var disableMessage bool
	if ok, err := m.Extension(sanitize.E_DisableMessage, &disableMessage); ok && err == nil && disableMessage {
		p.Debug("Skipping: ", m.Name())
		return ""
	}

	for _, f := range m.Fields() {
		var rules sanitize.FieldRules

		if ok, err := f.Extension(sanitize.E_Rules, &rules); ok && err == nil {
			switch rules.Kind {
			case sanitize.Sanitization_HTML:
				html = true
			case sanitize.Sanitization_TEXT:
				text = true
			}
		}
	}
	out := make([]string, 0)
	if html {
		out = append(out, "htmlSanitize := bluemonday.UGCPolicy()")
	}
	if text {
		out = append(out, "textSanitize := bluemonday.StrictPolicy()")
	}
	return strings.Join(out, "\n	")
}

func formatSanitizeCalls(varName string, argName string, sanitizeKind string, isRepeated bool, trim bool) string {

	sanitizeCalls := []string{}

	if isRepeated {
		varName = fmt.Sprintf("m.%s[i]", varName)
	} else {
		varName = fmt.Sprintf("m.%s", varName)
		argName = fmt.Sprintf("m.%s", argName)
	}

	sanitizeCalls = append(sanitizeCalls, fmt.Sprintf("%s = %sSanitize.Sanitize(%s)", varName, strings.ToLower(sanitizeKind), argName))

	if trim {
		sanitizeCalls = append(sanitizeCalls, fmt.Sprintf("%[1]s = strings.TrimSpace(%[1]s)", varName))
	}

	return strings.Join(sanitizeCalls, "\n")
}

func (p *SanitizeModule) buildSanitizeCall(f pgs.Field, name string, sanitizeKind string, trim bool) string {
	prefix := ""
	suffix := ""
	indent := ""
	sanitizeCall := ""
	iter := "i"
	elementName := name

	if f.Type().IsRepeated() {
		indent = "	"
		suffix = "\n}"
		elementName = strings.ToLower(string(name[0:2]))
		if sanitizeKind == "" {
			iter = "_"
		}
		prefix = fmt.Sprintf("for %s, %s := range m.%s {\n",
			iter,
			elementName,
			name,
		)
	}
	var format string
	if sanitizeKind == "" {
		// building call for message
		if f.Type().IsRepeated() {
			format = "%s.Sanitize()"
		} else {
			format = "m.%s.Sanitize()"
		}
		sanitizeCall = fmt.Sprintf(format, elementName)
	} else {
		// building calls for string
		sanitizeCall = formatSanitizeCalls(name, elementName, sanitizeKind, f.Type().IsRepeated(), trim)
	}
	return fmt.Sprintf("%s%s%s%s", prefix, indent, sanitizeCall, suffix)
}

func (p *SanitizeModule) isDisabledMessage(m pgs.Message) bool {
	disableMessage := false
	_, _ = m.Extension(sanitize.E_DisableMessage, &disableMessage)
	return disableMessage
}

func (p *SanitizeModule) checkNoSanitize(f pgs.Field) string {
	var rules sanitize.FieldRules

	name := p.ctx.Name(f)
	ok, err := f.Extension(sanitize.E_Rules, &rules)

	if err == nil && ok {
		fmt.Fprintf(
			os.Stderr,
			"%v:%d: sanitize option defined on %v for disabled message %v\n",
			f.File().Name(),
			f.SourceCodeInfo().Location().Span[0]+1,
			name,
			f.Message().Name())
		p.hasErrors = true
		return fmt.Sprintf("// sanitize option defined on %v for disabled message", name)
	}
	return ""
}

func (p *SanitizeModule) sanitizer(f pgs.Field) string {
	name := p.ctx.Name(f)

	var disableField bool
	if ok, err := f.Extension(sanitize.E_DisableField, &disableField); ok && err == nil && disableField {
		p.Debug("Skipping disabled field:", name)
		return "// sanitize is disabled on field " + name.String()
	}

	switch f.Type().ProtoType() {
	case pgs.StringT:
		var rules sanitize.FieldRules

		ok, err := f.Extension(sanitize.E_Rules, &rules)
		if err != nil {
			p.Logf(
				"%v:%d: Error can't retrieve rules extension for message %s with error: %s",
				f.File().Name(),
				f.SourceCodeInfo().Location().Span[0]+1,
				f.FullyQualifiedName(),
				err,
			)
			p.hasErrors = true
		} else if ok {
			switch rules.Kind {
			case sanitize.Sanitization_HTML:
				return p.buildSanitizeCall(f, string(name), "html", rules.GetTrim())
			case sanitize.Sanitization_TEXT:
				result := p.buildSanitizeCall(f, string(name), "text", rules.GetTrim())
				return result
			}
		}

	case pgs.MessageT:
		return p.buildSanitizeCall(f, string(name), "", false)
	}
	return ""
}

const sanitizeTpl = `{{ leadingCommenter . }}

package {{ package . }}

import "github.com/microcosm-cc/bluemonday"
import "strings"

{{ range .AllMessages }}
func (m *{{ name . }}) Sanitize() {
{{ if isDisabledMessage . }}
    // sanitize is disabled for this message
{{ range .Fields }}
    {{ checkNoSanitize . }}
{{ end }}
{{ else }}
	if m == nil {
		return
	}

	{{ initializer . }}

{{ range .Fields }}
    {{ sanitizer . }}
{{ end }}
{{ end }}
}
{{ end }}
`