FSM: Check unhandled transitions

4 files changed, 30 insertions, 13 deletions

diff --git a/client/client.go b/client/client.go
index 103136b..63e4f91 100644
--- a/client/client.go
+++ b/client/client.go
@@ -73,7 +73,7 @@ func (client *Client) Register(
 	name string,
 	directory string,
 	language string,
-	stateCallback func(FSMStateID, FSMStateID, interface{}),
+	stateCallback func(FSMStateID, FSMStateID, interface{}) bool,
 	debug bool,
 ) error {
 	errorMessage := "failed to register with the GO library"
@@ -155,11 +155,15 @@ func (client *Client) Deregister() {
 
 // askProfile asks the user for a profile by moving the FSM to the ASK_PROFILE state.
 func (client *Client) askProfile(chosenServer server.Server) error {
+	errorMessage := "failed asking for profiles"
 	profiles, profilesErr := server.GetValidProfiles(chosenServer, client.SupportsWireguard)
 	if profilesErr != nil {
-		return types.NewWrappedError("failed asking for profiles", profilesErr)
+		return types.NewWrappedError(errorMessage, profilesErr)
+	}
+	goTransitionErr := client.FSM.GoTransitionRequired(STATE_ASK_PROFILE, profiles)
+	if goTransitionErr != nil {
+		return types.NewWrappedError(errorMessage, goTransitionErr)
 	}
-	client.FSM.GoTransitionWithData(STATE_ASK_PROFILE, profiles)
 	return nil
 }
 
diff --git a/client/client_test.go b/client/client_test.go
index f386c3c..adf5c75 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -81,8 +81,9 @@ func Test_server(t *testing.T) {
 		"org.letsconnect-vpn.app.linux",
 		"configstest",
 		"en",
-		func(old FSMStateID, new FSMStateID, data interface{}) {
+		func(old FSMStateID, new FSMStateID, data interface{}) bool {
 			stateCallback(t, old, new, data, state)
+			return true
 		},
 		false,
 	)
@@ -113,7 +114,7 @@ func test_connect_oauth_parameter(
 		"org.letsconnect-vpn.app.linux",
 		configDirectory,
 		"en",
-		func(oldState FSMStateID, newState FSMStateID, data interface{}) {
+		func(oldState FSMStateID, newState FSMStateID, data interface{}) bool {
 			if newState == STATE_OAUTH_STARTED {
 				server, serverErr := state.Servers.GetCustomServer(serverURI)
 				if serverErr != nil {
@@ -142,6 +143,7 @@ func test_connect_oauth_parameter(
 					}
 				}()
 			}
+			return true
 		},
 		false,
 	)
@@ -219,8 +221,9 @@ func Test_token_expired(t *testing.T) {
 		"org.letsconnect-vpn.app.linux",
 		"configsexpired",
 		"en",
-		func(old FSMStateID, new FSMStateID, data interface{}) {
+		func(old FSMStateID, new FSMStateID, data interface{}) bool {
 			stateCallback(t, old, new, data, state)
+			return true
 		},
 		false,
 	)
@@ -279,8 +282,9 @@ func Test_token_invalid(t *testing.T) {
 		"org.letsconnect-vpn.app.linux",
 		"configsinvalid",
 		"en",
-		func(old FSMStateID, new FSMStateID, data interface{}) {
+		func(old FSMStateID, new FSMStateID, data interface{}) bool {
 			stateCallback(t, old, new, data, state)
+			return true
 		},
 		false,
 	)
@@ -336,8 +340,9 @@ func Test_invalid_profile_corrected(t *testing.T) {
 		"org.letsconnect-vpn.app.linux",
 		"configscancelprofile",
 		"en",
-		func(old FSMStateID, new FSMStateID, data interface{}) {
+		func(old FSMStateID, new FSMStateID, data interface{}) bool {
 			stateCallback(t, old, new, data, state)
+			return true
 		},
 		false,
 	)
@@ -393,8 +398,9 @@ func Test_prefer_tcp(t *testing.T) {
 		"org.letsconnect-vpn.app.linux",
 		"configsprefertcp",
 		"en",
-		func(old FSMStateID, new FSMStateID, data interface{}) {
+		func(old FSMStateID, new FSMStateID, data interface{}) bool {
 			stateCallback(t, old, new, data, state)
+			return true
 		},
 		false,
 	)
diff --git a/client/fsm.go b/client/fsm.go
index 767ceaa..c93a9a8 100644
--- a/client/fsm.go
+++ b/client/fsm.go
@@ -96,7 +96,7 @@ func GetStateName(s FSMStateID) string {
 }
 
 func newFSM(
-	callback func(FSMStateID, FSMStateID, interface{}),
+	callback func(FSMStateID, FSMStateID, interface{}) bool,
 	directory string,
 	debug bool,
 ) fsm.FSM {
diff --git a/client/server.go b/client/server.go
index 3bcecc6..d22dc65 100644
--- a/client/server.go
+++ b/client/server.go
@@ -479,16 +479,20 @@ func (client *Client) GetConfigCustomServer(url string, preferTCP bool) (string,
 
 // askSecureLocation asks the user to choose a Secure Internet location by moving the FSM to the STATE_ASK_LOCATION state.
 func (client *Client) askSecureLocation() error {
+	errorMessage := "failed settings secure location"
 	locations := client.Discovery.GetSecureLocationList()
 
 	// Ask for the location in the callback
-	client.FSM.GoTransitionWithData(STATE_ASK_LOCATION, locations)
+	goTransitionErr := client.FSM.GoTransitionRequired(STATE_ASK_LOCATION, locations)
+	if goTransitionErr != nil {
+		return types.NewWrappedError(errorMessage, goTransitionErr)
+	}
 
 	// The state has changed, meaning setting the secure location was not successful
 	if client.FSM.Current != STATE_ASK_LOCATION {
 		// TODO: maybe a custom type for this errors.new?
 		return types.NewWrappedError(
-			"failed setting secure location",
+			errorMessage,
 			errors.New("failed loading secure location"),
 		)
 	}
@@ -579,7 +583,10 @@ func (client *Client) ensureLogin(chosenServer server.Server) error {
 	if server.NeedsRelogin(chosenServer) {
 		url, urlErr := server.GetOAuthURL(chosenServer, client.Name)
 
-		client.FSM.GoTransitionWithData(STATE_OAUTH_STARTED, url)
+		goTransitionErr := client.FSM.GoTransitionRequired(STATE_OAUTH_STARTED, url)
+		if goTransitionErr != nil {
+			return types.NewWrappedError(errorMessage, goTransitionErr)
+		}
 
 		if urlErr != nil {
 			client.goBackInternal()