// -*- Mode: Go; indent-tabs-mode: t -*-

/*
 * Copyright (C) 2016 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package main

import (
	"fmt"
	"os"
	"runtime"
	"sync"
	"syscall"
)

// NSType is a type of Linux Namespace.
//
// Linux suppots the following namespaces:
// IPC, Network, Mount, PID, User and UTS.
type NSType string

// The constants below match file names in /proc/self/ns/

const (
	// IPCNS is the Linux IPC namespace.
	IPCNS NSType = "ipc"
	// NetworkNS is the Linux network namespace.
	NetworkNS NSType = "net"
	// MountNS is the Linux mount namespace.
	MountNS NSType = "mnt"
	// PIDNS is the Linux process identifier namespace.
	PIDNS NSType = "pid"
	// UserNS is the Linux user namespace.
	UserNS NSType = "user"
	// UTSNS is the Linux UTS (hostname) namespace.
	UTSNS NSType = "uts"
	// CGroupNS is the Linux control group namespace.
	CGroupNS NSType = "cgroup"
)

// SetnsArg returns 2nd argument (nstype) for setns.
func (t NSType) SetnsArg() int {
	// XXX: this might be arch specific, needs checking.
	switch t {
	case IPCNS:
		return 0x08000000
	case NetworkNS:
		return 0x40000000
	case MountNS:
		return 0x00020000
	case PIDNS:
		return 0x20000000
	case UserNS:
		return 0x10000000
	case UTSNS:
		return 0x04000000
	case CGroupNS:
		return 0x02000000
	}
	return 0
}

// CurrentNamespaceID returns the namespace identifier the calling thread.
//
// Because of how golang schedules goroutines this should only be called from a
// goroutine with a fixed OS thread (e.g. runtime.LockOSThread) or the results
// will not be deterministic.
func CurrentNamespaceID(nsType NSType) (string, error) {
	return os.Readlink(fmt.Sprintf("/proc/self/ns/%s", nsType))
}

// Namespace aids in running functions in a specific Linux namespace.
type Namespace struct {
	m sync.RWMutex

	enterFd, leaveFd int
	nsType           NSType
}

// OpenNS opens a namespace associated with a given file.
//
// The fname must point to a namespace file (e.g. /proc/$PID/ns/mnt or
// /run/snapd/ns/$SNAP_NAME.mnt). The type must be the type of namespace to
// open, e.g. MountNS.
func OpenNS(enterFname string, nsType NSType) (*Namespace, error) {
	leaveFname := fmt.Sprintf("/proc/self/ns/%s", nsType)
	leaveFd, err := syscall.Open(leaveFname, syscall.O_RDONLY|syscall.O_CLOEXEC, 0644)
	if err != nil {
		return nil, err
	}
	fmt.Printf("leaveFd: %d (from %q)\n", leaveFd, leaveFname)
	enterFd, err := syscall.Open(enterFname, syscall.O_RDONLY|syscall.O_CLOEXEC, 0644)
	if err != nil {
		syscall.Close(leaveFd)
		return nil, err
	}
	fmt.Printf("enterFd: %d (from %q)\n", enterFd, enterFname)
	return &Namespace{enterFd: enterFd, leaveFd: leaveFd, nsType: nsType}, nil
}

// Close releases system resources associated with an open namespace.
//
// Close waits for all goroutines executing in the namespace to finish.
func (ns *Namespace) Close() {
	ns.m.Lock()
	defer ns.m.Unlock()

	syscall.Close(ns.enterFd)
	syscall.Close(ns.leaveFd)
	ns.enterFd = -1
	ns.leaveFd = -1
}

// Type returns the type of the given namespace.
func (ns *Namespace) Type() NSType {
	return ns.nsType
}

// Do calls the given function in the given namespace.
func (ns *Namespace) Do(fn func() error) error {
	ns.m.RLock()
	defer ns.m.RUnlock()

	if ns.isClosed() {
		return fmt.Errorf("cannot perform operations on a closed namespace")
	}

	ch := make(chan error)
	go func() {
		// Lock the calling goroutine to one OS thread.
		runtime.LockOSThread()
		defer runtime.UnlockOSThread()

		if err := ns.enter(); err != nil {
			ch <- fmt.Errorf("cannot enter namespace: %s", err)
			return
		}

		fnErr := fn()

		if err := ns.leave(); err != nil {
			// XXX: Here we ignore fnErr, maybe this ought to return both?
			ch <- fmt.Errorf("cannot leave namespace: %s", err)
			return
		}

		ch <- fnErr
	}()
	return <-ch
}

// isClosed checks if a given namespace is closed
func (ns *Namespace) isClosed() bool {
	return ns.enterFd == -1 || ns.leaveFd == -1
}

// enter enters another namespace.
//
// Note that only the calling OS thread is affected.
func (ns *Namespace) enter() error {
	return setns(ns.enterFd, ns.nsType.SetnsArg())
}

// leave leaves the given namespace and goes to the initial namespace.
//
// Note that only the calling OS thread is affected.
func (ns *Namespace) leave() error {
	return setns(ns.leaveFd, ns.nsType.SetnsArg())
}

// sysSetNs holds the system call number for setns for the current architecture.
var sysSetNs int

func init() {
	// Lookup SYS_setns for this architecture.
	// This was lifted from https://github.com/opencontainers/runc/blob/master/libcontainer/system/setns_linux.go
	m := map[string]int{
		"386":     346,
		"arm64":   268,
		"amd64":   308,
		"arm":     375,
		"ppc":     350,
		"ppc64":   350,
		"ppc64le": 350,
		"s390x":   339,
	}
	num, ok := m[runtime.GOARCH]
	if !ok {
		panic(fmt.Errorf("please define syscall number for architecture %q", runtime.GOARCH))
	}
	sysSetNs = num
}

// setns calls the setns system call
func setns(fd, nsType int) error {
	fmt.Printf("syscall(%d, %d, %#x)\n", sysSetNs, fd, nsType)
	_, _, errno := syscall.RawSyscall(uintptr(sysSetNs), uintptr(fd), uintptr(nsType), 0)
	if errno != 0 {
		return errno
	}
	return nil
}

func demo(fname string) error {
	nsID, err := CurrentNamespaceID(MountNS)
	if err != nil {
		return err
	}
	fmt.Printf("(before) mount namespace ID is: %s\n", nsID)

	ns, err := OpenNS(fname, MountNS)
	if err != nil {
		return err
	}
	defer ns.Close()

	doErr := ns.Do(func() error {
		fmt.Printf("My thread id is: %d\n", syscall.Gettid())
		nsID, err := CurrentNamespaceID(MountNS)
		if err != nil {
			return err
		}
		fmt.Printf("(go-routine) mount namespace ID is: %s\n", nsID)
		return nil
	})

	nsID, err = CurrentNamespaceID(MountNS)
	if err != nil {
		return err
	}
	fmt.Printf("(after) mount namespace ID is: %s\n", nsID)

	return doErr
}

func main() {
	if len(os.Args) != 2 {
		fmt.Printf("usage: ns /path/to/namespace/file\n")
		os.Exit(0)
	}
	if err := demo(os.Args[1]); err != nil {
		fmt.Printf("error: %s\n", err)
		os.Exit(1)
	}
}