From cabcc398f648dfab148aeedd28c49f640c6bf10a Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Wed, 11 Feb 2015 14:21:38 -0500
Subject: [PATCH] Allow setting ulimits for containers

Signed-off-by: Brian Goff <cpuguy83@gmail.com>
---
 ulimit/ulimit.go      | 106 ++++++++++++++++++++++++++++++++++++++++++
 ulimit/ulimit_test.go |  41 ++++++++++++++++
 2 files changed, 147 insertions(+)
 create mode 100644 ulimit/ulimit.go
 create mode 100644 ulimit/ulimit_test.go

diff --git a/ulimit/ulimit.go b/ulimit/ulimit.go
new file mode 100644
index 0000000..2375315
--- /dev/null
+++ b/ulimit/ulimit.go
@@ -0,0 +1,106 @@
+package ulimit
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+// Human friendly version of Rlimit
+type Ulimit struct {
+	Name string
+	Hard int64
+	Soft int64
+}
+
+type Rlimit struct {
+	Type int    `json:"type,omitempty"`
+	Hard uint64 `json:"hard,omitempty"`
+	Soft uint64 `json:"soft,omitempty"`
+}
+
+const (
+	// magic numbers for making the syscall
+	// some of these are defined in the syscall package, but not all.
+	// Also since Windows client doesn't get access to the syscall package, need to
+	//	define these here
+	RLIMIT_AS         = 9
+	RLIMIT_CORE       = 4
+	RLIMIT_CPU        = 0
+	RLIMIT_DATA       = 2
+	RLIMIT_FSIZE      = 1
+	RLIMIT_LOCKS      = 10
+	RLIMIT_MEMLOCK    = 8
+	RLIMIT_MSGQUEUE   = 12
+	RLIMIT_NICE       = 13
+	RLIMIT_NOFILE     = 7
+	RLIMIT_NPROC      = 6
+	RLIMIT_RSS        = 5
+	RLIMIT_RTPRIO     = 14
+	RLIMIT_RTTIME     = 15
+	RLIMIT_SIGPENDING = 11
+	RLIMIT_STACK      = 3
+)
+
+var ulimitNameMapping = map[string]int{
+	//"as":         RLIMIT_AS, // Disbaled since this doesn't seem usable with the way Docker inits a container.
+	"core":       RLIMIT_CORE,
+	"cpu":        RLIMIT_CPU,
+	"data":       RLIMIT_DATA,
+	"fsize":      RLIMIT_FSIZE,
+	"locks":      RLIMIT_LOCKS,
+	"memlock":    RLIMIT_MEMLOCK,
+	"msgqueue":   RLIMIT_MSGQUEUE,
+	"nice":       RLIMIT_NICE,
+	"nofile":     RLIMIT_NOFILE,
+	"nproc":      RLIMIT_NPROC,
+	"rss":        RLIMIT_RSS,
+	"rtprio":     RLIMIT_RTPRIO,
+	"rttime":     RLIMIT_RTTIME,
+	"sigpending": RLIMIT_SIGPENDING,
+	"stack":      RLIMIT_STACK,
+}
+
+func Parse(val string) (*Ulimit, error) {
+	parts := strings.SplitN(val, "=", 2)
+	if len(parts) != 2 {
+		return nil, fmt.Errorf("invalid ulimit argument: %s", val)
+	}
+
+	if _, exists := ulimitNameMapping[parts[0]]; !exists {
+		return nil, fmt.Errorf("invalid ulimit type: %s", parts[0])
+	}
+
+	limitVals := strings.SplitN(parts[1], ":", 2)
+	if len(limitVals) > 2 {
+		return nil, fmt.Errorf("too many limit value arguments - %s, can only have up to two, `soft[:hard]`", parts[1])
+	}
+
+	soft, err := strconv.ParseInt(limitVals[0], 10, 64)
+	if err != nil {
+		return nil, err
+	}
+
+	hard := soft // in case no hard was set
+	if len(limitVals) == 2 {
+		hard, err = strconv.ParseInt(limitVals[1], 10, 64)
+	}
+	if soft > hard {
+		return nil, fmt.Errorf("ulimit soft limit must be less than or equal to hard limit: %d > %d", soft, hard)
+	}
+
+	return &Ulimit{Name: parts[0], Soft: soft, Hard: hard}, nil
+}
+
+func (u *Ulimit) GetRlimit() (*Rlimit, error) {
+	t, exists := ulimitNameMapping[u.Name]
+	if !exists {
+		return nil, fmt.Errorf("invalid ulimit name %s", u.Name)
+	}
+
+	return &Rlimit{Type: t, Soft: uint64(u.Soft), Hard: uint64(u.Hard)}, nil
+}
+
+func (u *Ulimit) String() string {
+	return fmt.Sprintf("%s=%s:%s", u.Name, u.Soft, u.Hard)
+}
diff --git a/ulimit/ulimit_test.go b/ulimit/ulimit_test.go
new file mode 100644
index 0000000..419b5e0
--- /dev/null
+++ b/ulimit/ulimit_test.go
@@ -0,0 +1,41 @@
+package ulimit
+
+import "testing"
+
+func TestParseInvalidLimitType(t *testing.T) {
+	if _, err := Parse("notarealtype=1024:1024"); err == nil {
+		t.Fatalf("expected error on invalid ulimit type")
+	}
+}
+
+func TestParseBadFormat(t *testing.T) {
+	if _, err := Parse("nofile:1024:1024"); err == nil {
+		t.Fatal("expected error on bad syntax")
+	}
+
+	if _, err := Parse("nofile"); err == nil {
+		t.Fatal("expected error on bad syntax")
+	}
+
+	if _, err := Parse("nofile="); err == nil {
+		t.Fatal("expected error on bad syntax")
+	}
+	if _, err := Parse("nofile=:"); err == nil {
+		t.Fatal("expected error on bad syntax")
+	}
+	if _, err := Parse("nofile=:1024"); err == nil {
+		t.Fatal("expected error on bad syntax")
+	}
+}
+
+func TestParseHardLessThanSoft(t *testing.T) {
+	if _, err := Parse("nofile:1024:1"); err == nil {
+		t.Fatal("expected error on hard limit less than soft limit")
+	}
+}
+
+func TestParseInvalidValueType(t *testing.T) {
+	if _, err := Parse("nofile:asdf"); err == nil {
+		t.Fatal("expected error on bad value type")
+	}
+}