From 154cb1d1f046ddc48d9ff76183698ef3a72bf01e Mon Sep 17 00:00:00 2001
From: Kr328 <kr328app@outlook.com>
Date: Sun, 18 Jun 2023 11:19:35 +0800
Subject: [PATCH] Improve: alloc using make if alloc size > 65536 (#2796)

---
 common/pool/alloc.go                 | 25 +++++++++++++++++--------
 common/pool/alloc_test.go            |  6 +++---
 component/process/process_windows.go | 11 ++++++++---
 3 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/common/pool/alloc.go b/common/pool/alloc.go
index 25f7989..f06b1a3 100644
--- a/common/pool/alloc.go
+++ b/common/pool/alloc.go
@@ -32,23 +32,32 @@ func NewAllocator() *Allocator {
 
 // Get a []byte from pool with most appropriate cap
 func (alloc *Allocator) Get(size int) []byte {
-	if size <= 0 || size > 65536 {
+	switch {
+	case size < 0:
+		panic("alloc.Get: len out of range")
+	case size == 0:
 		return nil
-	}
+	case size > 65536:
+		return make([]byte, size)
+	default:
+		bits := msb(size)
+		if size == 1<<bits {
+			return alloc.buffers[bits].Get().([]byte)[:size]
+		}
 
-	bits := msb(size)
-	if size == 1<<bits {
-		return alloc.buffers[bits].Get().([]byte)[:size]
+		return alloc.buffers[bits+1].Get().([]byte)[:size]
 	}
-
-	return alloc.buffers[bits+1].Get().([]byte)[:size]
 }
 
 // Put returns a []byte to pool for future use,
 // which the cap must be exactly 2^n
 func (alloc *Allocator) Put(buf []byte) error {
+	if cap(buf) == 0 || cap(buf) > 65536 {
+		return nil
+	}
+
 	bits := msb(cap(buf))
-	if cap(buf) == 0 || cap(buf) > 65536 || cap(buf) != 1<<bits {
+	if cap(buf) != 1<<bits {
 		return errors.New("allocator Put() incorrect buffer size")
 	}
 
diff --git a/common/pool/alloc_test.go b/common/pool/alloc_test.go
index 3d06331..0b7c3cd 100644
--- a/common/pool/alloc_test.go
+++ b/common/pool/alloc_test.go
@@ -19,17 +19,17 @@ func TestAllocGet(t *testing.T) {
 	assert.Equal(t, 1024, cap(alloc.Get(1023)))
 	assert.Equal(t, 1024, len(alloc.Get(1024)))
 	assert.Equal(t, 65536, len(alloc.Get(65536)))
-	assert.Nil(t, alloc.Get(65537))
+	assert.Equal(t, 65537, len(alloc.Get(65537)))
 }
 
 func TestAllocPut(t *testing.T) {
 	alloc := NewAllocator()
-	assert.NotNil(t, alloc.Put(nil), "put nil misbehavior")
+	assert.Nil(t, alloc.Put(nil), "put nil misbehavior")
 	assert.NotNil(t, alloc.Put(make([]byte, 3)), "put elem:3 []bytes misbehavior")
 	assert.Nil(t, alloc.Put(make([]byte, 4)), "put elem:4 []bytes misbehavior")
 	assert.Nil(t, alloc.Put(make([]byte, 1023, 1024)), "put elem:1024 []bytes misbehavior")
 	assert.Nil(t, alloc.Put(make([]byte, 65536)), "put elem:65536 []bytes misbehavior")
-	assert.NotNil(t, alloc.Put(make([]byte, 65537)), "put elem:65537 []bytes misbehavior")
+	assert.Nil(t, alloc.Put(make([]byte, 65537)), "put elem:65537 []bytes misbehavior")
 }
 
 func TestAllocPutThenGet(t *testing.T) {
diff --git a/component/process/process_windows.go b/component/process/process_windows.go
index 6ccb158..806adba 100644
--- a/component/process/process_windows.go
+++ b/component/process/process_windows.go
@@ -7,6 +7,8 @@ import (
 	"unsafe"
 
 	"golang.org/x/sys/windows"
+
+	"github.com/Dreamacro/clash/common/pool"
 )
 
 var (
@@ -41,8 +43,10 @@ func findProcessPath(network string, from netip.AddrPort, to netip.AddrPort) (st
 }
 
 func findPidByConnectionEndpoint(family uint32, protocol uint32, from netip.AddrPort, to netip.AddrPort) (uint32, error) {
-	buf := []byte(nil)
-	bufSize := uint32(0)
+	buf := pool.Get(0)
+	defer pool.Put(buf)
+
+	bufSize := uint32(len(buf))
 
 loop:
 	for {
@@ -77,7 +81,8 @@ loop:
 
 			break loop
 		case uintptr(windows.ERROR_INSUFFICIENT_BUFFER):
-			buf = make([]byte, bufSize)
+			pool.Put(buf)
+			buf = pool.Get(int(bufSize))
 
 			continue loop
 		default: