Fix potential deadlocks

internal/multiplex/session.go

@@ -56,6 +56,7 @@ type Session struct {
 	// Switchboard manages all connections to remote
 	sb *switchboard
 
+	// Used for LocalAddr() and RemoteAddr() etc.
 	addrs atomic.Value
 
 	// For accepting new streams
@@ -150,8 +151,8 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
 	}
 
 	sesh.streamsM.Lock()
-	defer sesh.streamsM.Unlock()
 	stream, existing := sesh.streams[frame.StreamID]
+	sesh.streamsM.Unlock()
 	if existing {
 		return stream.writeFrame(*frame)
 	} else {

internal/multiplex/stream.go

@@ -112,6 +112,9 @@ func (s *Stream) Write(in []byte) (n int, err error) {
 	n, err = s.session.sb.send(s.obfsBuf[:i], &s.assignedConnId)
 	log.Tracef("%v sent to remote through stream %v with err %v", len(in), s.id, err)
 	if err != nil {
+		if err == errBrokenSwitchboard {
+			s.session.Close()
+		}
 		return
 	}
 	return len(in), nil