// Copyright 2018 The gVisor Authors.
//
// 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
//
//     https://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 runsc implements Containerd Shim v2 interface.
package runsc

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"

	taskServer "gvisor.dev/gvisor/pkg/shim/v1/taskserver"
	pb "gvisor.dev/gvisor/pkg/shim/v1/taskserver/task_server_go_proto"

	cgroups "github.com/containerd/cgroups/v3"
	cgroup1 "github.com/containerd/cgroups/v3/cgroup1"
	cgroupsv2 "github.com/containerd/cgroups/v3/cgroup2"
	"github.com/containerd/containerd/api/events"
	task "github.com/containerd/containerd/api/runtime/task/v2"
	runctypes "github.com/containerd/containerd/api/types/runc/options"
	tasktypes "github.com/containerd/containerd/api/types/task"
	"github.com/containerd/containerd/v2/cmd/containerd-shim-runc-v2/process"
	"github.com/containerd/containerd/v2/core/runtime"
	"github.com/containerd/containerd/v2/pkg/protobuf"
	types "github.com/containerd/containerd/v2/pkg/protobuf/types"
	"github.com/containerd/containerd/v2/pkg/shim"
	"github.com/containerd/containerd/v2/pkg/shutdown"
	"github.com/containerd/containerd/v2/pkg/stdio"
	"github.com/containerd/containerd/v2/pkg/sys/reaper"
	"github.com/containerd/errdefs"
	errgrpc "github.com/containerd/errdefs/pkg/errgrpc"
	"github.com/containerd/log"
	typeurl "github.com/containerd/typeurl/v2"
	specs "github.com/opencontainers/runtime-spec/specs-go"

	"gvisor.dev/gvisor/pkg/shim/v1/extension"
	"gvisor.dev/gvisor/pkg/shim/v1/proc"
	"gvisor.dev/gvisor/pkg/shim/v1/runsccmd"
	"gvisor.dev/gvisor/pkg/shim/v1/utils"
	"gvisor.dev/gvisor/runsc/specutils"
)

var (
	empty   = &types.Empty{}
	bufPool = sync.Pool{
		New: func() any {
			buffer := make([]byte, 32<<10)
			return &buffer
		},
	}
)

const (
	// configFile is the default config file name. For containerd 1.2,
	// we assume that a config.toml should exist in the runtime root.
	configFile = "config.toml"

	cgroupParentAnnotation = "dev.gvisor.spec.cgroup-parent"
)

type oomPoller interface {
	io.Closer
	// add adds `cg` cgroup to oom poller. `cg` is cgroups.Cgroup in v1 and
	// `cgroupsv2.Manager` in v2
	add(id string, cg any) error
	// run monitors oom event and notifies the shim about them
	run(ctx context.Context)
	// isOOM synchronously checks if the container was OOM-killed by reading
	// the cgroup's memory events. This is used at container exit time to
	// ensure the TaskOOM event is published before TaskExit, even if the
	// async notification has not yet arrived. This is a one-shot operation:
	// the cgroup reference is consumed on the first call, and subsequent
	// calls for the same id return false.
	isOOM(id string) bool
}

// runscService is the shim implementation of a remote shim over gRPC. It converts
// shim calls into `runsc` commands. It runs in 2 different modes:
//  1. Service: process runs for the life time of the container and receives
//     calls described in shimapi.TaskService interface.
//  2. Tool: process is short lived and runs only to perform the requested
//     operations and then exits. It implements the direct functions in
//     shim.Shim interface.
//
// When the service is running, it saves a json file with state information so
// that commands sent to the tool can load the state and perform the operation
// with the required context.
type runscService struct {
	mu sync.Mutex

	events chan any

	// platform handles operations related to the console.
	platform stdio.Platform

	// ex gets notified whenever the container init process or an exec'd process
	// exits from inside the sandbox.
	ec chan proc.Exit

	// oomPoller monitors the sandbox's cgroup for OOM notifications.
	oomPoller oomPoller

	// containers maps container id to a container.
	containers map[string]*Container

	// root is the runsc root directory.
	root string

	// runtime is runsc runtime configured for sandbox.
	runtime *runsccmd.Runsc

	shutdown shutdown.Service
}

