From 63790c069fbca0c3548fdaa58148e6cefdf477d5 Mon Sep 17 00:00:00 2001
From: Kyren223 <Kyren223@proton.me>
Date: Mon, 6 Jan 2025 18:23:41 +0200
Subject: Server and client now properly handle the same user logging in from
 different locations, fixed the issue where it'd only work for the first time
 by checking if the addresses match before remove the session

---
 internal/client/ui/auth/auth.go    | 14 +++++++++++---
 internal/client/ui/core/core.go    |  6 ++++++
 internal/client/ui/ui.go           |  2 ++
 internal/server/server.go          |  5 ++++-
 internal/server/session/session.go | 11 ++++++++++-
 5 files changed, 33 insertions(+), 5 deletions(-)

(limited to 'internal')

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()
 }
-- 
cgit v1.3.1