Client: Use a mutex for state transitions

5 files changed, 165 insertions, 65 deletions

diff --git a/client/client.go b/client/client.go
index f0add42..74aea98 100644
--- a/client/client.go
+++ b/client/client.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/eduvpn/eduvpn-common/internal/config"
@@ -114,6 +115,8 @@ type Client struct {
 
 	// TokenGetter gets the tokens from the client
 	TokenGetter func(srv srvtypes.Current) *srvtypes.Tokens `json:"-"`
+
+	mu sync.Mutex
 }
 
 func (c *Client) updateTokens(srv server.Server) error {
@@ -134,8 +137,8 @@ func (c *Client) updateTokens(srv server.Server) error {
 	}
 
 	server.UpdateTokens(srv, oauth.Token{
-		Access: tokens.Access,
-		Refresh: tokens.Refresh,
+		Access:           tokens.Access,
+		Refresh:          tokens.Refresh,
 		ExpiredTimestamp: time.Unix(tokens.Expires, 0),
 	})
 
@@ -457,7 +460,8 @@ func (c *Client) profileCallback(ck *cookie.Cookie, srv server.Server) error {
 
 // AddServer adds a server with identifier and type
 func (c *Client) AddServer(ck *cookie.Cookie, identifier string, _type srvtypes.Type, ni bool) (err error) {
-
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	// If we have failed to add the server, we remove it again
 	// We add the server because we can then obtain it in other callback functions
 	previousState := c.FSM.Current
@@ -588,6 +592,8 @@ func (c *Client) server(identifier string, _type srvtypes.Type) (srv server.Serv
 
 // GetConfig gets a VPN configuration
 func (c *Client) GetConfig(ck *cookie.Cookie, identifier string, _type srvtypes.Type, pTCP bool) (cfg *srvtypes.Configuration, err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	previousState := c.FSM.Current
 	defer func() {
 		if err == nil {
@@ -848,6 +854,8 @@ func (c *Client) SetSecureLocation(ck *cookie.Cookie, countryCode string) (err e
 }
 
 func (c *Client) RenewSession(ck *cookie.Cookie) (err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	srv, err := c.Servers.Current()
 	if err != nil {
 		return err
@@ -877,13 +885,23 @@ func (c *Client) StartFailover(ck *cookie.Cookie, gateway string, mtu int, readR
 	return f.Start(ck.Context(), gateway, mtu)
 }
 
-
 func (c *Client) SetState(state FSMStateID) error {
-	err := c.FSM.CheckTransition(state)
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	_, err := c.FSM.GoTransition(state)
 	if err != nil {
+		// self-transitions are only debug errors
+		if c.FSM.InState(state) {
+			log.Logger.Debugf("attempt an invalid self-transition: %s", c.FSM.GetStateName(state))
+			return nil
+		}
 		return err
 	}
-	// TODO: Now we don't pass any data :/
-	c.FSM.GoTransition(state)
 	return nil
 }
+
+func (c *Client) InState(state FSMStateID) bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.FSM.InState(state)
+}
diff --git a/exports/exports.go b/exports/exports.go
index 5e62e4d..0e557fa 100644
--- a/exports/exports.go
+++ b/exports/exports.go
@@ -1,10 +1,10 @@
 // package main implements the main exported API to be used by other languages
 // Some notes:
-//  - Errors are returned as c strings, free them using FreeString. Same is the case for other string types, you should also free them
-//  - Types are converted from the Go representation to C using JSON strings
-//  - Cookies are used for cancellation, just fancy contexts. Create a cookie using `CookieNew`, pass it to the function that needs one as the first argument. To cancel the function, call `CookieCancel`, passing in the same cookie as argument
-//    - Cookies must also be freed, by using the CookieDelete function if the cookie is no longer needed
-//  - The state machine is used to track the state of a client. It is mainly used for asking for certain data from the client, e.g. asking for profiles and locations. But a client may also wish to build upon this state machine to build the whole UI around it
+//   - Errors are returned as c strings, free them using FreeString. Same is the case for other string types, you should also free them
+//   - Types are converted from the Go representation to C using JSON strings
+//   - Cookies are used for cancellation, just fancy contexts. Create a cookie using `CookieNew`, pass it to the function that needs one as the first argument. To cancel the function, call `CookieCancel`, passing in the same cookie as argument
+//   - Cookies must also be freed, by using the CookieDelete function if the cookie is no longer needed
+//   - The state machine is used to track the state of a client. It is mainly used for asking for certain data from the client, e.g. asking for profiles and locations. But a client may also wish to build upon this state machine to build the whole UI around it
 package main
 
 /*
@@ -104,13 +104,16 @@ func getVPNState() (*client.Client, error) {
 // Name is the name of the client, must be a valid client ID
 // Version is the version of the client. This version field is used for the user agent in all HTTP requests
 // cb is the state callback. It takes three arguments: The old state, the new state and the data for the state as JSON
-//  - Note that the states are defined in client/fsm.go, e.g. NO_SERVER (in Go: StateNoServer), ASK_PROFILE (in Go: StateAskProfile)
-//  - This callback returns non-zero if the state transition is handled. This is used to check if the client handles the needed transitions
+//   - Note that the states are defined in client/fsm.go, e.g. NO_SERVER (in Go: StateNoServer), ASK_PROFILE (in Go: StateAskProfile)
+//   - This callback returns non-zero if the state transition is handled. This is used to check if the client handles the needed transitions
+//
 // debug, if non-zero, enables debugging mode for the library, this means:
-//  - Log everything in debug mode, so you can get more detail of what is going on
-//  - Write the state graph to a file in the configDirectory. This can be used to create a FSM png file with mermaid https://mermaid.js.org/
+//   - Log everything in debug mode, so you can get more detail of what is going on
+//   - Write the state graph to a file in the configDirectory. This can be used to create a FSM png file with mermaid https://mermaid.js.org/
+//
 // After registering, the FSM is initialized and the state transition NO_SERVER should have been completed
 // If some error occurs during registering, it is returned as a C string
+//
 //export Register
 func Register(
 	name *C.char,
@@ -155,6 +158,7 @@ func Register(
 // e.g. when to show the renew button, when to show expiry notifications
 // The expiry times structure is defined in types/server/server.go `Expiry`
 // If some error occurs during registering, it is returned as a C string
+//
 //export ExpiryTimes
 func ExpiryTimes() (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -177,6 +181,7 @@ func ExpiryTimes() (*C.char, *C.char) {
 // This function SHOULD be called when the application exits such that the configuration file is saved correctly
 // Note that saving of the configuration file also happens in other cases, such as after getting a VPN configuration
 // Thus it is often not problematic if this function cannot be called due to a client crash
+//
 //export Deregister
 func Deregister() *C.char {
 	state, stateErr := getVPNState()
@@ -192,16 +197,18 @@ func Deregister() *C.char {
 // c is the cookie that is used for cancellation. Create a cookie first with CookieNew. This same cookie is also used for replying to state transitions
 // _type is the type of server that needs to be added. This type is defined in types/server/server.go Type
 // id is the identifier of the string
-//  - In case of secure internet: The organization ID
-//  - In case of custom server: The base URL
-//  - In case of institute access: The base URL
+//   - In case of secure internet: The organization ID
+//   - In case of custom server: The base URL
+//   - In case of institute access: The base URL
+//
 // ni stands for non-interactive. If non-zero, any state transitions will not be run.
 // This ni flag is useful for preprovisioned servers. For normal usage, you want to set this to zero (meaning: False)
 // If the server cannot be added it returns the error as a string
 // Note that the server is removed when an error has occured
 // The following state callbacks are mandatory to handle:
-//  - OAUTH_STARTED: This indicates that the OAuth procedure has been started, it returns the URL as the data.
-//                   The client should open the webbrowser with this URL and continue the authorization process.
+//   - OAUTH_STARTED: This indicates that the OAuth procedure has been started, it returns the URL as the data.
+//     The client should open the webbrowser with this URL and continue the authorization process.
+//
 //export AddServer
 func AddServer(c C.uintptr_t, _type C.int, id *C.char, ni C.int) *C.char {
 	// TODO: type
@@ -225,6 +232,7 @@ func AddServer(c C.uintptr_t, _type C.int, id *C.char, ni C.int) *C.char {
 // - In case of institute access: The base URL
 // If the server cannot be removed it returns the error as a string
 // Note that the server is not removed when an error has occured
+//
 //export RemoveServer
 func RemoveServer(_type C.int, id *C.char) *C.char {
 	state, stateErr := getVPNState()
@@ -239,6 +247,7 @@ func RemoveServer(_type C.int, id *C.char) *C.char {
 // In eduvpn-common, a server is marked as 'current' if you have gotten a VPN configuration for it
 // It returns the server as JSON, defined in types/server/server.go Current
 // If there is no current server or some other, e.g. there is no current state, an error is returned with a nil string
+//
 //export CurrentServer
 func CurrentServer() (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -260,6 +269,7 @@ func CurrentServer() (*C.char, *C.char) {
 // This is NOT the discovery list, but the servers that have previously been added with `AddServer`
 // It returns the server list as a JSON string defined in types/server/server.go List
 // If the server list cannot be retrieved it returns a nil string and an error
+//
 //export ServerList
 func ServerList() (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -288,37 +298,39 @@ func ServerList() (*C.char, *C.char) {
 // If the server cannot be added it returns the error as a string
 // Note that the server is removed when an error has occured
 // The current state callbacks MUST be handled
-//  - ASK_PROFILE:   This asks the client for profile.
-//                   This is called when the user/client has not set a profile for this server before, or the current profile is invalid
-//                   when the user has selected a profile. Reply with the choice using the `CookieReply` function and the profile ID
-//                   e.g. CookieReply(cookie, "wireguard")
-//                   CookieReply can be done in the background as the Go library waits for a reply
-//                   The data for this transition is defined in types/server/server.go RequiredAskTransition with embedded data Profiles in types/server/server.go
-//                   Note that RequiredTransition contains the cookie to be used for the CookieReply
-//                   so a client would:
-//                   - Parse the data to get the cookie and data
-//                   - get the cookie,
-//                   - get the profiles from the data
-//                   - show it in the UI and then reply with CookieReply using the choice
-//  - ASK_LOCATION:  This asks the client for a location. Note that under normal circumstances,
-//                   this callback is not actually called as the home organization for the secure internet server is set as the current
-//                   if for some reason, an invalid location has been configured, the library will ask the client for a new one
-//                   when the user has selected al ocation. Reply with the choice using the `CookieReply` function and the location ID 
-//                   e.g. CookieReply(cookie, "nl")
-//                   CookieReply can be done in the background as the Go library waits for a reply
-//                   The data for this transition is defined in types/server/server.go RequiredAskTransition with embedded data a list of strings ([]string)
-//                   Note that RequiredTransition contains the cookie to be used for the CookieReply,
-//                   so a client would:
-//                   - Parse the data to get the cookie and data
-//                   - get the cookie,
-//                   - get the list of locations from the data
-//                   - show it in the UI and then reply with CookieReply using the choice
-//  - OAUTH_STARTED: This indicates that the OAuth procedure has been started, it returns the URL as the data.
-//                   The client should open the webbrowser with this URL and continue the authorization process.
-//                   This is only called if authorization needs to be retriggered
+//   - ASK_PROFILE:   This asks the client for profile.
+//     This is called when the user/client has not set a profile for this server before, or the current profile is invalid
+//     when the user has selected a profile. Reply with the choice using the `CookieReply` function and the profile ID
+//     e.g. CookieReply(cookie, "wireguard")
+//     CookieReply can be done in the background as the Go library waits for a reply
+//     The data for this transition is defined in types/server/server.go RequiredAskTransition with embedded data Profiles in types/server/server.go
+//     Note that RequiredTransition contains the cookie to be used for the CookieReply
+//     so a client would:
+//   - Parse the data to get the cookie and data
+//   - get the cookie,
+//   - get the profiles from the data
+//   - show it in the UI and then reply with CookieReply using the choice
+//   - ASK_LOCATION:  This asks the client for a location. Note that under normal circumstances,
+//     this callback is not actually called as the home organization for the secure internet server is set as the current
+//     if for some reason, an invalid location has been configured, the library will ask the client for a new one
+//     when the user has selected al ocation. Reply with the choice using the `CookieReply` function and the location ID
+//     e.g. CookieReply(cookie, "nl")
+//     CookieReply can be done in the background as the Go library waits for a reply
+//     The data for this transition is defined in types/server/server.go RequiredAskTransition with embedded data a list of strings ([]string)
+//     Note that RequiredTransition contains the cookie to be used for the CookieReply,
+//     so a client would:
+//   - Parse the data to get the cookie and data
+//   - get the cookie,
+//   - get the list of locations from the data
+//   - show it in the UI and then reply with CookieReply using the choice
+//   - OAUTH_STARTED: This indicates that the OAuth procedure has been started, it returns the URL as the data.
+//     The client should open the webbrowser with this URL and continue the authorization process.
+//     This is only called if authorization needs to be retriggered
+//
 // After getting a configuration, the FSM moves to the GOT_CONFIG state
 // The return data is the configuration, marshalled as JSON and defined in types/server/server.go Configuration
 // This is nil if an error is returned as a string
+//
 //export GetConfig
 func GetConfig(c C.uintptr_t, _type C.int, id *C.char, pTCP C.int) (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -344,6 +356,7 @@ func GetConfig(c C.uintptr_t, _type C.int, id *C.char, pTCP C.int) (*C.char, *C.
 // This MUST only be called if the user/client wishes to manually set a profile instead of the common lib asking for one using a transition
 // Data is the profile ID
 // It returns an error if unsuccessful
+//
 //export SetProfileID
 func SetProfileID(data *C.char) *C.char {
 	state, stateErr := getVPNState()
@@ -357,9 +370,10 @@ func SetProfileID(data *C.char) *C.char {
 // SetSecureLocation sets the location for the secure internet server if it exists
 // This MUST only be called if the user/client wishes to manually set a location instead of the common lib asking for one using a transition
 // Because this does network requests to initialize the location, there is a cookie again :)
-// c is the Cookie that needs to be passed. To create a cookie, first call `CookieNew` 
+// c is the Cookie that needs to be passed. To create a cookie, first call `CookieNew`
 // Data is the location ID
 // It returns an error if unsuccessful
+//
 //export SetSecureLocation
 func SetSecureLocation(c C.uintptr_t, data *C.char) *C.char {
 	state, stateErr := getVPNState()
@@ -378,6 +392,7 @@ func SetSecureLocation(c C.uintptr_t, data *C.char) *C.char {
 // c is the Cookie that needs to be passed. Create a new Cookie using `CookieNew`
 // If it was unsuccessful, it returns an error. Note that when the lib was built in release mode the data is almost always non-nil, even when an error has occurred
 // This means it has just returned the cached list
+//
 //export DiscoServers
 func DiscoServers(c C.uintptr_t) (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -403,6 +418,7 @@ func DiscoServers(c C.uintptr_t) (*C.char, *C.char) {
 // c is the Cookie that needs to be passed. Create a new Cookie using `CookieNew`
 // If it was unsuccessful, it returns an error. Note that when the lib was built in release mode the data is almost always non-nil, even when an error has occurred
 // This means it has just returned the cached list
+//
 //export DiscoOrganizations
 func DiscoOrganizations(c C.uintptr_t) (*C.char, *C.char) {
 	state, stateErr := getVPNState()
@@ -428,6 +444,7 @@ func DiscoOrganizations(c C.uintptr_t) (*C.char, *C.char) {
 // This MUST be called when disconnecting, see https://docs.eduvpn.org/server/v3/api.html#application-flow
 // c is the Cookie that needs to be passed. Create a new Cookie using `CookieNew`
 // If it was unsuccessful, it returns an error
+//
 //export Cleanup
 func Cleanup(c C.uintptr_t) *C.char {
 	state, stateErr := getVPNState()
@@ -447,6 +464,7 @@ func Cleanup(c C.uintptr_t) *C.char {
 // And it also possibly re-runs every state callback you need when getting a config
 // At least you MUST handle the OAuth started transition
 // It returns an error if unsuccessful
+//
 //export RenewSession
 func RenewSession(c C.uintptr_t) *C.char {
 	state, stateErr := getVPNState()
@@ -466,6 +484,7 @@ func RenewSession(c C.uintptr_t) *C.char {
 // To disable it you can pass a 0 int to this
 // support thus indicates whether or not to enable WireGuard
 // An error is returned if this is not possible
+//
 //export SetSupportWireguard
 func SetSupportWireguard(support C.int) *C.char {
 	state, stateErr := getVPNState()
@@ -485,6 +504,7 @@ func SetSupportWireguard(support C.int) *C.char {
 // readRxBytes is a function that returns the current rx bytes of the VPN interface, this should return a `long long int` in c
 // It returns a boolean whether or not the common lib has determined that it cannot reach the gateway. Non-zero=dropped, zero=not dropped
 // It also returns an error, if it fails to indicate if it has dropped or not. In this case, dropped is also set to zero
+//
 //export StartFailover
 func StartFailover(c C.uintptr_t, gateway *C.char, mtu C.int, readRxBytes C.ReadRxBytes) (C.int, *C.char) {
 	state, stateErr := getVPNState()
@@ -514,6 +534,7 @@ func StartFailover(c C.uintptr_t, gateway *C.char, mtu C.int, readRxBytes C.Read
 
 // SetState sets the state of the statemachine
 // Note that this transitions the FSM into the new state without passing any data to it
+//
 //export SetState
 func SetState(fsmState C.int) *C.char {
 	state, stateErr := getVPNState()
@@ -523,10 +544,26 @@ func SetState(fsmState C.int) *C.char {
 	return getCError(state.SetState(client.FSMStateID(fsmState)))
 }
 
+// InState checks if the FSM is in `fsmState`
+//
+//export InState
+func InState(fsmState C.int) (C.int, *C.char) {
+	state, stateErr := getVPNState()
+	if stateErr != nil {
+		return 0, getCError(stateErr)
+	}
+
+	if yes := state.InState(client.FSMStateID(fsmState)); yes {
+		return 1, nil
+	}
+	return 0, nil
+}
+
 // FreeString frees a string that was allocated by the eduvpn-common Go library
 // This happens when we return strings, such as errors from the Go lib back to the client
 // The client MUST thus ensure that this memory is freed using this function
 // Simply pass the pointer to the string in here
+//
 //export FreeString
 func FreeString(addr *C.char) {
 	C.free(unsafe.Pointer(addr))
@@ -556,11 +593,14 @@ func getCookie(c C.uintptr_t) (*cookie.Cookie, error) {
 //   - The server for which to get the tokens for, marshalled as JSON and defined in types/server/server.go `Current`
 //   - The output buffer
 //   - The length of the output buffer
+//
 // This 'output buffer' must contain the tokens, marshalled as JSON that is defined in types/server/server.go `Tokens`
 // setter is the void function that sets tokens. It takes two arguments:
 //   - The server for which to get the tokens for, marshalled as JSON and defined in types/server/server.go `Current`
 //   - The tokens, defined in types/server/server.go `Tokens` marshalled as JSON
+//
 // It returns an error when the tokens cannot be set
+//
 //export SetTokenHandler
 func SetTokenHandler(getter C.TokenGetter, setter C.TokenSetter) *C.char {
 	state, stateErr := getVPNState()
@@ -615,7 +655,6 @@ func SetTokenHandler(getter C.TokenGetter, setter C.TokenSetter) *C.char {
 		return &gotT
 	}
 
-
 	return nil
 }
 
@@ -623,9 +662,11 @@ func SetTokenHandler(getter C.TokenGetter, setter C.TokenSetter) *C.char {
 // This value should not be parsed or converted somehow by the client
 // This value is simply to pass back to the Go library
 // This value has two purposes:
-//  - Cancel a long running function
-//  - Send a reply to a state transition (ASK_PROFILE and ASK_LOCATION)
+//   - Cancel a long running function
+//   - Send a reply to a state transition (ASK_PROFILE and ASK_LOCATION)
+//
 // Functions that take a cookie have it as the first argument
+//
 //export CookieNew
 func CookieNew() C.uintptr_t {
 	c := cookie.NewWithContext(context.Background())
@@ -636,6 +677,7 @@ func CookieNew() C.uintptr_t {
 // The data that is sent to the Go library is the second argument of this function
 // c is the Cookie
 // data is the data to send
+//
 //export CookieReply
 func CookieReply(c C.uintptr_t, data *C.char) *C.char {
 	v, err := getCookie(c)
@@ -647,7 +689,8 @@ func CookieReply(c C.uintptr_t, data *C.char) *C.char {
 }
 
 // CookieDelete deletes the cookie by cancelling it and deleting the underlying cgo handle
-// This function MUST be called when the cookie that is created using `CookieNew` is no longer needed 
+// This function MUST be called when the cookie that is created using `CookieNew` is no longer needed
+//
 //export CookieDelete
 func CookieDelete(c C.uintptr_t) *C.char {
 	v, err := getCookie(c)
@@ -664,6 +707,7 @@ func CookieDelete(c C.uintptr_t) *C.char {
 // This means that functions which take this as first argument, return if they're still running
 // The error cause is always context.Canceled for that cancelled function: https://pkg.go.dev/context#pkg-variables
 // This CookieCancel function can also return an error if cancelling was unsuccessful
+//
 //export CookieCancel
 func CookieCancel(c C.uintptr_t) *C.char {
 	v, err := getCookie(c)
diff --git a/wrappers/python/eduvpn_common/loader.py b/wrappers/python/eduvpn_common/loader.py
index afad569..75a6a0a 100644
--- a/wrappers/python/eduvpn_common/loader.py
+++ b/wrappers/python/eduvpn_common/loader.py
@@ -4,7 +4,14 @@ from collections import defaultdict
 from ctypes import CDLL, c_char_p, c_int, c_void_p, cdll
 
 from eduvpn_common import __version__
-from eduvpn_common.types import BoolError, DataError, ReadRxBytes, TokenGetter, TokenSetter, VPNStateChange
+from eduvpn_common.types import (
+    BoolError,
+    DataError,
+    ReadRxBytes,
+    TokenGetter,
+    TokenSetter,
+    VPNStateChange,
+)
 
 
 def load_lib() -> CDLL:
@@ -88,7 +95,10 @@ def initialize_functions(lib: CDLL) -> None:
         c_int,
     ], c_void_p
     lib.RenewSession.argtypes, lib.RenewSession.restype = [c_int], c_void_p
-    lib.SetTokenHandler.argtypes, lib.SetTokenHandler.restype = [TokenGetter, TokenSetter], c_void_p
+    lib.SetTokenHandler.argtypes, lib.SetTokenHandler.restype = [
+        TokenGetter,
+        TokenSetter,
+    ], c_void_p
     lib.Cleanup.argtypes, lib.Cleanup.restype = [c_int], c_void_p
     lib.SetProfileID.argtypes, lib.SetProfileID.restype = [c_char_p], c_void_p
     lib.CookieNew.argtypes, lib.CookieNew.restype = [], c_int
@@ -105,6 +115,9 @@ def initialize_functions(lib: CDLL) -> None:
     lib.SetState.argtypes, lib.SetState.restype = [
         c_int,
     ], c_void_p
+    lib.InState.argtypes, lib.InState.restype = [
+        c_int,
+    ], BoolError
     lib.StartFailover.argtypes, lib.StartFailover.restype = [
         c_int,
         c_char_p,
diff --git a/wrappers/python/eduvpn_common/main.py b/wrappers/python/eduvpn_common/main.py
index 64fa6ac..f3639e0 100644
--- a/wrappers/python/eduvpn_common/main.py
+++ b/wrappers/python/eduvpn_common/main.py
@@ -3,7 +3,17 @@ from enum import IntEnum
 from typing import Any, Callable, Iterator, Optional
 
 from eduvpn_common.loader import initialize_functions, load_lib
-from eduvpn_common.types import ReadRxBytes, TokenGetter, TokenSetter, VPNStateChange, decode_res, encode_args
+from eduvpn_common.types import (
+    ReadRxBytes,
+    TokenGetter,
+    TokenSetter,
+    VPNStateChange,
+    decode_res,
+    encode_args,
+)
+
+from eduvpn_common.event import EventHandler
+from eduvpn_common.state import State
 
 
 class WrappedError(Exception):
@@ -57,14 +67,20 @@ class EduVPN(object):
         self.version = version
         self.config_directory = config_directory
         self.jar = Jar(lambda x: self.go_function(self.lib.CookieCancel, x))
-        self.callback = None
         self.token_setter = None
         self.token_getter = None
+        self.event_handler = EventHandler()
 
         # Load the library
         self.lib = load_lib()
         initialize_functions(self.lib)
 
+    def register_class_callbacks(self, _class):
+        self.event_handler.change_class_callbacks(_class, add=True)
+
+    def deregister_class_callbacks(self, _class):
+        self.event_handler.change_class_callbacks(_class, add=False)
+
     def go_cookie_function(self, func: Any, *args: Iterator) -> Any:
         cookie = self.lib.CookieNew()
         self.jar.add(cookie)
@@ -95,7 +111,7 @@ class EduVPN(object):
         global global_object
         global_object = None
 
-    def register(self, handler: Optional[Callable] = None, debug: bool = False) -> None:
+    def register(self, debug: bool = False) -> None:
         """Register the Go shared library.
         This makes sure the FSM is initialized and that we can call Go functions
 
@@ -106,7 +122,6 @@ class EduVPN(object):
         global global_object
         if global_object is not None:
             raise Exception("Already registered")
-        self.callback = handler
         global_object = self
         register_err = self.go_function(
             self.lib.Register,
@@ -175,11 +190,17 @@ class EduVPN(object):
         if remove_err:
             forwardError(remove_err)
 
-    def set_state(self, state: int):
+    def set_state(self, state: State):
         state_err = self.go_function(self.lib.SetState, state)
         if state_err:
             forwardError(state_err)
 
+    def in_state(self, state: State) -> bool:
+        yes, state_err = self.go_function(self.lib.InState, state)
+        if state_err:
+            forwardError(state_err)
+        return yes
+
     def get_config(
         self, _type: ServerType, identifier: str, prefer_tcp: bool = False
     ) -> str:
@@ -257,7 +278,9 @@ class EduVPN(object):
     def set_token_handler(self, getter: Callable, setter: Callable) -> None:
         self.token_setter = setter
         self.token_getter = getter
-        handler_err = self.go_function(self.lib.SetTokenHandler, token_getter, token_setter)
+        handler_err = self.go_function(
+            self.lib.SetTokenHandler, token_getter, token_setter
+        )
 
         if handler_err:
             forwardError(handler_err)
@@ -309,6 +332,7 @@ class EduVPN(object):
 
 global_object: Optional[EduVPN] = None
 
+
 @TokenSetter
 def token_setter(server: ctypes.c_char_p, tokens: ctypes.c_char_p):
     global global_object
@@ -318,6 +342,7 @@ def token_setter(server: ctypes.c_char_p, tokens: ctypes.c_char_p):
         return 0
     global_object.token_setter(server.decode(), tokens.decode())
 
+
 @TokenGetter
 def token_getter(server: ctypes.c_char_p, buf: ctypes.c_char_p, size: ctypes.c_size_t):
     global global_object
@@ -332,6 +357,7 @@ def token_getter(server: ctypes.c_char_p, buf: ctypes.c_char_p, size: ctypes.c_s
     outbuf = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char * size))
     outbuf.contents.value = got.encode("utf-8")
 
+
 @VPNStateChange
 def state_callback(old_state: int, new_state: int, data: str) -> int:
     """The internal callback that is passed to the Go library
@@ -345,9 +371,7 @@ def state_callback(old_state: int, new_state: int, data: str) -> int:
     global global_object
     if global_object is None:
         return 0
-    if global_object.callback is None:
-        return 0
-    handled = global_object.callback(old_state, new_state, data.decode("utf-8"))
+    handled = global_object.event_handler.run(State(old_state), State(new_state), data.decode("utf-8"))
     if handled:
         return 1
     return 0
diff --git a/wrappers/python/eduvpn_common/types.py b/wrappers/python/eduvpn_common/types.py
index f83e710..f1be42c 100644
--- a/wrappers/python/eduvpn_common/types.py
+++ b/wrappers/python/eduvpn_common/types.py
@@ -38,6 +38,7 @@ ReadRxBytes = CFUNCTYPE(c_ulonglong)
 TokenGetter = CFUNCTYPE(c_void_p, c_char_p, POINTER(c_char), c_size_t)
 TokenSetter = CFUNCTYPE(c_void_p, c_char_p, c_char_p)
 
+
 def encode_args(args: List[Any], types: List[Any]) -> Iterator[Any]:
     """Encode the arguments ready to be used by the Go library