5 files changed, 33 insertions, 5 deletions
diff --git a/internal/client/ui/auth/auth.go b/internal/client/ui/auth/auth.go
index 42ea3e5..3f89298 100644
--- a/internal/client/ui/auth/auth.go
+++ b/internal/client/ui/auth/auth.go
@@ -80,14 +80,22 @@ ___] | |__] | \|    | | \|`)
 	centerStyle = lipgloss.NewStyle().Width(authWidth).AlignHorizontal(0.5)
 )
 
+func init() {
+	// HACK: to avoid a circular dependency, so core can transition to this
+	// I don't like how go has this issue, I would rather slower compilations
+	ui.NewAuth = func() tea.Model {
+		return New()
+	}
+}
+
 type Model struct {
-	focusIndex int
+	popup *choicepopup.Model
+
 	fields     []authfield.Model
+	focusIndex int
 	remember   bool
 
 	signup bool
-
-	popup *choicepopup.Model
 }
 
 func New() Model {
diff --git a/internal/client/ui/core/core.go b/internal/client/ui/core/core.go
index 50abbe6..1e2bc6a 100644
--- a/internal/client/ui/core/core.go
+++ b/internal/client/ui/core/core.go
@@ -179,6 +179,12 @@ func (m *Model) updateConnected(msg tea.Msg) tea.Cmd {
 		m.timeout = initialTimeout
 		return tea.Batch(gateway.Connect(m.privKey, connectionTimeout), m.loading.Init())
 
+	case *packet.Error:
+		err := "new connection from another location, closing this one"
+		if msg.PktType == packet.PacketError && err == msg.Error {
+			return ui.Transition(ui.NewAuth())
+		}
+
 	case *packet.NetworksInfo:
 		if msg.Set {
 			state.State.Networks = msg.Networks
diff --git a/internal/client/ui/ui.go b/internal/client/ui/ui.go
index 3ee0124..ab8b273 100644
--- a/internal/client/ui/ui.go
+++ b/internal/client/ui/ui.go
@@ -23,6 +23,8 @@ var Width int
 var Height int
 var Program *tea.Program
 
+var NewAuth func() tea.Model
+
 type ModelTransition struct {
 	Model tea.Model
 }
diff --git a/internal/server/server.go b/internal/server/server.go
index 495b2a5..f510821 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -196,7 +196,10 @@ func (server *server) handleConnection(conn net.Conn) {
 	}()
 	defer func() {
 		conn.Close()
-		server.RemoveSession(sess.ID())
+		sameAddress := addr.String() == server.Session(sess.ID()).Addr().String()
+		if sameAddress {
+			server.RemoveSession(sess.ID())
+		}
 		log.Println(addr, "disconnected")
 	}()
 
diff --git a/internal/server/session/session.go b/internal/server/session/session.go
index 88a19e7..a904d39 100644
--- a/internal/server/session/session.go
+++ b/internal/server/session/session.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/ed25519"
 	"crypto/rand"
+	"log"
 	"net"
 	"sync"
 	"time"
@@ -99,7 +100,7 @@ func (s *Session) Read(ctx context.Context) (packet.Packet, bool) {
 }
 
 func (s *Session) Close() {
-	timeout := 1 * time.Second
+	timeout := 10 * time.Millisecond
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	payload := &packet.Error{
 		Error:   "new connection from another location, closing this one",
@@ -109,5 +110,13 @@ func (s *Session) Close() {
 	s.Write(ctx, pkt)
 	cancel()
 
+	// Add some delay before canceling to let the writer enough time to
+	// actually write that into the connection
+	// HACK: consider just giving session the private connection and to
+	// write directly so we don't have to wait
+	time.Sleep(100 * time.Millisecond)
+
+	log.Println(s.addr, "closed due to new connection from another location")
+
 	s.cancel()
 }