github.com/metacubex/mihomo@v1.18.5/common/net/deadline/conn.go (about)

     1  package deadline
     2  
     3  import (
     4  	"net"
     5  	"os"
     6  	"time"
     7  
     8  	"github.com/metacubex/mihomo/common/atomic"
     9  
    10  	"github.com/sagernet/sing/common/buf"
    11  	"github.com/sagernet/sing/common/bufio"
    12  	"github.com/sagernet/sing/common/network"
    13  )
    14  
    15  type connReadResult struct {
    16  	buffer []byte
    17  	err    error
    18  }
    19  
    20  type Conn struct {
    21  	network.ExtendedConn
    22  	deadline     atomic.TypedValue[time.Time]
    23  	pipeDeadline pipeDeadline
    24  	disablePipe  atomic.Bool
    25  	inRead       atomic.Bool
    26  	resultCh     chan *connReadResult
    27  }
    28  
    29  func IsConn(conn any) bool {
    30  	_, ok := conn.(*Conn)
    31  	return ok
    32  }
    33  
    34  func NewConn(conn net.Conn) *Conn {
    35  	c := &Conn{
    36  		ExtendedConn: bufio.NewExtendedConn(conn),
    37  		pipeDeadline: makePipeDeadline(),
    38  		resultCh:     make(chan *connReadResult, 1),
    39  	}
    40  	c.resultCh <- nil
    41  	return c
    42  }
    43  
    44  func (c *Conn) Read(p []byte) (n int, err error) {
    45  	select {
    46  	case result := <-c.resultCh:
    47  		if result != nil {
    48  			n = copy(p, result.buffer)
    49  			err = result.err
    50  			if n >= len(result.buffer) {
    51  				c.resultCh <- nil // finish cache read
    52  			} else {
    53  				result.buffer = result.buffer[n:]
    54  				c.resultCh <- result // push back for next call
    55  			}
    56  			return
    57  		} else {
    58  			c.resultCh <- nil
    59  			break
    60  		}
    61  	case <-c.pipeDeadline.wait():
    62  		return 0, os.ErrDeadlineExceeded
    63  	}
    64  
    65  	if c.disablePipe.Load() {
    66  		return c.ExtendedConn.Read(p)
    67  	} else if c.deadline.Load().IsZero() {
    68  		c.inRead.Store(true)
    69  		defer c.inRead.Store(false)
    70  		return c.ExtendedConn.Read(p)
    71  	}
    72  
    73  	<-c.resultCh
    74  	go c.pipeRead(len(p))
    75  
    76  	return c.Read(p)
    77  }
    78  
    79  func (c *Conn) pipeRead(size int) {
    80  	buffer := make([]byte, size)
    81  	n, err := c.ExtendedConn.Read(buffer)
    82  	buffer = buffer[:n]
    83  	c.resultCh <- &connReadResult{
    84  		buffer: buffer,
    85  		err:    err,
    86  	}
    87  }
    88  
    89  func (c *Conn) ReadBuffer(buffer *buf.Buffer) (err error) {
    90  	select {
    91  	case result := <-c.resultCh:
    92  		if result != nil {
    93  			n, _ := buffer.Write(result.buffer)
    94  			err = result.err
    95  
    96  			if n >= len(result.buffer) {
    97  				c.resultCh <- nil // finish cache read
    98  			} else {
    99  				result.buffer = result.buffer[n:]
   100  				c.resultCh <- result // push back for next call
   101  			}
   102  			return
   103  		} else {
   104  			c.resultCh <- nil
   105  			break
   106  		}
   107  	case <-c.pipeDeadline.wait():
   108  		return os.ErrDeadlineExceeded
   109  	}
   110  
   111  	if c.disablePipe.Load() {
   112  		return c.ExtendedConn.ReadBuffer(buffer)
   113  	} else if c.deadline.Load().IsZero() {
   114  		c.inRead.Store(true)
   115  		defer c.inRead.Store(false)
   116  		return c.ExtendedConn.ReadBuffer(buffer)
   117  	}
   118  
   119  	<-c.resultCh
   120  	go c.pipeRead(buffer.FreeLen())
   121  
   122  	return c.ReadBuffer(buffer)
   123  }
   124  
   125  func (c *Conn) SetReadDeadline(t time.Time) error {
   126  	if c.disablePipe.Load() {
   127  		return c.ExtendedConn.SetReadDeadline(t)
   128  	} else if c.inRead.Load() {
   129  		c.disablePipe.Store(true)
   130  		return c.ExtendedConn.SetReadDeadline(t)
   131  	}
   132  	c.deadline.Store(t)
   133  	c.pipeDeadline.set(t)
   134  	return nil
   135  }
   136  
   137  func (c *Conn) ReaderReplaceable() bool {
   138  	select {
   139  	case result := <-c.resultCh:
   140  		c.resultCh <- result
   141  		if result != nil {
   142  			return false // cache reading
   143  		} else {
   144  			break
   145  		}
   146  	default:
   147  		return false // pipe reading
   148  	}
   149  	return c.disablePipe.Load() || c.deadline.Load().IsZero()
   150  }
   151  
   152  func (c *Conn) Upstream() any {
   153  	return c.ExtendedConn
   154  }