github.com/kelleygo/clashcore@v1.0.2/transport/tuic/v5/packet.go (about)

     1  package v5
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/kelleygo/clashcore/common/atomic"
    10  	N "github.com/kelleygo/clashcore/common/net"
    11  	"github.com/kelleygo/clashcore/common/pool"
    12  	"github.com/kelleygo/clashcore/transport/tuic/common"
    13  
    14  	"github.com/metacubex/quic-go"
    15  	"github.com/zhangyunhao116/fastrand"
    16  )
    17  
    18  type quicStreamPacketConn struct {
    19  	connId    uint16
    20  	quicConn  quic.Connection
    21  	inputConn *N.BufferedConn
    22  
    23  	udpRelayMode          common.UdpRelayMode
    24  	maxUdpRelayPacketSize int
    25  
    26  	deferQuicConnFn func(quicConn quic.Connection, err error)
    27  	closeDeferFn    func()
    28  	writeClosed     *atomic.Bool
    29  
    30  	closeOnce sync.Once
    31  	closeErr  error
    32  	closed    bool
    33  
    34  	deFragger
    35  }
    36  
    37  func (q *quicStreamPacketConn) Close() error {
    38  	q.closeOnce.Do(func() {
    39  		q.closed = true
    40  		q.closeErr = q.close()
    41  	})
    42  	return q.closeErr
    43  }
    44  
    45  func (q *quicStreamPacketConn) close() (err error) {
    46  	if q.closeDeferFn != nil {
    47  		defer q.closeDeferFn()
    48  	}
    49  	if q.deferQuicConnFn != nil {
    50  		defer func() {
    51  			q.deferQuicConnFn(q.quicConn, err)
    52  		}()
    53  	}
    54  	if q.inputConn != nil {
    55  		_ = q.inputConn.Close()
    56  		q.inputConn = nil
    57  
    58  		buf := pool.GetBuffer()
    59  		defer pool.PutBuffer(buf)
    60  		err = NewDissociate(q.connId).WriteTo(buf)
    61  		if err != nil {
    62  			return
    63  		}
    64  		var stream quic.SendStream
    65  		stream, err = q.quicConn.OpenUniStream()
    66  		if err != nil {
    67  			return
    68  		}
    69  		_, err = buf.WriteTo(stream)
    70  		if err != nil {
    71  			return
    72  		}
    73  		err = stream.Close()
    74  		if err != nil {
    75  			return
    76  		}
    77  	}
    78  	return
    79  }
    80  
    81  func (q *quicStreamPacketConn) SetDeadline(t time.Time) error {
    82  	//TODO implement me
    83  	return nil
    84  }
    85  
    86  func (q *quicStreamPacketConn) SetReadDeadline(t time.Time) error {
    87  	if q.inputConn != nil {
    88  		return q.inputConn.SetReadDeadline(t)
    89  	}
    90  	return nil
    91  }
    92  
    93  func (q *quicStreamPacketConn) SetWriteDeadline(t time.Time) error {
    94  	//TODO implement me
    95  	return nil
    96  }
    97  
    98  func (q *quicStreamPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    99  	if inputConn := q.inputConn; inputConn != nil { // copy inputConn avoid be nil in for loop
   100  		for {
   101  			var packet Packet
   102  			packet, err = ReadPacket(inputConn)
   103  			if err != nil {
   104  				return
   105  			}
   106  			if packetPtr := q.deFragger.Feed(&packet); packetPtr != nil {
   107  				n = copy(p, packet.DATA)
   108  				addr = packetPtr.ADDR.UDPAddr()
   109  				return
   110  			}
   111  		}
   112  	} else {
   113  		err = net.ErrClosed
   114  	}
   115  	return
   116  }
   117  
   118  func (q *quicStreamPacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) {
   119  	if inputConn := q.inputConn; inputConn != nil { // copy inputConn avoid be nil in for loop
   120  		for {
   121  			var packet Packet
   122  			packet, err = ReadPacket(inputConn)
   123  			if err != nil {
   124  				return
   125  			}
   126  			if packetPtr := q.deFragger.Feed(&packet); packetPtr != nil {
   127  				data = packetPtr.DATA
   128  				addr = packetPtr.ADDR.UDPAddr()
   129  				return
   130  			}
   131  		}
   132  	} else {
   133  		err = net.ErrClosed
   134  	}
   135  	return
   136  }
   137  
   138  func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
   139  	if len(p) > 0xffff { // uint16 max
   140  		return 0, &quic.DatagramTooLargeError{PeerMaxDatagramFrameSize: 0xffff}
   141  	}
   142  	if q.closed {
   143  		return 0, net.ErrClosed
   144  	}
   145  	if q.writeClosed != nil && q.writeClosed.Load() {
   146  		_ = q.Close()
   147  		return 0, net.ErrClosed
   148  	}
   149  	if q.deferQuicConnFn != nil {
   150  		defer func() {
   151  			q.deferQuicConnFn(q.quicConn, err)
   152  		}()
   153  	}
   154  	buf := pool.GetBuffer()
   155  	defer pool.PutBuffer(buf)
   156  	address, err := NewAddressNetAddr(addr)
   157  	if err != nil {
   158  		return
   159  	}
   160  	pktId := uint16(fastrand.Uint32())
   161  	packet := NewPacket(q.connId, pktId, 1, 0, uint16(len(p)), address, p)
   162  	switch q.udpRelayMode {
   163  	case common.QUIC:
   164  		err = packet.WriteTo(buf)
   165  		if err != nil {
   166  			return
   167  		}
   168  		var stream quic.SendStream
   169  		stream, err = q.quicConn.OpenUniStream()
   170  		if err != nil {
   171  			return
   172  		}
   173  		defer stream.Close()
   174  		_, err = buf.WriteTo(stream)
   175  		if err != nil {
   176  			return
   177  		}
   178  	default: // native
   179  		if len(p) > q.maxUdpRelayPacketSize {
   180  			err = fragWriteNative(q.quicConn, packet, buf, q.maxUdpRelayPacketSize)
   181  		} else {
   182  			err = packet.WriteTo(buf)
   183  			if err != nil {
   184  				return
   185  			}
   186  			data := buf.Bytes()
   187  			err = q.quicConn.SendDatagram(data)
   188  		}
   189  
   190  		var tooLarge *quic.DatagramTooLargeError
   191  		if errors.As(err, &tooLarge) {
   192  			err = fragWriteNative(q.quicConn, packet, buf, int(tooLarge.PeerMaxDatagramFrameSize)-PacketOverHead)
   193  		}
   194  		if err != nil {
   195  			return
   196  		}
   197  	}
   198  	n = len(p)
   199  
   200  	return
   201  }
   202  
   203  func (q *quicStreamPacketConn) LocalAddr() net.Addr {
   204  	return q.quicConn.LocalAddr()
   205  }
   206  
   207  var _ net.PacketConn = (*quicStreamPacketConn)(nil)