All: Implement a token handler

This implements a token handler for OAuth tokens. Clients can use the SetTokenHandler
function in exports to set a token handler.

It needs two arguments, a getter and a setter. The getter is a callback with three arguments:
- The server to get the tokens for, in types.server.current as JSON
- The output buffer
- The output buffer maximum length

The tokens should be written to the output buffer with maximum
length. The type should be types.server.Tokens and be marshalled as
JSON. If no tokens are available, leave the output buffer intact

The token setter is a callback with two arguments:
- The server for which to set the tokens for, in types.server.Current as JSON
- The tokens, defined in types.server.Tokens as JSON

Breaking changes:
- No more tokens as arguments, was already deprecated in previous commits
- Tokens are no longer returned in types.server.Configuration

1 files changed, 101 insertions, 7 deletions

diff --git a/client/client.go b/client/client.go
index 813f6dc..28080b0 100644
--- a/client/client.go
+++ b/client/client.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/eduvpn/eduvpn-common/internal/config"
 	"github.com/eduvpn/eduvpn-common/internal/discovery"
@@ -109,7 +110,59 @@ type Client struct {
 	Debug bool `json:"-"`
 
 	// The Failover monitor for the current VPN connection
-	Failover *failover.DroppedConMon
+	Failover *failover.DroppedConMon `json:"-"`
+
+	// TokenSetter sets the tokens in the client
+	TokenSetter func(srv srvtypes.Current, tok srvtypes.Tokens) `json:"-"`
+
+	// TokenGetter gets the tokens from the client
+	TokenGetter func(srv srvtypes.Current) *srvtypes.Tokens `json:"-"`
+}
+
+func (c *Client) updateTokens(srv server.Server) error {
+	if c.TokenGetter == nil {
+		return errors.New("no tokken getter defined")
+	}
+	pSrv, err := c.pubCurrentServer(srv)
+	if err != nil {
+		return err
+	}
+	// shouldn't happen
+	if pSrv == nil {
+		return errors.New("public server is nil when getting tokens")
+	}
+	tokens := c.TokenGetter(*pSrv)
+	if tokens == nil {
+		return errors.New("client returned nil for tokens")
+	}
+
+	server.UpdateTokens(srv, oauth.Token{
+		Access: tokens.Access,
+		Refresh: tokens.Refresh,
+		ExpiredTimestamp: time.Unix(tokens.Expires, 0),
+	})
+
+	return nil
+}
+
+func (c *Client) forwardTokens(srv server.Server) error {
+	if c.TokenSetter == nil {
+		return errors.New("no token setter defined")
+	}
+	pSrv, err := c.pubCurrentServer(srv)
+	if err != nil {
+		return err
+	}
+	if pSrv == nil {
+		return errors.New("public server is nil when updating tokens")
+	}
+	o := srv.OAuth()
+	if o == nil {
+		return errors.New("oauth was nil when forwarding tokens")
+	}
+	t := o.Token()
+	c.TokenSetter(*pSrv, t.Public())
+	return nil
 }
 
 // New creates a new client with the following parameters:
@@ -329,7 +382,18 @@ func (c *Client) callbacks(ck *cookie.Cookie, srv server.Server, forceauth bool)
 	// oauth
 	// TODO: This should be ck.Context()
 	// But needsrelogin needs a rewrite to support this properly
+
+	// first make sure we get the most up to date tokens from the client
+	err := c.updateTokens(srv)
+	if err != nil {
+		log.Logger.Debugf("failed to get tokens from client: %v", err)
+	}
 	if server.NeedsRelogin(context.Background(), srv) || forceauth {
+		// mark organizations as expired if the server is a secure internet server
+		b, berr := srv.Base()
+		if berr == nil && b.Type == srvtypes.TypeSecureInternet {
+			c.Discovery.MarkOrganizationsExpired()
+		}
 		err := c.loginCallback(ck, srv)
 		if err != nil {
 			return err
@@ -421,6 +485,10 @@ func (c *Client) AddServer(ck *cookie.Cookie, identifier string, _type srvtypes.
 	case srvtypes.TypeSecureInternet:
 		dOrg, dSrv, err := c.Discovery.SecureHomeArgs(identifier)
 		if err != nil {
+			// We mark the organizations as expired because we got an error
+			// Note that in the docs it states that it only should happen when the Org ID doesn't exist
+			// However, this is nice as well because it also catches the error where the SecureInternetHome server is not found
+			c.Discovery.MarkOrganizationsExpired()
 			return err
 		}
 		srv, err = c.Servers.AddSecureInternet(ck.Context(), dOrg, dSrv)
@@ -442,7 +510,15 @@ func (c *Client) AddServer(ck *cookie.Cookie, identifier string, _type srvtypes.
 	}
 
 	// callbacks
-	return c.callbacks(ck, srv, false)
+	err = c.callbacks(ck, srv, false)
+	if err != nil {
+		return err
+	}
+	terr := c.forwardTokens(srv)
+	if terr != nil {
+		log.Logger.Debugf("failed to forward tokens after adding: %v", terr)
+	}
+	return nil
 }
 
 func (c *Client) config(ck *cookie.Cookie, srv server.Server, pTCP bool, forceAuth bool) (cfg *srvtypes.Configuration, err error) {
@@ -529,6 +605,14 @@ func (c *Client) GetConfig(ck *cookie.Cookie, identifier string, _type srvtypes.
 		cfg, err = c.config(ck, srv, pTCP, true)
 	}
 
+	// tokens might be updated, forward them
+	defer func() {
+		terr := c.forwardTokens(srv)
+		if terr != nil {
+			log.Logger.Debugf("failed to forward tokens after get config: %v", terr)
+		}
+	}()
+
 	// still an error, return nil with the error
 	if err != nil {
 		return nil, err
@@ -700,15 +784,18 @@ func (c *Client) Cleanup(ck *cookie.Cookie) (err error) {
 	if err != nil {
 		return err
 	}
-	// TODO: Support cookie context here
-	// if server.NeedsRelogin(context.Background(), srv) {
-	// 	// TODO: ask client for tokens
-	// }
+	err = c.updateTokens(srv)
+	if err != nil {
+		log.Logger.Debugf("failed to update tokens for disconnect: %v", err)
+	}
 	err = server.Disconnect(ck.Context(), srv)
 	if err != nil {
 		return err
 	}
-	// TODO: Set tokens with callback
+	err = c.forwardTokens(srv)
+	if err != nil {
+		log.Logger.Debugf("failed to forward tokens after disconnect: %v", err)
+	}
 	return nil
 }
 
@@ -740,6 +827,13 @@ func (c *Client) RenewSession(ck *cookie.Cookie) (err error) {
 		c.FSM.GoTransition(StateLoadingServer)
 		c.FSM.GoTransition(StateChosenServer)
 	}
+	// update tokens in the end
+	defer func() {
+		terr := c.forwardTokens(srv)
+		if terr != nil {
+			log.Logger.Debugf("failed to forward tokens after renew: %v", terr)
+		}
+	}()
 	// TODO: Maybe this can be deleted because we force auth now
 	server.MarkTokensForRenew(srv)
 	// run the callbacks by forcing auth