server: add connection to transport context

internal/transport/http2_server.go

@@ -217,7 +217,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 	}
 	done := make(chan struct{})
 	t := &http2Server{
-		ctx:               context.Background(),
+		ctx:               setConnection(context.Background(), conn),
 		done:              done,
 		conn:              conn,
 		remoteAddr:        conn.RemoteAddr(),
@@ -1345,3 +1345,18 @@ func getJitter(v time.Duration) time.Duration {
 	j := grpcrand.Int63n(2*r) - r
 	return time.Duration(j)
 }
+
+type connectionKey struct{}
+
+// GetConnection gets the connection from the context.
+func GetConnection(ctx context.Context) net.Conn {
+	conn, _ := ctx.Value(connectionKey{}).(net.Conn)
+	return conn
+}
+
+// SetConnection adds the connection to the context to be able to get
+// information about the destination ip and port for an incoming RPC. This also
+// allows any unary or streaming interceptors to see the connection.
+func setConnection(ctx context.Context, conn net.Conn) context.Context {
+	return context.WithValue(ctx, connectionKey{}, conn)
+}

internal/xds/rbac/rbac_engine.go

@@ -32,11 +32,14 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/internal/transport"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 )
 
+var getConnection = transport.GetConnection
+
 // ChainEngine represents a chain of RBAC Engines, used to make authorization
 // decisions on incoming RPCs.
 type ChainEngine struct {
@@ -206,16 +209,3 @@ type rpcData struct {
 	// handshake.
 	certs []*x509.Certificate
 }
-
-type connectionKey struct{}
-
-func getConnection(ctx context.Context) net.Conn {
-	conn, _ := ctx.Value(connectionKey{}).(net.Conn)
-	return conn
-}
-
-// SetConnection adds the connection to the context to be able to get
-// information about the destination ip and port for an incoming RPC.
-func SetConnection(ctx context.Context, conn net.Conn) context.Context {
-	return context.WithValue(ctx, connectionKey{}, conn)
-}

internal/xds/rbac/rbac_engine_test.go

@@ -440,6 +440,9 @@ func (s) TestNewChainEngine(t *testing.T) {
 // different types of data representing incoming RPC's (piped into a context),
 // and verifies that it works as expected.
 func (s) TestChainEngine(t *testing.T) {
+	defer func(gc func(ctx context.Context) net.Conn) {
+		getConnection = gc
+	}(getConnection)
 	tests := []struct {
 		name        string
 		rbacConfigs []*v3rbacpb.RBAC
@@ -882,7 +885,9 @@ func (s) TestChainEngine(t *testing.T) {
 					}
 					conn := <-connCh
 					defer conn.Close()
-					ctx = SetConnection(ctx, conn)
+					getConnection = func(context.Context) net.Conn {
+						return conn
+					}
 					ctx = peer.NewContext(ctx, data.rpcData.peerInfo)
 					stream := &ServerTransportStreamWithMethod{
 						method: data.rpcData.fullMethod,

test/end2end_test.go

@@ -7766,3 +7766,57 @@ func (cvd *credentialsVerifyDeadline) Clone() credentials.TransportCredentials {
 func (cvd *credentialsVerifyDeadline) OverrideServerName(s string) error {
 	return nil
 }
+
+func unaryInterceptorVerifyConn(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	conn := transport.GetConnection(ctx)
+	if conn == nil {
+		return nil, status.Error(codes.NotFound, "connection was not in context")
+	}
+	return nil, status.Error(codes.OK, "")
+}
+
+// TestUnaryServerInterceptorGetsConnection tests whether the accepted conn on
+// the server gets to any unary interceptors on the server side.
+func (s) TestUnaryServerInterceptorGetsConnection(t *testing.T) {
+	ss := &stubserver.StubServer{}
+	if err := ss.Start([]grpc.ServerOption{grpc.UnaryInterceptor(unaryInterceptorVerifyConn)}); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	if _, err := ss.Client.EmptyCall(ctx, &testpb.Empty{}); status.Code(err) != codes.OK {
+		t.Fatalf("ss.Client.EmptyCall(_, _) = _, %v, want _, error code %s", err, codes.OK)
+	}
+}
+
+func streamingInterceptorVerifyConn(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	conn := transport.GetConnection(ss.Context())
+	if conn == nil {
+		return status.Error(codes.NotFound, "connection was not in context")
+	}
+	return status.Error(codes.OK, "")
+}
+
+// TestStreamingServerInterceptorGetsConnection tests whether the accepted conn on
+// the server gets to any streaming interceptors on the server side.
+func (s) TestStreamingServerInterceptorGetsConnection(t *testing.T) {
+	ss := &stubserver.StubServer{}
+	if err := ss.Start([]grpc.ServerOption{grpc.StreamInterceptor(streamingInterceptorVerifyConn)}); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	s, err := ss.Client.StreamingOutputCall(ctx, &testpb.StreamingOutputCallRequest{})
+	if err != nil {
+		t.Fatalf("ss.Client.StreamingOutputCall(_) = _, %v, want _, <nil>", err)
+	}
+	if _, err := s.Recv(); err != io.EOF {
+		t.Fatalf("ss.Client.StreamingInputCall(_) = _, %v, want _, %v", err, io.EOF)
+	}
+}