From af777f44f4f28ec099cdd77f4dfe9de3b4c967e5 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@dalby.cc>
Date: Wed, 18 Feb 2026 15:27:20 +0000
Subject: [PATCH] cmd/headscale/cli: extract bypassDatabase helper and simplify
 policy file reads

Add bypassDatabase() to consolidate the repeated LoadServerConfig +
NewHeadscaleDatabase pattern in getPolicy and setPolicy.

Replace os.Open + io.ReadAll with os.ReadFile in setPolicy and
checkPolicy, removing the manual file-handle management.
---
 cmd/headscale/cli/policy.go | 86 +++++++++++++++++++++++++++++++------
 1 file changed, 74 insertions(+), 12 deletions(-)

diff --git a/cmd/headscale/cli/policy.go b/cmd/headscale/cli/policy.go
index 270b75db..e5dfba3d 100644
--- a/cmd/headscale/cli/policy.go
+++ b/cmd/headscale/cli/policy.go
@@ -3,7 +3,6 @@ package cli
 import (
 	"errors"
 	"fmt"
-	"io"
 	"os"
 
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
@@ -20,6 +19,23 @@ const (
 
 var errAborted = errors.New("command aborted by user")
 
+// bypassDatabase loads the server config and opens the database directly,
+// bypassing the gRPC server. The caller is responsible for closing the
+// returned database handle.
+func bypassDatabase() (*db.HSDatabase, error) {
+	cfg, err := types.LoadServerConfig()
+	if err != nil {
+		return nil, fmt.Errorf("loading config: %w", err)
+	}
+
+	d, err := db.NewHeadscaleDatabase(cfg, nil)
+	if err != nil {
+		return nil, fmt.Errorf("opening database: %w", err)
+	}
+
+	return d, nil
+}
+
 func init() {
 	rootCmd.AddCommand(policyCmd)
 
@@ -52,15 +68,67 @@ var getPolicy = &cobra.Command{
 				return errAborted
 			}
 
-			cfg, err := types.LoadServerConfig()
+			d, err := bypassDatabase()
 			if err != nil {
-				return fmt.Errorf("loading config: %w", err)
+				return err
+			}
+			defer d.Close()
+
+			pol, err := d.GetPolicy()
+			if err != nil {
+				return fmt.Errorf("loading policy from database: %w", err)
 			}
 
-			d, err := db.NewHeadscaleDatabase(cfg, nil)
+			policyData = pol.Data
+		} else {
+			ctx, client, conn, cancel, err := newHeadscaleCLIWithConfig()
 			if err != nil {
-				return fmt.Errorf("opening database: %w", err)
+				return fmt.Errorf("connecting to headscale: %w", err)
 			}
+			defer cancel()
+			defer conn.Close()
+
+			response, err := client.GetPolicy(ctx, &v1.GetPolicyRequest{})
+			if err != nil {
+				return fmt.Errorf("loading ACL policy: %w", err)
+			}
+
+			policyData = response.GetPolicy()
+		}
+
+		// This does not pass output format as we don't support yaml, json or
+		// json-line output for this command. It is HuJSON already.
+		fmt.Println(policyData)
+
+		return nil
+	},
+}
+
+var setPolicy = &cobra.Command{
+	Use:   "set",
+	Short: "Updates the ACL Policy",
+	Long: `
+	Updates the existing ACL Policy with the provided policy. The policy must be a valid HuJSON object.
+	This command only works when the acl.policy_mode is set to "db", and the policy will be stored in the database.`,
+	Aliases: []string{"put", "update"},
+	RunE: func(cmd *cobra.Command, args []string) error {
+		policyPath, _ := cmd.Flags().GetString("file")
+
+		policyBytes, err := os.ReadFile(policyPath)
+		if err != nil {
+			return fmt.Errorf("reading policy file: %w", err)
+		}
+
+		if bypass, _ := cmd.Flags().GetBool(bypassFlag); bypass {
+			if !confirmAction(cmd, "DO NOT run this command if an instance of headscale is running, are you sure headscale is not running?") {
+				return errAborted
+			}
+
+			d, err := bypassDatabase()
+			if err != nil {
+				return err
+			}
+			defer d.Close()
 
 			users, err := d.ListUsers()
 			if err != nil {
@@ -104,13 +172,7 @@ var checkPolicy = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		policyPath, _ := cmd.Flags().GetString("file")
 
-		f, err := os.Open(policyPath)
-		if err != nil {
-			return fmt.Errorf("opening policy file: %w", err)
-		}
-		defer f.Close()
-
-		policyBytes, err := io.ReadAll(f)
+		policyBytes, err := os.ReadFile(policyPath)
 		if err != nil {
 			return fmt.Errorf("reading policy file: %w", err)
 		}