/
main_forklimits.go
146 lines (121 loc) · 3.25 KB
/
main_forklimits.go
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
package main
import (
"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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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 fmt.Errorf("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())
}