Client: Refresh server endpoints on renew and cleanup

1 files changed, 25 insertions, 0 deletions

diff --git a/client/client.go b/client/client.go
index eab704a..c54dd0e 100644
--- a/client/client.go
+++ b/client/client.go
@@ -835,6 +835,19 @@ func (c *Client) Cleanup(ck *cookie.Cookie) (err error) {
 	if err != nil {
 		return i18nerr.Wrap(err, "Failed to get the current server to cleanup the connection")
 	}
+
+	err = srv.RefreshEndpoints(ck.Context(), &c.Discovery)
+
+	// If we get a canceled error, return that, otherwise just log the error
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return i18nerr.Wrap(err, "Cleanup was canceled")
+		}
+
+		log.Logger.Warningf("failed to refresh server endpoints: %v", err)
+	}
+
+
 	defer c.SaveState()
 	err = c.updateTokens(srv)
 	if err != nil {
@@ -884,6 +897,18 @@ func (c *Client) RenewSession(ck *cookie.Cookie) (err error) {
 	if !c.FSM.InState(StateLoadingServer) {
 		c.FSM.GoTransition(StateLoadingServer) //nolint:errcheck
 	}
+	err = srv.RefreshEndpoints(ck.Context(), &c.Discovery)
+
+	// If we get a canceled error, return that, otherwise just log the error
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			return i18nerr.Wrap(err, "Renewing was canceled")
+		}
+
+		log.Logger.Warningf("failed to refresh server endpoints: %v", err)
+	}
+
+
 	// update tokens in the end
 	defer func() {
 		terr := c.forwardTokens(srv)