Let applications contribute template functions

internal/templates/registry.go

@@ -0,0 +1,109 @@
+package templates
+
+import (
+	"fmt"
+	"reflect"
+	"sync"
+	"text/template"
+)
+
+// Registry holds template functions added by the application, for the
+// functions a template needs that this library does not provide.
+type Registry struct {
+	// reserved is a function so a package can name its own built-ins without
+	// this file knowing them. It is resolved once, on first use.
+	reserved func() map[string]struct{}
+	once     sync.Once
+	names    map[string]struct{}
+
+	mu    sync.RWMutex
+	funcs map[string]any
+}
+
+// NewRegistry returns a Registry that refuses any name returned by reserved.
+func NewRegistry(reserved func() map[string]struct{}) *Registry {
+	return &Registry{reserved: reserved, funcs: map[string]any{}}
+}
+
+// Register adds fn to the registry. It returns an error if name is already
+// registered or reserved.
+func (r *Registry) Register(name string, fn any) error {
+	if err := validate(name, fn); err != nil {
+		return err
+	}
+
+	r.once.Do(func() { r.names = r.reserved() })
+	if _, ok := r.names[name]; ok {
+		return fmt.Errorf("template function %q is built in and cannot be replaced", name)
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if _, ok := r.funcs[name]; ok {
+		return fmt.Errorf("template function %q is already registered", name)
+	}
+	r.funcs[name] = fn
+	return nil
+}
+
+// Replace adds fn to the registry, replacing any reserved or previously
+// registered function with the same name.
+func (r *Registry) Replace(name string, fn any) error {
+	if err := validate(name, fn); err != nil {
+		return err
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.funcs[name] = fn
+	return nil
+}
+
+func validate(name string, fn any) error {
+	switch {
+	case name == "":
+		return fmt.Errorf("template function name is required")
+	case fn == nil:
+		return fmt.Errorf("template function %q is nil", name)
+	case reflect.TypeOf(fn).Kind() != reflect.Func:
+		return fmt.Errorf("template function %q is a %s, not a function", name, reflect.TypeOf(fn).Kind())
+	case !validName(name):
+		return fmt.Errorf("template function name %q is not a valid identifier", name)
+	}
+	return nil
+}
+
+// Unregister removes a function from the registry. It returns true if a
+// function was removed.
+func (r *Registry) Unregister(name string) bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if _, ok := r.funcs[name]; !ok {
+		return false
+	}
+	delete(r.funcs, name)
+	return true
+}
+
+// Apply adds the registered functions to funcMap.
+func (r *Registry) Apply(funcMap template.FuncMap) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	for name, fn := range r.funcs {
+		funcMap[name] = fn
+	}
+}
+
+// validName reports whether name is a valid Go identifier, as required by
+// "text/template".
+func validName(name string) bool {
+	for i, c := range name {
+		switch {
+		case c == '_':
+		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z':
+		case c >= '0' && c <= '9' && i > 0:
+		default:
+			return false
+		}
+	}
+	return true
+}

sshutil/funcs.go

@@ -0,0 +1,46 @@
+package sshutil
+
+import (
+	"text/template"
+
+	"go.step.sm/crypto/internal/templates"
+)
+
+// templateFuncs holds the functions registered by the application. It is
+// separate from the X.509 registry, so registering for one kind of certificate
+// does not affect the other.
+var templateFuncs = templates.NewRegistry(func() map[string]struct{} {
+	names := map[string]struct{}{}
+	for name := range builtinFuncMap(new(TemplateError)) {
+		names[name] = struct{}{}
+	}
+	return names
+})
+
+// RegisterTemplateFunc adds fn to the functions available to SSH certificate
+// templates. It returns an error if name is already registered or built in.
+//
+// It behaves as [go.step.sm/crypto/x509util.RegisterTemplateFunc] does, over a
+// separate registry; an application that wants a function in both calls both.
+func RegisterTemplateFunc(name string, fn any) error {
+	return templateFuncs.Register(name, fn)
+}
+
+// ReplaceTemplateFunc adds fn to the functions available to SSH certificate
+// templates, replacing a built-in or previously registered function with the
+// same name. Use [RegisterTemplateFunc] unless the replacement is intended.
+func ReplaceTemplateFunc(name string, fn any) error {
+	return templateFuncs.Replace(name, fn)
+}
+
+// UnregisterTemplateFunc removes a registered function. It returns true if a
+// function was removed.
+func UnregisterTemplateFunc(name string) bool {
+	return templateFuncs.Unregister(name)
+}
+
+// builtinFuncMap returns the functions provided by this package, excluding
+// those registered by the application.
+func builtinFuncMap(err *TemplateError) template.FuncMap {
+	return templates.GetFuncMap(&err.Message)
+}

