From e67c6be052c744fc6b2f917cd4960b036701646a Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Mon, 12 Jun 2017 10:14:40 -0700
Subject: [PATCH] Add default namespace to client

Adds a default namespace to the client.
Intercepts client requests to add namespace to context.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>
---
 client.go | 43 ++++++++++++++++++++++++++++++++-----------
 grpc.go   | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 11 deletions(-)
 create mode 100644 grpc.go

diff --git a/client.go b/client.go
index 08095d5df..a60588d0f 100644
--- a/client.go
+++ b/client.go
@@ -41,30 +41,50 @@ func init() {
 	grpclog.SetLogger(log.New(ioutil.Discard, "", log.LstdFlags))
 }
 
-type NewClientOpts func(c *Client) error
+type clientOpts struct {
+	defaultns string
+}
+
+type ClientOpt func(c *clientOpts) error
+
+func WithDefaultNamespace(ns string) ClientOpt {
+	return func(c *clientOpts) error {
+		c.defaultns = ns
+		return nil
+	}
+}
 
 // New returns a new containerd client that is connected to the containerd
 // instance provided by address
-func New(address string, opts ...NewClientOpts) (*Client, error) {
+func New(address string, opts ...ClientOpt) (*Client, error) {
+	var copts clientOpts
+	for _, o := range opts {
+		if err := o(&copts); err != nil {
+			return nil, err
+		}
+	}
+
 	gopts := []grpc.DialOption{
 		grpc.WithInsecure(),
 		grpc.WithTimeout(100 * time.Second),
 		grpc.WithDialer(dialer),
 	}
+	if copts.defaultns != "" {
+		unary, stream := newNSInterceptors(copts.defaultns)
+		gopts = append(gopts,
+			grpc.WithUnaryInterceptor(unary),
+			grpc.WithStreamInterceptor(stream),
+		)
+	}
+
 	conn, err := grpc.Dial(dialAddress(address), gopts...)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to dial %q", address)
 	}
-	c := &Client{
+	return &Client{
 		conn:    conn,
 		runtime: runtime.GOOS,
-	}
-	for _, o := range opts {
-		if err := o(c); err != nil {
-			return nil, err
-		}
-	}
-	return c, nil
+	}, nil
 }
 
 // Client is the client to interact with containerd and its various services
@@ -72,7 +92,8 @@ func New(address string, opts ...NewClientOpts) (*Client, error) {
 type Client struct {
 	conn *grpc.ClientConn
 
-	runtime string
+	defaultns string
+	runtime   string
 }
 
 func (c *Client) IsServing(ctx context.Context) (bool, error) {
diff --git a/grpc.go b/grpc.go
new file mode 100644
index 000000000..c56eeada3
--- /dev/null
+++ b/grpc.go
@@ -0,0 +1,35 @@
+package containerd
+
+import (
+	"github.com/containerd/containerd/namespaces"
+	"golang.org/x/net/context"
+	"google.golang.org/grpc"
+)
+
+type namespaceInterceptor struct {
+	namespace string
+}
+
+func (ni namespaceInterceptor) unary(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	_, ok := namespaces.Namespace(ctx)
+	if !ok {
+		ctx = namespaces.WithNamespace(ctx, ni.namespace)
+	}
+	return invoker(ctx, method, req, reply, cc, opts...)
+}
+
+func (ni namespaceInterceptor) stream(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	_, ok := namespaces.Namespace(ctx)
+	if !ok {
+		ctx = namespaces.WithNamespace(ctx, ni.namespace)
+	}
+
+	return streamer(ctx, desc, cc, method, opts...)
+}
+
+func newNSInterceptors(ns string) (grpc.UnaryClientInterceptor, grpc.StreamClientInterceptor) {
+	ni := namespaceInterceptor{
+		namespace: ns,
+	}
+	return grpc.UnaryClientInterceptor(ni.unary), grpc.StreamClientInterceptor(ni.stream)
+}