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())
}
|