From f700f4b6a380b2714f4c25cd64df03d0073f0d4b Mon Sep 17 00:00:00 2001
From: yaling888 <73897884+yaling888@users.noreply.github.com>
Date: Wed, 15 Jun 2022 04:35:06 +0800
Subject: [PATCH] Chore: upgrade gVisor

---
 go.mod                                   |  2 +-
 go.sum                                   |  4 ++--
 listener/tun/device/iobased/endpoint.go  | 23 +++++++++++++++--------
 listener/tun/device/tun/tun_wireguard.go |  6 +++++-
 4 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/go.mod b/go.mod
index 68d98212..eeffb82a 100644
--- a/go.mod
+++ b/go.mod
@@ -28,7 +28,7 @@ require (
 	golang.zx2c4.com/wireguard/windows v0.5.4-0.20220328111914-004c22c5647e
 	google.golang.org/protobuf v1.28.0
 	gopkg.in/yaml.v3 v3.0.1
-	gvisor.dev/gvisor v0.0.0-20220601233344-46e478629075
+	gvisor.dev/gvisor v0.0.0-20220614011939-d89c08b0f364
 )
 
 require (
diff --git a/go.sum b/go.sum
index 01a449a3..9e5cc5fc 100644
--- a/go.sum
+++ b/go.sum
@@ -158,5 +158,5 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-gvisor.dev/gvisor v0.0.0-20220601233344-46e478629075 h1:ucwwit0X39HmMQ1iSeNCIXw4g/B8bi+O6TSxxDbPs9E=
-gvisor.dev/gvisor v0.0.0-20220601233344-46e478629075/go.mod h1:TIvkJD0sxe8pIob3p6T8IzxXunlp6yfgktvTNp+DGNM=
+gvisor.dev/gvisor v0.0.0-20220614011939-d89c08b0f364 h1:D+X2kUQINrsyxdwX71CUWxxQCGRU4tS6gni5WRYLCjs=
+gvisor.dev/gvisor v0.0.0-20220614011939-d89c08b0f364/go.mod h1:TIvkJD0sxe8pIob3p6T8IzxXunlp6yfgktvTNp+DGNM=
diff --git a/listener/tun/device/iobased/endpoint.go b/listener/tun/device/iobased/endpoint.go
index e0a4583a..377bba76 100644
--- a/listener/tun/device/iobased/endpoint.go
+++ b/listener/tun/device/iobased/endpoint.go
@@ -8,8 +8,10 @@ import (
 	"io"
 	"sync"
 
+	"github.com/Dreamacro/clash/common/pool"
+
+	"gvisor.dev/gvisor/pkg/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip"
-	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
@@ -93,23 +95,29 @@ func (e *Endpoint) dispatchLoop(cancel context.CancelFunc) {
 
 	mtu := int(e.mtu)
 	for {
-		data := make([]byte, mtu)
+		data := pool.Get(mtu)
 
 		n, err := e.rw.Read(data)
 		if err != nil {
+			_ = pool.Put(data)
 			break
 		}
 
 		if n == 0 || n > mtu {
+			_ = pool.Put(data)
 			continue
 		}
 
 		if !e.IsAttached() {
+			_ = pool.Put(data)
 			continue /* unattached, drop packet */
 		}
 
 		pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-			Data: buffer.View(data[:n]).ToVectorisedView(),
+			Payload: buffer.NewWithData(data[:n]),
+			OnRelease: func() {
+				_ = pool.Put(data)
+			},
 		})
 
 		switch header.IPVersion(data) {
@@ -117,6 +125,8 @@ func (e *Endpoint) dispatchLoop(cancel context.CancelFunc) {
 			e.InjectInbound(header.IPv4ProtocolNumber, pkt)
 		case header.IPv6Version:
 			e.InjectInbound(header.IPv6ProtocolNumber, pkt)
+		default:
+			_ = pool.Put(data)
 		}
 		pkt.DecRef()
 	}
@@ -138,11 +148,8 @@ func (e *Endpoint) outboundLoop(ctx context.Context) {
 func (e *Endpoint) writePacket(pkt *stack.PacketBuffer) tcpip.Error {
 	defer pkt.DecRef()
 
-	size := pkt.Size()
-	views := pkt.Views()
-
-	vView := buffer.NewVectorisedView(size, views)
-	if _, err := e.rw.Write(vView.ToView()); err != nil {
+	buf := pkt.Buffer()
+	if _, err := e.rw.Write(buf.Flatten()); err != nil {
 		return &tcpip.ErrInvalidEndpointState{}
 	}
 	return nil
diff --git a/listener/tun/device/tun/tun_wireguard.go b/listener/tun/device/tun/tun_wireguard.go
index 1aafae54..04f86d4b 100644
--- a/listener/tun/device/tun/tun_wireguard.go
+++ b/listener/tun/device/tun/tun_wireguard.go
@@ -87,7 +87,11 @@ func (t *TUN) Write(packet []byte) (int, error) {
 
 	packet = append(t.cache[:t.offset], packet...)
 
-	return t.nt.Write(packet, t.offset)
+	n, err := t.nt.Write(packet, t.offset)
+	if n < t.offset {
+		return 0, err
+	}
+	return n - t.offset, err
 }
 
 func (t *TUN) Close() error {