var _ extension.TaskServiceExt = (*runscService)(nil)

// NewTaskService returns a runsc task service.
func NewTaskService(ctx context.Context, publisher shim.Publisher, sd shutdown.Service) (extension.TaskServiceExt, error) {
	var (
		ep  oomPoller
		err error
	)
	if cgroups.Mode() == cgroups.Unified {
		ep, err = newOOMv2Poller(publisher)
	} else {
		ep, err = newOOMEpoller(publisher)
	}
	if err != nil {
		return nil, err
	}
	go ep.run(ctx)
	s := &runscService{
		events:     make(chan any, 128),
		containers: make(map[string]*Container),
		ec:         proc.ExitCh,
		oomPoller:  ep,
		shutdown:   sd,
		root:       proc.RunscRoot,
		runtime:    &runsccmd.Runsc{Root: proc.RunscRoot},
	}
	go s.processExits(ctx)
	runsccmd.Monitor = &runsccmd.LogMonitor{Next: reaper.Default}
	if err := s.initPlatform(); err != nil {
		return nil, fmt.Errorf("failed to initialized platform behavior: %w", err)
	}
	go s.forward(ctx, publisher)
	sd.RegisterCallback(func(context.Context) error {
		close(s.events)
		return nil
	})
	return s, nil
}

// getContainer returns the container by id.
func (s *runscService) getContainer(id string) (*Container, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := s.containers[id]
	if c == nil {
		return nil, errgrpc.ToGRPCf(errdefs.ErrNotFound, "container not created")
	}
	return c, nil
}

// Create creates a new initial process and container with the underlying OCI
// runtime.
func (s *runscService) Create(ctx context.Context, r *task.CreateTaskRequest) (*task.CreateTaskResponse, error) {
	return s.CreateWithFSRestore(ctx, &extension.CreateWithFSRestoreRequest{
		Create: r,
	})
}

// CreateWithFSRestore is the same as Create, but it additionally restores the
// container's filesystem from a snapshot.
func (s *runscService) CreateWithFSRestore(ctx context.Context, rfs *extension.CreateWithFSRestoreRequest) (*task.CreateTaskResponse, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	c, err := NewContainer(ctx, s.platform, &ContainerConfig{
		ID:                 rfs.Create.ID,
		Bundle:             rfs.Create.Bundle,
		Rootfs:             rfs.Create.Rootfs,
		Options:            rfs.Create.Options,
		Terminal:           rfs.Create.Terminal,
		Stdin:              rfs.Create.Stdin,
		Stdout:             rfs.Create.Stdout,
		Stderr:             rfs.Create.Stderr,
		FSRestoreImagePath: rfs.Conf.ImagePath,
		FSRestoreDirect:    rfs.Conf.Direct,
	})
	if err != nil {
		return nil, err
	}

	s.containers[rfs.Create.ID] = c

	p, err := c.Process("")
	if err != nil {
		return nil, err
	}
	if initProcess, ok := p.(*proc.Init); ok && initProcess.Sandbox {
		s.root = initProcess.Runtime().Root
		s.runtime = initProcess.Runtime()
	}

	// Set up OOM notification on the sandbox's cgroup. This is done on
	// sandbox create since the sandbox process will be created here.
	pid := p.Pid()
	if pid > 0 {
		var (
			cg  any
			err error
		)
		if cgroups.Mode() == cgroups.Unified {
			var cgPath string
			cgPath, err = cgroupsv2.PidGroupPath(pid)
			if err == nil {
				cg, err = cgroupsv2.Load(cgPath)
			}
		} else {
			cg, err = cgroup1.Load(cgroup1.PidPath(pid))
		}
		if err != nil {
			return nil, fmt.Errorf("loading cgroup for %d: %w", pid, err)
		}
		// Register the sandbox cgroup with the OOM poller using this
		// container's id
		if err := s.oomPoller.add(rfs.Create.ID, cg); err != nil {
			return nil, fmt.Errorf("add cg to OOM monitor: %w", err)
		}
	}
	return &task.CreateTaskResponse{
		Pid: uint32(pid),
	}, nil
}

