From 00ff288f0cae53d25f4c40355d3bd77cb3b06f4b Mon Sep 17 00:00:00 2001
From: Rubens Peculis <rubenspeculis@users.noreply.github.com>
Date: Wed, 17 Jul 2024 21:12:16 +1000
Subject: [PATCH] fix(1996): Implement register method enum converter (#2013)

Added a new function `RegisterMethodToV1Enum()` to Node, converting the internal register method string to the corresponding V1 Enum value. Included corresponding unit test in `node_test.go` to ensure correct conversion for various register methods.
---
 hscontrol/types/node.go      | 16 +++++++++--
 hscontrol/types/node_test.go | 51 ++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/hscontrol/types/node.go b/hscontrol/types/node.go
index 6bee5c4..19b287a 100644
--- a/hscontrol/types/node.go
+++ b/hscontrol/types/node.go
@@ -373,8 +373,7 @@ func (node *Node) Proto() *v1.Node {
 		User:        node.User.Proto(),
 		ForcedTags:  node.ForcedTags,
 
-		// TODO(kradalby): Implement register method enum converter
-		// RegisterMethod: ,
+		RegisterMethod: node.RegisterMethodToV1Enum(),
 
 		CreatedAt: timestamppb.New(node.CreatedAt),
 	}
@@ -489,6 +488,19 @@ func (node *Node) PeerChangeFromMapRequest(req tailcfg.MapRequest) tailcfg.PeerC
 	return ret
 }
 
+func (node *Node) RegisterMethodToV1Enum() v1.RegisterMethod {
+	switch node.RegisterMethod {
+	case "authkey":
+		return v1.RegisterMethod_REGISTER_METHOD_AUTH_KEY
+	case "oidc":
+		return v1.RegisterMethod_REGISTER_METHOD_OIDC
+	case "cli":
+		return v1.RegisterMethod_REGISTER_METHOD_CLI
+	default:
+		return v1.RegisterMethod_REGISTER_METHOD_UNSPECIFIED
+	}
+}
+
 // ApplyPeerChange takes a PeerChange struct and updates the node.
 func (node *Node) ApplyPeerChange(change *tailcfg.PeerChange) {
 	if change.Key != nil {
diff --git a/hscontrol/types/node_test.go b/hscontrol/types/node_test.go
index 85857c3..798a54d 100644
--- a/hscontrol/types/node_test.go
+++ b/hscontrol/types/node_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
@@ -540,3 +541,53 @@ func TestApplyPeerChange(t *testing.T) {
 		})
 	}
 }
+
+func TestNodeRegisterMethodToV1Enum(t *testing.T) {
+	tests := []struct {
+		name string
+		node Node
+		want v1.RegisterMethod
+	}{
+		{
+			name: "authkey",
+			node: Node{
+				ID:             1,
+				RegisterMethod: util.RegisterMethodAuthKey,
+			},
+			want: v1.RegisterMethod_REGISTER_METHOD_AUTH_KEY,
+		},
+		{
+			name: "oidc",
+			node: Node{
+				ID:             1,
+				RegisterMethod: util.RegisterMethodOIDC,
+			},
+			want: v1.RegisterMethod_REGISTER_METHOD_OIDC,
+		},
+		{
+			name: "cli",
+			node: Node{
+				ID:             1,
+				RegisterMethod: util.RegisterMethodCLI,
+			},
+			want: v1.RegisterMethod_REGISTER_METHOD_CLI,
+		},
+		{
+			name: "unknown",
+			node: Node{
+				ID: 0,
+			},
+			want: v1.RegisterMethod_REGISTER_METHOD_UNSPECIFIED,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := tt.node.RegisterMethodToV1Enum()
+
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("RegisterMethodToV1Enum() unexpected result (-want +got):\n%s", diff)
+			}
+		})
+	}
+}