Skip to content

Commit eae7820

Browse files
joshdrakejoshdrake
authored
Merge pull request #1117 from smallstep/josh/register-template-func
Let applications contribute template functions
2 parents 38e55c7 + 88702aa commit eae7820

8 files changed

Lines changed: 452 additions & 11 deletions

File tree

‎internal/templates/registry.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
package templates
2+
3+
import (
4+
"fmt"
5+
"reflect"
6+
"sync"
7+
"text/template"
8+
)
9+
10+
// Registry holds template functions added by the application, for the
11+
// functions a template needs that this library does not provide.
12+
type Registry struct {
13+
// reserved is a function so a package can name its own built-ins without
14+
// this file knowing them. It is resolved once, on first use.
15+
reserved func() map[string]struct{}
16+
once sync.Once
17+
names map[string]struct{}
18+
19+
mu sync.RWMutex
20+
funcs map[string]any
21+
}
22+
23+
// NewRegistry returns a Registry that refuses any name returned by reserved.
24+
func NewRegistry(reserved func() map[string]struct{}) *Registry {
25+
return &Registry{reserved: reserved, funcs: map[string]any{}}
26+
}
27+
28+
// Register adds fn to the registry. It returns an error if name is already
29+
// registered or reserved.
30+
func (r *Registry) Register(name string, fn any) error {
31+
if err := validate(name, fn); err != nil {
32+
return err
33+
}
34+
35+
r.once.Do(func() { r.names = r.reserved() })
36+
if _, ok := r.names[name]; ok {
37+
return fmt.Errorf("template function %q is built in and cannot be replaced", name)
38+
}
39+
40+
r.mu.Lock()
41+
defer r.mu.Unlock()
42+
if _, ok := r.funcs[name]; ok {
43+
return fmt.Errorf("template function %q is already registered", name)
44+
}
45+
r.funcs[name] = fn
46+
return nil
47+
}
48+
49+
// Replace adds fn to the registry, replacing any reserved or previously
50+
// registered function with the same name.
51+
func (r *Registry) Replace(name string, fn any) error {
52+
if err := validate(name, fn); err != nil {
53+
return err
54+
}
55+
r.mu.Lock()
56+
defer r.mu.Unlock()
57+
r.funcs[name] = fn
58+
return nil
59+
}
60+
61+
func validate(name string, fn any) error {
62+
switch {
63+
case name == "":
64+
return fmt.Errorf("template function name is required")
65+
case fn == nil:
66+
return fmt.Errorf("template function %q is nil", name)
67+
case reflect.TypeOf(fn).Kind() != reflect.Func:
68+
return fmt.Errorf("template function %q is a %s, not a function", name, reflect.TypeOf(fn).Kind())
69+
case !validName(name):
70+
return fmt.Errorf("template function name %q is not a valid identifier", name)
71+
}
72+
return nil
73+
}
74+
75+
// Unregister removes a function from the registry. It returns true if a
76+
// function was removed.
77+
func (r *Registry) Unregister(name string) bool {
78+
r.mu.Lock()
79+
defer r.mu.Unlock()
80+
if _, ok := r.funcs[name]; !ok {
81+
return false
82+
}
83+
delete(r.funcs, name)
84+
return true
85+
}
86+
87+
// Apply adds the registered functions to funcMap.
88+
func (r *Registry) Apply(funcMap template.FuncMap) {
89+
r.mu.RLock()
90+
defer r.mu.RUnlock()
91+
for name, fn := range r.funcs {
92+
funcMap[name] = fn
93+
}
94+
}
95+
96+
// validName reports whether name is a valid Go identifier, as required by
97+
// "text/template".
98+
func validName(name string) bool {
99+
for i, c := range name {
100+
switch {
101+
case c == '_':
102+
case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z':
103+
case c >= '0' && c <= '9' && i > 0:
104+
default:
105+
return false
106+
}
107+
}
108+
return true
109+
}

‎sshutil/funcs.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package sshutil
2+
3+
import (
4+
"text/template"
5+
6+
"go.step.sm/crypto/internal/templates"
7+
)
8+
9+
// templateFuncs holds the functions registered by the application. It is
10+
// separate from the X.509 registry, so registering for one kind of certificate
11+
// does not affect the other.
12+
var templateFuncs = templates.NewRegistry(func() map[string]struct{} {
13+
names := map[string]struct{}{}
14+
for name := range builtinFuncMap(new(TemplateError)) {
15+
names[name] = struct{}{}
16+
}
17+
return names
18+
})
19+
20+
// RegisterTemplateFunc adds fn to the functions available to SSH certificate
21+
// templates. It returns an error if name is already registered or built in.
22+
//
23+
// It behaves as [go.step.sm/crypto/x509util.RegisterTemplateFunc] does, over a
24+
// separate registry; an application that wants a function in both calls both.
25+
func RegisterTemplateFunc(name string, fn any) error {
26+
return templateFuncs.Register(name, fn)
27+
}
28+
29+
// ReplaceTemplateFunc adds fn to the functions available to SSH certificate
30+
// templates, replacing a built-in or previously registered function with the
31+
// same name. Use [RegisterTemplateFunc] unless the replacement is intended.
32+
func ReplaceTemplateFunc(name string, fn any) error {
33+
return templateFuncs.Replace(name, fn)
34+
}
35+
36+
// UnregisterTemplateFunc removes a registered function. It returns true if a
37+
// function was removed.
38+
func UnregisterTemplateFunc(name string) bool {
39+
return templateFuncs.Unregister(name)
40+
}
41+
42+
// builtinFuncMap returns the functions provided by this package, excluding
43+
// those registered by the application.
44+
func builtinFuncMap(err *TemplateError) template.FuncMap {
45+
return templates.GetFuncMap(&err.Message)
46+
}

