github.com/kelleygo/clashcore@v1.0.2/transport/tuic/common/stream.go (about)

     1  package common
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/metacubex/quic-go"
     9  )
    10  
    11  type quicStreamConn struct {
    12  	quic.Stream
    13  	lock  sync.Mutex
    14  	lAddr net.Addr
    15  	rAddr net.Addr
    16  
    17  	closeDeferFn func()
    18  
    19  	closeOnce sync.Once
    20  	closeErr  error
    21  }
    22  
    23  func (q *quicStreamConn) Write(p []byte) (n int, err error) {
    24  	q.lock.Lock()
    25  	defer q.lock.Unlock()
    26  	return q.Stream.Write(p)
    27  }
    28  
    29  func (q *quicStreamConn) Close() error {
    30  	q.closeOnce.Do(func() {
    31  		q.closeErr = q.close()
    32  	})
    33  	return q.closeErr
    34  }
    35  
    36  func (q *quicStreamConn) close() error {
    37  	if q.closeDeferFn != nil {
    38  		defer q.closeDeferFn()
    39  	}
    40  
    41  	// https://github.com/cloudflare/cloudflared/commit/ed2bac026db46b239699ac5ce4fcf122d7cab2cd
    42  	// Make sure a possible writer does not block the lock forever. We need it, so we can close the writer
    43  	// side of the stream safely.
    44  	_ = q.Stream.SetWriteDeadline(time.Now())
    45  
    46  	// This lock is eventually acquired despite Write also acquiring it, because we set a deadline to writes.
    47  	q.lock.Lock()
    48  	defer q.lock.Unlock()
    49  
    50  	// We have to clean up the receiving stream ourselves since the Close in the bottom does not handle that.
    51  	q.Stream.CancelRead(0)
    52  	return q.Stream.Close()
    53  }
    54  
    55  func (q *quicStreamConn) LocalAddr() net.Addr {
    56  	return q.lAddr
    57  }
    58  
    59  func (q *quicStreamConn) RemoteAddr() net.Addr {
    60  	return q.rAddr
    61  }
    62  
    63  var _ net.Conn = (*quicStreamConn)(nil)
    64  
    65  func NewQuicStreamConn(stream quic.Stream, lAddr, rAddr net.Addr, closeDeferFn func()) net.Conn {
    66  	return &quicStreamConn{Stream: stream, lAddr: lAddr, rAddr: rAddr, closeDeferFn: closeDeferFn}
    67  }