State + FSM: Properly handle the disconnect flow

- /disconnect is now called
- A new state is added (DISCONNECTING) that waits for the disconnect to complete
- A helper function is exposed (InFSMState) that can be used by clients to see in which state they are in

7 files changed, 92 insertions, 19 deletions

diff --git a/exports/exports.go b/exports/exports.go
index e0e8464..3ebd7f0 100644
--- a/exports/exports.go
+++ b/exports/exports.go
@@ -26,7 +26,7 @@ var P_StateCallbacks map[string]C.PythonCB
 
 var VPNStates map[string]*eduvpn.VPNState
 
-func StateCallback(name string, old_state eduvpn.VPNStateID, new_state eduvpn.VPNStateID, data interface{}) {
+func StateCallback(name string, old_state eduvpn.StateID, new_state eduvpn.StateID, data interface{}) {
 	P_StateCallback, exists := P_StateCallbacks[name]
 	if !exists || P_StateCallback == nil {
 		return
@@ -73,7 +73,7 @@ func Register(name *C.char, config_directory *C.char, stateCallback C.PythonCB,
 	}
 	VPNStates[nameStr] = state
 	P_StateCallbacks[nameStr] = stateCallback
-	registerErr := state.Register(nameStr, C.GoString(config_directory), func(old eduvpn.VPNStateID, new eduvpn.VPNStateID, data interface{}) {
+	registerErr := state.Register(nameStr, C.GoString(config_directory), func(old eduvpn.StateID, new eduvpn.StateID, data interface{}) {
 		StateCallback(nameStr, old, new, data)
 	}, debug != 0)
 
@@ -254,6 +254,17 @@ func SetDisconnected(name *C.char) *C.char {
 	return C.CString(ErrorToString(setDisconnectedErr))
 }
 
+//export SetDisconnecting
+func SetDisconnecting(name *C.char) *C.char {
+	nameStr := C.GoString(name)
+	state, stateErr := GetVPNState(nameStr)
+	if stateErr != nil {
+		return C.CString(ErrorToString(stateErr))
+	}
+	setDisconnectingErr := state.SetDisconnecting()
+	return C.CString(ErrorToString(setDisconnectingErr))
+}
+
 //export SetConnecting
 func SetConnecting(name *C.char) *C.char {
 	nameStr := C.GoString(name)
@@ -301,6 +312,20 @@ func ShouldRenewButton(name *C.char) C.int {
 	return C.int(0)
 }
 
+//export InFSMState
+func InFSMState(name *C.char, checkState C.int) C.int {
+	nameStr := C.GoString(name)
+	state, stateErr := GetVPNState(nameStr)
+	if stateErr != nil {
+		return C.int(0)
+	}
+	inStateBool := state.InFSMState(eduvpn.StateID(checkState))
+	if inStateBool {
+		return C.int(1)
+	}
+	return C.int(0)
+}
+
 //export FreeString
 func FreeString(addr *C.char) {
 	C.free(unsafe.Pointer(addr))
diff --git a/internal/fsm/fsm.go b/internal/fsm/fsm.go
index 3c51ce7..f5b1507 100644
--- a/internal/fsm/fsm.go
+++ b/internal/fsm/fsm.go
@@ -62,6 +62,9 @@ const (
 	// Has config means the user has gotten a config
 	HAS_CONFIG
 
+	// Disconnecting means the OS is disconnecting and the Go code is doing the /disconnect
+	DISCONNECTING
+
 	// Connecting means the OS is establishing a connection to the server
 	CONNECTING
 
@@ -93,6 +96,8 @@ func (s FSMStateID) String() string {
 		return "Ask_Profile"
 	case AUTHORIZED:
 		return "Authorized"
+	case DISCONNECTING:
+		return "Disconnecting"
 	case CONNECTING:
 		return "Connecting"
 	case CONNECTED:
@@ -142,8 +147,9 @@ func (fsm *FSM) Init(name string, callback func(FSMStateID, FSMStateID, interfac
 		REQUEST_CONFIG: FSMState{Transitions: []FSMTransition{{ASK_PROFILE, "Multiple profiles found and no profile chosen"}, {HAS_CONFIG, "Only one profile or profile already chosen"}, {NO_SERVER, "Cancel or Error"}, {OAUTH_STARTED, "Re-authorize"}}},
 		ASK_PROFILE:    FSMState{Transitions: []FSMTransition{{HAS_CONFIG, "User chooses profile"}, {NO_SERVER, "Cancel or Error"}, {SEARCH_SERVER, "Cancel or Error"}}},
 		HAS_CONFIG:     FSMState{Transitions: []FSMTransition{{CONNECTING, "OS reports it is trying to connect"}, {REQUEST_CONFIG, "User reconnects"}, {NO_SERVER, "User wants to choose a new server"}, {OAUTH_STARTED, "Re-authorize with OAuth"}}, BackState: NO_SERVER},
+		DISCONNECTING:     FSMState{Transitions: []FSMTransition{{HAS_CONFIG, "Cancel or Error"}, {HAS_CONFIG, "Done disconnecting"}}},
 		CONNECTING:     FSMState{Transitions: []FSMTransition{{HAS_CONFIG, "Cancel or Error"}, {CONNECTED, "Done connecting"}}},
-		CONNECTED:      FSMState{Transitions: []FSMTransition{{HAS_CONFIG, "OS reports disconnected"}}},
+		CONNECTED:      FSMState{Transitions: []FSMTransition{{DISCONNECTING, "App wants to disconnect"}}},
 	}
 	fsm.Current = DEREGISTERED
 	fsm.Name = name
diff --git a/internal/server/common.go b/internal/server/common.go
index 7ee8be6..56c8af0 100644
--- a/internal/server/common.go
+++ b/internal/server/common.go
@@ -523,6 +523,10 @@ func GetConfig(server Server, forceTCP bool) (string, string, error) {
 	return getConfigWithProfile(server, forceTCP)
 }
 
+func Disconnect(server Server) {
+	APIDisconnect(server)
+}
+
 type ServerGetCurrentProfileNotFoundError struct {
 	ProfileID string
 }
diff --git a/state.go b/state.go
index 3ad979d..139f67a 100644
--- a/state.go
+++ b/state.go
@@ -14,7 +14,7 @@ import (
 )
 
 type ServerInfo = server.ServerInfoScreen
-type VPNStateID = fsm.FSMStateID
+type StateID = fsm.FSMStateID
 
 type VPNState struct {
 	// The chosen server
@@ -40,9 +40,9 @@ func (state *VPNState) GetSavedServers() *server.ServersConfiguredScreen {
 	return state.Servers.GetServersConfigured()
 }
 
-func (state *VPNState) Register(name string, directory string, stateCallback func(VPNStateID, VPNStateID, interface{}), debug bool) error {
+func (state *VPNState) Register(name string, directory string, stateCallback func(StateID, StateID, interface{}), debug bool) error {
 	errorMessage := "failed to register with the GO library"
-	if !state.FSM.InState(fsm.DEREGISTERED) {
+	if !state.InFSMState(fsm.DEREGISTERED) {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.DeregisteredError{}.CustomError()}
 	}
 	// Initialize the logger
@@ -92,7 +92,7 @@ func (state *VPNState) Deregister() error {
 
 func (state *VPNState) GoBack() error {
 	errorMessage := "failed to go back"
-	if state.FSM.InState(fsm.DEREGISTERED) {
+	if state.InFSMState(fsm.DEREGISTERED) {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.DeregisteredError{}.CustomError()}
 	}
 
@@ -104,7 +104,7 @@ func (state *VPNState) GoBack() error {
 
 func (state *VPNState) getConfig(chosenServer server.Server, forceTCP bool) (string, string, error) {
 	errorMessage := "failed to get a configuration for OpenVPN/Wireguard"
-	if state.FSM.InState(fsm.DEREGISTERED) {
+	if state.InFSMState(fsm.DEREGISTERED) {
 		return "", "", &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.DeregisteredError{}.CustomError()}
 	}
 
@@ -273,7 +273,7 @@ func (state *VPNState) GetConfigCustomServer(url string, forceTCP bool) (string,
 
 func (state *VPNState) CancelOAuth() error {
 	errorMessage := "failed to cancel OAuth"
-	if !state.FSM.InState(fsm.OAUTH_STARTED) {
+	if !state.InFSMState(fsm.OAUTH_STARTED) {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.WrongStateError{Got: state.FSM.Current, Want: fsm.OAUTH_STARTED}.CustomError()}
 	}
 
@@ -289,7 +289,7 @@ func (state *VPNState) CancelOAuth() error {
 func (state *VPNState) ChangeSecureLocation() error {
 	errorMessage := "failed to change location from the main screen"
 
-	if !state.FSM.InState(fsm.NO_SERVER) {
+	if !state.InFSMState(fsm.NO_SERVER) {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.WrongStateError{Got: state.FSM.Current, Want: fsm.NO_SERVER}.CustomError()}
 	}
 
@@ -306,14 +306,14 @@ func (state *VPNState) ChangeSecureLocation() error {
 }
 
 func (state *VPNState) GetDiscoOrganizations() (string, error) {
-	if state.FSM.InState(fsm.DEREGISTERED) {
+	if state.InFSMState(fsm.DEREGISTERED) {
 		return "", &types.WrappedErrorMessage{Message: "failed to get the organizations with Discovery", Err: fsm.DeregisteredError{}.CustomError()}
 	}
 	return state.Discovery.GetOrganizationsList()
 }
 
 func (state *VPNState) GetDiscoServers() (string, error) {
-	if state.FSM.InState(fsm.DEREGISTERED) {
+	if state.InFSMState(fsm.DEREGISTERED) {
 		return "", &types.WrappedErrorMessage{Message: "failed to get the servers with Discovery", Err: fsm.DeregisteredError{}.CustomError()}
 	}
 	return state.Discovery.GetServersList()
@@ -351,7 +351,7 @@ func (state *VPNState) getServerInfoData() *server.ServerInfoScreen {
 }
 
 func (state *VPNState) SetConnected() error {
-	if state.FSM.InState(fsm.CONNECTED) {
+	if state.InFSMState(fsm.CONNECTED) {
 		// already connected, show no error
 		return nil
 	}
@@ -364,7 +364,7 @@ func (state *VPNState) SetConnected() error {
 }
 
 func (state *VPNState) SetConnecting() error {
-	if state.FSM.InState(fsm.CONNECTING) {
+	if state.InFSMState(fsm.CONNECTING) {
 		// already loading connection, show no error
 		return nil
 	}
@@ -376,15 +376,37 @@ func (state *VPNState) SetConnecting() error {
 	return nil
 }
 
+func (state *VPNState) SetDisconnecting() error {
+	if state.InFSMState(fsm.DISCONNECTING) {
+		// already disconnecting, show no error
+		return nil
+	}
+	if !state.FSM.HasTransition(fsm.DISCONNECTING) {
+		return &types.WrappedErrorMessage{Message: "failed to set disconnecting", Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.DISCONNECTING}.CustomError()}
+	}
+
+
+	state.FSM.GoTransitionWithData(fsm.DISCONNECTING, state.getServerInfoData(), false)
+	return nil
+}
+
 func (state *VPNState) SetDisconnected() error {
-	if state.FSM.InState(fsm.HAS_CONFIG) {
+	errorMessage := "failed to set disconnected"
+	if state.InFSMState(fsm.HAS_CONFIG) {
 		// already disconnected, show no error
 		return nil
 	}
 	if !state.FSM.HasTransition(fsm.HAS_CONFIG) {
-		return &types.WrappedErrorMessage{Message: "failed to set disconnected", Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.HAS_CONFIG}.CustomError()}
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.HAS_CONFIG}.CustomError()}
 	}
 
+	// Do the /disconnect API call and go to disconnected after...
+	currentServer, currentServerErr := state.Servers.GetCurrentServer()
+	if currentServerErr != nil {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: currentServerErr}
+	}
+	server.Disconnect(currentServer)
+
 	state.FSM.GoTransitionWithData(fsm.HAS_CONFIG, state.getServerInfoData(), false)
 
 	return nil
@@ -417,7 +439,7 @@ func (state *VPNState) RenewSession() error {
 }
 
 func (state *VPNState) ShouldRenewButton() bool {
-	if !state.FSM.InState(fsm.CONNECTED) {
+	if !state.InFSMState(fsm.CONNECTED) {
 		return false
 	}
 
@@ -431,6 +453,10 @@ func (state *VPNState) ShouldRenewButton() bool {
 	return server.ShouldRenewButton(currentServer)
 }
 
+func (state *VPNState) InFSMState(checkState StateID) bool {
+	return state.FSM.InState(checkState)
+}
+
 func GetErrorCause(err error) error {
 	return types.GetErrorCause(err)
 }
diff --git a/wrappers/python/src/__init__.py b/wrappers/python/src/__init__.py
index c0b6679..761adf5 100644
--- a/wrappers/python/src/__init__.py
+++ b/wrappers/python/src/__init__.py
@@ -83,12 +83,14 @@ lib.SetSecureLocation.argtypes, lib.SetSecureLocation.restype = [
     c_char_p,
 ], c_void_p
 lib.SetConnected.argtypes, lib.SetConnected.restype = [c_char_p], c_void_p
+lib.SetDisconnecting.argtypes, lib.SetDisconnecting.restype = [c_char_p], c_void_p
 lib.SetConnecting.argtypes, lib.SetConnecting.restype = [c_char_p], c_void_p
 lib.SetDisconnected.argtypes, lib.SetDisconnected.restype = [c_char_p], c_void_p
 lib.SetSearchServer.argtypes, lib.SetSearchServer.restype = [c_char_p], c_void_p
 lib.ShouldRenewButton.argtypes, lib.ShouldRenewButton.restype = [], int
 lib.RenewSession.argtypes, lib.RenewSession.restype = [c_char_p], c_void_p
 lib.FreeString.argtypes, lib.FreeString.restype = [c_void_p], None
+lib.InFSMState.argtypes, lib.InFSMState.restype = [c_void_p, c_int], int
 
 
 class WrappedError:
diff --git a/wrappers/python/src/main.py b/wrappers/python/src/main.py
index 8440c7d..574d0e2 100644
--- a/wrappers/python/src/main.py
+++ b/wrappers/python/src/main.py
@@ -145,6 +145,12 @@ class EduVPN(object):
         if connect_err:
             raise Exception(connect_err)
 
+    def set_disconnecting(self) -> None:
+        disconnecting_err = self.go_function(lib.SetDisconnecting)
+
+        if disconnecting_err:
+            raise Exception(disconnecting_err)
+
     def set_connecting(self) -> None:
         connecting_err = self.go_function(lib.SetConnecting)
 
@@ -216,3 +222,6 @@ class EduVPN(object):
 
     def should_renew_button(self) -> bool:
         return self.go_function(lib.ShouldRenewButton)
+
+    def in_fsm_state(self, state_id: State) -> bool:
+        return self.go_function(lib.InFSMState, state_id)
diff --git a/wrappers/python/src/state.py b/wrappers/python/src/state.py
index cd5bd90..a4b11a8 100644
--- a/wrappers/python/src/state.py
+++ b/wrappers/python/src/state.py
@@ -19,5 +19,6 @@ class State(IntEnum):
     REQUEST_CONFIG = 8
     ASK_PROFILE = 9
     HAS_CONFIG = 10
-    CONNECTING = 11
-    CONNECTED = 12
+    DISCONNECTING = 11
+    CONNECTING = 12
+    CONNECTED = 13