github.com/metacubex/mihomo@v1.18.5/transport/tuic/v4/packet.go (about)

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