From 8a3a0b64037bea8c15700943d37a16f02375cf28 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Sun, 27 Feb 2022 09:04:48 +0100
Subject: [PATCH] Add YAML support to ACLs

---
 acls.go       | 40 +++++++++++++++++++++++++++++++---------
 acls_types.go | 44 +++++++++++++++++++++++++++++++++-----------
 2 files changed, 64 insertions(+), 20 deletions(-)

diff --git a/acls.go b/acls.go
index ce14a89a..1a597c24 100644
--- a/acls.go
+++ b/acls.go
@@ -5,11 +5,13 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 
 	"github.com/rs/zerolog/log"
 	"github.com/tailscale/hujson"
+	"gopkg.in/yaml.v3"
 	"inet.af/netaddr"
 	"tailscale.com/tailcfg"
 )
@@ -53,16 +55,36 @@ func (h *Headscale) LoadACLPolicy(path string) error {
 		return err
 	}
 
-	ast, err := hujson.Parse(policyBytes)
-	if err != nil {
-		return err
-	}
-	ast.Standardize()
-	policyBytes = ast.Pack()
-	err = json.Unmarshal(policyBytes, &policy)
-	if err != nil {
-		return err
+	switch filepath.Ext(path) {
+	case ".yml", ".yaml":
+		log.Debug().
+			Str("path", path).
+			Bytes("file", policyBytes).
+			Msg("Loading ACLs from YAML")
+
+		err := yaml.Unmarshal(policyBytes, &policy)
+		if err != nil {
+			return err
+		}
+
+		log.Trace().
+			Interface("policy", policy).
+			Msg("Loaded policy from YAML")
+
+	default:
+		ast, err := hujson.Parse(policyBytes)
+		if err != nil {
+			return err
+		}
+
+		ast.Standardize()
+		policyBytes = ast.Pack()
+		err = json.Unmarshal(policyBytes, &policy)
+		if err != nil {
+			return err
+		}
 	}
+
 	if policy.IsZero() {
 		return errEmptyPolicy
 	}
diff --git a/acls_types.go b/acls_types.go
index 08e650ff..fb869821 100644
--- a/acls_types.go
+++ b/acls_types.go
@@ -5,23 +5,24 @@ import (
 	"strings"
 
 	"github.com/tailscale/hujson"
+	"gopkg.in/yaml.v3"
 	"inet.af/netaddr"
 )
 
 // ACLPolicy represents a Tailscale ACL Policy.
 type ACLPolicy struct {
-	Groups    Groups    `json:"Groups"`
-	Hosts     Hosts     `json:"Hosts"`
-	TagOwners TagOwners `json:"TagOwners"`
-	ACLs      []ACL     `json:"ACLs"`
-	Tests     []ACLTest `json:"Tests"`
+	Groups    Groups    `json:"Groups"    yaml:"Groups"`
+	Hosts     Hosts     `json:"Hosts"     yaml:"Hosts"`
+	TagOwners TagOwners `json:"TagOwners" yaml:"TagOwners"`
+	ACLs      []ACL     `json:"ACLs"      yaml:"ACLs"`
+	Tests     []ACLTest `json:"Tests"     yaml:"Tests"`
 }
 
 // ACL is a basic rule for the ACL Policy.
 type ACL struct {
-	Action string   `json:"Action"`
-	Users  []string `json:"Users"`
-	Ports  []string `json:"Ports"`
+	Action string   `json:"Action" yaml:"Action"`
+	Users  []string `json:"Users"  yaml:"Users"`
+	Ports  []string `json:"Ports"  yaml:"Ports"`
 }
 
 // Groups references a series of alias in the ACL rules.
@@ -35,9 +36,9 @@ type TagOwners map[string][]string
 
 // ACLTest is not implemented, but should be use to check if a certain rule is allowed.
 type ACLTest struct {
-	User  string   `json:"User"`
-	Allow []string `json:"Allow"`
-	Deny  []string `json:"Deny,omitempty"`
+	User  string   `json:"User"           yaml:"User"`
+	Allow []string `json:"Allow"          yaml:"Allow"`
+	Deny  []string `json:"Deny,omitempty" yaml:"Deny,omitempty"`
 }
 
 // UnmarshalJSON allows to parse the Hosts directly into netaddr objects.
@@ -69,6 +70,27 @@ func (hosts *Hosts) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
+// UnmarshalYAML allows to parse the Hosts directly into netaddr objects.
+func (hosts *Hosts) UnmarshalYAML(data []byte) error {
+	newHosts := Hosts{}
+	hostIPPrefixMap := make(map[string]string)
+
+	err := yaml.Unmarshal(data, &hostIPPrefixMap)
+	if err != nil {
+		return err
+	}
+	for host, prefixStr := range hostIPPrefixMap {
+		prefix, err := netaddr.ParseIPPrefix(prefixStr)
+		if err != nil {
+			return err
+		}
+		newHosts[host] = prefix
+	}
+	*hosts = newHosts
+
+	return nil
+}
+
 // IsZero is perhaps a bit naive here.
 func (policy ACLPolicy) IsZero() bool {
 	if len(policy.Groups) == 0 && len(policy.Hosts) == 0 && len(policy.ACLs) == 0 {