From 5024c58388b2dedfd59f08c5127d232def298cac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tadej=20Jane=C5=BE?= <tadej.j@nez.si>
Date: Thu, 2 Sep 2021 13:50:32 +0200
Subject: [PATCH] go/common/node: Add custom text (un)marshaler for RolesMask
 type

This will result in easy to understand "roles" fields in various CLI
commands that output JSON.
---
 go/common/node/node.go      | 44 ++++++++++++++++++++++++++++-
 go/common/node/node_test.go | 56 +++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/go/common/node/node.go b/go/common/node/node.go
index a9baee46ff8..60fd0c3ba51 100644
--- a/go/common/node/node.go
+++ b/go/common/node/node.go
@@ -24,6 +24,9 @@ import (
 )
 
 var (
+	// ErrInvalidRole is the error returned when a node role is invalid.
+	ErrInvalidRole = errors.New("node: invalid role")
+
 	// ErrInvalidTEEHardware is the error returned when a TEE hardware
 	// implementation is invalid.
 	ErrInvalidTEEHardware = errors.New("node: invalid TEE implementation")
@@ -109,6 +112,15 @@ const (
 	// RoleReserved are all the bits of the Oasis node roles bitmask
 	// that are reserved and must not be used.
 	RoleReserved RolesMask = ((1 << 32) - 1) & ^((RoleStorageRPC << 1) - 1)
+
+	RoleComputeWorkerName = "compute"
+	RoleStorageWorkerName = "storage"
+	RoleKeyManagerName    = "key-manager"
+	RoleValidatorName     = "validator"
+	RoleConsensusRPCName  = "consensus-rpc"
+	RoleStorageRPCName    = "storage-rpc"
+
+	RolesMaskStringSep = ","
 )
 
 // Roles returns a list of available valid roles.
@@ -154,7 +166,37 @@ func (m RolesMask) String() string {
 		ret = append(ret, "storage-rpc")
 	}
 
-	return strings.Join(ret, ",")
+	return strings.Join(ret, RolesMaskStringSep)
+}
+
+// MarshalText encodes a RolesMask into text form.
+func (m RolesMask) MarshalText() ([]byte, error) {
+	return []byte(m.String()), nil
+}
+
+// UnmarshalText decodes a text slice into a RolesMask.
+func (k *RolesMask) UnmarshalText(text []byte) error {
+	*k = 0
+	roles := strings.Split(string(text), RolesMaskStringSep)
+	for _, role := range roles {
+		switch role {
+		case RoleComputeWorkerName:
+			*k |= RoleComputeWorker
+		case RoleStorageWorkerName:
+			*k |= RoleStorageWorker
+		case RoleKeyManagerName:
+			*k |= RoleKeyManager
+		case RoleValidatorName:
+			*k |= RoleValidator
+		case RoleConsensusRPCName:
+			*k |= RoleConsensusRPC
+		case RoleStorageRPCName:
+			*k |= RoleStorageRPC
+		default:
+			return fmt.Errorf("%w: '%s'", ErrInvalidRole, role)
+		}
+	}
+	return nil
 }
 
 // ValidateBasic performs basic descriptor validity checks.
diff --git a/go/common/node/node_test.go b/go/common/node/node_test.go
index ac6098df749..a67ef2c71f4 100644
--- a/go/common/node/node_test.go
+++ b/go/common/node/node_test.go
@@ -9,6 +9,62 @@ import (
 	"github.com/oasisprotocol/oasis-core/go/common/cbor"
 )
 
+func TestRolesMask(t *testing.T) {
+	require := require.New(t)
+
+	testVectors := []struct {
+		rolesMaskString            string
+		rolesMaskStringUnmarshable bool
+		rolesMask                  RolesMask
+		rolesMaskStringCanonical   bool
+		errMsg                     string
+	}{
+		// Valid single roles.
+		{"compute", true, 1, true, ""},
+		{"storage", true, 2, true, ""},
+		{"key-manager", true, 4, true, ""},
+		{"validator", true, 8, true, ""},
+		{"consensus-rpc", true, 16, true, ""},
+		{"storage-rpc", true, 32, true, ""},
+		// Valid multiple roles.
+		{"compute,storage", true, 3, true, ""},
+		{"compute,storage,validator", true, 11, true, ""},
+		{"compute,storage,validator,consensus-rpc", true, 27, true, ""},
+		{"validator,consensus-rpc", true, 24, true, ""},
+		{"storage,storage-rpc", true, 34, true, ""},
+
+		// Invalid.
+		{"compute ", false, 1, false, "node: invalid role: 'compute '"},
+		{"compute ,", false, 1, false, "node: invalid role: 'compute '"},
+		{" validator", false, 1, false, "node: invalid role: ' validator'"},
+		{"compute, storage", false, 1, false, "node: invalid role: ' storage'"},
+		{"master", false, 1, false, "node: invalid role: 'master'"},
+		{"storage-rpc,storage", true, 34, false, ""},
+	}
+
+	for _, v := range testVectors {
+		var unmarshaledRolesMask RolesMask
+		err := unmarshaledRolesMask.UnmarshalText([]byte(v.rolesMaskString))
+		if !v.rolesMaskStringUnmarshable {
+			require.EqualErrorf(
+				err,
+				v.errMsg,
+				"Unmarshaling invalid roles mask: '%s' should fail with expected error message",
+				v.rolesMaskString,
+			)
+			continue
+		}
+		require.NoErrorf(err, "Failed to unmarshal a valid roles mask: '%s'", v.rolesMaskString)
+		require.Equal(v.rolesMask, unmarshaledRolesMask, "Unmarshaled roles mask doesn't equal expected roles mask")
+
+		textRolesMask, err := v.rolesMask.MarshalText()
+		require.NoError(err, "Failed to marshal a valid roles mask: '%s'", v.rolesMask)
+		if v.rolesMaskStringCanonical {
+			require.Equal(v.rolesMaskString, string(textRolesMask), "Marshaled roles mask doesn't equal expected text roles mask")
+		}
+	}
+}
+
 func TestNodeDescriptor(t *testing.T) {
 	require := require.New(t)