Ensure namespace is proxied to grpc/ttrpc plugins

cpuguy83 · cpuguy83 · commit 130a9c7ddbdd · 2021-10-16T00:09:50.000Z
Before this change we only ever had the grpc incoming medata set so when
we make a request to a shim or a grpc plugin the namespace is not sent
over the RPC.

I need this for github.com/cpuguy83/systemdshim, and I am sure there are
other use-cases where this would be needed.

Signed-off-by: Brian Goff &lt;cpuguy83@gmail.com&gt;
diff --git a/services/server/namespace.go b/services/server/namespace.go
@@ -0,0 +1,52 @@
+/*
+   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 server
+
+import (
+	"context"
+
+	"github.com/containerd/containerd/namespaces"
+	"google.golang.org/grpc"
+)
+
+func unaryNamespaceInterceptor(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	if ns, ok := namespaces.Namespace(ctx); ok {
+		// The above call checks the *incoming* metadata, this makes sure the outgoing metadata is also set
+		ctx = namespaces.WithNamespace(ctx, ns)
+	}
+	return handler(ctx, req)
+}
+
+func streamNamespaceInterceptor(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	ctx := ss.Context()
+	if ns, ok := namespaces.Namespace(ctx); ok {
+		// The above call checks the *incoming* metadata, this makes sure the outgoing metadata is also set
+		ctx = namespaces.WithNamespace(ctx, ns)
+		ss = &wrappedSSWithContext{ctx: ctx, ServerStream: ss}
+	}
+
+	return handler(srv, ss)
+}
+
+type wrappedSSWithContext struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (w *wrappedSSWithContext) Context() context.Context {
+	return w.ctx
+}
diff --git a/services/server/server.go b/services/server/server.go
@@ -102,10 +102,12 @@ func New(ctx context.Context, config *srvconfig.Config) (*Server, error) {
 		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
 			otelgrpc.StreamServerInterceptor(),
 			grpc.StreamServerInterceptor(grpc_prometheus.StreamServerInterceptor),
+			streamNamespaceInterceptor,
 		)),
 		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
 			otelgrpc.UnaryServerInterceptor(),
 			grpc.UnaryServerInterceptor(grpc_prometheus.UnaryServerInterceptor),
+			unaryNamespaceInterceptor,
 		)),
 	}
 	if config.GRPC.MaxRecvMsgSize > 0 {
