kirillius
/
Cloak
mirror of https://github.com/cbeuw/Cloak
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use Compare-And-Swap for atomic booleans indicating session and switchboard closed

This commit is contained in:
Andy Wang 2020-12-25 23:16:57 +00:00
parent 4f34e69006
commit 2f17841f85
No known key found for this signature in database
GPG Key ID: 181B49F9F38F3374
2 changed files with 21 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
internal/multiplex/session.go
View File

@ -66,6 +66,8 @@ type Session struct {
streamsM sync.Mutex
streams map[uint32]*Stream
// For accepting new streams
acceptCh chan *Stream
// a pool of heap allocated frame objects so we don't have to allocate a new one each time we receive a frame
recvFramePool sync.Pool
@ -78,9 +80,6 @@ type Session struct {
// Used for LocalAddr() and RemoteAddr() etc.
addrs atomic.Value
// For accepting new streams
acceptCh chan *Stream
closed uint32
terminalMsg atomic.Value
@ -181,7 +180,7 @@ func (sesh *Session) Accept() (net.Conn, error) {
}
func (sesh *Session) closeStream(s *Stream, active bool) error {
if atomic.SwapUint32(&s.closed, 1) == 1 {
if !atomic.CompareAndSwapUint32(&s.closed, 0, 1) {
return fmt.Errorf("closing stream %v: %w", s.id, errRepeatStreamClosing)
}
_ = s.getRecvBuf().Close() // recvBuf.Close should not return error
@ -244,6 +243,10 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
}
sesh.streamsM.Lock()
if sesh.IsClosed() {
sesh.streamsM.Unlock()
return ErrBrokenSession
}
existingStream, existing := sesh.streams[frame.StreamID]
if existing {
sesh.streamsM.Unlock()
@ -255,10 +258,10 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
} else {
newStream := makeStream(sesh, frame.StreamID)
sesh.streams[frame.StreamID] = newStream
sesh.acceptCh <- newStream
sesh.streamsM.Unlock()
// new stream
sesh.streamCountIncr()
sesh.acceptCh <- newStream
return newStream.recvFrame(frame)
}
}
@ -276,14 +279,14 @@ func (sesh *Session) TerminalMsg() string {
}
}
func (sesh *Session) closeSession(closeSwitchboard bool) error {
if atomic.SwapUint32(&sesh.closed, 1) == 1 {
func (sesh *Session) closeSession() error {
if !atomic.CompareAndSwapUint32(&sesh.closed, 0, 1) {
log.Debugf("session %v has already been closed", sesh.id)
return errRepeatSessionClosing
}
sesh.acceptCh <- nil
sesh.streamsM.Lock()
close(sesh.acceptCh)
for id, stream := range sesh.streams {
if stream == nil {
continue
@ -294,26 +297,23 @@ func (sesh *Session) closeSession(closeSwitchboard bool) error {
sesh.streamCountDecr()
}
sesh.streamsM.Unlock()
if closeSwitchboard {
sesh.sb.closeAll()
}
return nil
}
func (sesh *Session) passiveClose() error {
log.Debugf("attempting to passively close session %v", sesh.id)
err := sesh.closeSession(true)
err := sesh.closeSession()
if err != nil {
return err
}
sesh.sb.closeAll()
log.Debugf("session %v closed gracefully", sesh.id)
return nil
}
func (sesh *Session) Close() error {
log.Debugf("attempting to actively close session %v", sesh.id)
err := sesh.closeSession(false)
err := sesh.closeSession()
if err != nil {
return err
}
@ -339,7 +339,6 @@ func (sesh *Session) Close() error {
if err != nil {
return err
}
sesh.sb.closeAll()
log.Debugf("session %v closed gracefully", sesh.id)
return nil

17
internal/multiplex/switchboard.go
View File

@ -71,7 +71,8 @@ func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
n, err = conn.Write(d)
if err != nil {
sb.conns.Delete(*connId)
sb.close("failed to write to remote " + err.Error())
sb.session.SetTerminalMsg("failed to write to remote " + err.Error())
sb.session.passiveClose()
return n, err
}
sb.valve.AddTx(int64(n))
@ -138,16 +139,11 @@ func (sb *switchboard) pickRandConn() (uint32, net.Conn, error) {
return id, conn, nil
}
func (sb *switchboard) close(terminalMsg string) {
atomic.StoreUint32(&sb.broken, 1)
if !sb.session.IsClosed() {
sb.session.SetTerminalMsg(terminalMsg)
sb.session.passiveClose()
}
}
// actively triggered by session.Close()
func (sb *switchboard) closeAll() {
if !atomic.CompareAndSwapUint32(&sb.broken, 0, 1) {
return
}
sb.conns.Range(func(key, connI interface{}) bool {
conn := connI.(net.Conn)
conn.Close()
@ -168,7 +164,8 @@ func (sb *switchboard) deplex(connId uint32, conn net.Conn) {
log.Debugf("a connection for session %v has closed: %v", sb.session.id, err)
sb.conns.Delete(connId)
atomic.AddUint32(&sb.numConns, ^uint32(0))
sb.close("a connection has dropped unexpectedly")
sb.session.SetTerminalMsg("a connection has dropped unexpectedly")
sb.session.passiveClose()
return
}