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 <[email protected]>

Author: mrolla

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
 		}

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
+}

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
 		}
 	}