From e84e19a04eefccc743f0d397efe49ff42626f4b3 Mon Sep 17 00:00:00 2001
From: Tyler Kropp <tyler@tylerkropp.xyz>
Date: Mon, 2 May 2022 16:55:34 -0400
Subject: [PATCH] templates: Add custom template function registration (#4757)

* Add custom template function registration

* Rename TemplateFunctions to CustomFunctions

* Add documentation

* Document CustomFunctions interface

* Preallocate custom functions map list

* Fix interface name in error message
---
 modules/caddyhttp/templates/templates.go  | 32 ++++++++++++++++++++---
 modules/caddyhttp/templates/tplcontext.go | 14 +++++++---
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/modules/caddyhttp/templates/templates.go b/modules/caddyhttp/templates/templates.go
index 8f2ae9dcfc3..5aff27f5b68 100644
--- a/modules/caddyhttp/templates/templates.go
+++ b/modules/caddyhttp/templates/templates.go
@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"text/template"
 
 	"github.com/caddyserver/caddy/v2"
 	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
@@ -36,6 +37,8 @@ func init() {
 //
 // ⚠️ Template functions/actions are still experimental, so they are subject to change.
 //
+// Custom template functions can be registered by creating a plugin module under the `http.handlers.templates.functions.*` namespace that implements the `CustomFunctions` interface.
+//
 // [All Sprig functions](https://masterminds.github.io/sprig/) are supported.
 //
 // In addition to the standard functions and the Sprig library, Caddy adds
@@ -249,6 +252,14 @@ type Templates struct {
 	// The template action delimiters. If set, must be precisely two elements:
 	// the opening and closing delimiters. Default: `["{{", "}}"]`
 	Delimiters []string `json:"delimiters,omitempty"`
+
+	customFuncs []template.FuncMap
+}
+
+// Customfunctions is the interface for registering custom template functions.
+type CustomFunctions interface {
+	// CustomTemplateFunctions should return the mapping from custom function names to implementations.
+	CustomTemplateFunctions() template.FuncMap
 }
 
 // CaddyModule returns the Caddy module information.
@@ -261,6 +272,18 @@ func (Templates) CaddyModule() caddy.ModuleInfo {
 
 // Provision provisions t.
 func (t *Templates) Provision(ctx caddy.Context) error {
+	fnModInfos := caddy.GetModules("http.handlers.templates.functions")
+	customFuncs := make([]template.FuncMap, len(fnModInfos), 0)
+	for _, modInfo := range fnModInfos {
+		mod := modInfo.New()
+		fnMod, ok := mod.(CustomFunctions)
+		if !ok {
+			return fmt.Errorf("module %q does not satisfy the CustomFunctions interface", modInfo.ID)
+		}
+		customFuncs = append(customFuncs, fnMod.CustomTemplateFunctions())
+	}
+	t.customFuncs = customFuncs
+
 	if t.MIMETypes == nil {
 		t.MIMETypes = defaultMIMETypes
 	}
@@ -331,10 +354,11 @@ func (t *Templates) executeTemplate(rr caddyhttp.ResponseRecorder, r *http.Reque
 	}
 
 	ctx := &TemplateContext{
-		Root:       fs,
-		Req:        r,
-		RespHeader: WrappedHeader{rr.Header()},
-		config:     t,
+		Root:        fs,
+		Req:         r,
+		RespHeader:  WrappedHeader{rr.Header()},
+		config:      t,
+		CustomFuncs: t.customFuncs,
 	}
 
 	err := ctx.executeTemplateInBuffer(r.URL.Path, rr.Buffer())
diff --git a/modules/caddyhttp/templates/tplcontext.go b/modules/caddyhttp/templates/tplcontext.go
index 4f3cbf50a80..7843455a5a3 100644
--- a/modules/caddyhttp/templates/tplcontext.go
+++ b/modules/caddyhttp/templates/tplcontext.go
@@ -40,10 +40,11 @@ import (
 
 // TemplateContext is the TemplateContext with which HTTP templates are executed.
 type TemplateContext struct {
-	Root       http.FileSystem
-	Req        *http.Request
-	Args       []interface{} // defined by arguments to funcInclude
-	RespHeader WrappedHeader
+	Root        http.FileSystem
+	Req         *http.Request
+	Args        []interface{} // defined by arguments to funcInclude
+	RespHeader  WrappedHeader
+	CustomFuncs []template.FuncMap // functions added by plugins
 
 	config *Templates
 	tpl    *template.Template
@@ -62,6 +63,11 @@ func (c *TemplateContext) NewTemplate(tplName string) *template.Template {
 	// add sprig library
 	c.tpl.Funcs(sprigFuncMap)
 
+	// add all custom functions
+	for _, funcMap := range c.CustomFuncs {
+		c.tpl.Funcs(funcMap)
+	}
+
 	// add our own library
 	c.tpl.Funcs(template.FuncMap{
 		"include":          c.funcInclude,