// Copyright (C) 2018 Marius Schellenberger

// +build go1.11

// Package godrop provides a simple privileges dropping library
package godrop

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
)

func errf(i interface{}) error {
	return fmt.Errorf("godrop: %s", i)
}

// Config represents the drop config
type Config struct {
	// User is the user to drop privileges to.
	User string
	// Group is the group to drop privileges to.
	Group string
	// Chroot is the directory to chroot into. Leave this emptry for no chroot.
	// When compiling without cgo, make sure the chroot directory contains the /etc/passwd and /etc/group files.
	Chroot string
	// Set to true, to run the process in the foreground.
	Foreground bool
}

// Drop will spawn a new process and hand over the listening socket file descriptor
func Drop(c Config, f func() (net.Listener, error)) error {
	return MultiDrop(c, func() ([]net.Listener, error) {
		l, err := f()
		return []net.Listener{l}, err
	})
}

// MultiDrop will spawn a new process and hand over the all listening sockets
func MultiDrop(c Config, f func() ([]net.Listener, error)) error {
	uid, err := userID(c.User)
	if err != nil {
		return errf(err)
	}
	if uid == 0 {
		return errf("unable to drop privileges to uid 0 (root)")
	}
	gid, err := groupID(c.Group)
	if err != nil {
		return errf(err)
	}
	switch os.Getuid() {
	case 0:
		cmd := exec.Command(os.Args[0], os.Args[1:]...)
		ln, err := f()
		if err != nil {
			return errf(err)
		}
		for i, v := range ln {
			var f *os.File
			switch l := v.(type) {
			case *net.TCPListener:
				f, err = l.File()
				l.Close()
			case *net.UnixListener:
				f, err = l.File()
				l.Close()
			default:
				return errf(fmt.Sprintf("index %d listener is not type of either *net.TCPListener or *net.UnixListener", i))
			}
			if err != nil {
				return errf(fmt.Sprintf("index %d %s", i, err))
			}
			cmd.ExtraFiles = append(cmd.ExtraFiles, f)
		}

		cmd.SysProcAttr = &syscall.SysProcAttr{
			Chroot: c.Chroot,
			Credential: &syscall.Credential{
				Uid: uint32(uid),
				Gid: uint32(gid),
			},
			Setsid: true,
		}

		if c.Foreground {
			cmd.Stdin = os.Stdin
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
		}

		if err := cmd.Start(); err != nil {
			return errf(err)
		}

		if c.Foreground {
			go func() {
				term := make(chan os.Signal)
				signal.Notify(term, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
				sigs := make(chan os.Signal)
				signal.Notify(sigs, syscall.SIGSTOP, syscall.SIGCONT, syscall.SIGUSR1, syscall.SIGUSR2)
				for {
					select {
					case sig := <-term:
						cmd.Process.Signal(sig)
						return
					case sig := <-sigs:
						cmd.Process.Signal(sig)
					}
				}
			}()
			_ = cmd.Wait()
			os.Exit(int(cmd.ProcessState.Sys().(syscall.WaitStatus)))
		}
		cmd.Process.Release()
		os.Exit(0)
	case uid:
		return nil
	}
	return errf("dropping priviledges failed")
}

// GetListener returns the listener socket of file descriptor 3
func GetListener() (net.Listener, error) {
	return GetListenerFd(3)
}

// GetListenerFd returns the listener socket of the given file descriptor
func GetListenerFd(fd int) (net.Listener, error) {
	if fd < 3 {
		return nil, errf("fd is less than 3")
	}
	f := os.NewFile(uintptr(fd), "")
	defer f.Close()
	return net.FileListener(f)
}