feat(sasl): enrich authentication context with metadata (#725)

* feat: enrich sasl context with metadata * feedback from first review * second round of feedback * remove unnecessary blank line * it has been a long day * Update dialer.go * Update sasl/sasl.go * Update sasl/sasl.go * Update transport.go Co-authored-by: Achille <achille.roussel@gmail.com>
segmentio · Nov 5, 2021 · f7dd036 · f7dd036
1 parent ea83b29
commit f7dd036
Showing 3 changed files with 28 additions and 2 deletions.
diff --git a/dialer.go b/dialer.go
@@ -281,7 +281,10 @@ func (d *Dialer) connect(ctx context.Context, network, address string, connCfg C
 	conn := NewConnWith(c, connCfg)
 
 	if d.SASLMechanism != nil {
-		if err := d.authenticateSASL(ctx, conn); err != nil {
+		metadata := &sasl.Metadata{
+			Host: address,
+		}
+		if err := d.authenticateSASL(sasl.WithMetadata(ctx, metadata), conn); err != nil {
 			_ = conn.Close()
 			return nil, err
 		}

diff --git a/sasl/sasl.go b/sasl/sasl.go
@@ -2,6 +2,8 @@ package sasl
 
 import "context"
 
+type ctxKey struct{}
+
 // Mechanism implements the SASL state machine for a particular mode of
 // authentication.  It is used by the kafka.Dialer to perform the SASL
 // handshake.
@@ -42,3 +44,21 @@ type StateMachine interface {
 	// value will be true.
 	Next(ctx context.Context, challenge []byte) (done bool, response []byte, err error)
 }
+
+// Metadata contains additional data for performing SASL authentication.
+type Metadata struct {
+	// Host is the address of the broker the authentication will be
+	// performed on.
+	Host string
+}
+
+// WithMetadata returns a copy of the context with associated Metadata.
+func WithMetadata(ctx context.Context, m *Metadata) context.Context {
+	return context.WithValue(ctx, ctxKey{}, m)
+}
+
+// MetadataFromContext retrieves the Metadata from the context.
+func MetadataFromContext(ctx context.Context) *Metadata {
+	m, _ := ctx.Value(ctxKey{}).(*Metadata)
+	return m
+}
diff --git a/transport.go b/transport.go
@@ -1197,7 +1197,10 @@ func (g *connGroup) connect(ctx context.Context, addr net.Addr) (*conn, error) {
 	pc.SetDeadline(time.Time{})
 
 	if g.pool.sasl != nil {
-		if err := authenticateSASL(ctx, pc, g.pool.sasl); err != nil {
+		metadata := &sasl.Metadata{
+			Host: netAddr.String(),
+		}
+		if err := authenticateSASL(sasl.WithMetadata(ctx, metadata), pc, g.pool.sasl); err != nil {
 			return nil, err
 		}
 	}