From fa2cbd14b1fd763c564e2479447e1febebc262ba Mon Sep 17 00:00:00 2001 From: Jianhui Zhao Date: Thu, 9 Sep 2021 08:14:39 +0000 Subject: [PATCH] refactor: change the login process to improve stability Signed-off-by: Jianhui Zhao --- broker.go | 141 ++++++++++++++++++++++-------------------- device.go | 52 +++++++--------- ui/src/views/Rtty.vue | 26 ++++---- user.go | 33 +++++----- 4 files changed, 121 insertions(+), 131 deletions(-) diff --git a/broker.go b/broker.go index 6b218c3..2b5f6ea 100644 --- a/broker.go +++ b/broker.go @@ -3,10 +3,12 @@ package main import ( "crypto/x509" "encoding/binary" + "sync/atomic" "time" "rttys/client" "rttys/config" + "rttys/utils" "github.com/gorilla/websocket" jsoniter "github.com/json-iterator/go" @@ -14,47 +16,45 @@ import ( ) type session struct { - devid string - devsid byte - u client.Client + dev client.Client + user client.Client + confirmed uint32 } type broker struct { - cfg *config.Config - devices map[string]client.Client - loginAck chan *loginAckMsg - logout chan string - register chan client.Client - unregister chan client.Client - waitLoginUsers map[string]client.Client - sessions map[string]*session - cmdReq chan *commandReq - webCon chan *webNewCon - webReq chan *webReq - termMessage chan *termMessage - userMessage chan *usrMessage - cmdMessage chan []byte - webMessage chan *webResp - devCertPool *x509.CertPool + cfg *config.Config + devices map[string]client.Client + loginAck chan *loginAckMsg + logout chan string + register chan client.Client + unregister chan client.Client + sessions map[string]*session + cmdReq chan *commandReq + webCon chan *webNewCon + webReq chan *webReq + termMessage chan *termMessage + userMessage chan *usrMessage + cmdMessage chan []byte + webMessage chan *webResp + devCertPool *x509.CertPool } func newBroker(cfg *config.Config) *broker { return &broker{ - cfg: cfg, - loginAck: make(chan *loginAckMsg, 1000), - logout: make(chan string, 1000), - register: make(chan client.Client, 1000), - unregister: make(chan client.Client, 1000), - devices: make(map[string]client.Client), - waitLoginUsers: make(map[string]client.Client), - sessions: make(map[string]*session), - cmdReq: make(chan *commandReq, 1000), - webCon: make(chan *webNewCon, 1000), - webReq: make(chan *webReq, 1000), - termMessage: make(chan *termMessage, 1000), - userMessage: make(chan *usrMessage, 1000), - cmdMessage: make(chan []byte, 1000), - webMessage: make(chan *webResp, 1000), + cfg: cfg, + loginAck: make(chan *loginAckMsg, 1000), + logout: make(chan string, 1000), + register: make(chan client.Client, 1000), + unregister: make(chan client.Client, 1000), + devices: make(map[string]client.Client), + sessions: make(map[string]*session), + cmdReq: make(chan *commandReq, 1000), + webCon: make(chan *webNewCon, 1000), + webReq: make(chan *webReq, 1000), + termMessage: make(chan *termMessage, 1000), + userMessage: make(chan *usrMessage, 1000), + cmdMessage: make(chan []byte, 1000), + webMessage: make(chan *webResp, 1000), } } @@ -87,15 +87,23 @@ func (br *broker) run() { c.WriteMsg(msgTypeRegister, append([]byte{err}, msg...)) } else { if dev, ok := br.devices[devid]; ok { - if _, ok := br.waitLoginUsers[devid]; ok { - log.Error().Msg("Another user is logining the device, wait...") - time.AfterFunc(time.Millisecond*10, func() { - br.register <- c - }) - } else { - br.waitLoginUsers[devid] = c - dev.WriteMsg(msgTypeLogin, []byte{}) + sid := utils.GenUniqueID("sid") + + s := &session{ + dev: dev, + user: c, } + + time.AfterFunc(time.Second*3, func() { + if atomic.LoadUint32(&s.confirmed) == 0 { + c.Close() + } + }) + + br.sessions[sid] = s + + dev.WriteMsg(msgTypeLogin, []byte(sid)) + log.Info().Msg("New session: " + sid) } else { userLoginAck(loginErrorOffline, c) log.Error().Msgf("Not found the device '%s'", devid) @@ -103,29 +111,29 @@ func (br *broker) run() { } case c := <-br.unregister: - id := c.DeviceID() + devid := c.DeviceID() if c.IsDevice() { - delete(br.devices, id) + delete(br.devices, devid) for sid, s := range br.sessions { - if s.devid == id { - s.u.Close() + if s.dev == c { + s.user.Close() delete(br.sessions, sid) log.Info().Msg("Delete session: " + sid) } } - log.Info().Msgf("Device '%s' unregistered", id) + log.Info().Msgf("Device '%s' unregistered", devid) } else { sid := c.(*user).sid - if s, ok := br.sessions[sid]; ok { + if _, ok := br.sessions[sid]; ok { delete(br.sessions, sid) c.Close() - if dev, ok := br.devices[s.devid]; ok { - dev.WriteMsg(msgTypeLogout, []byte{sid[len(sid)-1] - '0'}) + if dev, ok := br.devices[devid]; ok { + dev.WriteMsg(msgTypeLogout, []byte(sid)) } log.Info().Msg("Delete session: " + sid) @@ -133,22 +141,18 @@ func (br *broker) run() { } case msg := <-br.loginAck: - if c, ok := br.waitLoginUsers[msg.devid]; ok { + if s, ok := br.sessions[msg.sid]; ok { if msg.isBusy { - userLoginAck(loginErrorBusy, c) + userLoginAck(loginErrorBusy, s.user) log.Error().Msg("login fail, device busy") } else { - sid := msg.devid + string(msg.sid+'0') - br.sessions[sid] = &session{msg.devid, msg.sid, c} + atomic.StoreUint32(&s.confirmed, 1) - u := c.(*user) - u.sid = sid + u := s.user.(*user) + u.sid = msg.sid - userLoginAck(loginErrorNone, c) - - log.Info().Msg("New session: " + sid) + userLoginAck(loginErrorNone, s.user) } - delete(br.waitLoginUsers, msg.devid) } // device active logout @@ -156,7 +160,7 @@ func (br *broker) run() { case sid := <-br.logout: if s, ok := br.sessions[sid]; ok { delete(br.sessions, sid) - s.u.Close() + s.user.Close() log.Info().Msg("Delete session: " + sid) } @@ -164,33 +168,34 @@ func (br *broker) run() { // from device, includes terminal data and file data case msg := <-br.termMessage: if s, ok := br.sessions[msg.sid]; ok { - s.u.WriteMsg(websocket.BinaryMessage, msg.data) + s.user.WriteMsg(websocket.BinaryMessage, msg.data) } case msg := <-br.userMessage: if s, ok := br.sessions[msg.sid]; ok { - if dev, ok := br.devices[s.devid]; ok { - devsid := msg.sid[len(msg.sid)-1] - '0' + if dev, ok := br.devices[s.dev.DeviceID()]; ok { data := msg.data if msg.typ == websocket.BinaryMessage { if data[0] == 1 { dev.WriteMsg(msgTypeFile, data[1:]) } else { - dev.WriteMsg(msgTypeTermData, append([]byte{devsid}, data[1:]...)) + dev.WriteMsg(msgTypeTermData, append([]byte(msg.sid), data[1:]...)) } } else { typ := jsoniter.Get(data, "type").ToString() switch typ { case "winsize": - b := [5]byte{devsid} + b := [32 + 4]byte{} + + copy(b[:], msg.sid) cols := jsoniter.Get(data, "cols").ToUint() rows := jsoniter.Get(data, "rows").ToUint() - binary.BigEndian.PutUint16(b[1:], uint16(cols)) - binary.BigEndian.PutUint16(b[3:], uint16(rows)) + binary.BigEndian.PutUint16(b[32:], uint16(cols)) + binary.BigEndian.PutUint16(b[34:], uint16(rows)) dev.WriteMsg(msgTypeWinsize, b[:]) } diff --git a/device.go b/device.go index c80f30a..f9a0cb9 100644 --- a/device.go +++ b/device.go @@ -12,7 +12,7 @@ import ( "io/ioutil" "net" "strings" - "sync" + "sync/atomic" "time" "github.com/rs/zerolog/log" @@ -43,8 +43,7 @@ type device struct { conn net.Conn active time.Time registered bool - closeMutex sync.Mutex - closed bool + closed uint32 cancel context.CancelFunc send chan []byte // Buffered channel of outbound messages. } @@ -56,7 +55,7 @@ type termMessage struct { type loginAckMsg struct { devid string - sid byte + sid string isBusy bool } @@ -77,22 +76,19 @@ func (dev *device) WriteMsg(typ int, data []byte) { } func (dev *device) Close() { - defer dev.closeMutex.Unlock() + if atomic.LoadUint32(&dev.closed) == 1 { + return + } + atomic.StoreUint32(&dev.closed, 1) - dev.closeMutex.Lock() + log.Debug().Msgf("Device '%s' disconnected", dev.conn.RemoteAddr()) - if !dev.closed { - log.Debug().Msgf("Device '%s' disconnected", dev.conn.RemoteAddr()) + dev.conn.Close() - dev.closed = true + dev.cancel() - dev.conn.Close() - - dev.cancel() - - if dev.registered { - dev.br.unregister <- dev - } + if dev.registered { + dev.br.unregister <- dev } } @@ -224,41 +220,35 @@ func (dev *device) readLoop() { dev.br.register <- dev case msgTypeLogin: - if msgLen < 1 { + if msgLen < 33 { log.Error().Msg("msgTypeLogin: invalid") return } - code := b[0] - sid := byte(0) + sid := string(b[:32]) + code := b[32] - if code == 0 { - if msgLen < 2 { - log.Error().Msg("msgTypeLogin: invalid") - return - } - - sid = b[1] - } dev.br.loginAck <- &loginAckMsg{dev.id, sid, code == 1} case msgTypeLogout: - if msgLen < 1 { + if msgLen < 32 { log.Error().Msg("msgTypeLogout: invalid") return } - dev.br.logout <- dev.id + string(b[0]+'0') + dev.br.logout <- string(b[:32]) case msgTypeTermData: fallthrough case msgTypeFile: - if msgLen < 1 { + if msgLen < 32 { log.Error().Msg("msgTypeTermData|msgTypeFile: invalid") return } - sid := dev.id + string(b[0]+'0') + sid := string(b[:32]) + + b = b[31:] if typ == msgTypeFile { b[0] = 1 diff --git a/ui/src/views/Rtty.vue b/ui/src/views/Rtty.vue index 4fe418b..145b2f6 100644 --- a/ui/src/views/Rtty.vue +++ b/ui/src/views/Rtty.vue @@ -272,23 +272,11 @@ const overlayAddon = new OverlayAddon(); term.loadAddon(overlayAddon); - term.open(this.$refs['terminal'] as HTMLElement); - term.focus(); - - window.addEventListener('resize', this.fitTerm); - const socket = new WebSocket(protocol + location.host + `/connect/${this.devid}`); this.disposables.push({dispose: () => socket.close()}); socket.binaryType = 'arraybuffer'; this.socket = socket; - socket.addEventListener('open', () => { - this.axios.get('/fontsize').then(r => { - this.term?.setOption('fontSize', r.data.size); - this.fitTerm(); - }); - }); - socket.addEventListener('close', () => this.dispose()); socket.addEventListener('error', () => this.dispose()); @@ -300,13 +288,23 @@ if (msg.type === 'login') { if (msg.err === LoginErrorOffline) { this.$message.error(this.$t('Device offline').toString()); - this.dispose(); + this.$router.push('/'); return; } else if (msg.err === LoginErrorBusy) { this.$message.error(this.$t('Sessions is full').toString()); - this.dispose(); + this.$router.push('/'); return; } + + window.addEventListener('resize', this.fitTerm); + + term.open(this.$refs['terminal'] as HTMLElement); + term.focus(); + + this.axios.get('/fontsize').then(r => { + this.term?.setOption('fontSize', r.data.size); + this.fitTerm(); + }); } else if (msg.type === 'logout') { this.dispose(); } diff --git a/user.go b/user.go index 0dc295d..e6e7bc3 100644 --- a/user.go +++ b/user.go @@ -3,7 +3,7 @@ package main import ( "context" "fmt" - "sync" + "sync/atomic" "time" "rttys/client" @@ -22,14 +22,13 @@ const ( ) type user struct { - br *broker - sid string - devid string - conn *websocket.Conn - closeMutex sync.Mutex - closed bool - cancel context.CancelFunc - send chan *usrMessage // Buffered channel of outbound messages. + br *broker + sid string + devid string + conn *websocket.Conn + closed uint32 + cancel context.CancelFunc + send chan *usrMessage // Buffered channel of outbound messages. } type usrMessage struct { @@ -60,16 +59,14 @@ func (u *user) WriteMsg(typ int, data []byte) { } func (u *user) Close() { - defer u.closeMutex.Unlock() - - u.closeMutex.Lock() - - if !u.closed { - u.closed = true - u.cancel() - u.conn.Close() - u.br.unregister <- u + if atomic.LoadUint32(&u.closed) == 1 { + return } + atomic.StoreUint32(&u.closed, 1) + + u.cancel() + u.conn.Close() + u.br.unregister <- u } func userLoginAck(code int, c client.Client) {