From bdd84abf056b153b3e37883884e91c52f98e15b9 Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Fri, 15 Mar 2019 12:23:51 -0400
Subject: [PATCH] Add additional capability handling opts

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 oci/spec_opts.go      | 63 ++++++++++++++++++++++++++++++++++++++++---
 oci/spec_opts_test.go | 44 ++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/oci/spec_opts.go b/oci/spec_opts.go
index dcffc0787..0769d511f 100644
--- a/oci/spec_opts.go
+++ b/oci/spec_opts.go
@@ -33,7 +33,7 @@ import (
 	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/continuity/fs"
-	"github.com/opencontainers/image-spec/specs-go/v1"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/opencontainers/runc/libcontainer/user"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/pkg/errors"
@@ -741,9 +741,11 @@ func WithCapabilities(caps []string) SpecOpts {
 }
 
 // WithAllCapabilities sets all linux capabilities for the process
-var WithAllCapabilities = WithCapabilities(getAllCapabilities())
+var WithAllCapabilities = WithCapabilities(GetAllCapabilities())
 
-func getAllCapabilities() []string {
+// GetAllCapabilities returns all caps up to CAP_LAST_CAP
+// or CAP_BLOCK_SUSPEND on RHEL6
+func GetAllCapabilities() []string {
 	last := capability.CAP_LAST_CAP
 	// hack for RHEL6 which has no /proc/sys/kernel/cap_last_cap
 	if last == capability.Cap(63) {
@@ -759,6 +761,61 @@ func getAllCapabilities() []string {
 	return caps
 }
 
+func capsContain(caps []string, s string) bool {
+	for _, c := range caps {
+		if c == s {
+			return true
+		}
+	}
+	return false
+}
+
+func removeCap(caps *[]string, s string) {
+	for i, c := range *caps {
+		if c == s {
+			*caps = append((*caps)[:i], (*caps)[i+1:]...)
+		}
+	}
+}
+
+// WithAddedCapabilities adds the provided capabilities
+func WithAddedCapabilities(caps []string) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		setCapabilities(s)
+		for _, c := range caps {
+			for _, cl := range []*[]string{
+				&s.Process.Capabilities.Bounding,
+				&s.Process.Capabilities.Effective,
+				&s.Process.Capabilities.Permitted,
+				&s.Process.Capabilities.Inheritable,
+			} {
+				if !capsContain(*cl, c) {
+					*cl = append(*cl, c)
+				}
+			}
+		}
+		return nil
+	}
+}
+
+// WithDroppedCapabilities removes the provided capabilities
+func WithDroppedCapabilities(caps []string) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		setCapabilities(s)
+		for _, c := range caps {
+			for _, cl := range []*[]string{
+				&s.Process.Capabilities.Bounding,
+				&s.Process.Capabilities.Effective,
+				&s.Process.Capabilities.Permitted,
+				&s.Process.Capabilities.Inheritable,
+			} {
+				removeCap(cl, c)
+			}
+		}
+		return nil
+	}
+}
+
 // WithAmbientCapabilities set the Linux ambient capabilities for the process
 // Ambient capabilities should only be set for non-root users or the caller should
 // understand how these capabilities are used and set
diff --git a/oci/spec_opts_test.go b/oci/spec_opts_test.go
index 74b246566..7cc65af94 100644
--- a/oci/spec_opts_test.go
+++ b/oci/spec_opts_test.go
@@ -422,3 +422,47 @@ func TestWithImageConfigArgs(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestAddCaps(t *testing.T) {
+	t.Parallel()
+
+	var s specs.Spec
+
+	if err := WithAddedCapabilities([]string{"CAP_CHOWN"})(nil, nil, nil, &s); err != nil {
+		t.Fatal(err)
+	}
+	for i, cl := range [][]string{
+		s.Process.Capabilities.Bounding,
+		s.Process.Capabilities.Effective,
+		s.Process.Capabilities.Permitted,
+		s.Process.Capabilities.Inheritable,
+	} {
+		if !capsContain(cl, "CAP_CHOWN") {
+			t.Errorf("cap list %d does not contain added cap", i)
+		}
+	}
+}
+
+func TestDropCaps(t *testing.T) {
+	t.Parallel()
+
+	var s specs.Spec
+
+	if err := WithAllCapabilities(nil, nil, nil, &s); err != nil {
+		t.Fatal(err)
+	}
+	if err := WithDroppedCapabilities([]string{"CAP_CHOWN"})(nil, nil, nil, &s); err != nil {
+		t.Fatal(err)
+	}
+
+	for i, cl := range [][]string{
+		s.Process.Capabilities.Bounding,
+		s.Process.Capabilities.Effective,
+		s.Process.Capabilities.Permitted,
+		s.Process.Capabilities.Inheritable,
+	} {
+		if capsContain(cl, "CAP_CHOWN") {
+			t.Errorf("cap list %d contains dropped cap", i)
+		}
+	}
+}