sshutil/funcs_test.go

@@ -0,0 +1,69 @@
+package sshutil
+
+import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.step.sm/crypto/x509util"
+)
+
+func TestRegisterTemplateFunc(t *testing.T) {
+	cr := CertificateRequest{Key: mustGeneratePublicKey(t), Type: UserCert.String()}
+	data := CreateTemplateData(UserCert, "jane@example.com", []string{"jane"})
+
+	require.NoError(t, RegisterTemplateFunc("testPrincipals", func(data any) (any, error) {
+		m, _ := data.(TemplateData)
+		return m[PrincipalsKey], nil
+	}))
+	t.Cleanup(func() { UnregisterTemplateFunc("testPrincipals") })
+
+	var o Options
+	require.NoError(t, WithTemplate(`{{ testPrincipals $ | toJson }}`, data)(cr, &o))
+	assert.Equal(t, `["jane"]`, o.CertBuffer.String())
+}
+
+func TestRegisterTemplateFuncErrors(t *testing.T) {
+	require.Error(t, RegisterTemplateFunc("", func() string { return "" }))
+	require.Error(t, RegisterTemplateFunc("notfn", "a string"))
+
+	err := RegisterTemplateFunc("toJson", func() string { return "" })
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "built in and cannot be replaced")
+}
+
+// TestRegistriesAreSeparate checks that a function registered for one kind of
+// certificate is not available to the other.
+func TestRegistriesAreSeparate(t *testing.T) {
+	require.NoError(t, RegisterTemplateFunc("testSSHOnly", func() string { return "ssh" }))
+	t.Cleanup(func() { UnregisterTemplateFunc("testSSHOnly") })
+
+	cr := CertificateRequest{Key: mustGeneratePublicKey(t), Type: UserCert.String()}
+	data := CreateTemplateData(UserCert, "jane@example.com", []string{"jane"})
+
+	var o Options
+	require.NoError(t, WithTemplate(`{{ testSSHOnly }}`, data)(cr, &o))
+	assert.Equal(t, "ssh", o.CertBuffer.String())
+
+	// The same name is undefined for X.509 templates.
+	signer, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+	xcr, err := x509util.CreateCertificateRequest("foo", []string{"foo.com"}, signer)
+	require.NoError(t, err)
+	var xo x509util.Options
+	err = x509util.WithTemplate(`{{ testSSHOnly }}`, x509util.TemplateData{})(xcr, &xo)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), `function "testSSHOnly" not defined`)
+
+	// An application that wants it in both registers with both.
+	require.NoError(t, x509util.RegisterTemplateFunc("testSSHOnly", func() string { return "x509" }))
+	t.Cleanup(func() { x509util.UnregisterTemplateFunc("testSSHOnly") })
+
+	var xo2 x509util.Options
+	require.NoError(t, x509util.WithTemplate(`{{ testSSHOnly }}`, x509util.TemplateData{})(xcr, &xo2))
+	assert.Equal(t, "x509", xo2.CertBuffer.String())
+}

sshutil/options.go

@@ -7,8 +7,6 @@ import (
 	"text/template"
 
 	"github.com/pkg/errors"
-
-	"go.step.sm/crypto/internal/templates"
 )
 
 // Options are the options that can be passed to NewCertificate.
@@ -36,7 +34,9 @@ func GetFuncMap() template.FuncMap {
 }
 
 func getFuncMap(err *TemplateError) template.FuncMap {
-	return templates.GetFuncMap(&err.Message)
+	funcMap := builtinFuncMap(err)
+	templateFuncs.Apply(funcMap)
+	return funcMap
 }
 
 // WithTemplate is an options that executes the given template text with the

x509util/funcs.go

