From e3ba82caac1e1c3f3c9768404b819819cb541d97 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 23 Dec 2020 00:38:42 +0000
Subject: [PATCH] Code style changes and unlock mutex when failed to open
 stream

---
 internal/client/piper.go | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/internal/client/piper.go b/internal/client/piper.go
index 811a8f5..4bb9fc0 100644
--- a/internal/client/piper.go
+++ b/internal/client/piper.go
@@ -46,13 +46,14 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
 					sesh.Close()
 				}
 				log.Errorf("Failed to open stream: %v", err)
+				streamsMutex.Unlock()
 				continue
 			}
-			_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
-
 			streams[addr.String()] = stream
 			streamsMutex.Unlock()
 
+			_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
+
 			proxyAddr := addr
 			go func(stream *mux.Stream, localConn *net.UDPConn) {
 				buf := make([]byte, 8192)
@@ -60,24 +61,21 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration
 					n, err := stream.Read(buf)
 					if err != nil {
 						log.Tracef("copying stream to proxy client: %v", err)
-						streamsMutex.Lock()
-						delete(streams, addr.String())
-						streamsMutex.Unlock()
-						stream.Close()
-						return
+						break
 					}
 					_ = stream.SetReadDeadline(time.Now().Add(streamTimeout))
 
 					_, err = localConn.WriteTo(buf[:n], proxyAddr)
 					if err != nil {
 						log.Tracef("copying stream to proxy client: %v", err)
-						streamsMutex.Lock()
-						delete(streams, addr.String())
-						streamsMutex.Unlock()
-						stream.Close()
-						return
+						break
 					}
 				}
+				streamsMutex.Lock()
+				delete(streams, addr.String())
+				streamsMutex.Unlock()
+				stream.Close()
+				return
 			}(stream, localConn)
 		} else {
 			streamsMutex.Unlock()