package instancectl

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"unsafe"

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

type claimBPFInsn struct {
	Code, Registers uint8
	Offset          int16
	Immediate       int32
}

func claimConnectFilter(ipv6 bool) []claimBPFInsn {
	// bpf_sock_addr has network-order user_ip4 at 4 and user_ip6 at 8.
	load := func(offset int16) []claimBPFInsn { return []claimBPFInsn{{0x61, 0x10, offset, 0}, {0xdc, 0, 0, 32}} }
	var p []claimBPFInsn
	if !ipv6 {
		p = append(p, load(4)...)
		p = append(p, claimBPFInsn{0x54, 0, 0, -16777216}, claimBPFInsn{0x15, 0, 2, 0x7f000000})
	} else {
		p = append(p, load(8)...)
		p = append(p, claimBPFInsn{0x55, 0, 14, 0})
		p = append(p, load(12)...)
		p = append(p, claimBPFInsn{0x55, 0, 11, 0})
		p = append(p, load(16)...)
		p = append(p, claimBPFInsn{0x15, 0, 4, 0xffff}, claimBPFInsn{0x55, 0, 7, 0})
		p = append(p, load(20)...)
		p = append(p, claimBPFInsn{0x15, 0, 6, 1})
		p = append(p, load(20)...)
		p = append(p, claimBPFInsn{0x54, 0, 0, -16777216}, claimBPFInsn{0x15, 0, 2, 0x7f000000})
	}
	return append(p, claimBPFInsn{0xb7, 0, 0, 1}, claimBPFInsn{0x95, 0, 0, 0}, claimBPFInsn{0xb7, 0, 0, 0}, claimBPFInsn{0x95, 0, 0, 0})
}

func claimAttachConnectFilter(group int, kind uint32, ipv6 bool) error {
	p := claimConnectFilter(ipv6)
	license := []byte("GPL\x00")
	attr := struct {
		Type, Count           uint32
		Instructions, License uint64
		LogLevel, LogSize     uint32
		LogBuffer             uint64
		Version, Flags        uint32
		Name                  [16]byte
		IfIndex, Attach       uint32
	}{Type: unix.BPF_PROG_TYPE_CGROUP_SOCK_ADDR, Count: uint32(len(p)), Instructions: uint64(uintptr(unsafe.Pointer(&p[0]))), License: uint64(uintptr(unsafe.Pointer(&license[0]))), Attach: kind}
	fd, _, errno := unix.Syscall(unix.SYS_BPF, unix.BPF_PROG_LOAD, uintptr(unsafe.Pointer(&attr)), unsafe.Sizeof(attr))
	runtime.KeepAlive(p)
	runtime.KeepAlive(license)
	if errno != 0 {
		return errno
	}
	defer unix.Close(int(fd))
	attach := struct{ Target, Program, Kind, Flags uint32 }{uint32(group), uint32(fd), kind, unix.BPF_F_ALLOW_MULTI}
	_, _, errno = unix.Syscall(unix.SYS_BPF, unix.BPF_PROG_ATTACH, uintptr(unsafe.Pointer(&attach)), unsafe.Sizeof(attach))
	if errno != 0 {
		return errno
	}
	return nil
}

func (u ClaimUnits) restrictBoundNetwork(s ClaimUnitStarted) error {
	_, group, err := u.Placement(u.BoundHost)
	if err != nil || s.ControlGroup != group || s.Unit != u.BoundUnit || u.CgroupRoot != "/sys/fs/cgroup" {
		return os.ErrPermission
	}
	// Pin each directory and check every ancestor before attaching. No program
	// is pinned on disk; the kernel retains it until this cgroup is destroyed.
	fd, err := unix.Open("/", unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0)
	if err != nil {
		return err
	}
	defer func() { unix.Close(fd) }()
	for _, part := range strings.Split(strings.TrimPrefix(filepath.Join(u.CgroupRoot, group), "/"), "/") {
		next, e := unix.Openat(fd, part, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
		if e != nil {
			return e
		}
		unix.Close(fd)
		fd = next
		var st unix.Stat_t
		if e = unix.Fstat(fd, &st); e != nil {
			return e
		}
		if st.Uid != 0 || st.Mode&0022 != 0 {
			return os.ErrPermission
		}
	}
	var fs unix.Statfs_t
	if err = unix.Fstatfs(fd, &fs); err != nil {
		return err
	}
	if fs.Type != unix.CGROUP2_SUPER_MAGIC {
		return os.ErrPermission
	}
	for _, hook := range []struct {
		kind uint32
		ipv6 bool
	}{
		{unix.BPF_CGROUP_INET4_CONNECT, false}, {unix.BPF_CGROUP_INET6_CONNECT, true},
		{unix.BPF_CGROUP_UDP4_SENDMSG, false}, {unix.BPF_CGROUP_UDP6_SENDMSG, true},
	} {
		if err = claimAttachConnectFilter(fd, hook.kind, hook.ipv6); err != nil {
			return err
		}
	}
	return nil
}