Client + Server: Implement a token updater callback

2 files changed, 48 insertions, 1 deletions

diff --git a/client/client.go b/client/client.go
index 2bb0cc2..89db8f5 100644
--- a/client/client.go
+++ b/client/client.go
@@ -11,6 +11,7 @@ import (
 	"github.com/eduvpn/eduvpn-common/internal/fsm"
 	"github.com/eduvpn/eduvpn-common/internal/http"
 	"github.com/eduvpn/eduvpn-common/internal/log"
+	"github.com/eduvpn/eduvpn-common/internal/oauth"
 	"github.com/eduvpn/eduvpn-common/internal/server"
 	"github.com/eduvpn/eduvpn-common/internal/util"
 	"github.com/eduvpn/eduvpn-common/types"
@@ -116,6 +117,30 @@ type Client struct {
 
 	// The Failover monitor for the current VPN connection
 	Failover *failover.DroppedConMon `json:"-"`
+
+	// tokenCB is the token callback
+	tokenCB func(srv server.Server, tok oauth.Token)
+}
+
+func (c *Client) ForwardTokenUpdate(srv server.Server) {
+	if c.tokenCB == nil {
+		log.Logger.Debugf("No token update callback available")
+		return
+	}
+	t := oauth.Token{}
+	o := srv.OAuth()
+	if o != nil {
+		t = o.Token()
+	} else {
+		log.Logger.Debugf("OAuth was nil when forwarding token callback")
+	}
+	log.Logger.Debugf("Running token callback")
+	c.tokenCB(srv, t)
+}
+
+// SetTokenUpdater sets the token updater callback
+func (c *Client) SetTokenUpdater(updater func(srv server.Server, tok oauth.Token)) {
+	c.tokenCB = updater
 }
 
 // Register initializes the clientwith the following parameters:
diff --git a/client/server.go b/client/server.go
index 55a3b9e..5578f04 100644
--- a/client/server.go
+++ b/client/server.go
@@ -108,6 +108,8 @@ func (c *Client) Cleanup(ct oauth.Token) error {
 	if server.NeedsRelogin(srv) {
 		server.UpdateTokens(srv, ct)
 	}
+	// update tokens to client
+	defer c.ForwardTokenUpdate(srv)
 	// Do the /disconnect API call
 	err = server.Disconnect(srv)
 	if err != nil {
@@ -260,6 +262,9 @@ func (c *Client) AddInstituteServer(url string) (srv server.Server, err error) {
 	}
 
 	c.FSM.GoTransitionWithData(StateNoServer, c.Servers)
+
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
 	return srv, nil
 }
 
@@ -322,6 +327,9 @@ func (c *Client) AddSecureInternetHomeServer(orgID string) (srv server.Server, e
 		return nil, err
 	}
 	c.FSM.GoTransitionWithData(StateNoServer, c.Servers)
+
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
 	return srv, nil
 }
 
@@ -369,6 +377,9 @@ func (c *Client) AddCustomServer(url string) (srv server.Server, err error) {
 	}
 
 	c.FSM.GoTransitionWithData(StateNoServer, c.Servers)
+
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
 	return srv, nil
 }
 
@@ -408,6 +419,9 @@ func (c *Client) GetConfigInstituteAccess(url string, preferTCP bool, t oauth.To
 		c.goBackInternal()
 	}
 
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
+
 	return cfg, err
 }
 
@@ -446,6 +460,9 @@ func (c *Client) GetConfigSecureInternet(orgID string, preferTCP bool, t oauth.T
 		c.goBackInternal()
 	}
 
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
+
 	return cfg, err
 }
 
@@ -483,6 +500,9 @@ func (c *Client) GetConfigCustomServer(url string, preferTCP bool, t oauth.Token
 		c.goBackInternal()
 	}
 
+	// Also forward tokens using the callback
+	c.ForwardTokenUpdate(srv)
+
 	return cfg, err
 }
 
@@ -543,7 +563,9 @@ func (c *Client) RenewSession() (err error) {
 	}
 
 	server.MarkTokensForRenew(srv)
-	return c.ensureLogin(srv, oauth.Token{})
+	c.ForwardTokenUpdate(srv)
+	err = c.ensureLogin(srv, oauth.Token{})
+	return err
 }
 
 // ShouldRenewButton returns true if the renew button should be shown