4 files changed, 60 insertions, 3 deletions
diff --git a/internal/server/common.go b/internal/server/common.go
index 6f57c7f..fcba07f 100644
--- a/internal/server/common.go
+++ b/internal/server/common.go
@@ -5,9 +5,9 @@ import (
 	"time"
 
 	"github.com/eduvpn/eduvpn-common/internal/oauth"
-	"github.com/eduvpn/eduvpn-common/types"
 	"github.com/eduvpn/eduvpn-common/internal/util"
 	"github.com/eduvpn/eduvpn-common/internal/wireguard"
+	"github.com/eduvpn/eduvpn-common/types"
 )
 
 // The base type for servers
@@ -155,8 +155,6 @@ func (servers *Servers) addInstituteAndCustom(
 		server = &InstituteAccessServer{}
 	}
 
-	// Set the current server
-	toAddServers.CurrentURL = url
 	instituteInitErr := server.init(
 		url,
 		discoServer.DisplayName,
diff --git a/internal/server/custom.go b/internal/server/custom.go
index 52a0094..4d3c91c 100644
--- a/internal/server/custom.go
+++ b/internal/server/custom.go
@@ -1,10 +1,31 @@
 package server
 
 import (
+	"errors"
 	"fmt"
 	"github.com/eduvpn/eduvpn-common/types"
 )
 
+func (servers *Servers) SetCustomServer(server Server) error {
+	errorMessage := "failed setting custom server"
+	base, baseErr := server.GetBase()
+	if baseErr != nil {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+	}
+
+	if base.Type != "custom_server" {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: errors.New("Not a custom server")}
+	}
+
+	if _, ok := servers.CustomServers.Map[base.URL]; ok {
+		servers.CustomServers.CurrentURL = base.URL
+		servers.IsType = CustomServerType
+	} else {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: errors.New("Not a custom server")}
+	}
+	return nil
+}
+
 func (servers *Servers) GetCustomServer(url string) (*InstituteAccessServer, error) {
 	if server, ok := servers.CustomServers.Map[url]; ok {
 		return server, nil
diff --git a/internal/server/instituteaccess.go b/internal/server/instituteaccess.go
index f2669b8..bf0e2bc 100644
--- a/internal/server/instituteaccess.go
+++ b/internal/server/instituteaccess.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/eduvpn/eduvpn-common/internal/oauth"
@@ -21,6 +22,26 @@ type InstituteAccessServers struct {
 	CurrentURL string                            `json:"current_url"`
 }
 
+func (servers *Servers) SetInstituteAccess(server Server) error {
+	errorMessage := "failed setting institute access server"
+	base, baseErr := server.GetBase()
+	if baseErr != nil {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+	}
+
+	if base.Type != "institute_access" {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: errors.New("Not an institute access server")}
+	}
+
+	if _, ok := servers.InstituteServers.Map[base.URL]; ok {
+		servers.InstituteServers.CurrentURL = base.URL
+		servers.IsType = InstituteAccessServerType
+	} else {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: errors.New("No such institute access server")}
+	}
+	return nil
+}
+
 func (servers *Servers) GetInstituteAccess(url string) (*InstituteAccessServer, error) {
 	if server, ok := servers.InstituteServers.Map[url]; ok {
 		return server, nil
diff --git a/internal/server/secureinternet.go b/internal/server/secureinternet.go
index 776bb72..27d48a5 100644
--- a/internal/server/secureinternet.go
+++ b/internal/server/secureinternet.go
@@ -31,6 +31,23 @@ func (servers *Servers) GetSecureInternetHomeServer() (*SecureInternetHomeServer
 	return &servers.SecureInternetHomeServer, nil
 }
 
+func (servers *Servers) SetSecureInternet(server Server) error {
+	errorMessage := "failed setting secure internet server"
+	base, baseErr := server.GetBase()
+	if baseErr != nil {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+	}
+
+	if base.Type != "secure_internet" {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: errors.New("Not a secure internet server")}
+	}
+
+	// The location should already be configured
+	// TODO: check for location?
+	servers.IsType = SecureInternetServerType
+	return nil
+}
+
 func (servers *Servers) RemoveSecureInternet() {
 	// Empty out the struct
 	servers.SecureInternetHomeServer = SecureInternetHomeServer{}