chore: add tuic-server listener

2022-11-28 17:09:25 +08:00
parent cd53e2d4a7
commit 551283c16e
13 changed files with 1063 additions and 330 deletions
--- a/transport/tuic/client.go
+++ b/transport/tuic/client.go
@@ -6,10 +6,8 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
-	"fmt"
 	"math/rand"
 	"net"
-	"net/netip"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -21,7 +19,6 @@ import (
 	"github.com/Dreamacro/clash/common/pool"
 	"github.com/Dreamacro/clash/component/dialer"
 	C "github.com/Dreamacro/clash/constant"
-	"github.com/Dreamacro/clash/transport/tuic/congestion"
 )

 var (
@@ -39,7 +36,7 @@ type ClientOption struct {
 	UdpRelayMode          string
 	CongestionController  string
 	ReduceRtt             bool
-	RequestTimeout        int
+	RequestTimeout        time.Duration
 	MaxUdpRelayPacketSize int
 	FastOpen              bool
 	MaxOpenStreams        int64
@@ -53,6 +50,7 @@ type Client struct {
 	connMutex sync.Mutex

 	openStreams atomic.Int64
+	closed      atomic.Bool

 	udpInputMap sync.Map

@@ -82,160 +80,141 @@ func (t *Client) getQuicConn(ctx context.Context) (quic.Connection, error) {
 		return nil, err
 	}

-	switch t.CongestionController {
-	case "cubic":
-		quicConn.SetCongestionControl(
-			congestion.NewCubicSender(
-				congestion.DefaultClock{},
-				congestion.GetMaxPacketSize(quicConn.RemoteAddr()),
-				false,
-				nil,
-			),
-		)
-	case "new_reno":
-		quicConn.SetCongestionControl(
-			congestion.NewCubicSender(
-				congestion.DefaultClock{},
-				congestion.GetMaxPacketSize(quicConn.RemoteAddr()),
-				true,
-				nil,
-			),
-		)
-	case "bbr":
-		quicConn.SetCongestionControl(
-			congestion.NewBBRSender(
-				congestion.DefaultClock{},
-				congestion.GetMaxPacketSize(quicConn.RemoteAddr()),
-				congestion.InitialCongestionWindow,
-				congestion.DefaultBBRMaxCongestionWindow,
-			),
-		)
-	}
+	SetCongestionController(quicConn, t.CongestionController)

-	sendAuthentication := func(quicConn quic.Connection) (err error) {
-		defer func() {
-			t.deferQuicConn(quicConn, err)
-		}()
-		stream, err := quicConn.OpenUniStream()
-		if err != nil {
-			return err
-		}
-		buf := pool.GetBuffer()
-		defer pool.PutBuffer(buf)
-		err = NewAuthenticate(t.Token).WriteTo(buf)
-		if err != nil {
-			return err
-		}
-		_, err = buf.WriteTo(stream)
-		if err != nil {
-			return err
-		}
-		err = stream.Close()
-		if err != nil {
-			return
-		}
-		return nil
-	}
-
-	parseUDP := func(quicConn quic.Connection) (err error) {
-		defer func() {
-			t.deferQuicConn(quicConn, err)
-		}()
-		switch t.UdpRelayMode {
-		case "quic":
-			for {
-				var stream quic.ReceiveStream
-				stream, err = quicConn.AcceptUniStream(context.Background())
-				if err != nil {
-					return err
-				}
-				go func() (err error) {
-					var assocId uint32
-					defer func() {
-						t.deferQuicConn(quicConn, err)
-						if err != nil && assocId != 0 {
-							if val, ok := t.udpInputMap.LoadAndDelete(assocId); ok {
-								if conn, ok := val.(net.Conn); ok {
-									_ = conn.Close()
-								}
-							}
-						}
-						stream.CancelRead(0)
-					}()
-					reader := bufio.NewReader(stream)
-					packet, err := ReadPacket(reader)
-					if err != nil {
-						return
-					}
-					assocId = packet.ASSOC_ID
-					if val, ok := t.udpInputMap.Load(assocId); ok {
-						if conn, ok := val.(net.Conn); ok {
-							writer := bufio.NewWriterSize(conn, packet.BytesLen())
-							_ = packet.WriteTo(writer)
-							_ = writer.Flush()
-						}
-					}
-					return
-				}()
-			}
-		default: // native
-			for {
-				var message []byte
-				message, err = quicConn.ReceiveMessage()
-				if err != nil {
-					return err
-				}
-				go func() (err error) {
-					var assocId uint32
-					defer func() {
-						t.deferQuicConn(quicConn, err)
-						if err != nil && assocId != 0 {
-							if val, ok := t.udpInputMap.LoadAndDelete(assocId); ok {
-								if conn, ok := val.(net.Conn); ok {
-									_ = conn.Close()
-								}
-							}
-						}
-					}()
-					buffer := bytes.NewBuffer(message)
-					packet, err := ReadPacket(buffer)
-					if err != nil {
-						return
-					}
-					assocId = packet.ASSOC_ID
-					if val, ok := t.udpInputMap.Load(assocId); ok {
-						if conn, ok := val.(net.Conn); ok {
-							_, _ = conn.Write(message)
-						}
-					}
-					return
-				}()
-			}
-		}
-	}
-
-	go sendAuthentication(quicConn)
+	go func() {
+		_ = t.sendAuthentication(quicConn)
+	}()

 	if t.udp {
-		go parseUDP(quicConn)
+		go func() {
+			_ = t.parseUDP(quicConn)
+		}()
 	}

 	t.quicConn = quicConn
+	t.openStreams.Store(0)
 	return quicConn, nil
 }

+func (t *Client) sendAuthentication(quicConn quic.Connection) (err error) {
+	defer func() {
+		t.deferQuicConn(quicConn, err)
+	}()
+	stream, err := quicConn.OpenUniStream()
+	if err != nil {
+		return err
+	}
+	buf := pool.GetBuffer()
+	defer pool.PutBuffer(buf)
+	err = NewAuthenticate(t.Token).WriteTo(buf)
+	if err != nil {
+		return err
+	}
+	_, err = buf.WriteTo(stream)
+	if err != nil {
+		return err
+	}
+	err = stream.Close()
+	if err != nil {
+		return
+	}
+	return nil
+}
+
+func (t *Client) parseUDP(quicConn quic.Connection) (err error) {
+	defer func() {
+		t.deferQuicConn(quicConn, err)
+	}()
+	switch t.UdpRelayMode {
+	case "quic":
+		for {
+			var stream quic.ReceiveStream
+			stream, err = quicConn.AcceptUniStream(context.Background())
+			if err != nil {
+				return err
+			}
+			go func() (err error) {
+				var assocId uint32
+				defer func() {
+					t.deferQuicConn(quicConn, err)
+					if err != nil && assocId != 0 {
+						if val, ok := t.udpInputMap.LoadAndDelete(assocId); ok {
+							if conn, ok := val.(net.Conn); ok {
+								_ = conn.Close()
+							}
+						}
+					}
+					stream.CancelRead(0)
+				}()
+				reader := bufio.NewReader(stream)
+				packet, err := ReadPacket(reader)
+				if err != nil {
+					return
+				}
+				assocId = packet.ASSOC_ID
+				if val, ok := t.udpInputMap.Load(assocId); ok {
+					if conn, ok := val.(net.Conn); ok {
+						writer := bufio.NewWriterSize(conn, packet.BytesLen())
+						_ = packet.WriteTo(writer)
+						_ = writer.Flush()
+					}
+				}
+				return
+			}()
+		}
+	default: // native
+		for {
+			var message []byte
+			message, err = quicConn.ReceiveMessage()
+			if err != nil {
+				return err
+			}
+			go func() (err error) {
+				var assocId uint32
+				defer func() {
+					t.deferQuicConn(quicConn, err)
+					if err != nil && assocId != 0 {
+						if val, ok := t.udpInputMap.LoadAndDelete(assocId); ok {
+							if conn, ok := val.(net.Conn); ok {
+								_ = conn.Close()
+							}
+						}
+					}
+				}()
+				buffer := bytes.NewBuffer(message)
+				packet, err := ReadPacket(buffer)
+				if err != nil {
+					return
+				}
+				assocId = packet.ASSOC_ID
+				if val, ok := t.udpInputMap.Load(assocId); ok {
+					if conn, ok := val.(net.Conn); ok {
+						_, _ = conn.Write(message)
+					}
+				}
+				return
+			}()
+		}
+	}
+}
+
 func (t *Client) deferQuicConn(quicConn quic.Connection, err error) {
 	var netError net.Error
 	if err != nil && errors.As(err, &netError) {
 		t.connMutex.Lock()
 		defer t.connMutex.Unlock()
 		if t.quicConn == quicConn {
-			t.Close(err)
+			t.forceClose(err, true)
 		}
 	}
 }

-func (t *Client) Close(err error) {
+func (t *Client) forceClose(err error, locked bool) {
+	if !locked {
+		t.connMutex.Lock()
+		defer t.connMutex.Unlock()
+	}
 	quicConn := t.quicConn
 	if quicConn != nil {
 		_ = quicConn.CloseWithError(ProtocolError, err.Error())
@@ -250,6 +229,13 @@ func (t *Client) Close(err error) {
 	}
 }

+func (t *Client) Close() {
+	t.closed.Store(true)
+	if t.openStreams.Load() == 0 {
+		t.forceClose(ClientClosed, false)
+	}
+}
+
 func (t *Client) DialContext(ctx context.Context, metadata *C.Metadata) (net.Conn, error) {
 	quicConn, err := t.getQuicConn(ctx)
 	if err != nil {
@@ -278,7 +264,15 @@ func (t *Client) DialContext(ctx context.Context, metadata *C.Metadata) (net.Con
 			Stream: quicStream,
 			lAddr:  quicConn.LocalAddr(),
 			rAddr:  quicConn.RemoteAddr(),
-			client: t,
+			ref:    t,
+			closeDeferFn: func() {
+				time.AfterFunc(C.DefaultTCPTimeout, func() {
+					openStreams := t.openStreams.Add(-1)
+					if openStreams == 0 && t.closed.Load() {
+						t.forceClose(ClientClosed, false)
+					}
+				})
+			},
 		}
 		_, err = buf.WriteTo(stream)
 		if err != nil {
@@ -306,12 +300,12 @@ type earlyConn struct {
 	resOnce sync.Once
 	resErr  error

-	RequestTimeout int
+	RequestTimeout time.Duration
 }

 func (conn *earlyConn) response() error {
 	if conn.RequestTimeout > 0 {
-		_ = conn.SetReadDeadline(time.Now().Add(time.Duration(conn.RequestTimeout) * time.Millisecond))
+		_ = conn.SetReadDeadline(time.Now().Add(conn.RequestTimeout))
 	}
 	response, err := ReadResponse(conn)
 	if err != nil {
@@ -341,59 +335,6 @@ func (conn *earlyConn) Read(b []byte) (n int, err error) {
 	return conn.BufferedConn.Read(b)
 }

-type quicStreamConn struct {
-	quic.Stream
-	lock   sync.Mutex
-	lAddr  net.Addr
-	rAddr  net.Addr
-	client *Client
-
-	closeOnce sync.Once
-	closeErr  error
-}
-
-func (q *quicStreamConn) Write(p []byte) (n int, err error) {
-	q.lock.Lock()
-	defer q.lock.Unlock()
-	return q.Stream.Write(p)
-}
-
-func (q *quicStreamConn) Close() error {
-	q.closeOnce.Do(func() {
-		q.closeErr = q.close()
-	})
-	return q.closeErr
-}
-
-func (q *quicStreamConn) close() error {
-	defer time.AfterFunc(C.DefaultTCPTimeout, func() {
-		q.client.openStreams.Add(-1)
-	})
-
-	// https://github.com/cloudflare/cloudflared/commit/ed2bac026db46b239699ac5ce4fcf122d7cab2cd
-	// Make sure a possible writer does not block the lock forever. We need it, so we can close the writer
-	// side of the stream safely.
-	_ = q.Stream.SetWriteDeadline(time.Now())
-
-	// This lock is eventually acquired despite Write also acquiring it, because we set a deadline to writes.
-	q.lock.Lock()
-	defer q.lock.Unlock()
-
-	// We have to clean up the receiving stream ourselves since the Close in the bottom does not handle that.
-	q.Stream.CancelRead(0)
-	return q.Stream.Close()
-}
-
-func (q *quicStreamConn) LocalAddr() net.Addr {
-	return q.lAddr
-}
-
-func (q *quicStreamConn) RemoteAddr() net.Addr {
-	return q.rAddr
-}
-
-var _ net.Conn = &quicStreamConn{}
-
 func (t *Client) ListenPacketContext(ctx context.Context, metadata *C.Metadata) (net.PacketConn, error) {
 	quicConn, err := t.getQuicConn(ctx)
 	if err != nil {
@@ -415,139 +356,27 @@ func (t *Client) ListenPacketContext(ctx context.Context, metadata *C.Metadata)
 		}
 	}
 	pc := &quicStreamPacketConn{
-		connId:    connId,
-		quicConn:  quicConn,
-		lAddr:     quicConn.LocalAddr(),
-		client:    t,
-		inputConn: N.NewBufferedConn(pipe2),
+		connId:                connId,
+		quicConn:              quicConn,
+		lAddr:                 quicConn.LocalAddr(),
+		inputConn:             N.NewBufferedConn(pipe2),
+		udpRelayMode:          t.UdpRelayMode,
+		maxUdpRelayPacketSize: t.MaxUdpRelayPacketSize,
+		ref:                   t,
+		deferQuicConnFn:       t.deferQuicConn,
+		closeDeferFn: func() {
+			t.udpInputMap.Delete(connId)
+			time.AfterFunc(C.DefaultUDPTimeout, func() {
+				openStreams := t.openStreams.Add(-1)
+				if openStreams == 0 && t.closed.Load() {
+					t.forceClose(ClientClosed, false)
+				}
+			})
+		},
 	}
 	return pc, nil
 }

-type quicStreamPacketConn struct {
-	connId    uint32
-	quicConn  quic.Connection
-	lAddr     net.Addr
-	client    *Client
-	inputConn *N.BufferedConn
-
-	closeOnce sync.Once
-	closeErr  error
-	closed    bool
-}
-
-func (q *quicStreamPacketConn) Close() error {
-	q.closeOnce.Do(func() {
-		q.closed = true
-		q.closeErr = q.close()
-	})
-	return q.closeErr
-}
-
-func (q *quicStreamPacketConn) close() (err error) {
-	defer time.AfterFunc(C.DefaultTCPTimeout, func() {
-		q.client.openStreams.Add(-1)
-	})
-	defer func() {
-		q.client.deferQuicConn(q.quicConn, err)
-	}()
-	q.client.udpInputMap.Delete(q.connId)
-	_ = q.inputConn.Close()
-	buf := pool.GetBuffer()
-	defer pool.PutBuffer(buf)
-	err = NewDissociate(q.connId).WriteTo(buf)
-	if err != nil {
-		return
-	}
-	stream, err := q.quicConn.OpenUniStream()
-	if err != nil {
-		return
-	}
-	_, err = buf.WriteTo(stream)
-	if err != nil {
-		return
-	}
-	err = stream.Close()
-	if err != nil {
-		return
-	}
-	return
-}
-
-func (q *quicStreamPacketConn) SetDeadline(t time.Time) error {
-	//TODO implement me
-	return nil
-}
-
-func (q *quicStreamPacketConn) SetReadDeadline(t time.Time) error {
-	return q.inputConn.SetReadDeadline(t)
-}
-
-func (q *quicStreamPacketConn) SetWriteDeadline(t time.Time) error {
-	//TODO implement me
-	return nil
-}
-
-func (q *quicStreamPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
-	packet, err := ReadPacket(q.inputConn)
-	if err != nil {
-		return
-	}
-	n = copy(p, packet.DATA)
-	addr = packet.ADDR.UDPAddr()
-	return
-}
-
-func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if len(p) > q.client.MaxUdpRelayPacketSize {
-		return 0, fmt.Errorf("udp packet too large(%d > %d)", len(p), q.client.MaxUdpRelayPacketSize)
-	}
-	if q.closed {
-		return 0, net.ErrClosed
-	}
-	defer func() {
-		q.client.deferQuicConn(q.quicConn, err)
-	}()
-	addr.String()
-	buf := pool.GetBuffer()
-	defer pool.PutBuffer(buf)
-	addrPort, err := netip.ParseAddrPort(addr.String())
-	if err != nil {
-		return
-	}
-	err = NewPacket(q.connId, uint16(len(p)), NewAddressAddrPort(addrPort), p).WriteTo(buf)
-	if err != nil {
-		return
-	}
-	switch q.client.UdpRelayMode {
-	case "quic":
-		var stream quic.SendStream
-		stream, err = q.quicConn.OpenUniStream()
-		if err != nil {
-			return
-		}
-		defer stream.Close()
-		_, err = buf.WriteTo(stream)
-		if err != nil {
-			return
-		}
-	default: // native
-		err = q.quicConn.SendMessage(buf.Bytes())
-		if err != nil {
-			return
-		}
-	}
-	n = len(p)
-
-	return
-}
-
-func (q *quicStreamPacketConn) LocalAddr() net.Addr {
-	return q.lAddr
-}
-
-var _ net.PacketConn = &quicStreamPacketConn{}
-
 func NewClient(clientOption *ClientOption, udp bool) *Client {
 	c := &Client{
 		ClientOption: clientOption,
@@ -558,5 +387,5 @@ func NewClient(clientOption *ClientOption, udp bool) *Client {
 }

 func closeClient(client *Client) {
-	client.Close(ClientClosed)
+	client.forceClose(ClientClosed, false)
 }