‎sshutil/funcs_test.go‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package sshutil
2+
3+
import (
4+
"crypto/ecdsa"
5+
"crypto/elliptic"
6+
"crypto/rand"
7+
"testing"
8+
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
12+
"go.step.sm/crypto/x509util"
13+
)
14+
15+
func TestRegisterTemplateFunc(t *testing.T) {
16+
cr := CertificateRequest{Key: mustGeneratePublicKey(t), Type: UserCert.String()}
17+
data := CreateTemplateData(UserCert, "jane@example.com", []string{"jane"})
18+
19+
require.NoError(t, RegisterTemplateFunc("testPrincipals", func(data any) (any, error) {
20+
m, _ := data.(TemplateData)
21+
return m[PrincipalsKey], nil
22+
}))
23+
t.Cleanup(func() { UnregisterTemplateFunc("testPrincipals") })
24+
25+
var o Options
26+
require.NoError(t, WithTemplate(`{{ testPrincipals $ | toJson }}`, data)(cr, &o))
27+
assert.Equal(t, `["jane"]`, o.CertBuffer.String())
28+
}
29+
30+
func TestRegisterTemplateFuncErrors(t *testing.T) {
31+
require.Error(t, RegisterTemplateFunc("", func() string { return "" }))
32+
require.Error(t, RegisterTemplateFunc("notfn", "a string"))
33+
34+
err := RegisterTemplateFunc("toJson", func() string { return "" })
35+
require.Error(t, err)
36+
assert.Contains(t, err.Error(), "built in and cannot be replaced")
37+
}
38+
39+
// TestRegistriesAreSeparate checks that a function registered for one kind of
40+
// certificate is not available to the other.
41+
func TestRegistriesAreSeparate(t *testing.T) {
42+
require.NoError(t, RegisterTemplateFunc("testSSHOnly", func() string { return "ssh" }))
43+
t.Cleanup(func() { UnregisterTemplateFunc("testSSHOnly") })
44+
45+
cr := CertificateRequest{Key: mustGeneratePublicKey(t), Type: UserCert.String()}
46+
data := CreateTemplateData(UserCert, "jane@example.com", []string{"jane"})
47+
48+
var o Options
49+
require.NoError(t, WithTemplate(`{{ testSSHOnly }}`, data)(cr, &o))
50+
assert.Equal(t, "ssh", o.CertBuffer.String())
51+
52+
// The same name is undefined for X.509 templates.
53+
signer, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
54+
require.NoError(t, err)
55+
xcr, err := x509util.CreateCertificateRequest("foo", []string{"foo.com"}, signer)
56+
require.NoError(t, err)
57+
var xo x509util.Options
58+
err = x509util.WithTemplate(`{{ testSSHOnly }}`, x509util.TemplateData{})(xcr, &xo)
59+
require.Error(t, err)
60+
assert.Contains(t, err.Error(), `function "testSSHOnly" not defined`)
61+
62+
// An application that wants it in both registers with both.
63+
require.NoError(t, x509util.RegisterTemplateFunc("testSSHOnly", func() string { return "x509" }))
64+
t.Cleanup(func() { x509util.UnregisterTemplateFunc("testSSHOnly") })
65+
66+
var xo2 x509util.Options
67+
require.NoError(t, x509util.WithTemplate(`{{ testSSHOnly }}`, x509util.TemplateData{})(xcr, &xo2))
68+
assert.Equal(t, "x509", xo2.CertBuffer.String())
69+
}

‎sshutil/options.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@ import (
77
"text/template"
88

99
"github.com/pkg/errors"
10-
11-
"go.step.sm/crypto/internal/templates"
1210
)
1311

1412
// Options are the options that can be passed to NewCertificate.
@@ -36,7 +34,9 @@ func GetFuncMap() template.FuncMap {
3634
}
3735

3836
func getFuncMap(err *TemplateError) template.FuncMap {
39-
return templates.GetFuncMap(&err.Message)
37+
funcMap := builtinFuncMap(err)
38+
templateFuncs.Apply(funcMap)
39+
return funcMap
4040
}
4141

