diff options
Diffstat (limited to 'internal/server')
| -rw-r--r-- | internal/server/server.go | 5 | ||||
| -rw-r--r-- | internal/server/session/session.go | 11 |
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() } |