From 6c6a804736170dc8a35dc349b7083d57410dff46 Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Sat, 19 Nov 2022 16:13:13 -0800
Subject: [PATCH] perf: set buffers support for udp hop conns

---
 go.mod                  |  2 +-
 go.sum                  |  4 +--
 pkg/pktconns/udp/hop.go | 62 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index ab458c6..1242e91 100644
--- a/go.mod
+++ b/go.mod
@@ -89,6 +89,6 @@ require (
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
-replace github.com/lucas-clemente/quic-go => github.com/apernet/quic-go v0.31.1-0.20221119042740-af3f46bd335f
+replace github.com/lucas-clemente/quic-go => github.com/apernet/quic-go v0.31.1-0.20221119235156-55bf700f2dd4
 
 replace github.com/LiamHaworth/go-tproxy => github.com/apernet/go-tproxy v0.0.0-20221025153553-ed04a2935f88
diff --git a/go.sum b/go.sum
index d28789c..b72a8a5 100644
--- a/go.sum
+++ b/go.sum
@@ -49,8 +49,8 @@ github.com/antonfisher/nested-logrus-formatter v1.3.1 h1:NFJIr+pzwv5QLHTPyKz9UME
 github.com/antonfisher/nested-logrus-formatter v1.3.1/go.mod h1:6WTfyWFkBc9+zyBaKIqRrg/KwMqBbodBjgbHjDz7zjA=
 github.com/apernet/go-tproxy v0.0.0-20221025153553-ed04a2935f88 h1:YNsl7PMiU9x/0CleMHJ7GUdS8y1aRTFwTxdSmLLEijQ=
 github.com/apernet/go-tproxy v0.0.0-20221025153553-ed04a2935f88/go.mod h1:uxH+nFzlJug5OHjPYmzKwvVVb9wOToeGuLNVeerwWtc=
-github.com/apernet/quic-go v0.31.1-0.20221119042740-af3f46bd335f h1:pE353n6QvdHxUk6Lt0614VMmvxHjqQzsGiboZWzlndk=
-github.com/apernet/quic-go v0.31.1-0.20221119042740-af3f46bd335f/go.mod h1:0wFbizLgYzqHqtlyxyCaJKlE7bYgE6JQ+54TLd/Dq2g=
+github.com/apernet/quic-go v0.31.1-0.20221119235156-55bf700f2dd4 h1:rNk86XSaAK/nPyab0ZxI2uRWwYqse9JehOVG+ijOh0I=
+github.com/apernet/quic-go v0.31.1-0.20221119235156-55bf700f2dd4/go.mod h1:0wFbizLgYzqHqtlyxyCaJKlE7bYgE6JQ+54TLd/Dq2g=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
diff --git a/pkg/pktconns/udp/hop.go b/pkg/pktconns/udp/hop.go
index 259599d..e1e4f9e 100644
--- a/pkg/pktconns/udp/hop.go
+++ b/pkg/pktconns/udp/hop.go
@@ -1,11 +1,13 @@
 package udp
 
 import (
+	"errors"
 	"math/rand"
 	"net"
 	"strconv"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/apernet/hysteria/pkg/pktconns/obfs"
@@ -29,6 +31,9 @@ type ObfsUDPHopClientPacketConn struct {
 	currentConn net.PacketConn
 	addrIndex   int
 
+	readBufferSize  int
+	writeBufferSize int
+
 	recvQueue chan *udpPacket
 	closeChan chan struct{}
 	closed    bool
@@ -155,6 +160,13 @@ func (c *ObfsUDPHopClientPacketConn) hop() {
 	} else {
 		c.currentConn = newConn
 	}
+	// Set buffer sizes if previously set
+	if c.readBufferSize > 0 {
+		_ = trySetPacketConnReadBuffer(c.currentConn, c.readBufferSize)
+	}
+	if c.writeBufferSize > 0 {
+		_ = trySetPacketConnWriteBuffer(c.currentConn, c.writeBufferSize)
+	}
 	go c.recvRoutine(c.currentConn)
 	c.addrIndex = rand.Intn(len(c.serverAddrs))
 }
@@ -244,6 +256,56 @@ func (c *ObfsUDPHopClientPacketConn) SetDeadline(t time.Time) error {
 	return c.SetWriteDeadline(t)
 }
 
+func (c *ObfsUDPHopClientPacketConn) SetReadBuffer(bytes int) error {
+	c.connMutex.Lock()
+	defer c.connMutex.Unlock()
+	c.readBufferSize = bytes
+	if c.prevConn != nil {
+		_ = trySetPacketConnReadBuffer(c.prevConn, bytes)
+	}
+	return trySetPacketConnReadBuffer(c.currentConn, bytes)
+}
+
+func (c *ObfsUDPHopClientPacketConn) SetWriteBuffer(bytes int) error {
+	c.connMutex.Lock()
+	defer c.connMutex.Unlock()
+	c.writeBufferSize = bytes
+	if c.prevConn != nil {
+		_ = trySetPacketConnWriteBuffer(c.prevConn, bytes)
+	}
+	return trySetPacketConnWriteBuffer(c.currentConn, bytes)
+}
+
+func (c *ObfsUDPHopClientPacketConn) SyscallConn() (syscall.RawConn, error) {
+	c.connMutex.RLock()
+	defer c.connMutex.RUnlock()
+	sc, ok := c.currentConn.(syscall.Conn)
+	if !ok {
+		return nil, errors.New("not supported")
+	}
+	return sc.SyscallConn()
+}
+
+func trySetPacketConnReadBuffer(pc net.PacketConn, bytes int) error {
+	sc, ok := pc.(interface {
+		SetReadBuffer(bytes int) error
+	})
+	if ok {
+		return sc.SetReadBuffer(bytes)
+	}
+	return nil
+}
+
+func trySetPacketConnWriteBuffer(pc net.PacketConn, bytes int) error {
+	sc, ok := pc.(interface {
+		SetWriteBuffer(bytes int) error
+	})
+	if ok {
+		return sc.SetWriteBuffer(bytes)
+	}
+	return nil
+}
+
 // parseAddr parses the multi-port server address and returns the host and ports.
 // Supports both comma-separated single ports and dash-separated port ranges.
 // Format: "host:port1,port2-port3,port4"