@@ -0,0 +1,57 @@
+package x509util
+
+import (
+	"text/template"
+
+	"go.step.sm/crypto/internal/templates"
+)
+
+// templateFuncs holds the functions registered by the application. The reserved
+// set is this package's own function map, so a registration cannot shadow one.
+var templateFuncs = templates.NewRegistry(func() map[string]struct{} {
+	names := map[string]struct{}{}
+	for name := range builtinFuncMap(new(TemplateError)) {
+		names[name] = struct{}{}
+	}
+	return names
+})
+
+// RegisterTemplateFunc adds fn to the functions available to X.509 certificate
+// templates. It returns an error if name is already registered or built in.
+//
+// Register during start-up. "text/template" resolves function names when it
+// parses, so a template rendered before the call will fail to parse.
+//
+// A function receives only its own arguments. One that needs the template data
+// takes it as a parameter, which the template passes as "$" rather than ".",
+// as the dot is rebound inside a range block:
+//
+//	{{ cel "device.serial" $ | toJson }}
+func RegisterTemplateFunc(name string, fn any) error {
+	return templateFuncs.Register(name, fn)
+}
+
+// ReplaceTemplateFunc adds fn to the functions available to X.509 certificate
+// templates, replacing a built-in or previously registered function with the
+// same name. Use [RegisterTemplateFunc] unless the replacement is intended.
+func ReplaceTemplateFunc(name string, fn any) error {
+	return templateFuncs.Replace(name, fn)
+}
+
+// UnregisterTemplateFunc removes a registered function. It returns true if a
+// function was removed.
+func UnregisterTemplateFunc(name string) bool {
+	return templateFuncs.Unregister(name)
+}
+
+// builtinFuncMap returns the functions provided by this package, excluding
+// those registered by the application.
+func builtinFuncMap(err *TemplateError) template.FuncMap {
+	funcMap := templates.GetFuncMap(&err.Message)
+	// asn1 methods
+	funcMap["asn1Enc"] = asn1Encode
+	funcMap["asn1Marshal"] = asn1Marshal
+	funcMap["asn1Seq"] = asn1Sequence
+	funcMap["asn1Set"] = asn1Set
+	return funcMap
+}

x509util/funcs_replace_test.go

@@ -0,0 +1,48 @@
+package x509util
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReplaceTemplateFunc checks that a built-in can be replaced deliberately,
+// but not by [RegisterTemplateFunc].
+func TestReplaceTemplateFunc(t *testing.T) {
+	cr, _ := createCertificateRequest(t, "foo", []string{"foo.com"})
+
+	require.Error(t, RegisterTemplateFunc("toJson", func(any) string { return "replaced" }))
+
+	require.NoError(t, ReplaceTemplateFunc("toJson", func(any) string { return "replaced" }))
+	t.Cleanup(func() { UnregisterTemplateFunc("toJson") })
+
+	var o Options
+	require.NoError(t, WithTemplate(`{{ toJson .Subject }}`, TemplateData{})(cr, &o))
+	assert.Equal(t, "replaced", o.CertBuffer.String())
+
+	// Removing it restores the built-in.
+	UnregisterTemplateFunc("toJson")
+	var o2 Options
+	require.NoError(t, WithTemplate(`{{ toJson "x" }}`, TemplateData{})(cr, &o2))
+	assert.Equal(t, `"x"`, o2.CertBuffer.String())
+}
+
+func TestReplaceTemplateFuncOverridesARegistration(t *testing.T) {
+	require.NoError(t, RegisterTemplateFunc("testReplaceMe", func() string { return "first" }))
+	t.Cleanup(func() { UnregisterTemplateFunc("testReplaceMe") })
+
+	require.NoError(t, ReplaceTemplateFunc("testReplaceMe", func() string { return "second" }))
+
+	cr, _ := createCertificateRequest(t, "foo", []string{"foo.com"})
+	var o Options
+	require.NoError(t, WithTemplate(`{{ testReplaceMe }}`, TemplateData{})(cr, &o))
+	assert.Equal(t, "second", o.CertBuffer.String())
+}
+
+func TestReplaceTemplateFuncStillValidates(t *testing.T) {
+	require.Error(t, ReplaceTemplateFunc("", func() string { return "" }))
+	require.Error(t, ReplaceTemplateFunc("bad-name", func() string { return "" }))
+	require.Error(t, ReplaceTemplateFunc("notfn", "a string"))
+	require.Error(t, ReplaceTemplateFunc("nilfn", nil))
+}