4242
// WithTemplate is an options that executes the given template text with the

‎x509util/funcs.go‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package x509util
2+
3+
import (
4+
"text/template"
5+
6+
"go.step.sm/crypto/internal/templates"
7+
)
8+
9+
// templateFuncs holds the functions registered by the application. The reserved
10+
// set is this package's own function map, so a registration cannot shadow one.
11+
var templateFuncs = templates.NewRegistry(func() map[string]struct{} {
12+
names := map[string]struct{}{}
13+
for name := range builtinFuncMap(new(TemplateError)) {
14+
names[name] = struct{}{}
15+
}
16+
return names
17+
})
18+
19+
// RegisterTemplateFunc adds fn to the functions available to X.509 certificate
20+
// templates. It returns an error if name is already registered or built in.
21+
//
22+
// Register during start-up. "text/template" resolves function names when it
23+
// parses, so a template rendered before the call will fail to parse.
24+
//
25+
// A function receives only its own arguments. One that needs the template data
26+
// takes it as a parameter, which the template passes as "$" rather than ".",
27+
// as the dot is rebound inside a range block:
28+
//
29+
// {{ cel "device.serial" $ | toJson }}
30+
func RegisterTemplateFunc(name string, fn any) error {
31+
return templateFuncs.Register(name, fn)
32+
}
33+
34+
// ReplaceTemplateFunc adds fn to the functions available to X.509 certificate
35+
// templates, replacing a built-in or previously registered function with the
36+
// same name. Use [RegisterTemplateFunc] unless the replacement is intended.
37+
func ReplaceTemplateFunc(name string, fn any) error {
38+
return templateFuncs.Replace(name, fn)
39+
}
40+
41+
// UnregisterTemplateFunc removes a registered function. It returns true if a
42+
// function was removed.
43+
func UnregisterTemplateFunc(name string) bool {
44+
return templateFuncs.Unregister(name)
45+
}
46+
47+
// builtinFuncMap returns the functions provided by this package, excluding
48+
// those registered by the application.
49+
func builtinFuncMap(err *TemplateError) template.FuncMap {
50+
funcMap := templates.GetFuncMap(&err.Message)
51+
// asn1 methods
52+
funcMap["asn1Enc"] = asn1Encode
53+
funcMap["asn1Marshal"] = asn1Marshal
54+
funcMap["asn1Seq"] = asn1Sequence
55+
funcMap["asn1Set"] = asn1Set
56+
return funcMap
57+
}

‎x509util/funcs_replace_test.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
package x509util
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/assert"
7+
"github.com/stretchr/testify/require"
8+
)
9+
10+
// TestReplaceTemplateFunc checks that a built-in can be replaced deliberately,
11+
// but not by [RegisterTemplateFunc].
12+
func TestReplaceTemplateFunc(t *testing.T) {
13+
cr, _ := createCertificateRequest(t, "foo", []string{"foo.com"})
14+
15+
require.Error(t, RegisterTemplateFunc("toJson", func(any) string { return "replaced" }))
16+
17+
require.NoError(t, ReplaceTemplateFunc("toJson", func(any) string { return "replaced" }))
18+
t.Cleanup(func() { UnregisterTemplateFunc("toJson") })
19+
20+
var o Options
21+
require.NoError(t, WithTemplate(`{{ toJson .Subject }}`, TemplateData{})(cr, &o))
22+
assert.Equal(t, "replaced", o.CertBuffer.String())
23+
24+
// Removing it restores the built-in.
25+
UnregisterTemplateFunc("toJson")
26+
var o2 Options
27+
require.NoError(t, WithTemplate(`{{ toJson "x" }}`, TemplateData{})(cr, &o2))
28+
assert.Equal(t, `"x"`, o2.CertBuffer.String())
29+
}
30+
31+
func TestReplaceTemplateFuncOverridesARegistration(t *testing.T) {
32+
require.NoError(t, RegisterTemplateFunc("testReplaceMe", func() string { return "first" }))
33+
t.Cleanup(func() { UnregisterTemplateFunc("testReplaceMe") })
34+
35+
require.NoError(t, ReplaceTemplateFunc("testReplaceMe", func() string { return "second" }))
36+
37+
cr, _ := createCertificateRequest(t, "foo", []string{"foo.com"})
38+
var o Options
39+
require.NoError(t, WithTemplate(`{{ testReplaceMe }}`, TemplateData{})(cr, &o))
40+
assert.Equal(t, "second", o.CertBuffer.String())
41+
}
42+
43+
func TestReplaceTemplateFuncStillValidates(t *testing.T) {
44+
require.Error(t, ReplaceTemplateFunc("", func() string { return "" }))
45+
require.Error(t, ReplaceTemplateFunc("bad-name", func() string { return "" }))
46+
require.Error(t, ReplaceTemplateFunc("notfn", "a string"))
47+
require.Error(t, ReplaceTemplateFunc("nilfn", nil))
48+
}

0 commit comments

Comments
 (0)