From c7592da63fcd2ff85cb8bc01d01e1af08e1ad59e Mon Sep 17 00:00:00 2001
From: Jianhui Zhao <zhaojh329@gmail.com>
Date: Mon, 16 Jun 2025 15:29:14 +0800
Subject: [PATCH] fix: Ensure the client is only close once

Signed-off-by: Jianhui Zhao <zhaojh329@gmail.com>
---
 device.go | 21 ++++++++++-----------
 user.go   | 21 +++++++++------------
 2 files changed, 19 insertions(+), 23 deletions(-)

diff --git a/device.go b/device.go
index b62b77f..2128f8d 100644
--- a/device.go
+++ b/device.go
@@ -12,7 +12,7 @@ import (
 	"os"
 	"rttys/client"
 	"strings"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	"github.com/rs/zerolog/log"
@@ -81,7 +81,8 @@ type device struct {
 	token      string
 	conn       net.Conn
 	registered bool
-	closed     uint32
+	closed     bool
+	close      sync.Once
 	err        byte
 	send       chan []byte // Buffered channel of outbound messages.
 }
@@ -158,7 +159,7 @@ func (dev *device) WriteMsg(typ int, data []byte) {
 }
 
 func (dev *device) Closed() bool {
-	return atomic.LoadUint32(&dev.closed) == 1
+	return dev.closed
 }
 
 func (dev *device) CloseConn() {
@@ -166,17 +167,15 @@ func (dev *device) CloseConn() {
 }
 
 func (dev *device) Close() {
-	if dev.Closed() {
-		return
-	}
+	dev.close.Do(func() {
+		dev.closed = true
 
-	atomic.StoreUint32(&dev.closed, 1)
+		log.Debug().Msgf("Device '%s' disconnected", dev.conn.RemoteAddr())
 
-	log.Debug().Msgf("Device '%s' disconnected", dev.conn.RemoteAddr())
+		dev.CloseConn()
 
-	dev.CloseConn()
-
-	close(dev.send)
+		close(dev.send)
+	})
 }
 
 func parseDeviceInfo(dev *device, b []byte) bool {
diff --git a/user.go b/user.go
index f1e3501..eb7febe 100644
--- a/user.go
+++ b/user.go
@@ -3,7 +3,7 @@ package main
 import (
 	"fmt"
 	"net/http"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	"rttys/client"
@@ -24,7 +24,8 @@ type user struct {
 	sid    string
 	devid  string
 	conn   *websocket.Conn
-	closed uint32
+	closed bool
+	close  sync.Once
 	send   chan *usrMessage // Buffered channel of outbound messages.
 }
 
@@ -56,7 +57,7 @@ func (u *user) WriteMsg(typ int, data []byte) {
 }
 
 func (u *user) Closed() bool {
-	return atomic.LoadUint32(&u.closed) == 1
+	return u.closed
 }
 
 func (u *user) CloseConn() {
@@ -64,15 +65,11 @@ func (u *user) CloseConn() {
 }
 
 func (u *user) Close() {
-	if u.Closed() {
-		return
-	}
-
-	atomic.StoreUint32(&u.closed, 1)
-
-	u.CloseConn()
-
-	close(u.send)
+	u.close.Do(func() {
+		u.closed = true
+		u.CloseConn()
+		close(u.send)
+	})
 }
 
 func userLoginAck(code int, c client.Client) {