From 059a222394a39556adb8247c7cbb3bbece03bf13 Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Sun, 28 Jul 2019 12:52:57 +0100
Subject: [PATCH] Use defer to unlock mutexes

---
 internal/client/state.go      | 3 +--
 internal/multiplex/session.go | 4 +---
 internal/multiplex/stream.go  | 8 ++------
 internal/server/activeuser.go | 9 +++------
 4 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/internal/client/state.go b/internal/client/state.go
index 9208180..f2d7e61 100644
--- a/internal/client/state.go
+++ b/internal/client/state.go
@@ -74,9 +74,9 @@ func InitState(localHost, localPort, remoteHost, remotePort string, nowFunc func
 
 func (sta *State) UpdateIntervalKeys() {
 	sta.intervalDataM.Lock()
+	defer sta.intervalDataM.Unlock()
 	currentInterval := sta.Now().Unix() / int64(sta.TicketTimeHint)
 	if currentInterval == sta.intervalData.interval {
-		sta.intervalDataM.Unlock()
 		return
 	}
 	sta.intervalData.interval = currentInterval
@@ -84,7 +84,6 @@ func (sta *State) UpdateIntervalKeys() {
 	intervalKey := ecdh.GenerateSharedSecret(ephPv, sta.staticPub)
 	seed := int64(binary.BigEndian.Uint64(ephPv.(*[32]byte)[0:8]))
 	sta.intervalData.ephPv, sta.intervalData.ephPub, sta.intervalData.intervalKey, sta.intervalData.seed = ephPv, ephPub, intervalKey, seed
-	sta.intervalDataM.Unlock()
 }
 
 func (sta *State) GetIntervalKeys() (crypto.PublicKey, []byte, int64) {
diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index f30717f..6ecc6f2 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -110,22 +110,20 @@ func (sesh *Session) getStream(id uint32, closingFrame bool) *Stream {
 	// it would have been neater to use defer Unlock(), however it gives
 	// non-negligable overhead and this function is performance critical
 	sesh.streamsM.Lock()
+	defer sesh.streamsM.Unlock()
 	stream := sesh.streams[id]
 	if stream != nil {
-		sesh.streamsM.Unlock()
 		return stream
 	} else {
 		if closingFrame {
 			// If the stream has been closed and the current frame is a closing frame,
 			// we return nil
-			sesh.streamsM.Unlock()
 			return nil
 		} else {
 			stream = makeStream(id, sesh)
 			sesh.streams[id] = stream
 			sesh.acceptCh <- stream
 			//log.Printf("Adding stream %v\n", id)
-			sesh.streamsM.Unlock()
 			return stream
 		}
 	}
diff --git a/internal/multiplex/stream.go b/internal/multiplex/stream.go
index 4620562..7b3c679 100644
--- a/internal/multiplex/stream.go
+++ b/internal/multiplex/stream.go
@@ -80,8 +80,8 @@ func (s *Stream) Write(in []byte) (n int, err error) {
 	// in the middle of the execution of Write. This may cause the closing frame
 	// to be sent before the data frame and cause loss of packet.
 	s.writingM.RLock()
+	defer s.writingM.RUnlock()
 	if s.isClosed() {
-		s.writingM.RUnlock()
 		return 0, ErrBrokenStream
 	}
 
@@ -94,12 +94,9 @@ func (s *Stream) Write(in []byte) (n int, err error) {
 
 	tlsRecord, err := s.session.obfs(f)
 	if err != nil {
-		s.writingM.RUnlock()
 		return 0, err
 	}
 	n, err = s.session.sb.send(tlsRecord)
-	s.writingM.RUnlock()
-
 	return
 
 }
@@ -122,8 +119,8 @@ func (s *Stream) passiveClose() {
 func (s *Stream) Close() error {
 
 	s.writingM.Lock()
+	defer s.writingM.Unlock()
 	if s.isClosed() {
-		s.writingM.Unlock()
 		return errors.New("Already Closed")
 	}
 
@@ -144,7 +141,6 @@ func (s *Stream) Close() error {
 	s._close()
 	s.session.delStream(s.id)
 	//log.Printf("%v actively closed\n", stream.id)
-	s.writingM.Unlock()
 	return nil
 }
 
diff --git a/internal/server/activeuser.go b/internal/server/activeuser.go
index 282f42f..f9bd405 100644
--- a/internal/server/activeuser.go
+++ b/internal/server/activeuser.go
@@ -32,18 +32,16 @@ func (u *ActiveUser) DelSession(sessionID uint32) {
 
 func (u *ActiveUser) GetSession(sessionID uint32, obfs mux.Obfser, deobfs mux.Deobfser, obfsedRead func(net.Conn, []byte) (int, error)) (sesh *mux.Session, existing bool, err error) {
 	u.sessionsM.Lock()
+	defer u.sessionsM.Unlock()
 	if sesh = u.sessions[sessionID]; sesh != nil {
-		u.sessionsM.Unlock()
 		return sesh, true, nil
 	} else {
 		err := u.panel.Manager.authoriseNewSession(u)
 		if err != nil {
-			u.sessionsM.Unlock()
 			return nil, false, err
 		}
 		sesh = mux.MakeSession(sessionID, u.valve, obfs, deobfs, obfsedRead)
 		u.sessions[sessionID] = sesh
-		u.sessionsM.Unlock()
 		return sesh, false, nil
 	}
 }
@@ -64,7 +62,6 @@ func (u *ActiveUser) Terminate(reason string) {
 
 func (u *ActiveUser) NumSession() int {
 	u.sessionsM.RLock()
-	l := len(u.sessions)
-	u.sessionsM.RUnlock()
-	return l
+	defer u.sessionsM.RUnlock()
+	return len(u.sessions)
 }