From fa6a9f063586e30833bdb9842e84289f4ef11020 Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Tue, 4 Jun 2019 15:21:10 +0000
Subject: [PATCH] Add ttrpc namespace support

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 namespaces/context.go | 12 +++++-------
 namespaces/ttrpc.go   | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 namespaces/ttrpc.go

diff --git a/namespaces/context.go b/namespaces/context.go
index cc5621a68..b4e988e7b 100644
--- a/namespaces/context.go
+++ b/namespaces/context.go
@@ -36,10 +36,9 @@ type namespaceKey struct{}
 // WithNamespace sets a given namespace on the context
 func WithNamespace(ctx context.Context, namespace string) context.Context {
 	ctx = context.WithValue(ctx, namespaceKey{}, namespace) // set our key for namespace
-
-	// also store on the grpc headers so it gets picked up by any clients that
+	// also store on the grpc and ttrpc headers so it gets picked up by any clients that
 	// are using this.
-	return withGRPCNamespaceHeader(ctx, namespace)
+	return withTTRPCNamespaceHeader(withGRPCNamespaceHeader(ctx, namespace), namespace)
 }
 
 // NamespaceFromEnv uses the namespace defined in CONTAINERD_NAMESPACE or
@@ -58,9 +57,10 @@ func NamespaceFromEnv(ctx context.Context) context.Context {
 func Namespace(ctx context.Context) (string, bool) {
 	namespace, ok := ctx.Value(namespaceKey{}).(string)
 	if !ok {
-		return fromGRPCHeader(ctx)
+		if namespace, ok = fromGRPCHeader(ctx); !ok {
+			return fromTTRPCHeader(ctx)
+		}
 	}
-
 	return namespace, ok
 }
 
@@ -70,10 +70,8 @@ func NamespaceRequired(ctx context.Context) (string, error) {
 	if !ok || namespace == "" {
 		return "", errors.Wrapf(errdefs.ErrFailedPrecondition, "namespace is required")
 	}
-
 	if err := Validate(namespace); err != nil {
 		return "", errors.Wrap(err, "namespace validation")
 	}
-
 	return namespace, nil
 }
diff --git a/namespaces/ttrpc.go b/namespaces/ttrpc.go
new file mode 100644
index 000000000..28e039417
--- /dev/null
+++ b/namespaces/ttrpc.go
@@ -0,0 +1,41 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package namespaces
+
+import (
+	"context"
+
+	"github.com/containerd/ttrpc"
+)
+
+const (
+	// TTRPCHeader defines the header name for specifying a containerd namespace
+	TTRPCHeader = "containerd-namespace-ttrpc"
+)
+
+func withTTRPCNamespaceHeader(ctx context.Context, namespace string) context.Context {
+	md, ok := ttrpc.GetMetadata(ctx)
+	if !ok {
+		md = ttrpc.Metadata{}
+	}
+	md.Set(TTRPCHeader, namespace)
+	return ttrpc.WithMetadata(ctx, md)
+}
+
+func fromTTRPCHeader(ctx context.Context) (string, bool) {
+	return ttrpc.GetMetadataValue(ctx, TTRPCHeader)
+}