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.
This commit is contained in:
Rubens Peculis 2024-07-17 21:12:16 +10:00 committed by GitHub
parent 8823778d05
commit 00ff288f0c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 65 additions and 2 deletions

View file

@ -373,8 +373,7 @@ func (node *Node) Proto() *v1.Node {
User: node.User.Proto(), User: node.User.Proto(),
ForcedTags: node.ForcedTags, ForcedTags: node.ForcedTags,
// TODO(kradalby): Implement register method enum converter RegisterMethod: node.RegisterMethodToV1Enum(),
// RegisterMethod: ,
CreatedAt: timestamppb.New(node.CreatedAt), CreatedAt: timestamppb.New(node.CreatedAt),
} }
@ -489,6 +488,19 @@ func (node *Node) PeerChangeFromMapRequest(req tailcfg.MapRequest) tailcfg.PeerC
return ret 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. // ApplyPeerChange takes a PeerChange struct and updates the node.
func (node *Node) ApplyPeerChange(change *tailcfg.PeerChange) { func (node *Node) ApplyPeerChange(change *tailcfg.PeerChange) {
if change.Key != nil { if change.Key != nil {

View file

@ -6,6 +6,7 @@ import (
"github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts" "github.com/google/go-cmp/cmp/cmpopts"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/juanfont/headscale/hscontrol/util" "github.com/juanfont/headscale/hscontrol/util"
"tailscale.com/tailcfg" "tailscale.com/tailcfg"
"tailscale.com/types/key" "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)
}
})
}
}