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

2 files changed, 14 insertions, 2 deletions

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()
 }