From ca51bd91bd7a1df8914c4caf618ece5bb5ecf803 Mon Sep 17 00:00:00 2001
From: Chris Baker <1675087+cgbaker@users.noreply.github.com>
Date: Mon, 22 Mar 2021 01:49:21 +0000
Subject: [PATCH] refactor?

---
 nomad/structs/json_encoding.go | 43 ++++++++++++++++++++++++++--------
 nomad/structs/structs.go       |  4 ++--
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/nomad/structs/json_encoding.go b/nomad/structs/json_encoding.go
index 3f1a6e2b0..2a2ccaa66 100644
--- a/nomad/structs/json_encoding.go
+++ b/nomad/structs/json_encoding.go
@@ -6,13 +6,11 @@ import (
 	"github.com/hashicorp/go-msgpack/codec"
 )
 
-// Special encoding for structs.Node, to perform the following:
-// 1. provide backwards compatibility for the following fields:
-//   * Node.Drain
-type nodeExt struct{}
+func init() {
+	registerExtension(reflect.TypeOf(Node{}), nodeExt)
+}
 
-// ConvertExt converts a structs.Node to a struct with the extra field, Drain
-func (n nodeExt) ConvertExt(v interface{}) interface{} {
+func nodeExt(v interface{}) interface{} {
 	node := v.(*Node)
 	if node == nil {
 		return nil
@@ -27,10 +25,35 @@ func (n nodeExt) ConvertExt(v interface{}) interface{} {
 	}
 }
 
-// UpdateExt is not used
-func (n nodeExt) UpdateExt(_ interface{}, _ interface{}) {}
+// BOILERPLATE GOES HERE
 
-func RegisterJSONEncodingExtensions(h *codec.JsonHandle) *codec.JsonHandle {
-	h.SetInterfaceExt(reflect.TypeOf(Node{}), 1, nodeExt{})
+type extendFunc func(interface{}) interface{}
+
+var (
+	extendedTypes = map[reflect.Type]extendFunc{}
+)
+
+func registerExtension(tpe reflect.Type, ext extendFunc) {
+	extendedTypes[tpe] = ext
+}
+
+type nomadJsonEncodingExtensions struct{}
+
+// ConvertExt calls the registered conversions functions
+func (n nomadJsonEncodingExtensions) ConvertExt(v interface{}) interface{} {
+	if fn, ok := extendedTypes[reflect.TypeOf(v)]; ok {
+		return fn(v)
+	} else {
+		return nil
+	}
+}
+
+// UpdateExt is not used
+func (n nomadJsonEncodingExtensions) UpdateExt(_ interface{}, _ interface{}) {}
+
+func NomadJsonEncodingExtensions(h *codec.JsonHandle) *codec.JsonHandle {
+	for tpe, _ := range extendedTypes {
+		h.SetInterfaceExt(tpe, 1, nomadJsonEncodingExtensions{})
+	}
 	return h
 }
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 02d6a2648..222c7b86a 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -10589,10 +10589,10 @@ var (
 	JsonHandle = &codec.JsonHandle{
 		HTMLCharsAsIs: true,
 	}
-	JsonHandleWithExtensions = RegisterJSONEncodingExtensions(&codec.JsonHandle{
+	JsonHandleWithExtensions = NomadJsonEncodingExtensions(&codec.JsonHandle{
 		HTMLCharsAsIs: true,
 	})
-	JsonHandlePretty = RegisterJSONEncodingExtensions(&codec.JsonHandle{
+	JsonHandlePretty = NomadJsonEncodingExtensions(&codec.JsonHandle{
 		HTMLCharsAsIs: true,
 		Indent:        4,
 	})