File: main_forklimits.go

package info (click to toggle)
incus 6.0.5-6
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 25,788 kB
  • sloc: sh: 16,313; ansic: 3,121; python: 457; makefile: 337; ruby: 51; sql: 50; lisp: 6
file content (147 lines) | stat: -rw-r--r-- 3,340 bytes parent folder | download | duplicates (4)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
package main

import (
	"errors"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
)

var reLimitsArg = regexp.MustCompile(`^limit=(\w+):(\w+):(\w+)$`)

type cmdForklimits struct {
	global *cmdGlobal
}

func (c *cmdForklimits) command() *cobra.Command {
	// Main subcommand
	cmd := &cobra.Command{}
	cmd.Use = "forklimits [fd=<number>...] [limit=<name>:<softlimit>:<hardlimit>...] -- <command> [<arg>...]"
	cmd.Short = "Execute a task inside the container"
	cmd.Long = `Description:
  Execute a command with specific limits set.

  This internal command is used to spawn a command with limits set. It can also pass through one or more filed escriptors specified by fd=n arguments.
  These are passed through in the order they are specified.
`
	cmd.RunE = c.run
	cmd.Hidden = true

	return cmd
}

func (c *cmdForklimits) run(cmd *cobra.Command, _ []string) error {
	// Use raw args instead of cobra passed args, as we need to access the "--" argument.
	args := c.global.rawArgs(cmd)

	if len(args) == 0 {
		_ = cmd.Help()
		return nil
	}

	// Only root should run this
	if os.Geteuid() != 0 {
		return errors.New("This must be run as root")
	}

	type limit struct {
		name string
		soft string
		hard string
	}

	var limits []limit
	var fds []uintptr
	var cmdParts []string

	for i, arg := range args {
		matches := reLimitsArg.FindStringSubmatch(arg)
		if len(matches) == 4 {
			limits = append(limits, limit{
				name: matches[1],
				soft: matches[2],
				hard: matches[3],
			})
		} else if strings.HasPrefix(arg, "fd=") {
			fdParts := strings.SplitN(arg, "=", 2)
			fdNum, err := strconv.Atoi(fdParts[1])
			if err != nil {
				_ = cmd.Help()
				return errors.New("Invalid file descriptor number")
			}

			fds = append(fds, uintptr(fdNum))
		} else if arg == "--" {
			if len(args)-1 > i {
				cmdParts = args[i+1:]
			}

			break // No more passing of arguments needed.
		} else {
			_ = cmd.Help()
			return errors.New("Unrecognised argument")
		}
	}

	// Setup rlimits.
	for _, limit := range limits {
		var resource int
		var rLimit unix.Rlimit

		if limit.name == "memlock" {
			resource = unix.RLIMIT_MEMLOCK
		} else {
			return fmt.Errorf("Unsupported limit type: %q", limit.name)
		}

		if limit.soft == "unlimited" {
			rLimit.Cur = unix.RLIM_INFINITY
		} else {
			softLimit, err := strconv.ParseUint(limit.soft, 10, 64)
			if err != nil {
				return fmt.Errorf("Invalid soft limit for %q", limit.name)
			}

			rLimit.Cur = softLimit
		}

		if limit.hard == "unlimited" {
			rLimit.Max = unix.RLIM_INFINITY
		} else {
			hardLimit, err := strconv.ParseUint(limit.hard, 10, 64)
			if err != nil {
				return fmt.Errorf("Invalid hard limit for %q", limit.name)
			}

			rLimit.Max = hardLimit
		}

		err := unix.Setrlimit(resource, &rLimit)
		if err != nil {
			return err
		}
	}

	if len(cmdParts) == 0 {
		_ = cmd.Help()
		return errors.New("Missing required command argument")
	}

	// Clear the cloexec flag on the file descriptors we are passing through.
	for _, fd := range fds {
		_, _, syscallErr := unix.Syscall(unix.SYS_FCNTL, fd, unix.F_SETFD, uintptr(0))
		if syscallErr != 0 {
			err := os.NewSyscallError(fmt.Sprintf("fcntl failed on FD %d", fd), syscallErr)
			if err != nil {
				return err
			}
		}
	}

	return unix.Exec(cmdParts[0], cmdParts, os.Environ())
}