fix: panic: send on closed channel

Signed-off-by: Jianhui Zhao <zhaojh329@gmail.com>

broker.go

@@ -95,6 +95,13 @@ func (br *broker) run() {
 				}
 
 				c.WriteMsg(msgTypeRegister, append([]byte{err}, msg...))
+
+				if err > 0 {
+					// ensure the last packet was sent
+					time.AfterFunc(time.Millisecond*100, func() {
+						dev.Close()
+					})
+				}
 			} else {
 				if dev, ok := br.devices[devid]; ok {
 					sid := utils.GenUniqueID("sid")
@@ -125,19 +132,11 @@ func (br *broker) run() {
 
 			c.Close()
 
-			if !c.IsDevice() {
-				sid := c.(*user).sid
-
-				if _, ok := br.sessions[sid]; ok {
-					delete(br.sessions, sid)
-
-					if dev, ok := br.devices[devid]; ok {
-						dev.WriteMsg(msgTypeLogout, []byte(sid))
+			if c.IsDevice() {
+				if !c.(*device).registered {
+					break
 				}
 
-					log.Info().Msg("Delete session: " + sid)
-				}
-			} else if c.(*device).registered {
 				delete(br.devices, devid)
 
 				for sid, s := range br.sessions {
@@ -149,6 +148,18 @@ func (br *broker) run() {
 				}
 
 				log.Info().Msgf("Device '%s' unregistered", devid)
+			} else {
+				sid := c.(*user).sid
+
+				if _, ok := br.sessions[sid]; ok {
+					delete(br.sessions, sid)
+
+					if dev, ok := br.devices[devid]; ok {
+						dev.WriteMsg(msgTypeLogout, []byte(sid))
+					}
+
+					log.Info().Msg("Delete session: " + sid)
+				}
 			}
 
 		case msg := <-br.loginAck:

device.go

@@ -191,7 +191,6 @@ func parseHeartbeat(dev *device, b []byte) {
 func (dev *device) readLoop() {
 	defer func() {
 		dev.br.unregister <- dev
-		dev.conn.Close()
 	}()
 
 	br := bufio.NewReader(dev.conn)
@@ -303,7 +302,7 @@ func (dev *device) writeLoop() {
 
 	defer func() {
 		ticker.Stop()
-		dev.conn.Close()
+		dev.br.unregister <- dev
 	}()
 
 	ninactive := 0
@@ -325,15 +324,13 @@ func (dev *device) writeLoop() {
 		case <-ticker.C:
 			now := time.Now()
 			if now.Sub(dev.active) > heartbeatInterval*3/2 {
-				if !dev.registered {
-					dev.Close()
+				if dev.id == "" {
 					return
 				}
 
 				log.Error().Msgf("Inactive device in long time: %s", dev.id)
 				if ninactive > 1 {
 					log.Error().Msgf("Inactive 3 times, now kill it: %s", dev.id)
-					dev.Close()
 					return
 				}
 				ninactive = ninactive + 1