Skip to content

Commit

Permalink
[ADDED] UserInfoHandler for dynamically setting user/password
Browse files Browse the repository at this point in the history
Signed-off-by: Piotr Piotrowski <piotr@synadia.com>
  • Loading branch information
piotrpio committed Sep 10, 2024
1 parent a06b6a9 commit 82acab9
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 0 deletions.
21 changes: 21 additions & 0 deletions nats.go
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ var (
ErrNkeysNotSupported = errors.New("nats: nkeys not supported by the server")
ErrStaleConnection = errors.New("nats: " + STALE_CONNECTION)
ErrTokenAlreadySet = errors.New("nats: token and token handler both set")
ErrUserInfoAlreadySet = errors.New("nats: cannot set user info callback and user/pass")
ErrMsgNotBound = errors.New("nats: message is not bound to subscription/connection")
ErrMsgNoReply = errors.New("nats: message does not have a reply")
ErrClientIPNotSupported = errors.New("nats: client IP not supported by this server")
Expand Down Expand Up @@ -230,6 +231,9 @@ type SignatureHandler func([]byte) ([]byte, error)
// AuthTokenHandler is used to generate a new token.
type AuthTokenHandler func() string

// UserInfoCB is used to pass the username and password when establishing connection.
type UserInfoCB func() (string, string)

// ReconnectDelayHandler is used to get from the user the desired
// delay the library should pause before attempting to reconnect
// again. Note that this is invoked after the library tried the
Expand Down Expand Up @@ -443,6 +447,9 @@ type Options struct {
// Password sets the password to be used when connecting to a server.
Password string

// UserInfo sets the callback handler that will fetch the username and password.
UserInfo UserInfoCB

// Token sets the token to be used when connecting to a server.
Token string

Expand Down Expand Up @@ -1166,6 +1173,13 @@ func UserInfo(user, password string) Option {
}
}

func UserInfoHandler(cb UserInfoCB) Option {
return func(o *Options) error {
o.UserInfo = cb
return nil
}
}

// Token is an Option to set the token to use
// when a token is not included directly in the URLs
// and when a token handler is not provided.
Expand Down Expand Up @@ -2563,6 +2577,13 @@ func (nc *Conn) connectProto() (string, error) {
pass = o.Password
token = o.Token
nkey = o.Nkey

if nc.Opts.UserInfo != nil {
if user != _EMPTY_ || pass != _EMPTY_ {
return _EMPTY_, ErrUserInfoAlreadySet
}
user, pass = nc.Opts.UserInfo()
}
}

// Look for user jwt.
Expand Down
74 changes: 74 additions & 0 deletions test/auth_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ import (
"errors"
"fmt"
"io/fs"
"net"
"os"
"strings"
"sync/atomic"
"testing"
Expand Down Expand Up @@ -377,3 +379,75 @@ func TestConnectMissingCreds(t *testing.T) {
t.Fatalf("Expected not exists error, got: %v", err)
}
}

func TestUserInfoHandler(t *testing.T) {
conf := createConfFile(t, []byte(`
listen: 127.0.0.1:-1
accounts: {
A {
users: [{ user: "pp", password: "foo" }]
}
}
`))
defer os.Remove(conf)

s, _ := RunServerWithConfig(conf)
defer s.Shutdown()

user, pass := "pp", "foo"
userInfoCB := func() (string, string) {
return user, pass
}

// check that we cannot set the user info twice
_, err := nats.Connect(s.ClientURL(), nats.UserInfo("pp", "foo"), nats.UserInfoHandler(userInfoCB))
if !errors.Is(err, nats.ErrUserInfoAlreadySet) {
t.Fatalf("Expected ErrUserInfoAlreadySet, got: %v", err)
}

addr, ok := s.Addr().(*net.TCPAddr)
if !ok {
t.Fatalf("Expected a TCP address, got %T", addr)
}

// check that user/pass from url takes precedence
_, err = nats.Connect(fmt.Sprintf("nats://bad:bad@localhost:%d", addr.Port),
nats.UserInfoHandler(userInfoCB))
if !errors.Is(err, nats.ErrAuthorization) {
t.Fatalf("Expected ErrAuthorization, got: %v", err)
}

// connect using the handler
nc, err := nats.Connect(s.ClientURL(),
nats.ReconnectWait(100*time.Millisecond),
nats.UserInfoHandler(userInfoCB))
if err != nil {
t.Fatalf("Error on connect: %v", err)
}
defer nc.Close()

// now change the password and reload the server
newConfig := []byte(`
listen: 127.0.0.1:-1
accounts: {
A {
users: [{ user: "dd", password: "bar" }]
}
}
`)
if err := os.WriteFile(conf, newConfig, 0666); err != nil {
t.Fatalf("Error writing conf file: %v", err)
}

// update the user info used by the callback
user, pass = "dd", "bar"

status := nc.StatusChanged(nats.CONNECTED)

if err := s.Reload(); err != nil {
t.Fatalf("Error on reload: %v", err)
}

// we should get a reconnected event meaning the new credentials were used
WaitOnChannel(t, status, nats.CONNECTED)
}

0 comments on commit 82acab9

Please sign in to comment.