From 15c6e832aafd445653f5ed1faff1629a3072d3bb Mon Sep 17 00:00:00 2001
From: Samuel Karp <skarp@amazon.com>
Date: Fri, 16 Jun 2017 14:29:00 -0700
Subject: [PATCH] containerd: Add WithNamespace to replace namespace

Signed-off-by: Samuel Karp <skarp@amazon.com>
---
 spec_unix.go      | 18 ++++++++++++++++++
 spec_unix_test.go | 26 +++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/spec_unix.go b/spec_unix.go
index c74d88e5f..9eb337710 100644
--- a/spec_unix.go
+++ b/spec_unix.go
@@ -188,6 +188,24 @@ func WithHostNamespace(ns specs.LinuxNamespaceType) SpecOpts {
 	}
 }
 
+// WithLinuxNamespace uses the passed in namespace for the spec. If a namespace of the same type already exists in the
+// spec, the existing namespace is replaced by the one provided.
+func WithLinuxNamespace(ns specs.LinuxNamespace) SpecOpts {
+	return func(s *specs.Spec) error {
+		for i, n := range s.Linux.Namespaces {
+			if n.Type == ns.Type {
+				before := s.Linux.Namespaces[:i]
+				after := s.Linux.Namespaces[i+1:]
+				s.Linux.Namespaces = append(before, ns)
+				s.Linux.Namespaces = append(s.Linux.Namespaces, after...)
+				return nil
+			}
+		}
+		s.Linux.Namespaces = append(s.Linux.Namespaces, ns)
+		return nil
+	}
+}
+
 func WithImageConfig(ctx context.Context, i Image) SpecOpts {
 	return func(s *specs.Spec) error {
 		var (
diff --git a/spec_unix_test.go b/spec_unix_test.go
index 74fe49179..394771ea3 100644
--- a/spec_unix_test.go
+++ b/spec_unix_test.go
@@ -2,7 +2,11 @@
 
 package containerd
 
-import "testing"
+import (
+	"testing"
+
+	specs "github.com/opencontainers/runtime-spec/specs-go"
+)
 
 func TestGenerateSpec(t *testing.T) {
 	s, err := GenerateSpec()
@@ -56,3 +60,23 @@ func TestSpecWithTTY(t *testing.T) {
 		t.Errorf("xterm not set in env for TTY")
 	}
 }
+
+func TestWithLinuxNamespace(t *testing.T) {
+	replacedNS := specs.LinuxNamespace{Type: specs.NetworkNamespace, Path: "/var/run/netns/test"}
+	s, err := GenerateSpec(WithLinuxNamespace(replacedNS))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defaultNS := defaultNamespaces()
+	found := false
+	for i, ns := range s.Linux.Namespaces {
+		if ns == replacedNS && !found {
+			found = true
+			continue
+		}
+		if defaultNS[i] != ns {
+			t.Errorf("ns at %d does not match set %q != %q", i, defaultNS[i], ns)
+		}
+	}
+}