github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/transport/internet/request/assembler/simple/client.go (about)

     1  package simple
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"crypto/rand"
     7  	"io"
     8  	"time"
     9  
    10  	"github.com/v2fly/v2ray-core/v5/common"
    11  
    12  	"github.com/v2fly/v2ray-core/v5/transport/internet/request"
    13  )
    14  
    15  func newClient(config *ClientConfig) request.SessionAssemblerClient {
    16  	return &simpleAssemblerClient{config: config}
    17  }
    18  
    19  type simpleAssemblerClient struct {
    20  	assembly request.TransportClientAssembly
    21  	config   *ClientConfig
    22  }
    23  
    24  func (s *simpleAssemblerClient) OnTransportClientAssemblyReady(assembly request.TransportClientAssembly) {
    25  	s.assembly = assembly
    26  }
    27  
    28  func (s *simpleAssemblerClient) NewSession(ctx context.Context, opts ...request.SessionOption) (request.Session, error) {
    29  	sessionID := make([]byte, 16)
    30  	_, err := io.ReadFull(rand.Reader, sessionID)
    31  	if err != nil {
    32  		return nil, err
    33  	}
    34  	sessionContext, finish := context.WithCancel(ctx)
    35  	session := &simpleAssemblerClientSession{
    36  		sessionID: sessionID, tripper: s.assembly.Tripper(), readBuffer: bytes.NewBuffer(nil),
    37  		ctx: sessionContext, finish: finish, writerChan: make(chan []byte), readerChan: make(chan []byte, 16), assembler: s,
    38  	}
    39  	go session.keepRunning()
    40  	return session, nil
    41  }
    42  
    43  type simpleAssemblerClientSession struct {
    44  	sessionID        []byte
    45  	currentWriteWait int
    46  
    47  	assembler  *simpleAssemblerClient
    48  	tripper    request.Tripper
    49  	readBuffer *bytes.Buffer
    50  	writerChan chan []byte
    51  	readerChan chan []byte
    52  	ctx        context.Context
    53  	finish     func()
    54  }
    55  
    56  func (s *simpleAssemblerClientSession) keepRunning() {
    57  	s.currentWriteWait = int(s.assembler.config.InitialPollingIntervalMs)
    58  	for s.ctx.Err() == nil {
    59  		s.runOnce()
    60  	}
    61  }
    62  
    63  func (s *simpleAssemblerClientSession) runOnce() {
    64  	sendBuffer := bytes.NewBuffer(nil)
    65  	if s.currentWriteWait != 0 {
    66  		waitTimer := time.NewTimer(time.Millisecond * time.Duration(s.currentWriteWait))
    67  		waitForFirstWrite := true
    68  	copyFromWriterLoop:
    69  		for {
    70  			select {
    71  			case <-s.ctx.Done():
    72  				return
    73  			case data := <-s.writerChan:
    74  				sendBuffer.Write(data)
    75  				if sendBuffer.Len() >= int(s.assembler.config.MaxWriteSize) {
    76  					break copyFromWriterLoop
    77  				}
    78  				if waitForFirstWrite {
    79  					waitForFirstWrite = false
    80  					waitTimer.Reset(time.Millisecond * time.Duration(s.assembler.config.WaitSubsequentWriteMs))
    81  				}
    82  			case <-waitTimer.C:
    83  				break copyFromWriterLoop
    84  			}
    85  		}
    86  		waitTimer.Stop()
    87  	}
    88  
    89  	firstRound := true
    90  	pollConnection := true
    91  	for sendBuffer.Len() != 0 || firstRound {
    92  		firstRound = false
    93  		sendAmount := sendBuffer.Len()
    94  		if sendAmount > int(s.assembler.config.MaxWriteSize) {
    95  			sendAmount = int(s.assembler.config.MaxWriteSize)
    96  		}
    97  		data := sendBuffer.Next(sendAmount)
    98  		if len(data) != 0 {
    99  			pollConnection = false
   100  		}
   101  		for {
   102  			resp, err := s.tripper.RoundTrip(s.ctx, request.Request{Data: data, ConnectionTag: s.sessionID})
   103  			if err != nil {
   104  				newError("failed to send data").Base(err).WriteToLog()
   105  				if s.ctx.Err() != nil {
   106  					return
   107  				}
   108  				time.Sleep(time.Millisecond * time.Duration(s.assembler.config.FailedRetryIntervalMs))
   109  				continue
   110  			}
   111  			if len(resp.Data) != 0 {
   112  				s.readerChan <- resp.Data
   113  			}
   114  			if len(resp.Data) != 0 {
   115  				pollConnection = false
   116  			}
   117  			break
   118  		}
   119  	}
   120  	if pollConnection {
   121  		s.currentWriteWait = int(s.assembler.config.BackoffFactor * float32(s.currentWriteWait))
   122  		if s.currentWriteWait > int(s.assembler.config.MaxPollingIntervalMs) {
   123  			s.currentWriteWait = int(s.assembler.config.MaxPollingIntervalMs)
   124  		}
   125  		if s.currentWriteWait < int(s.assembler.config.MinPollingIntervalMs) {
   126  			s.currentWriteWait = int(s.assembler.config.MinPollingIntervalMs)
   127  		}
   128  	} else {
   129  		s.currentWriteWait = int(0)
   130  	}
   131  }
   132  
   133  func (s *simpleAssemblerClientSession) Read(p []byte) (n int, err error) {
   134  	if s.readBuffer.Len() == 0 {
   135  		select {
   136  		case <-s.ctx.Done():
   137  			return 0, s.ctx.Err()
   138  		case data := <-s.readerChan:
   139  			s.readBuffer.Write(data)
   140  		}
   141  	}
   142  	n, err = s.readBuffer.Read(p)
   143  	if err == io.EOF {
   144  		s.readBuffer.Reset()
   145  		return 0, nil
   146  	}
   147  	return
   148  }
   149  
   150  func (s *simpleAssemblerClientSession) Write(p []byte) (n int, err error) {
   151  	buf := make([]byte, len(p))
   152  	copy(buf, p)
   153  	select {
   154  	case <-s.ctx.Done():
   155  		return 0, s.ctx.Err()
   156  	case s.writerChan <- buf:
   157  		return len(p), nil
   158  	}
   159  }
   160  
   161  func (s *simpleAssemblerClientSession) Close() error {
   162  	s.finish()
   163  	return nil
   164  }
   165  
   166  func init() {
   167  	common.Must(common.RegisterConfig((*ClientConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
   168  		clientConfig, ok := config.(*ClientConfig)
   169  		if !ok {
   170  			return nil, newError("not a ClientConfig")
   171  		}
   172  		return newClient(clientConfig), nil
   173  	}))
   174  }