// Start starts the container.
func (s *runscService) Start(ctx context.Context, r *task.StartRequest) (*task.StartResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	p, err := c.Start(ctx, r)
	if err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	if len(r.ExecID) == 0 {
		s.send(&events.TaskStart{
			ContainerID: r.ID,
			Pid:         uint32(p.Pid()),
		})
	} else {
		s.send(&events.TaskExecStarted{
			ContainerID: r.ID,
			ExecID:      r.ExecID,
			Pid:         uint32(p.Pid()),
		})
	}
	// TODO: Set the cgroup and oom notifications on restore.
	return &task.StartResponse{
		Pid: uint32(p.Pid()),
	}, nil
}

// Delete deletes the initial process and container.
func (s *runscService) Delete(ctx context.Context, r *task.DeleteRequest) (*task.DeleteResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	p, err := c.Delete(ctx, r)
	if err != nil {
		return nil, errgrpc.ToGRPC(err)
	}

	// ExecID will be empty for init container process.
	if len(r.ExecID) == 0 {
		s.mu.Lock()
		delete(s.containers, r.ID)
		hasCont := len(s.containers) > 0
		s.mu.Unlock()

		if !hasCont && s.platform != nil {
			s.platform.Close()
		}
	}
	return &task.DeleteResponse{
		ExitStatus: uint32(p.ExitStatus()),
		ExitedAt:   protobuf.ToTimestamp(p.ExitedAt()),
		Pid:        uint32(p.Pid()),
	}, nil
}

// Exec spawns an additional process inside the container.
func (s *runscService) Exec(ctx context.Context, r *task.ExecProcessRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	// Check whether or not the process already exists in the container.
	p, _ := c.Process(r.ExecID)
	if p != nil {
		return nil, errgrpc.ToGRPCf(errdefs.ErrAlreadyExists, "id %s", r.ExecID)
	}
	p, err = c.Exec(ctx, r)
	if err != nil {
		return nil, err
	}
	s.send(&events.TaskExecAdded{
		ContainerID: r.ID,
		ExecID:      p.ID(),
	})
	return empty, nil
}

// ResizePty resizes the terminal of a process.
func (s *runscService) ResizePty(ctx context.Context, r *task.ResizePtyRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, errgrpc.ToGRPCf(errdefs.ErrFailedPrecondition, "container must be created")
	}
	if err := c.ResizePty(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	return empty, nil
}

// State returns runtime state information for the container.
func (s *runscService) State(ctx context.Context, r *task.StateRequest) (*task.StateResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	p, err := c.Process(r.ExecID)
	if err != nil {
		log.L.Debugf("State failed to find process: %v", err)
		return nil, errgrpc.ToGRPC(err)
	}
	st, err := p.Status(ctx)
	if err != nil {
		log.L.Debugf("State failed: %v", err)
		return nil, err
	}
	status := tasktypes.Status_UNKNOWN
	switch st {
	case "created":
		status = tasktypes.Status_CREATED
	case "running":
		status = tasktypes.Status_RUNNING
	case "stopped":
		status = tasktypes.Status_STOPPED
	case "pausing":
		status = tasktypes.Status_PAUSING
	case "paused":
		status = tasktypes.Status_PAUSED
	}
	sio := p.Stdio()
	return &task.StateResponse{
		ID:         p.ID(),
		Bundle:     c.Bundle,
		Pid:        uint32(p.Pid()),
		Status:     status,
		Stdin:      sio.Stdin,
		Stdout:     sio.Stdout,
		Stderr:     sio.Stderr,
		Terminal:   sio.Terminal,
		ExitStatus: uint32(p.ExitStatus()),
		ExitedAt:   protobuf.ToTimestamp(p.ExitedAt()),
	}, nil
}

