From 34f087de1c24e39fce7201bfe8d238f7e5bc85cf Mon Sep 17 00:00:00 2001
From: cathugger <cathugger@cock.li>
Date: Sun, 22 Sep 2024 15:46:54 +0000
Subject: [PATCH] argument to change uid/gid (#927)

different from
https://github.com/yggdrasil-network/yggdrasil-go/pull/817 in that it
can resolve user names, automatically use user's primary gid & allows
specifying gid in the same argument, with `:` eg `username:groupname`.
feel free to criticize & suggest different argument name & description
because i didn't put much of thought to that.

---------

Co-authored-by: Neil <git@neilalexander.dev>
Co-authored-by: VNAT <xepjk@protonmail.com>
Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>
---
 cmd/yggdrasil/chuser_other.go | 10 ++++
 cmd/yggdrasil/chuser_unix.go  | 87 +++++++++++++++++++++++++++++++++++
 cmd/yggdrasil/main.go         |  9 ++++
 3 files changed, 106 insertions(+)
 create mode 100644 cmd/yggdrasil/chuser_other.go
 create mode 100644 cmd/yggdrasil/chuser_unix.go

diff --git a/cmd/yggdrasil/chuser_other.go b/cmd/yggdrasil/chuser_other.go
new file mode 100644
index 00000000..702f3715
--- /dev/null
+++ b/cmd/yggdrasil/chuser_other.go
@@ -0,0 +1,10 @@
+//go:build !aix && !darwin && !dragonfly && !freebsd && !linux && !netbsd && !openbsd && !solaris
+// +build !aix,!darwin,!dragonfly,!freebsd,!linux,!netbsd,!openbsd,!solaris
+
+package main
+
+import "errors"
+
+func chuser(user string) error {
+	return errors.New("setting uid/gid is not supported on this platform")
+}
diff --git a/cmd/yggdrasil/chuser_unix.go b/cmd/yggdrasil/chuser_unix.go
new file mode 100644
index 00000000..6e802c69
--- /dev/null
+++ b/cmd/yggdrasil/chuser_unix.go
@@ -0,0 +1,87 @@
+//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
+// +build aix darwin dragonfly freebsd linux netbsd openbsd solaris
+
+package main
+
+import (
+	"errors"
+	"fmt"
+	"math"
+	osuser "os/user"
+	"strconv"
+	"strings"
+	"syscall"
+)
+
+func chuser(user string) error {
+	group := ""
+	if i := strings.IndexByte(user, ':'); i >= 0 {
+		user, group = user[:i], user[i+1:]
+	}
+
+	u := (*osuser.User)(nil)
+	g := (*osuser.Group)(nil)
+
+	if user != "" {
+		if _, err := strconv.ParseUint(user, 10, 32); err == nil {
+			u, err = osuser.LookupId(user)
+			if err != nil {
+				return fmt.Errorf("failed to lookup user by id %q: %v", user, err)
+			}
+		} else {
+			u, err = osuser.Lookup(user)
+			if err != nil {
+				return fmt.Errorf("failed to lookup user by name %q: %v", user, err)
+			}
+		}
+	}
+	if group != "" {
+		if _, err := strconv.ParseUint(group, 10, 32); err == nil {
+			g, err = osuser.LookupGroupId(group)
+			if err != nil {
+				return fmt.Errorf("failed to lookup group by id %q: %v", user, err)
+			}
+		} else {
+			g, err = osuser.LookupGroup(group)
+			if err != nil {
+				return fmt.Errorf("failed to lookup group by name %q: %v", user, err)
+			}
+		}
+	}
+
+	if g != nil {
+		gid, _ := strconv.ParseUint(g.Gid, 10, 32)
+		var err error
+		if gid < math.MaxInt {
+			err = syscall.Setgid(int(gid))
+		} else {
+			err = errors.New("gid too big")
+		}
+
+		if err != nil {
+			return fmt.Errorf("failed to setgid %d: %v", gid, err)
+		}
+	} else if u != nil {
+		gid, _ := strconv.ParseUint(u.Gid, 10, 32)
+		err := syscall.Setgid(int(uint32(gid)))
+		if err != nil {
+			return fmt.Errorf("failed to setgid %d: %v", gid, err)
+		}
+	}
+
+	if u != nil {
+		uid, _ := strconv.ParseUint(u.Uid, 10, 32)
+		var err error
+		if uid < math.MaxInt {
+			err = syscall.Setuid(int(uid))
+		} else {
+			err = errors.New("uid too big")
+		}
+
+		if err != nil {
+			return fmt.Errorf("failed to setuid %d: %v", uid, err)
+		}
+	}
+
+	return nil
+}
diff --git a/cmd/yggdrasil/main.go b/cmd/yggdrasil/main.go
index 29afdf5d..3a00b3fc 100644
--- a/cmd/yggdrasil/main.go
+++ b/cmd/yggdrasil/main.go
@@ -52,6 +52,7 @@ func main() {
 	getsnet := flag.Bool("subnet", false, "use in combination with either -useconf or -useconffile, outputs your IPv6 subnet")
 	getpkey := flag.Bool("publickey", false, "use in combination with either -useconf or -useconffile, outputs your public key")
 	loglevel := flag.String("loglevel", "info", "loglevel to enable")
+	chuserto := flag.String("user", "", "user (and, optionally, group) to set UID/GID to")
 	flag.Parse()
 	
 	done := make(chan struct{})
@@ -280,6 +281,14 @@ func main() {
 		<-done
 	})
 
+	// Change user if requested
+	if *chuserto != "" {
+		err = chuser(*chuserto)
+		if err != nil {
+			panic(err)
+		}
+	}
+
 	// Block until we are told to shut down.
 	<-ctx.Done()