// Copyright (c) 2026 Storm Dragon
// SPDX-License-Identifier: GPL-3.0-or-later

package tentacle

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
)

type launchReply struct {
	PID   int
	Errno uint32
}

// launch passes exactly four descriptors: stdin, stdout, stderr and the already
// limited cgroup. The service performs clone3 inside its delegated tree. No
// shell, environment update, or settings operation is exposed by this protocol.
func (delegation *Delegation) launch(ctx context.Context, spec Spec, files ...*os.File) (int, <-chan struct{}, func(), error) {
	connection, err := net.DialUnix("unixpacket", nil, &net.UnixAddr{Name: filepath.Join(delegation.directory, "launch"), Net: "unixpacket"})
	if err != nil {
		return 0, nil, nil, err
	}
	keep := false
	defer func() {
		if !keep {
			connection.Close()
		}
	}()
	connection.SetDeadline(time.Now().Add(3 * time.Second))
	stopCancel := context.AfterFunc(ctx, func() { connection.Close() })
	defer func() {
		if !keep {
			stopCancel()
		}
	}()
	data, err := json.Marshal(spec)
	if err != nil || len(data) > MaxFrameBytes || len(files) != 4 {
		return 0, nil, nil, ErrLimit
	}
	descriptors := make([]int, len(files))
	for index, file := range files {
		descriptors[index] = int(file.Fd())
	}
	if _, _, err := connection.WriteMsgUnix(data, syscall.UnixRights(descriptors...), nil); err != nil {
		return 0, nil, nil, err
	}
	buffer := make([]byte, 256)
	count, err := connection.Read(buffer)
	if err != nil {
		return 0, nil, nil, err
	}
	var reply launchReply
	if err := json.Unmarshal(buffer[:count], &reply); err != nil {
		return 0, nil, nil, ErrProtocol
	}
	if reply.Errno != 0 {
		return 0, nil, nil, syscall.Errno(reply.Errno)
	}
	if reply.PID <= 0 {
		return 0, nil, nil, ErrProtocol
	}
	connection.SetDeadline(time.Time{})
	done := make(chan struct{})
	go func() { defer stopCancel(); var end [1]byte; connection.Read(end[:]); connection.Close(); close(done) }()
	keep = true
	return reply.PID, done, func() { connection.Close() }, nil
}

func serve_launches(input io.Reader, output io.Writer, socketPath, parent string) error {
	listener, err := net.ListenUnix("unixpacket", &net.UnixAddr{Name: socketPath, Net: "unixpacket"})
	if err != nil {
		return err
	}
	defer listener.Close()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var workers sync.WaitGroup
	capacity := make(chan struct{}, DefaultLimits().Processes)
	accepted := make(chan struct{})
	go func() {
		defer close(accepted)
		for {
			connection, err := listener.AcceptUnix()
			if err != nil {
				return
			}
			select {
			case capacity <- struct{}{}:
				workers.Add(1)
				go func() { defer workers.Done(); defer func() { <-capacity }(); serve_launch(ctx, connection, parent) }()
			default:
				connection.Close()
			}
		}
	}()
	defer func() { cancel(); listener.Close(); <-accepted; workers.Wait() }()
	if _, err := fmt.Fprintln(output, parent); err != nil {
		return err
	}
	_, err = io.Copy(io.Discard, input)
	return err
}

func serve_launch(ctx context.Context, connection *net.UnixConn, parent string) {
	defer connection.Close()
	stop := context.AfterFunc(ctx, func() { connection.Close() })
	defer stop()
	connection.SetDeadline(time.Now().Add(3 * time.Second))
	data, ancillary := make([]byte, MaxFrameBytes), make([]byte, syscall.CmsgSpace(4*4))
	count, rightsCount, flags, _, err := connection.ReadMsgUnix(data, ancillary)
	if err != nil {
		return
	}
	// Always close received descriptors, even when framing or validation fails.
	var files []*os.File
	defer func() {
		for _, file := range files {
			file.Close()
		}
	}()
	messages, err := syscall.ParseSocketControlMessage(ancillary[:rightsCount])
	if err != nil {
		return
	}
	for _, message := range messages {
		descriptors, err := syscall.ParseUnixRights(&message)
		if err != nil {
			return
		}
		for _, descriptor := range descriptors {
			syscall.CloseOnExec(descriptor)
			files = append(files, os.NewFile(uintptr(descriptor), "tentacle-launch"))
		}
	}
	if flags&(syscall.MSG_TRUNC|syscall.MSG_CTRUNC) != 0 || len(files) != 4 {
		return
	}
	var spec Spec
	if decode_strict(data[:count], &spec) != nil || validate_spec(spec) != nil {
		return
	}
	var stat syscall.Statfs_t
	if syscall.Fstatfs(int(files[3].Fd()), &stat) != nil || stat.Type != 0x63677270 {
		return
	}
	groupPath, err := os.Readlink(fmt.Sprintf("/proc/self/fd/%d", files[3].Fd()))
	if err != nil {
		return
	}
	relative, err := filepath.Rel(parent, groupPath)
	parts := strings.Split(relative, string(filepath.Separator))
	if err != nil || len(parts) != 2 || !strings.HasPrefix(parts[0], "cthulhu-tentacles-") {
		return
	}
	for _, identity := range []string{strings.TrimPrefix(parts[0], "cthulhu-tentacles-"), parts[1]} {
		if decoded, err := hex.DecodeString(identity); err != nil || len(decoded) != 16 {
			return
		}
	}
	command := exec.Command(spec.Program[0], spec.Program[1:]...)
	command.Dir = spec.Directory
	command.Stdin, command.Stdout, command.Stderr = files[0], files[1], files[2]
	command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, UseCgroupFD: true, CgroupFD: int(files[3].Fd())}
	if err := command.Start(); err != nil {
		var errno syscall.Errno
		if !errors.As(err, &errno) {
			errno = syscall.EIO
		}
		reply, _ := json.Marshal(launchReply{Errno: uint32(errno)})
		connection.Write(reply)
		return
	}
	for _, file := range files {
		file.Close()
	}
	files = nil
	// A closed launch connection kills the direct child. The reader's cgroup
	// cleanup or systemd's guardian-exit cleanup covers every descendant.
	disconnected := make(chan struct{})
	go func() { defer close(disconnected); var end [1]byte; connection.Read(end[:]); command.Process.Kill() }()
	reply, _ := json.Marshal(launchReply{PID: command.Process.Pid})
	if _, err := connection.Write(reply); err != nil {
		command.Process.Kill()
	}
	connection.SetDeadline(time.Time{})
	_ = command.Wait()
	connection.Write([]byte{0})
	connection.Close()
	<-disconnected
}