// Pause the container.
func (s *runscService) Pause(ctx context.Context, r *task.PauseRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	if err := c.Pause(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	s.send(&events.TaskPaused{
		ContainerID: c.ID,
	})
	return empty, nil
}

// Resume the container.
func (s *runscService) Resume(ctx context.Context, r *task.ResumeRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	if err := c.Resume(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	s.send(&events.TaskResumed{
		ContainerID: c.ID,
	})
	return empty, nil
}

// Kill the container with the provided signal.
func (s *runscService) Kill(ctx context.Context, r *task.KillRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	if err := c.Kill(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	return empty, nil
}

// Pids returns all pids inside the container.
func (s *runscService) Pids(ctx context.Context, r *task.PidsRequest) (*task.PidsResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	pids, err := s.getContainerPids(ctx, c)
	if err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	var processes []*tasktypes.ProcessInfo
	for _, pid := range pids {
		pInfo := tasktypes.ProcessInfo{
			Pid: pid,
		}
		for _, p := range c.ExecdProcesses() {
			if p.Pid() == int(pid) {
				d := &runctypes.ProcessDetails{
					ExecID: p.ID(),
				}
				a, err := typeurl.MarshalAnyToProto(d)
				if err != nil {
					return nil, fmt.Errorf("failed to marshal process %d info: %w", pid, err)
				}
				pInfo.Info = a
				break
			}
		}
		processes = append(processes, &pInfo)
	}
	return &task.PidsResponse{
		Processes: processes,
	}, nil
}

// CloseIO closes the I/O context of the container.
func (s *runscService) CloseIO(ctx context.Context, r *task.CloseIORequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	if err := c.CloseIO(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	return empty, nil
}

// Checkpoint checkpoints the container.
func (s *runscService) Checkpoint(ctx context.Context, r *task.CheckpointTaskRequest) (*types.Empty, error) {
	return empty, errdefs.ErrNotImplemented
}

// Restore restores the container.
func (s *runscService) Restore(ctx context.Context, r *extension.RestoreRequest) (*task.StartResponse, error) {
	c, err := s.getContainer(r.Start.ID)
	if err != nil {
		return nil, err
	}
	p, err := c.Restore(ctx, r)
	if err != nil {
		return nil, err
	}
	return &task.StartResponse{
		Pid: uint32(p.Pid()),
	}, nil
}

// Connect returns shim information such as the shim's pid.
func (s *runscService) Connect(ctx context.Context, r *task.ConnectRequest) (*task.ConnectResponse, error) {
	var pid int
	if c, err := s.getContainer(r.ID); err == nil {
		pid = c.Pid()
	}
	return &task.ConnectResponse{
		ShimPid: uint32(os.Getpid()),
		TaskPid: uint32(pid),
	}, nil
}

func (s *runscService) Shutdown(ctx context.Context, r *task.ShutdownRequest) (*types.Empty, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if len(s.containers) > 0 {
		return nil, fmt.Errorf("shim still servicing %d containers", len(s.containers))
	}

	return empty, nil
}

func (s *runscService) Stats(ctx context.Context, r *task.StatsRequest) (*task.StatsResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}

	// gvisor currently (as of 2020-03-03) only returns the total memory
	// usage and current PID value[0]. However, we copy the common fields here
	// so that future updates will propagate correct information.  We're
	// using the cgroups.Metrics structure so we're returning the same type
	// as runc.
	//
	// [0]: https://github.com/google/gvisor/blob/277a0d5a1fbe8272d4729c01ee4c6e374d047ebc/runsc/boot/events.go#L61-L81
	return c.Stats(ctx)
}

// Update updates a running container.
func (s *runscService) Update(ctx context.Context, r *task.UpdateTaskRequest) (*types.Empty, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	if err := c.Update(ctx, r); err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	return empty, nil
}

// Wait waits for the container to exit.
func (s *runscService) Wait(ctx context.Context, r *task.WaitRequest) (*task.WaitResponse, error) {
	c, err := s.getContainer(r.ID)
	if err != nil {
		return nil, err
	}
	p, err := c.Process(r.ExecID)
	if err != nil {
		return nil, errgrpc.ToGRPC(err)
	}
	p.Wait()

	res := &task.WaitResponse{
		ExitStatus: uint32(p.ExitStatus()),
		ExitedAt:   protobuf.ToTimestamp(p.ExitedAt()),
	}
	log.L.Debugf("Wait succeeded, response: %+v", res)
	return res, nil
}

func (s *runscService) processExits(ctx context.Context) {
	for e := range s.ec {
		s.checkProcesses(ctx, e)
	}
}

func (s *runscService) checkProcesses(ctx context.Context, e proc.Exit) {
	// Find the container whose process exited.
	s.mu.Lock()
	var (
		containerID string
		exitingP    process.Process
	)
	for cid, c := range s.containers {
		for _, p := range c.All() {
			if p.ID() == e.ID {
				containerID = cid
				exitingP = p
				break
			}
		}
		if exitingP != nil {
			break
		}
	}
	s.mu.Unlock()
	if exitingP == nil {
		return
	}
	p := exitingP

	ip, isInit := p.(*proc.Init)
	if isInit {
		// Ensure all children are killed.
		log.L.Debugf("Container init process exited, killing all container processes")
		ip.KillAll(ctx)
	}
	p.SetExited(e.Status)
	// On init process exit, synchronously check the cgroup for OOM kills
	// before publishing the exit event. The async OOM notification via
	// EventChan (cgroups v2) can lose the race against the container exit
	// on some architectures (notably aarch64), causing kubelet to report
	// reason=Error instead of reason=OOMKilled. Only check on init exit —
	// exec processes share the sandbox cgroup and would produce spurious
	// events.
	// Use the per-container id for TaskOOM routing.
	if isInit && s.oomPoller.isOOM(containerID) {
		s.send(&events.TaskOOM{ContainerID: containerID})
	}
	s.send(&events.TaskExit{
		ContainerID: containerID,
		ID:          p.ID(),
		Pid:         uint32(p.Pid()),
		ExitStatus:  uint32(e.Status),
		ExitedAt:    protobuf.ToTimestamp(p.ExitedAt()),
	})
}

func (s *runscService) send(event any) {
	s.events <- event
}

func (s *runscService) allProcessesForAllContainers() (o []process.Process) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, p := range s.containers {
		for _, p := range p.All() {
			o = append(o, p)
		}
	}
	return o
}

func (s *runscService) getContainerPids(ctx context.Context, c *Container) ([]uint32, error) {
	p, err := c.Process("")
	if err != nil {
		return nil, err
	}
	ps, err := p.(*proc.Init).Runtime().Ps(ctx, c.ID)
	if err != nil {
		return nil, err
	}
	pids := make([]uint32, 0, len(ps))
	for _, pid := range ps {
		pids = append(pids, uint32(pid))
	}
	return pids, nil
}

func (s *runscService) forward(ctx context.Context, publisher shim.Publisher) {
	for e := range s.events {
		err := publisher.Publish(ctx, getTopic(e), e)
		if err != nil {
			// Should not happen.
			panic(fmt.Errorf("post event: %w", err))
		}
	}
}

func getTopic(e any) string {
	switch e.(type) {
	case *events.TaskCreate:
		return runtime.TaskCreateEventTopic
	case *events.TaskStart:
		return runtime.TaskStartEventTopic
	case *events.TaskOOM:
		return runtime.TaskOOMEventTopic
	case *events.TaskExit:
		return runtime.TaskExitEventTopic
	case *events.TaskDelete:
		return runtime.TaskDeleteEventTopic
	case *events.TaskExecAdded:
		return runtime.TaskExecAddedEventTopic
	case *events.TaskExecStarted:
		return runtime.TaskExecStartedEventTopic
	case *events.TaskPaused:
		return runtime.TaskPausedEventTopic
	case *events.TaskResumed:
		return runtime.TaskResumedEventTopic
	default:
		log.L.Infof("no topic for type %#v", e)
	}
	return runtime.TaskUnknownTopic
}

func newInit(workDir, namespace string, platform stdio.Platform, r *proc.CreateConfig, options *Options, rootfs string) (*proc.Init, error) {
	spec, err := utils.ReadSpec(r.Bundle)
	if err != nil {
		return nil, fmt.Errorf("read oci spec: %w", err)
	}

	updated, err := utils.UpdateVolumeAnnotations(spec)
	if err != nil {
		return nil, fmt.Errorf("update volume annotations: %w", err)
	}
	updated = setPodCgroup(spec) || updated

	if updated {
		if err := utils.WriteSpec(r.Bundle, spec); err != nil {
			return nil, err
		}
	}

	runtime := proc.NewRunsc(options.Root, r.Bundle, namespace, options.BinaryName, options.RunscConfig, spec)
	p := proc.New(r.ID, runtime, stdio.Stdio{
		Stdin:    r.Stdin,
		Stdout:   r.Stdout,
		Stderr:   r.Stderr,
		Terminal: r.Terminal,
	})
	p.Bundle = r.Bundle
	p.Platform = platform
	p.Rootfs = rootfs
	p.WorkDir = workDir
	p.IoUID = int(options.IoUID)
	p.IoGID = int(options.IoGID)
	p.Sandbox = specutils.SpecContainerType(spec) == specutils.ContainerTypeSandbox
	p.UserLog = utils.UserLogPath(spec)
	p.Monitor = reaper.Default
	return p, nil
}

// setPodCgroup searches for the pod cgroup path inside the container's cgroup
// path. If found, it's set as an annotation in the spec. This is done so that
// the sandbox joins the pod cgroup. Otherwise, the sandbox would join the pause
// container cgroup. Returns true if the spec was modified. Ex.:
// /kubepods/burstable/pod123/container123 => kubepods/burstable/pod123
func setPodCgroup(spec *specs.Spec) bool {
	if !utils.IsSandbox(spec) {
		return false
	}
	if spec.Linux == nil || len(spec.Linux.CgroupsPath) == 0 {
		return false
	}

	// Search backwards for the pod cgroup path to make the sandbox use it,
	// instead of the pause container's cgroup.
	parts := strings.Split(spec.Linux.CgroupsPath, string(filepath.Separator))
	for i := len(parts) - 1; i >= 0; i-- {
		if strings.HasPrefix(parts[i], "pod") {
			var path string
			for j := 0; j <= i; j++ {
				path = filepath.Join(path, parts[j])
			}
			// Add back the initial '/' that may have been lost above.
			if filepath.IsAbs(spec.Linux.CgroupsPath) {
				path = string(filepath.Separator) + path
			}
			if spec.Linux.CgroupsPath == path {
				return false
			}
			if spec.Annotations == nil {
				spec.Annotations = make(map[string]string)
			}
			spec.Annotations[cgroupParentAnnotation] = path
			return true
		}
	}
	return false
}

// GvisorTaskServer adapters runscService to taskServer.GvisorTaskServiceExt.
type GvisorTaskServer struct {
	s *runscService
}

var _ taskServer.GvisorTaskServiceExt = (*GvisorTaskServer)(nil)

// NewGvisorTaskServer creates a new GvisorTaskServer adapter.
func NewGvisorTaskServer(s extension.TaskServiceExt) *GvisorTaskServer {
	rs, ok := s.(*runscService)
	if !ok {
		return nil
	}
	return &GvisorTaskServer{s: rs}
}

// Checkpoint implements taskServer.GvisorTaskServiceExt.
func (g *GvisorTaskServer) Checkpoint(ctx context.Context, req *pb.CheckpointRequest) (*pb.CheckpointResponse, error) {
	c, err := g.s.getContainer(req.GetId())
	if err != nil {
		return nil, err
	}
	p, err := c.Process("")
	if err != nil {
		return nil, err
	}
	initProc, ok := p.(*proc.Init)
	if !ok {
		return nil, fmt.Errorf("process is not init process")
	}

	// Set debug options on the runtime for this execution
	runsc := setDebug(initProc.Runtime(), req.GetDebug())

	opts := &runsccmd.CheckpointOpts{
		ImagePath:                 req.GetImagePath(),
		LeaveRunning:              req.GetLeaveRunning(),
		Direct:                    req.GetDirect(),
		Compression:               req.GetCompression(),
		ExcludeCommittedZeroPages: req.GetExcludeCommittedZeroPages(),
		SaveRestoreExecArgv:       req.GetSaveRestoreExecArgv(),
		SaveRestoreExecTimeout:    req.GetSaveRestoreExecTimeout(),
		CudaCheckpointPath:        req.GetCudaCheckpointPath(),
		CudaCheckpointSequential:  req.GetCudaCheckpointSequential(),
		WorkPath:                  req.GetWorkPath(),
		FSPath:                    req.GetFsPath(),
	}

	if req.GetFsCheckpoint() {
		if err := runsc.FSCheckpoint(ctx, c.ID, opts); err != nil {
			return nil, err
		}
	} else {
		if err := runsc.Checkpoint(ctx, c.ID, opts); err != nil {
			return nil, err
		}
	}

	return &pb.CheckpointResponse{}, nil
}

// Wait implements taskServer.GvisorTaskServiceExt.
func (g *GvisorTaskServer) Wait(ctx context.Context, req *pb.WaitRequest) (*pb.WaitResponse, error) {
	c, err := g.s.getContainer(req.GetId())
	if err != nil {
		return nil, err
	}
	p, err := c.Process("")
	if err != nil {
		return nil, err
	}
	initProc, ok := p.(*proc.Init)
	if !ok {
		return nil, fmt.Errorf("process is not init process")
	}

	// Set debug options on the runtime for this execution
	runsc := setDebug(initProc.Runtime(), req.GetDebug())

	opts := &runsccmd.WaitOpts{
		PID:     int(req.GetPid()),
		RootPID: int(req.GetRootPid()),
	}
	switch req.GetWaitType() {
	case pb.WaitRequest_CHECKPOINT:
		opts.Checkpoint = true
	case pb.WaitRequest_RESTORE:
		opts.Restore = true
	case pb.WaitRequest_FSCHECKPOINT:
		opts.FSCheckpoint = true
	case pb.WaitRequest_FSRESTORE:
		opts.FSRestore = true
	}

	exitStatus, err := runsc.WaitWithOptions(ctx, c.ID, opts)
	if err != nil {
		return nil, err
	}

	return &pb.WaitResponse{ExitStatus: int32(exitStatus)}, nil
}

// State implements taskServer.GvisorTaskServiceExt.
func (g *GvisorTaskServer) State(ctx context.Context, req *pb.StateRequest) (*pb.StateResponse, error) {
	if c, err := g.s.getContainer(req.GetId()); err == nil {
		if p, err := c.Process(""); err == nil {
			if initProc, ok := p.(*proc.Init); ok {
				runsc := setDebug(initProc.Runtime(), req.GetDebug())
				if _, err := runsc.State(ctx, c.ID); err != nil {
					log.L.Debugf("State with debug failed: %v", err)
				}
			}
		}
	}

	resp, err := g.s.State(ctx, &task.StateRequest{ID: req.GetId()})
	if err != nil {
		return nil, err
	}
	stateStr := ""
	switch resp.Status {
	case tasktypes.Status_CREATED:
		stateStr = "created"
	case tasktypes.Status_RUNNING:
		stateStr = "running"
	case tasktypes.Status_STOPPED:
		stateStr = "stopped"
	case tasktypes.Status_PAUSING:
		stateStr = "pausing"
	case tasktypes.Status_PAUSED:
		stateStr = "paused"
	default:
		stateStr = "unknown"
	}
	id := resp.ID
	pid := int32(resp.Pid)
	return &pb.StateResponse{
		Id:    id,
		State: stateStr,
		Pid:   pid,
	}, nil
}

// Version implements taskServer.GvisorTaskServiceExt.
func (g *GvisorTaskServer) Version(ctx context.Context, req *pb.VersionRequest) (*pb.VersionResponse, error) {
	cmd := exec.Command("runsc", "-version")
	if g.s.runtime.Command != "" {
		cmd = exec.Command(g.s.runtime.Command, "-version")
	}
	out, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("failed to get runsc version: %w, output: %s", err, string(out))
	}
	versionStr := strings.TrimSpace(string(out))
	return &pb.VersionResponse{Version: versionStr}, nil
}

func setDebug(r *runsccmd.Runsc, d *pb.Debug) *runsccmd.Runsc {
	if r == nil {
		return nil
	}
	cp := *r
	if d == nil {
		cp.Debug = false
		cp.DebugLog = ""
		cp.DebugLogFD = nil
		return &cp
	}
	cp.Debug = true
	cp.DebugLog = d.GetDebugLog()
	cp.DebugLogFD = nil
	if fd := int(d.GetDebugLogFd()); fd != 0 {
		cp.DebugLogFD = &fd
	}
	return &cp
}
