From b2ae8082337e1b45e03bd6ee1cbad31a7470e924 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Mon, 23 Aug 2021 18:41:19 +0200
Subject: [PATCH] add merciful reconnect option.

---
 internal/config/session.go    | 11 +++++++++--
 internal/session/manager.go   |  4 ++++
 internal/types/session.go     |  1 +
 internal/websocket/manager.go | 33 +++++++++++++++++++++------------
 4 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/internal/config/session.go b/internal/config/session.go
index a990b3cb..1cff4b32 100644
--- a/internal/config/session.go
+++ b/internal/config/session.go
@@ -8,8 +8,9 @@ import (
 )
 
 type Session struct {
-	ImplicitHosting bool
-	APIToken        string
+	ImplicitHosting   bool
+	MercifulReconnect bool
+	APIToken          string
 
 	CookieEnabled    bool
 	CookieName       string
@@ -23,6 +24,11 @@ func (Session) Init(cmd *cobra.Command) error {
 		return err
 	}
 
+	cmd.PersistentFlags().Bool("session.merciful_reconnect", true, "allow reconnecting to websocket even if previous connection was not closed")
+	if err := viper.BindPFlag("session.merciful_reconnect", cmd.PersistentFlags().Lookup("session.merciful_reconnect")); err != nil {
+		return err
+	}
+
 	cmd.PersistentFlags().String("session.api_token", "", "API token for interacting with external services")
 	if err := viper.BindPFlag("session.api_token", cmd.PersistentFlags().Lookup("session.api_token")); err != nil {
 		return err
@@ -54,6 +60,7 @@ func (Session) Init(cmd *cobra.Command) error {
 
 func (s *Session) Set() {
 	s.ImplicitHosting = viper.GetBool("session.implicit_hosting")
+	s.MercifulReconnect = viper.GetBool("session.merciful_reconnect")
 	s.APIToken = viper.GetString("session.api_token")
 
 	s.CookieEnabled = viper.GetBool("session.cookie.enabled")
diff --git a/internal/session/manager.go b/internal/session/manager.go
index 50173467..1ec7f9f3 100644
--- a/internal/session/manager.go
+++ b/internal/session/manager.go
@@ -304,3 +304,7 @@ func (manager *SessionManagerCtx) ImplicitHosting() bool {
 func (manager *SessionManagerCtx) CookieEnabled() bool {
 	return manager.config.CookieEnabled
 }
+
+func (manager *SessionManagerCtx) MercifulReconnect() bool {
+	return manager.config.MercifulReconnect
+}
diff --git a/internal/types/session.go b/internal/types/session.go
index 1eb1f17e..6df70d87 100644
--- a/internal/types/session.go
+++ b/internal/types/session.go
@@ -50,6 +50,7 @@ type SessionManager interface {
 
 	ImplicitHosting() bool
 	CookieEnabled() bool
+	MercifulReconnect() bool
 
 	CookieSetToken(w http.ResponseWriter, token string)
 	CookieClearToken(w http.ResponseWriter, r *http.Request)
diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index 22489928..cbb90fbb 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -189,22 +189,31 @@ func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Reque
 	}
 
 	if session.State().IsConnected {
-		manager.logger.Debug().Str("session_id", session.ID()).Msg("already connected")
+		manager.logger.Warn().Str("session_id", session.ID()).Msg("already connected")
 
-		// TODO: Refactor, return error code.
-		if err = connection.WriteJSON(
-			message.SystemDisconnect{
-				Event:   event.SYSTEM_DISCONNECT,
-				Message: "already connected",
-			}); err != nil {
-			manager.logger.Error().Err(err).Msg("failed to send disconnect event")
+		if !manager.sessions.MercifulReconnect() {
+			// TODO: Refactor, return error code.
+			if err = connection.WriteJSON(
+				message.SystemDisconnect{
+					Event:   event.SYSTEM_DISCONNECT,
+					Message: "already connected",
+				}); err != nil {
+				manager.logger.Error().Err(err).Msg("failed to send disconnect event")
+			}
+
+			if err := connection.Close(); err != nil {
+				manager.logger.Warn().Err(err).Msg("connection closed with an error")
+			}
+
+			return
 		}
 
-		if err := connection.Close(); err != nil {
-			manager.logger.Warn().Err(err).Msg("connection closed with an error")
-		}
+		manager.logger.Debug().Str("session_id", session.ID()).Msg("replacing peer connection")
 
-		return
+		// replace peer connection
+		if err := session.GetWebSocketPeer().Destroy(); err != nil {
+			manager.logger.Warn().Err(err).Msg("previous connection closed with an error")
+		}
 	}
 
 	peer := &WebSocketPeerCtx{