github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/uot/lazy.go (about)

     1  package uot
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/sagernet/sing/common/buf"
     7  	"github.com/sagernet/sing/common/bufio"
     8  	N "github.com/sagernet/sing/common/network"
     9  )
    10  
    11  type LazyClientConn struct {
    12  	net.Conn
    13  	writer         N.VectorisedWriter
    14  	request        Request
    15  	requestWritten bool
    16  }
    17  
    18  func NewLazyClientConn(conn net.Conn, request Request) *LazyClientConn {
    19  	return &LazyClientConn{
    20  		Conn:    conn,
    21  		request: request,
    22  		writer:  bufio.NewVectorisedWriter(conn),
    23  	}
    24  }
    25  
    26  func NewLazyConn(conn net.Conn, request Request) *Conn {
    27  	clientConn := NewLazyClientConn(conn, request)
    28  	return NewConn(clientConn, request)
    29  }
    30  
    31  func (c *LazyClientConn) Write(p []byte) (n int, err error) {
    32  	if !c.requestWritten {
    33  		var request *buf.Buffer
    34  		request, err = EncodeRequest(c.request)
    35  		if err != nil {
    36  			return
    37  		}
    38  		err = c.writer.WriteVectorised([]*buf.Buffer{request, buf.As(p)})
    39  		if err != nil {
    40  			return
    41  		}
    42  		c.requestWritten = true
    43  		return len(p), nil
    44  	}
    45  	return c.Conn.Write(p)
    46  }
    47  
    48  func (c *LazyClientConn) WriteVectorised(buffers []*buf.Buffer) error {
    49  	if !c.requestWritten {
    50  		request, err := EncodeRequest(c.request)
    51  		if err != nil {
    52  			return err
    53  		}
    54  
    55  		err = c.writer.WriteVectorised(append([]*buf.Buffer{request}, buffers...))
    56  		c.requestWritten = true
    57  		return err
    58  	}
    59  	return c.writer.WriteVectorised(buffers)
    60  }
    61  
    62  func (c *LazyClientConn) NeedHandshake() bool {
    63  	return !c.requestWritten
    64  }
    65  
    66  func (c *LazyClientConn) ReaderReplaceable() bool {
    67  	return true
    68  }
    69  
    70  func (c *LazyClientConn) WriterReplaceable() bool {
    71  	return c.requestWritten
    72  }
    73  
    74  func (c *LazyClientConn) Upstream() any {
    75  	return c.Conn
    76  }