drpc: add TLS certificate handling and metadata infra for auth interceptors

This commit adds infrastructure needed for authentication interceptors:

1. New drpcctx/tlscert.go: Functions to store/retrieve TLS peer certificates in context
2. Server-side TLS certificate extraction in drpcserver
3. Improved metadata API with ClearContext, ClearContextExcept, and GetValue functions
4. Client-side per-RPC metadata support via WithPerRPCMetadata option

Author: shubhamdhama

drpcclient/clientconn.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"storj.io/drpc"
+	"storj.io/drpc/drpcmetadata"
 )
 
 // ClientConn represents a DRPC client connection, with support for configuring the
@@ -33,6 +34,9 @@ func finalInvoker(ctx context.Context, rpc string, enc drpc.Encoding, in, out dr
 }
 
 func (c *ClientConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) error {
+	if c.dopts.perRPCMetadata != nil {
+		ctx = drpcmetadata.AddPairs(ctx, c.dopts.perRPCMetadata)
+	}
 	if c.dopts.unaryInt != nil {
 		return c.dopts.unaryInt(ctx, rpc, enc, in, out, c, finalInvoker)
 	}
@@ -45,6 +49,9 @@ func finalStreamer(ctx context.Context, rpc string, enc drpc.Encoding, cc *Clien
 }
 
 func (c *ClientConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) {
+	if c.dopts.perRPCMetadata != nil {
+		ctx = drpcmetadata.AddPairs(ctx, c.dopts.perRPCMetadata)
+	}
 	if c.dopts.streamInt != nil {
 		return c.dopts.streamInt(ctx, rpc, enc, c, finalStreamer)
 	}

drpcclient/dialoptions.go

@@ -8,6 +8,8 @@ type dialOptions struct {
 
 	unaryInts  []UnaryClientInterceptor
 	streamInts []StreamClientInterceptor
+
+	perRPCMetadata map[string]string
 }
 
 // DialOption configures how we set up the client connection.
@@ -32,3 +34,9 @@ func WithChainStreamInterceptor(ints ...StreamClientInterceptor) DialOption {
 		opt.streamInts = append(opt.streamInts, ints...)
 	}
 }
+
+func WithPerRPCMetadata(metadata map[string]string) DialOption {
+	return func(opt *dialOptions) {
+		opt.perRPCMetadata = metadata
+	}
+}

drpcctx/tlscert.go

@@ -0,0 +1,27 @@
+package drpcctx
+
+import (
+	"context"
+	"crypto/x509"
+)
+
+// PeerConnectionInfo contains TLS peer connection information.
+type PeerConnectionInfo struct {
+	Certificates []*x509.Certificate
+}
+
+// TLSPeerCertKey is used to store TLS peer connection info in the context.
+type TLSPeerCertKey struct{}
+
+// WithPeerConnectionInfo associates the peer connection information of the TLS connection
+// with the context.
+func WithPeerConnectionInfo(ctx context.Context, info PeerConnectionInfo) context.Context {
+	return context.WithValue(ctx, TLSPeerCertKey{}, info)
+}
+
+// GetPeerConnectionInfo returns the TLS peer connection information associated with the
+// context and a bool indicating if it existed.
+func GetPeerConnectionInfo(ctx context.Context) (PeerConnectionInfo, bool) {
+	peerConnectionInfo, ok := ctx.Value(TLSPeerCertKey{}).(PeerConnectionInfo)
+	return peerConnectionInfo, ok
+}

drpcmetadata/metadata.go

@@ -50,6 +50,28 @@ func Decode(buf []byte) (map[string]string, error) {
 
 type metadataKey struct{}
 
+// ClearContext removes all metadata from the context and returns a new context
+// with no metadata attached.
+func ClearContext(ctx context.Context) context.Context {
+	return context.WithValue(ctx, metadataKey{}, nil)
+}
+
+// ClearContextExcept removes all metadata from the context except for the
+// specified key. If the specified key doesn't exist in the metadata, it clears
+// all metadata. Returns a new context with only the specified key-value pair
+// preserved.
+func ClearContextExcept(ctx context.Context, key string) context.Context {
+	md, ok := Get(ctx)
+	if !ok {
+		return ClearContext(ctx)
+	}
+	value, ok := md[key]
+	if !ok {
+		return ClearContext(ctx)
+	}
+	return context.WithValue(ctx, metadataKey{}, map[string]string{key: value})
+}
+
 // Add associates a key/value pair on the context.
 func Add(ctx context.Context, key, value string) context.Context {
 	metadata, ok := Get(ctx)
@@ -66,3 +88,13 @@ func Get(ctx context.Context) (map[string]string, bool) {
 	metadata, ok := ctx.Value(metadataKey{}).(map[string]string)
 	return metadata, ok
 }
+
+// GetValue retrieves a specific value by key from the context's metadata.
+func GetValue(ctx context.Context, key string) (string, bool) {
+	metadata, ok := Get(ctx)
+	if !ok {
+		return "", false
+	}
+	val, ok := metadata[key]
+	return val, ok
+}

drpcserver/server.go

@@ -5,12 +5,12 @@ package drpcserver
 
 import (
 	"context"
+	"crypto/tls"
 	"net"
 	"sync"
 	"time"
 
 	"github.com/zeebo/errs"
-
 	"storj.io/drpc"
 	"storj.io/drpc/drpccache"
 	"storj.io/drpc/drpcctx"
@@ -92,6 +92,32 @@ func (s *Server) getStats(rpc string) *drpcstats.Stats {
 
 // ServeOne serves a single set of rpcs on the provided transport.
 func (s *Server) ServeOne(ctx context.Context, tr drpc.Transport) (err error) {
+	// Check if the transport is a TLS connection
+	if tlsConn, ok := tr.(*tls.Conn); ok {
+		// Manually perform the TLS handshake to access peer certificate
+		// information. In Go's TLS implementation, the handshake is normally
+		// performed lazily on the first read/write operation. However, the
+		// transport received by ServeOne hasn't performed any I/O yet, so
+		// ConnectionState() would be empty. Only after the handshake completes
+		// is ConnectionState populated with peer certificates and other
+		// connection details that we need for authentication context.
+		//
+		// This explicit Handshake() call is safe and appropriate here. The
+		// connection hasn't started processing requests yet, so we're not
+		// interrupting any ongoing communication. Even if we didn't call it
+		// explicitly, the first read/write operation would call it internally
+		// anyway.
+		err := tlsConn.Handshake()
+		if err != nil {
+			return err
+		}
+		state := tlsConn.ConnectionState()
+		if len(state.PeerCertificates) > 0 {
+			ctx = drpcctx.WithPeerConnectionInfo(
+				ctx, drpcctx.PeerConnectionInfo{Certificates: state.PeerCertificates})
+		}
+	}
+
 	man := drpcmanager.NewWithOptions(tr, s.opts.Manager)
 	defer func() { err = errs.Combine(err, man.Close()) }()