drpc: add changes for auth interceptor changes

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,24 @@
+// Copyright (C) 2025 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package drpcctx
+
+import (
+	"context"
+	"crypto/x509"
+)
+
+// TLSPeerCertKey is used to store TLS info in the context.
+type TLSPeerCertKey struct{}
+
+// WithPeerCertificate associates the TLS info with the context.
+func WithPeerCertificate(ctx context.Context, certificate *x509.Certificate) context.Context {
+	return context.WithValue(ctx, TLSPeerCertKey{}, certificate)
+}
+
+// GetPeerCertificate returns the TLS info associated with the context and a bool
+// indicating if they existed.
+func GetPeerCertificate(ctx context.Context) (*x509.Certificate, bool) {
+	tlsInfo, ok := ctx.Value(TLSPeerCertKey{}).(*x509.Certificate)
+	return tlsInfo, ok
+}

drpcmetadata/metadata.go

@@ -88,3 +88,12 @@ func Get(ctx context.Context) (map[string]string, bool) {
 	metadata, ok := ctx.Value(metadataKey{}).(map[string]string)
 	return metadata, ok
 }
+
+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,6 +5,7 @@ package drpcserver
 
 import (
 	"context"
+	"crypto/tls"
 	"net"
 	"sync"
 	"time"
@@ -92,6 +93,18 @@ 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 {
+		err := tlsConn.Handshake()
+		if err != nil {
+			return err
+		}
+		state := tlsConn.ConnectionState()
+		if len(state.PeerCertificates) > 0 {
+			ctx = drpcctx.WithPeerCertificate(ctx, state.PeerCertificates[0])
+		}
+	}
+
 	man := drpcmanager.NewWithOptions(tr, s.opts.Manager)
 	defer func() { err = errs.Combine(err, man.Close()) }()