FSM: More states and fix graph order

4 files changed, 74 insertions, 28 deletions

diff --git a/src/fsm.go b/src/fsm.go
index 7a54fd8..a4e7be5 100644
--- a/src/fsm.go
+++ b/src/fsm.go
@@ -3,9 +3,25 @@ package eduvpn
 import (
 	"fmt"
 	"os"
+	"sort"
 )
 
-type FSMStateID int8
+type (
+	FSMStateID      int8
+	FSMStateIDSlice []FSMStateID
+)
+
+func (v FSMStateIDSlice) Len() int {
+	return len(v)
+}
+
+func (v FSMStateIDSlice) Less(i, j int) bool {
+	return v[i] < v[j]
+}
+
+func (v FSMStateIDSlice) Swap(i, j int) {
+	v[i], v[j] = v[j], v[i]
+}
 
 const (
 	// Deregistered means the app is not registered with the wrapper
@@ -23,6 +39,12 @@ const (
 	// Authenticated means the OAuth process has finished and the user is now authenticated with the server
 	AUTHENTICATED
 
+	// Requested config means the user has requested a config for connecting
+	REQUEST_CONFIG
+
+	// Has config means the user has gotten a config
+	HAS_CONFIG
+
 	// Ask profile means the go code is asking for a profile selection from the ui
 	ASK_PROFILE
 
@@ -40,6 +62,10 @@ func (s FSMStateID) String() string {
 		return "Chosen_Server"
 	case OAUTH_STARTED:
 		return "OAuth_Started"
+	case HAS_CONFIG:
+		return "Has_Config"
+	case REQUEST_CONFIG:
+		return "Request_Config"
 	case ASK_PROFILE:
 		return "Ask_Profile"
 	case AUTHENTICATED:
@@ -52,7 +78,7 @@ func (s FSMStateID) String() string {
 }
 
 type FSMTransition struct {
-	To FSMStateID
+	To          FSMStateID
 	Description string
 }
 
@@ -128,9 +154,19 @@ remincross = false;
 
 func (eduvpn *VPNState) generateMermaidGraph() string {
 	graph := "graph TD\n"
-	graph += "style " + eduvpn.FSM.Current.String() + " fill:cyan\n"
-	for state, transitions := range eduvpn.FSM.States {
+	sorted_fsm := make(FSMStateIDSlice, 0, len(eduvpn.FSM.States))
+	for state_id := range eduvpn.FSM.States {
+		sorted_fsm = append(sorted_fsm, state_id)
+	}
+	sort.Sort(sorted_fsm)
+	for _, state := range sorted_fsm {
+		transitions := eduvpn.FSM.States[state]
 		for _, transition := range transitions {
+			if state == eduvpn.FSM.Current {
+				graph += "\nstyle " + state.String() + " fill:cyan\n"
+			} else {
+				graph += "\nstyle " + state.String() + " fill:white\n"
+			}
 			graph += state.String() + "(" + state.String() + ") " + "-->|" + transition.Description + "| " + transition.To.String() + "\n"
 		}
 	}
@@ -143,14 +179,16 @@ func (eduvpn *VPNState) GenerateGraph() string {
 
 func (eduvpn *VPNState) InitializeFSM() {
 	eduvpn.FSM = &FSM{
-		States: FSMStates {
-			DEREGISTERED: {{NO_SERVER, "Client registers"}},
-			NO_SERVER: {{CHOSEN_SERVER, "User chooses a server"}},
-			CHOSEN_SERVER: {{AUTHENTICATED, "Found tokens in config"}, {OAUTH_STARTED, "No tokens found in config"}},
-			OAUTH_STARTED: {{AUTHENTICATED, "User authorizes with browser"}},
-			AUTHENTICATED: {{CONNECTED, "OS reports connected"}, {OAUTH_STARTED, "Re-authenticate with OAuth"}, {ASK_PROFILE, "Connect, multiple profiles detected"}},
-			ASK_PROFILE: {{CONNECTED, "OS reports connected"}},
-			CONNECTED: {{AUTHENTICATED, "OS reports disconnected"}},
+		States: FSMStates{
+			DEREGISTERED:   {{NO_SERVER, "Client registers"}},
+			NO_SERVER:      {{CHOSEN_SERVER, "User chooses a server"}},
+			CHOSEN_SERVER:  {{AUTHENTICATED, "Found tokens in config"}, {OAUTH_STARTED, "No tokens found in config"}},
+			OAUTH_STARTED:  {{AUTHENTICATED, "User authorizes with browser"}},
+			AUTHENTICATED:  {{OAUTH_STARTED, "Re-authenticate with OAuth"}, {REQUEST_CONFIG, "Client requests a config"}},
+			REQUEST_CONFIG: {{ASK_PROFILE, "Multiple profiles found"}, {HAS_CONFIG, "Success, only one profile"}},
+			ASK_PROFILE:    {{HAS_CONFIG, "User chooses profile and success"}},
+			HAS_CONFIG:     {{CONNECTED, "OS reports connected"}},
+			CONNECTED:      {{AUTHENTICATED, "OS reports disconnected"}},
 		},
 		Current: DEREGISTERED,
 	}
diff --git a/src/server.go b/src/server.go
index 20d9136..d8afb01 100644
--- a/src/server.go
+++ b/src/server.go
@@ -6,11 +6,11 @@ import (
 )
 
 type Server struct {
-	BaseURL   string             `json:"base_url"`
-	Endpoints *ServerEndpoints   `json:"endpoints"`
-	OAuth     *OAuth             `json:"oauth"`
-	Profiles  *ServerProfileInfo `json:"profiles"`
-	ProfilesRaw string           `json:"profiles_raw"`
+	BaseURL     string             `json:"base_url"`
+	Endpoints   *ServerEndpoints   `json:"endpoints"`
+	OAuth       *OAuth             `json:"oauth"`
+	Profiles    *ServerProfileInfo `json:"profiles"`
+	ProfilesRaw string             `json:"profiles_raw"`
 }
 
 type ServerProfile struct {
@@ -104,6 +104,9 @@ func (server *Server) getProfileForID(profile_id string) (*ServerProfile, error)
 }
 
 func (server *Server) getConfigWithProfile(profile_id string) (string, error) {
+	if !GetVPNState().HasTransition(HAS_CONFIG) {
+		return "", errors.New("cannot get a config with a profile, invalid state")
+	}
 	profile, profileErr := server.getProfileForID(profile_id)
 
 	if profileErr != nil {
@@ -117,11 +120,17 @@ func (server *Server) getConfigWithProfile(profile_id string) (string, error) {
 }
 
 func (server *Server) askForProfileID() (string, error) {
+	if !GetVPNState().HasTransition(ASK_PROFILE) {
+		return "", errors.New("cannot ask for a profile id, invalid state")
+	}
 	_, profile_id := GetVPNState().GoTransition(ASK_PROFILE, server.ProfilesRaw)
 	return profile_id, nil
 }
 
 func (server *Server) GetConfig() (string, error) {
+	if !GetVPNState().InState(REQUEST_CONFIG) {
+		return "", errors.New("cannot get a config, invalid state")
+	}
 	infoErr := server.APIInfo()
 
 	if infoErr != nil {
diff --git a/src/server_test.go b/src/server_test.go
index 5b6ec5a..3045983 100644
--- a/src/server_test.go
+++ b/src/server_test.go
@@ -40,8 +40,8 @@ func StateCallback(t *testing.T, oldState string, newState string, data string)
 		go LoginOAuthSelenium(t, data)
 	}
 
-    // We have no data to send back
-    return ""
+	// We have no data to send back
+	return ""
 }
 
 func Test_server(t *testing.T) {
@@ -133,7 +133,6 @@ func Test_token_expired(t *testing.T) {
 		return StateCallback(t, old, new, data)
 	}, false)
 
-
 	_, configErr := state.Connect("https://eduvpnserver")
 
 	if configErr != nil {
diff --git a/src/state.go b/src/state.go
index 95bb75c..b9bb920 100644
--- a/src/state.go
+++ b/src/state.go
@@ -6,10 +6,10 @@ import (
 
 type VPNState struct {
 	// Info passed by the client
-	ConfigDirectory string                       `json:"-"`
-	Name            string                       `json:"-"`
-	StateCallback   func(string, string, string) string `json:"-"`
-	StateCallbackData string `json:"-"`
+	ConfigDirectory   string                              `json:"-"`
+	Name              string                              `json:"-"`
+	StateCallback     func(string, string, string) string `json:"-"`
+	StateCallbackData string                              `json:"-"`
 
 	// The chosen server
 	Server *Server `json:"server"`
@@ -93,14 +93,14 @@ func (state *VPNState) Connect(url string) (string, error) {
 		state.GoTransition(AUTHENTICATED, "")
 	}
 
+	state.GoTransition(REQUEST_CONFIG, "")
+
 	config, configErr := state.Server.GetConfig()
 
 	if configErr != nil {
 		return "", configErr
-	}
-
-	if !state.HasTransition(CONNECTED) {
-		return "", errors.New("cannot connect to server, invalid state")
+	} else {
+		state.GoTransition(HAS_CONFIG, "")
 	}
 
 	return config, nil