fix locking scheme

xtaci · Dec 18, 2019 · 5ad2b95 · 5ad2b95
1 parent 361ba76
commit 5ad2b95
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/generic/rawcopy_unix.go b/generic/rawcopy_unix.go
@@ -16,17 +16,20 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 	}
 
 	buf := ctrl.Buffer
+	var locked bool
 	for {
 		var er error
 		var nr int
 		rr := c.Read(func(s uintptr) bool {
-			ctrl.Lock() // writelock will block reading
-			defer ctrl.Unlock()
+			ctrl.Lock() // acquire rights to read & write
+			locked = true
 			nr, er = syscall.Read(int(s), buf)
 			if er == syscall.EAGAIN {
+				ctrl.Unlock()
+				locked = false
 				return false
 			}
-			return true
+			return true // keep lock
 		})
 
 		// read EOF
@@ -35,9 +38,10 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 		}
 
 		if nr > 0 {
-			ctrl.Lock()
 			nw, ew := dst.Write(buf[0:nr])
 			ctrl.Unlock()
+			locked = false
+
 			if nw > 0 {
 				written += int64(nw)
 			}
@@ -64,5 +68,9 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 		}
 	}
 
+	if locked {
+		ctrl.Unlock()
+	}
+
 	return written, err
 }
diff --git a/generic/rawcopy_windows.go b/generic/rawcopy_windows.go
@@ -15,12 +15,13 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 	}
 
 	buf := ctrl.Buffer
+	var locked bool
 	for {
 		var er error
 		var nr int
 		rr := c.Read(func(s uintptr) bool {
 			ctrl.Lock()
-			defer ctrl.Unlock()
+			locked = true
 			var read uint32
 			var flags uint32
 			var wsabuf syscall.WSABuf
@@ -37,9 +38,10 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 		}
 
 		if nr > 0 {
-			ctrl.Lock()
 			nw, ew := dst.Write(buf[0:nr])
 			ctrl.Unlock()
+			locked = false
+
 			if nw > 0 {
 				written += int64(nw)
 			}
@@ -66,5 +68,9 @@ func rawCopy(dst io.Writer, src *net.TCPConn, ctrl *CopyControl) (written int64,
 		}
 	}
 
+	if locked {
+		ctrl.Unlock()
+	}
+
 	return written, err
 }