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

     1  package tls
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"sync"
     7  )
     8  
     9  type Listener struct {
    10  	net.Listener
    11  	config ServerConfig
    12  }
    13  
    14  func NewListener(inner net.Listener, config ServerConfig) net.Listener {
    15  	return &Listener{
    16  		Listener: inner,
    17  		config:   config,
    18  	}
    19  }
    20  
    21  func (l *Listener) Accept() (net.Conn, error) {
    22  	conn, err := l.Listener.Accept()
    23  	if err != nil {
    24  		return nil, err
    25  	}
    26  	return NewLazyConn(conn, l.config), nil
    27  }
    28  
    29  type LazyConn struct {
    30  	net.Conn
    31  	tlsConfig     ServerConfig
    32  	access        sync.Mutex
    33  	needHandshake bool
    34  }
    35  
    36  func NewLazyConn(conn net.Conn, config ServerConfig) *LazyConn {
    37  	return &LazyConn{
    38  		Conn:          conn,
    39  		tlsConfig:     config,
    40  		needHandshake: true,
    41  	}
    42  }
    43  
    44  func (c *LazyConn) HandshakeContext(ctx context.Context) error {
    45  	if !c.needHandshake {
    46  		return nil
    47  	}
    48  	c.access.Lock()
    49  	defer c.access.Unlock()
    50  	if c.needHandshake {
    51  		tlsConn, err := ServerHandshake(ctx, c.Conn, c.tlsConfig)
    52  		if err != nil {
    53  			return err
    54  		}
    55  		c.Conn = tlsConn
    56  		c.needHandshake = false
    57  	}
    58  	return nil
    59  }
    60  
    61  func (c *LazyConn) Read(p []byte) (n int, err error) {
    62  	err = c.HandshakeContext(context.Background())
    63  	if err != nil {
    64  		return
    65  	}
    66  	return c.Conn.Read(p)
    67  }
    68  
    69  func (c *LazyConn) Write(p []byte) (n int, err error) {
    70  	err = c.HandshakeContext(context.Background())
    71  	if err != nil {
    72  		return
    73  	}
    74  	return c.Conn.Write(p)
    75  }
    76  
    77  func (c *LazyConn) NeedHandshake() bool {
    78  	return c.needHandshake
    79  }
    80  
    81  func (c *LazyConn) ReaderReplaceable() bool {
    82  	return !c.needHandshake
    83  }
    84  
    85  func (c *LazyConn) WriterReplaceable() bool {
    86  	return !c.needHandshake
    87  }
    88  
    89  func (c *LazyConn) Upstream() any {
    90  	return c.Conn
    91  }