// Copyright (C) 2016-2017 Luke Shumaker <lukeshu@sbcglobal.net>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sd_login

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"

	"golang.org/x/sys/unix"
)

type _Cgroup string

type _CgroupTree string
type _CgroupLeaf string

var cgVersion_cache uint

func cgVersion() uint {
	if cgVersion_cache == 0 {
		fs, err := statfs("/sys/fs/cgroup/")
		if err != nil {
			return 0
		}
		if fs.Type == magic_CGROUP2_SUPER {
			cgVersion_cache = 2
		} else if fs.Type == magic_TMPFS {
			fs, err = statfs("/sys/fs/cgroup/systemd/")
			if err != nil {
				return 0
			}
			if fs.Type == magic_CGROUP2_SUPER {
				cgVersion_cache = 2
			} else {
				cgVersion_cache = 1
			}
		}
	}
	return cgVersion_cache
}

// GetCgroup returns the cgroup path of the process, relative to the
// root of the cgroup hierarchy.
//
// In cgroup v2, there is a only one cgroup hierarchy, so the behavior
// is obvious.
//
// However, in cgroup v1, there were multiple cgroup hierarchies, so
// this function must decide which hierarchy to use.  We choose the
// first hierarchy with either the "name=systemd" controller or the
// "name=elogind" controller attached to it[1].  If no such hierarchy
// exists, then an error is returned.
//
// However, it is possible to generally use cgroup v1, but use a
// single (named) v2 hierarchy alongside many v1 hierarchies.  In this
// case, we use the v2 hierarchy iff it is named "systemd", otherwise
// we use the cgroup v1 behavior.
//
// [1]: The "first" in that sentence is worrying; shouldn't the choice
// of hierarchy not depend on the undefined order that controllers are
// listed in?  Well, a controller may be attached to only one
// hierarchy at a time.  So there is only an ambiguity for "first" to
// come in if both "name=systemd" and "name=elogind" controllers
// exist.  Systemd and elogind cannot be used together, so this isn't
// a concern.
func (pid ProcessID) GetCgroup() (string, error) {
	cgVer := cgVersion()

	cgroupFilename := fmt.Sprintf("/proc/%d/cgroup", pid)

	f, err := os.Open(cgroupFilename)
	if err != nil {
		return "", err
	}
	defer f.Close()

	bf := bufio.NewReader(f)

	for {
		line, err := bf.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			return "", err
		}
		line = strings.TrimSuffix(line, "\n")

		parts := strings.SplitN(line, ":", 3)
		if len(parts) != 3 {
			continue
		}

		hierarchy := parts[0]
		controllers := parts[1]
		path := parts[2]

		switch cgVer {
		case 1:
			for _, controller := range strings.Split(controllers, ",") {
				if controller == "name=systemd" || controller == "name=elogind" {
					return path, nil
				}
			}
		case 2:
			if hierarchy != "0" {
				continue
			}
			return path, nil
		}
	}
	return "", unix.ENODATA
}

func (pid ProcessID) getCgroup() (_Cgroup, error) {
	cgroup, err := pid.GetCgroup()
	return _Cgroup(cgroup), err
}

// cgGetRootPath determines the cgroup that all other cgroups belong
// to.  The common case is just "/", but it could be something else if
// we are inside of a container, but have a view of the entier cgroup
// hierarchy.
func cgGetRootPath() (_CgroupTree, error) {
	cgroup, err := ProcessID(1).getCgroup()
	if err != nil {
		return "/", err
	}

	cgTree := _CgroupTree(trimOneSuffix(string(cgroup),
		"/init.scope",   // modern systemd
		"/system.slice", // legacy systemd
		"/system",       // even more legacy systemd
	))

	return cgTree, nil
}

func cgUnescape(s string) string {
	return strings.TrimPrefix(s, "_")
}

func cgDecodeUnit(unit string) string {
	if len(unit) < 3 {
		// because "X.Y" is at least 3 characters
		return ""
	}
	unit = cgUnescape(unit)
	if valid_unit_name(unit)&(unit_name_plain|unit_name_instance) == 0 {
		return ""
	}
	return unit
}

func (tree _CgroupTree) Parse(cgroup _Cgroup) (leaf _CgroupLeaf) {
	rest, ok := path_startswith(string(cgroup), string(tree))
	if ok {
		return _CgroupLeaf(rest)
	} else {
		return ""
	}
}

// Even if there is an invalid unit, a slice might still be returned.
func (leaf _CgroupLeaf) Parse() (slice, unit string, extra _CgroupLeaf) {
	strLeaf := string(leaf)

	slice = "-.slice"
	for {
		part, rest := split2(strings.TrimLeft(strLeaf, "/"), '/')
		if valid_slice_name(part) {
			slice = part
			strLeaf = rest
		} else {
			break
		}
	}
	slice = cgDecodeUnit(slice)
	if slice == "" {
		return
	}

	unit, extraStr := split2(strings.TrimLeft(strLeaf, "/"), '/')
	unit = cgDecodeUnit(unit)
	if unit == "" {
		return
	}

	extra = _CgroupLeaf(extraStr)
	return
}

func (leaf _CgroupLeaf) GetSlice() string {
	slice, _, _ := leaf.Parse()
	return slice
}

func (leaf _CgroupLeaf) GetUnit() string {
	_, unit, _ := leaf.Parse()
	return unit
}

func (leaf _CgroupLeaf) GetRest() _CgroupLeaf {
	_, _, extra := leaf.Parse()
	return extra
}