wip refactor entraid

Author: ndyakov

internal/auth/conn_reauth_credentials_listener.go

@@ -0,0 +1,77 @@
+package streaming
+
+import (
+	"github.com/redis/go-redis/v9/auth"
+	"github.com/redis/go-redis/v9/internal/pool"
+)
+
+// ConnReAuthCredentialsListener is a struct that implements the CredentialsListener interface.
+// It is used to re-authenticate the credentials when they are updated.
+// It holds reference to the connection to re-authenticate and will pass it to the reAuth and onErr callbacks.
+// It contains:
+// - reAuth: a function that takes the new credentials and returns an error if any.
+// - onErr: a function that takes an error and handles it.
+// - conn: the connection to re-authenticate.
+type ConnReAuthCredentialsListener struct {
+	// reAuth is called when the credentials are updated.
+	reAuth func(conn *pool.Conn, credentials auth.Credentials) error
+	// onErr is called when an error occurs.
+	onErr func(conn *pool.Conn, err error)
+	// conn is the connection to re-authenticate.
+	conn *pool.Conn
+
+	manager *Manager
+}
+
+// OnNext is called when the credentials are updated.
+// It calls the reAuth function with the new credentials.
+// If the reAuth function returns an error, it calls the onErr function with the error.
+func (c *ConnReAuthCredentialsListener) OnNext(credentials auth.Credentials) {
+	if c.conn.IsClosed() {
+		return
+	}
+
+	if c.reAuth == nil {
+		return
+	}
+
+	c.manager.MarkForReAuth(c.conn, func(err error) {
+		if err != nil {
+			c.OnError(err)
+			return
+		}
+		err = c.reAuth(c.conn, credentials)
+		if err != nil {
+			c.OnError(err)
+			return
+		}
+	})
+
+}
+
+// OnError is called when an error occurs.
+// It can be called from both the credentials provider and the reAuth function.
+func (c *ConnReAuthCredentialsListener) OnError(err error) {
+	if c.onErr == nil {
+		return
+	}
+
+	c.onErr(c.conn, err)
+}
+
+// newConnReAuthCredentialsListener creates a new ConnReAuthCredentialsListener.
+// Implements the auth.CredentialsListener interface.
+func newConnReAuthCredentialsListener(
+	conn *pool.Conn,
+	reAuth func(conn *pool.Conn, credentials auth.Credentials) error,
+	onErr func(conn *pool.Conn, err error),
+) *ConnReAuthCredentialsListener {
+	return &ConnReAuthCredentialsListener{
+		conn:   conn,
+		reAuth: reAuth,
+		onErr:  onErr,
+	}
+}
+
+// Ensure ConnReAuthCredentialsListener implements the CredentialsListener interface.
+var _ auth.CredentialsListener = (*ConnReAuthCredentialsListener)(nil)

internal/auth/cred_listeners.go

@@ -0,0 +1,44 @@
+package streaming
+
+import (
+	"sync"
+
+	"github.com/redis/go-redis/v9/auth"
+)
+
+type CredentialsListeners struct {
+	// connid -> listener
+	listeners map[uint64]auth.CredentialsListener
+	lock      sync.RWMutex
+}
+
+func NewCredentialsListeners() *CredentialsListeners {
+	return &CredentialsListeners{
+		listeners: make(map[uint64]auth.CredentialsListener),
+	}
+}
+
+func (c *CredentialsListeners) Add(connID uint64, listener auth.CredentialsListener) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.listeners == nil {
+		c.listeners = make(map[uint64]auth.CredentialsListener)
+	}
+	c.listeners[connID] = listener
+}
+
+func (c *CredentialsListeners) Get(connID uint64) (auth.CredentialsListener, bool) {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	if len(c.listeners) == 0 {
+		return nil, false
+	}
+	listener, ok := c.listeners[connID]
+	return listener, ok
+}
+
+func (c *CredentialsListeners) Remove(connID uint64) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	delete(c.listeners, connID)
+}

internal/auth/streaming/conn_reauth_credentials_listener.go

@@ -0,0 +1,50 @@
+package streaming
+
+import (
+	"time"
+
+	"github.com/redis/go-redis/v9/auth"
+	"github.com/redis/go-redis/v9/internal/pool"
+)
+
+type Manager struct {
+	credentialsListeners *CredentialsListeners
+	pool                 pool.Pooler
+	poolHookRef          *ReAuthPoolHook
+}
+
+func NewManager(pl pool.Pooler, reAuthTimeout time.Duration) *Manager {
+	return &Manager{
+		pool:                 pl,
+		poolHookRef:          NewReAuthPoolHook(pl.Size(), reAuthTimeout),
+		credentialsListeners: NewCredentialsListeners(),
+	}
+}
+
+func (m *Manager) PoolHook() pool.PoolHook {
+	return m.poolHookRef
+}
+
+func (m *Manager) Listener(
+	poolCn *pool.Conn,
+	reAuth func(*pool.Conn, auth.Credentials) error,
+	onErr func(*pool.Conn, error),
+) auth.CredentialsListener {
+	connID := poolCn.GetID()
+	listener, ok := m.credentialsListeners.Get(connID)
+	if !ok {
+		newCredListener := newConnReAuthCredentialsListener(
+			poolCn,
+			reAuth,
+			onErr,
+		)
+		newCredListener.manager = m
+		m.credentialsListeners.Add(connID, newCredListener)
+	}
+	return listener
+}
+
+func (m *Manager) MarkForReAuth(poolCn *pool.Conn, reAuthFn func(error)) {
+	connID := poolCn.GetID()
+	m.poolHookRef.MarkForReAuth(connID, reAuthFn)
+}