session: add session ID to grpc metadata via context

Add grpc interceptors that inject an LNC session's ID into the context
as gRPC metadata. By injecting it as such, it will be transported over
the wire in any outgoing gRPC calls.

This lets us be sure that any session call sent to the RPCMiddleware
interceptor in LND will continue to be grouped along with the
appropriate session ID. This gives LND a way to send the metadata we
include back to LiT meaning that we will later on be able to extract the
session ID again.

Author: ellemouton

session/context.go

@@ -0,0 +1,57 @@
+package session
+
+import (
+	"encoding/hex"
+	"fmt"
+
+	"github.com/lightningnetwork/lnd/fn"
+	"google.golang.org/grpc/metadata"
+)
+
+// contextKey is a struct that is used as a key for storing session IDs
+// in a context. Using this unexported type prevents collisions with other
+// context keys that may be used in the same context. However, this only
+// applies if the context is passed around in the same binary and not if the
+// value is converted to grpc metadata and sent over the wire. In that case,
+// we need to use a string key to avoid collisions with other metadata keys.
+type contextKey struct {
+	name string
+}
+
+// sessionIDCtxKey is the context key used to store the session ID in
+// a context. The key is a string to avoid collisions with other context values
+// that may also be included in grpc metadata which is why we add the 'lit'
+// prefix.
+var sessionIDCtxKey = contextKey{"lit_session_id"}
+
+// FromGRPCMetadata extracts the session ID from the given gRPC metadata kv
+// pairs if one is found.
+func FromGRPCMetadata(md metadata.MD) (fn.Option[ID], error) {
+	val := md.Get(sessionIDCtxKey.name)
+	if len(val) == 0 {
+		return fn.None[ID](), nil
+	}
+
+	if len(val) != 1 {
+		return fn.None[ID](), fmt.Errorf("more than one session ID "+
+			"found in gRPC metadata: %v", val)
+	}
+
+	b, err := hex.DecodeString(val[0])
+	if err != nil {
+		return fn.None[ID](), err
+	}
+
+	sessID, err := IDFromBytes(b)
+	if err != nil {
+		return fn.None[ID](), err
+	}
+
+	return fn.Some(sessID), nil
+}
+
+// AddToGRPCMetadata adds the session ID to the given gRPC metadata kv pairs.
+// The session ID is encoded as a hex string.
+func AddToGRPCMetadata(md metadata.MD, id ID) {
+	md.Set(sessionIDCtxKey.name, hex.EncodeToString(id[:]))
+}

session/server.go

@@ -18,7 +18,8 @@ import (
 
 type sessionID [33]byte
 
-type GRPCServerCreator func(opts ...grpc.ServerOption) *grpc.Server
+type GRPCServerCreator func(sessionID ID,
+	opts ...grpc.ServerOption) *grpc.Server
 
 type mailboxSession struct {
 	server *grpc.Server
@@ -70,7 +71,7 @@ func (m *mailboxSession) start(session *Session,
 	}
 
 	noiseConn := mailbox.NewNoiseGrpcConn(keys)
-	m.server = serverCreator(grpc.Creds(noiseConn))
+	m.server = serverCreator(session.ID, grpc.Creds(noiseConn))
 
 	m.wg.Add(1)
 	go m.run(mailboxServer)

session_rpcserver.go

@@ -26,6 +26,7 @@ import (
 	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/macaroons"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 	"gopkg.in/macaroon-bakery.v2/bakery"
 	"gopkg.in/macaroon-bakery.v2/bakery/checkers"
 	"gopkg.in/macaroon.v2"
@@ -77,10 +78,23 @@ func newSessionRPCServer(cfg *sessionRpcServerConfig) (*sessionRpcServer,
 	// actual mailbox server that spins up the Terminal Connect server
 	// interface.
 	server := session.NewServer(
-		func(opts ...grpc.ServerOption) *grpc.Server {
-			allOpts := append(cfg.grpcOptions, opts...)
+		func(id session.ID, opts ...grpc.ServerOption) *grpc.Server {
+			// Add the session ID injector interceptors first so
+			// that the session ID is available in the context of
+			// all interceptors that come after.
+			allOpts := []grpc.ServerOption{
+				addSessionIDToStreamCtx(id),
+				addSessionIDToUnaryCtx(id),
+			}
+
+			allOpts = append(allOpts, cfg.grpcOptions...)
+			allOpts = append(allOpts, opts...)
+
+			// Construct the gRPC server with the options.
 			grpcServer := grpc.NewServer(allOpts...)
 
+			// Register various grpc servers with the LNC session
+			// server.
 			cfg.registerGrpcServers(grpcServer)
 
 			return grpcServer
@@ -94,6 +108,62 @@ func newSessionRPCServer(cfg *sessionRpcServerConfig) (*sessionRpcServer,
 	}, nil
 }
 
+// wrappedServerStream is a wrapper around the grpc.ServerStream that allows us
+// to set a custom context. This is needed since the stream handler function
+// doesn't take a context as an argument, but rather has a Context method on the
+// handler itself. So we use this custom wrapper to override this method.
+type wrappedServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+// Context returns the context of the stream.
+//
+// NOTE: This implements the grpc.ServerStream Context method.
+func (w *wrappedServerStream) Context() context.Context {
+	return w.ctx
+}
+
+// addSessionIDToStreamCtx is a gRPC stream interceptor that adds the given
+// session ID to the context of the stream. This allows us to access the
+// session ID later on for any gRPC calls made through this stream.
+func addSessionIDToStreamCtx(id session.ID) grpc.ServerOption {
+	return grpc.StreamInterceptor(func(srv any, ss grpc.ServerStream,
+		info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler) error {
+
+		md, _ := metadata.FromIncomingContext(ss.Context())
+		mdCopy := md.Copy()
+		session.AddToGRPCMetadata(mdCopy, id)
+
+		// Wrap the original stream with our custom context.
+		wrapped := &wrappedServerStream{
+			ServerStream: ss,
+			ctx: metadata.NewIncomingContext(
+				ss.Context(), mdCopy,
+			),
+		}
+
+		return handler(srv, wrapped)
+	})
+}
+
+// addSessionIDToUnaryCtx is a gRPC unary interceptor that adds the given
+// session ID to the context of the unary call. This allows us to access the
+// session ID later on for any gRPC calls made through this context.
+func addSessionIDToUnaryCtx(id session.ID) grpc.ServerOption {
+	return grpc.UnaryInterceptor(func(ctx context.Context, req any,
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler) (resp any, err error) {
+
+		md, _ := metadata.FromIncomingContext(ctx)
+		mdCopy := md.Copy()
+		session.AddToGRPCMetadata(mdCopy, id)
+
+		return handler(metadata.NewIncomingContext(ctx, mdCopy), req)
+	})
+}
+
 // start all the components necessary for the sessionRpcServer to start serving
 // requests. This includes resuming all non-revoked sessions.
 func (s *sessionRpcServer) start(ctx context.Context) error {