github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/network/udp_agent.go (about)

     1  package network
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/15mga/kiwi"
     7  	"github.com/15mga/kiwi/ds"
     8  	"net"
     9  	"time"
    10  
    11  	"github.com/15mga/kiwi/util"
    12  )
    13  
    14  func NewUdpAgent(addr string, receiver kiwi.FnAgentBytes, options ...kiwi.AgentOption) *udpAgent {
    15  	return &udpAgent{
    16  		agent: newAgent(addr, receiver, options...),
    17  	}
    18  }
    19  
    20  type udpAgent struct {
    21  	agent
    22  	conn net.Conn
    23  }
    24  
    25  func (a *udpAgent) Start(ctx context.Context, conn net.Conn) {
    26  	a.conn = conn
    27  	a.onClose = a.conn.Close
    28  	a.start(ctx)
    29  	switch a.option.AgentMode {
    30  	case kiwi.AgentRW:
    31  		go a.read()
    32  		go a.write()
    33  	case kiwi.AgentR:
    34  		go a.read()
    35  	case kiwi.AgentW:
    36  		go a.write()
    37  	}
    38  }
    39  
    40  func (a *udpAgent) read() {
    41  	var (
    42  		newData [2048]byte
    43  		err     *util.Err
    44  	)
    45  	defer func() {
    46  		r := recover()
    47  		if r != nil {
    48  			kiwi.Error2(util.EcRecover, util.M{
    49  				"remote addr": a.conn.RemoteAddr().String(),
    50  				"recover":     fmt.Sprintf("%s", r),
    51  			})
    52  			a.read()
    53  			return
    54  		}
    55  		a.close(err)
    56  	}()
    57  
    58  	dur := time.Duration(a.option.DeadlineSecs)
    59  	for {
    60  		select {
    61  		case <-a.ctx.Done():
    62  			return
    63  		default:
    64  			if dur > 0 {
    65  				_ = a.conn.SetReadDeadline(time.Now().Add(time.Second * dur))
    66  			}
    67  			newLen, e := a.conn.Read(newData[:])
    68  			if e != nil {
    69  				err = util.WrapErr(util.EcIo, e)
    70  				return
    71  			}
    72  			a.receiver(a, newData[:newLen])
    73  		}
    74  	}
    75  }
    76  
    77  func (a *udpAgent) write() {
    78  	var (
    79  		err *util.Err
    80  	)
    81  	defer func() {
    82  		a.close(err)
    83  	}()
    84  
    85  	for {
    86  		select {
    87  		case <-a.ctx.Done():
    88  			return
    89  		case <-a.writeSignCh:
    90  			var elem *ds.LinkElem[[]byte]
    91  			a.enable.Mtx.Lock()
    92  			if a.enable.Disabled() {
    93  				a.enable.Mtx.Unlock()
    94  				return
    95  			}
    96  			elem = a.bytesLink.PopAll()
    97  			a.enable.Mtx.Unlock()
    98  			if elem == nil {
    99  				continue
   100  			}
   101  
   102  			for ; elem != nil; elem = elem.Next {
   103  				bytes := elem.Value
   104  				_, e := a.conn.Write(bytes)
   105  				util.RecycleBytes(bytes)
   106  				if e != nil {
   107  					err = util.WrapErr(util.EcIo, e)
   108  					return
   109  				}
   110  			}
   111  		}
   112  	}
   113  }