From 2f6338f5640126f9ac9814fef8988a1c8acaa5fc Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Thu, 10 Oct 2024 13:17:00 +0300 Subject: [PATCH] Add template mapping functions These are handy to deal with properties in protobufs within minder's templates. Signed-off-by: Juan Antonio Osorio --- internal/util/safe_template.go | 73 +++++++++++++- internal/util/safe_template_test.go | 150 ++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+), 4 deletions(-) diff --git a/internal/util/safe_template.go b/internal/util/safe_template.go index 207e3015cd..9f291ddfd4 100644 --- a/internal/util/safe_template.go +++ b/internal/util/safe_template.go @@ -22,6 +22,7 @@ import ( "fmt" htmltemplate "html/template" "io" + "reflect" "text/template" "github.com/rs/zerolog" @@ -32,6 +33,17 @@ var ( ErrExceededSizeLimit = errors.New("exceeded size limit") ) +var ( + // TemplateFuncs is a map of functions that can be used in templates + // It introduces two custom functions: + // - asMap: converts a structpb (or anything that implements the AsMap function call) to a map + // - mapGet: returns the value of a key in a map + TemplateFuncs = template.FuncMap{ + "asMap": asMap, + "mapGet": mapGet, + } +) + // SafeTemplate is a `template` wrapper that ensures that the template is // rendered in a safe and secure manner. That is, with memory limits // and timeouts. @@ -44,9 +56,56 @@ type templater interface { Name() string } +// This is a utility interface that allows us to accept any type +type asMapper interface { + AsMap() map[string]interface{} +} + +// asMap converts a structpb to a map +func asMap(s any) (reflect.Value, error) { + if s == nil { + return reflect.Value{}, fmt.Errorf("asMap called with nil") + } + + inspb, ok := s.(asMapper) + if !ok { + return reflect.Value{}, fmt.Errorf("invalid type: %T", s) + } + + return reflect.ValueOf(inspb.AsMap()), nil +} + +// mapGet returns the value of a key in a map +// The map could be a map[string]interface{} or a asMapper +// So we need to handle both cases +func mapGet(m any, key string) (reflect.Value, error) { + if m == nil { + return reflect.Value{}, fmt.Errorf("map is nil") + } + + // Check if the map is a map[string]interface{} + if mm, ok := m.(map[string]interface{}); ok { + return valueOfKey(mm, key) + } + + if mm, ok := m.(asMapper); ok { + mm := mm.AsMap() + return valueOfKey(mm, key) + } + + return reflect.Value{}, fmt.Errorf("invalid type: %T", m) +} + +func valueOfKey(m map[string]interface{}, key string) (reflect.Value, error) { + if v, ok := m[key]; ok { + return reflect.ValueOf(v), nil + } + return reflect.Value{}, fmt.Errorf("key not found: %s", key) +} + // NewSafeTextTemplate creates a new SafeTemplate for text templates func NewSafeTextTemplate(tmpl *string, name string) (*SafeTemplate, error) { - t, err := parseNewTextTemplate(tmpl, name) + t, err := parseNewTextTemplate(tmpl, name, TemplateFuncs) if err != nil { return nil, err } @@ -58,7 +117,7 @@ func NewSafeTextTemplate(tmpl *string, name string) (*SafeTemplate, error) { // NewSafeHTMLTemplate creates a new SafeTemplate for HTML templates func NewSafeHTMLTemplate(tmpl *string, name string) (*SafeTemplate, error) { - t, err := parseNewHtmlTemplate(tmpl, name) + t, err := parseNewHtmlTemplate(tmpl, name, TemplateFuncs) if err != nil { return nil, err } @@ -96,12 +155,15 @@ func (t *SafeTemplate) Execute(ctx context.Context, w io.Writer, data any, limit } // parseNewTextTemplate parses a named template from a string, ensuring it is not empty -func parseNewTextTemplate(tmpl *string, name string) (*template.Template, error) { +func parseNewTextTemplate(tmpl *string, name string, fnmap template.FuncMap) (*template.Template, error) { if tmpl == nil || len(*tmpl) == 0 { return nil, fmt.Errorf("missing template") } t := template.New(name).Option("missingkey=error") + if fnmap != nil { + t = t.Funcs(fnmap) + } t, err := t.Parse(*tmpl) if err != nil { return nil, fmt.Errorf("cannot parse template: %w", err) @@ -111,12 +173,15 @@ func parseNewTextTemplate(tmpl *string, name string) (*template.Template, error) } // parseNewHtmlTemplate parses a named template from a string, ensuring it is not empty -func parseNewHtmlTemplate(tmpl *string, name string) (*htmltemplate.Template, error) { +func parseNewHtmlTemplate(tmpl *string, name string, fnmap template.FuncMap) (*htmltemplate.Template, error) { if tmpl == nil || len(*tmpl) == 0 { return nil, fmt.Errorf("missing template") } t := htmltemplate.New(name).Option("missingkey=error") + if fnmap != nil { + t = t.Funcs(fnmap) + } t, err := t.Parse(*tmpl) if err != nil { return nil, fmt.Errorf("cannot parse template: %w", err) diff --git a/internal/util/safe_template_test.go b/internal/util/safe_template_test.go index ac74759a45..58ded6ff46 100644 --- a/internal/util/safe_template_test.go +++ b/internal/util/safe_template_test.go @@ -22,6 +22,8 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + structpb "google.golang.org/protobuf/types/known/structpb" "github.com/stacklok/minder/internal/util" ) @@ -208,3 +210,151 @@ func verybigstring(n int) string { } return s } + +func TestRenderStructPB(t *testing.T) { + t.Parallel() + + const limit = 1024 + + type args struct { + tmpl string + s any + } + tests := []struct { + name string + args args + expected string + wantErr bool + }{ + { + name: "asMap: valid template", + args: args{ + tmpl: "{{ with $m := asMap . }}{{ $m.name }}{{ end }}", + s: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "name": { + Kind: &structpb.Value_StringValue{ + StringValue: "test", + }, + }, + }, + }, + }, + expected: "test", + wantErr: false, + }, + { + name: "asMap: using wrong key", + args: args{ + tmpl: "{{ with $m := asMap . }}{{ $m.name2 }}{{ end }}", + s: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "name": { + Kind: &structpb.Value_StringValue{ + StringValue: "test", + }, + }, + }, + }, + }, + expected: "", + wantErr: true, + }, + { + name: "asMap: using wrong type", + args: args{ + tmpl: "{{ with $m := asMap . }}{{ $m.name }}{{ end }}", + s: "test", + }, + expected: "", + wantErr: true, + }, + { + name: "asMap: nil structpb", + args: args{ + tmpl: "{{ with $m := asMap . }}{{ $m.name }}{{ end }}", + s: nil, + }, + expected: "", + wantErr: true, + }, + { + name: "mapGet: valid with map[string]any", + args: args{ + tmpl: "{{ mapGet . \"name\" }}", + s: map[string]any{ + "name": "test", + }, + }, + expected: "test", + wantErr: false, + }, + { + name: "mapGet: valid with asMapper", + args: args{ + tmpl: "{{ mapGet . \"name\" }}", + s: &structpb.Struct{ + Fields: map[string]*structpb.Value{ + "name": { + Kind: &structpb.Value_StringValue{ + StringValue: "test", + }, + }, + }, + }, + }, + expected: "test", + wantErr: false, + }, + { + name: "mapGet: using wrong key", + args: args{ + tmpl: "{{ mapGet . \"name2\" }}", + s: map[string]any{ + "name": "test", + }, + }, + expected: "", + wantErr: true, + }, + { + name: "mapGet: using wrong type", + args: args{ + tmpl: "{{ mapGet . \"name\" }}", + s: "test", + }, + expected: "", + wantErr: true, + }, + { + name: "mapGet: nil map", + args: args{ + tmpl: "{{ mapGet . \"name\" }}", + s: nil, + }, + expected: "", + wantErr: true, + }, + } + + for _, tt := range tests { + tt := tt + + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tmpl, err := util.NewSafeTextTemplate(&tt.args.tmpl, "test") + // We're not testing the template parsing here + require.NoError(t, err, "unexpected error") + + out, err := tmpl.Render(context.Background(), tt.args.s, limit) + if tt.wantErr { + assert.Error(t, err, "expected error") + } else { + assert.NoError(t, err, "unexpected error") + assert.Equal(t, tt.expected, out, "expected output") + } + }) + } + +}