Fix missing UDP timeout

SagerNet · Dec 20, 2023 · 9525cab · 9525cab
1 parent 499e257
commit 9525cab
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 3 deletions.
diff --git a/hysteria/service.go b/hysteria/service.go
@@ -9,13 +9,15 @@ import (
 	"os"
 	"runtime"
 	"sync"
+	"time"
 
 	"github.com/sagernet/quic-go"
 	"github.com/sagernet/sing-quic"
 	hyCC "github.com/sagernet/sing-quic/hysteria/congestion"
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/auth"
 	"github.com/sagernet/sing/common/baderror"
+	"github.com/sagernet/sing/common/canceler"
 	E "github.com/sagernet/sing/common/exceptions"
 	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
@@ -32,6 +34,7 @@ type ServiceOptions struct {
 	XPlusPassword string
 	TLSConfig     aTLS.ServerConfig
 	UDPDisabled   bool
+	UDPTimeout    time.Duration
 	Handler       ServerHandler
 
 	// Legacy options
@@ -58,6 +61,7 @@ type Service[U comparable] struct {
 	quicConfig    *quic.Config
 	userMap       map[string]U
 	udpDisabled   bool
+	udpTimeout    time.Duration
 	handler       ServerHandler
 	quicListener  io.Closer
 }
@@ -109,6 +113,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) {
 		userMap:       make(map[string]U),
 		handler:       options.Handler,
 		udpDisabled:   options.UDPDisabled,
+		udpTimeout:    options.UDPTimeout,
 	}, nil
 }
 
@@ -272,7 +277,8 @@ func (s *serverSession[U]) handleStream(stream quic.Stream) error {
 			udpConn.closeWithError(E.Cause(err, "write server response"))
 			return err
 		}
-		go s.handler.NewPacketConnection(udpConn.ctx, udpConn, M.Metadata{
+		newCtx, newConn := canceler.NewPacketConn(udpConn.ctx, udpConn, s.udpTimeout)
+		go s.handler.NewPacketConnection(newCtx, newConn, M.Metadata{
 			Source:      s.source,
 			Destination: M.ParseSocksaddrHostPort(request.Host, request.Port),
 		})

diff --git a/hysteria2/service.go b/hysteria2/service.go
@@ -41,6 +41,7 @@ type ServiceOptions struct {
 	SalamanderPassword    string
 	TLSConfig             aTLS.ServerConfig
 	UDPDisabled           bool
+	UDPTimeout            time.Duration
 	Handler               ServerHandler
 	MasqueradeHandler     http.Handler
 }
@@ -62,6 +63,7 @@ type Service[U comparable] struct {
 	quicConfig            *quic.Config
 	userMap               map[string]U
 	udpDisabled           bool
+	udpTimeout            time.Duration
 	handler               ServerHandler
 	masqueradeHandler     http.Handler
 	quicListener          io.Closer
@@ -97,6 +99,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) {
 		quicConfig:            quicConfig,
 		userMap:               make(map[string]U),
 		udpDisabled:           options.UDPDisabled,
+		udpTimeout:            options.UDPTimeout,
 		handler:               options.Handler,
 		masqueradeHandler:     options.MasqueradeHandler,
 	}, nil

diff --git a/hysteria2/service_packet.go b/hysteria2/service_packet.go
@@ -3,6 +3,7 @@ package hysteria2
 import (
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/auth"
+	"github.com/sagernet/sing/common/canceler"
 	E "github.com/sagernet/sing/common/exceptions"
 	M "github.com/sagernet/sing/common/metadata"
 )
@@ -47,7 +48,8 @@ func (s *serverSession[U]) handleUDPMessage(message *udpMessage) {
 		s.udpAccess.Lock()
 		s.udpConnMap[message.sessionID] = udpConn
 		s.udpAccess.Unlock()
-		go s.handler.NewPacketConnection(udpConn.ctx, udpConn, M.Metadata{
+		newCtx, newConn := canceler.NewPacketConn(udpConn.ctx, udpConn, s.udpTimeout)
+		go s.handler.NewPacketConnection(newCtx, newConn, M.Metadata{
 			Source:      s.source,
 			Destination: M.ParseSocksaddr(message.destination),
 		})

diff --git a/tuic/service.go b/tuic/service.go
@@ -35,6 +35,7 @@ type ServiceOptions struct {
 	AuthTimeout       time.Duration
 	ZeroRTTHandshake  bool
 	Heartbeat         time.Duration
+	UDPTimeout        time.Duration
 	Handler           ServiceHandler
 }
 
@@ -53,6 +54,7 @@ type Service[U comparable] struct {
 	passwordMap       map[U]string
 	congestionControl string
 	authTimeout       time.Duration
+	udpTimeout        time.Duration
 	handler           ServiceHandler
 
 	quicListener io.Closer
@@ -88,6 +90,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) {
 		userMap:           make(map[[16]byte]U),
 		congestionControl: options.CongestionControl,
 		authTimeout:       options.AuthTimeout,
+		udpTimeout:        options.UDPTimeout,
 		handler:           options.Handler,
 	}, nil
 }

diff --git a/tuic/service_packet.go b/tuic/service_packet.go
@@ -3,6 +3,7 @@ package tuic
 import (
 	"github.com/sagernet/sing/common"
 	"github.com/sagernet/sing/common/auth"
+	"github.com/sagernet/sing/common/canceler"
 	E "github.com/sagernet/sing/common/exceptions"
 	M "github.com/sagernet/sing/common/metadata"
 )
@@ -65,7 +66,8 @@ func (s *serverSession[U]) handleUDPMessage(message *udpMessage, udpStream bool)
 		s.udpAccess.Lock()
 		s.udpConnMap[message.sessionID] = udpConn
 		s.udpAccess.Unlock()
-		go s.handler.NewPacketConnection(udpConn.ctx, udpConn, M.Metadata{
+		newCtx, newConn := canceler.NewPacketConn(udpConn.ctx, udpConn, s.udpTimeout)
+		go s.handler.NewPacketConnection(newCtx, newConn, M.Metadata{
 			Source:      s.source,
 			Destination: message.destination,
 		})