github.com/ipfans/trojan-go@v0.11.0/tunnel/router/conn.go (about)

     1  package router
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/ipfans/trojan-go/common"
     9  	"github.com/ipfans/trojan-go/log"
    10  	"github.com/ipfans/trojan-go/tunnel"
    11  )
    12  
    13  type packetInfo struct {
    14  	src     *tunnel.Metadata
    15  	payload []byte
    16  }
    17  
    18  type PacketConn struct {
    19  	proxy tunnel.PacketConn
    20  	net.PacketConn
    21  	packetChan chan *packetInfo
    22  	*Client
    23  	ctx    context.Context
    24  	cancel context.CancelFunc
    25  }
    26  
    27  func (c *PacketConn) packetLoop() {
    28  	go func() {
    29  		for {
    30  			buf := make([]byte, MaxPacketSize)
    31  			n, addr, err := c.proxy.ReadWithMetadata(buf)
    32  			if err != nil {
    33  				select {
    34  				case <-c.ctx.Done():
    35  					return
    36  				default:
    37  					log.Error("router packetConn error", err)
    38  					continue
    39  				}
    40  			}
    41  			c.packetChan <- &packetInfo{
    42  				src:     addr,
    43  				payload: buf[:n],
    44  			}
    45  		}
    46  	}()
    47  	for {
    48  		buf := make([]byte, MaxPacketSize)
    49  		n, addr, err := c.PacketConn.ReadFrom(buf)
    50  		if err != nil {
    51  			select {
    52  			case <-c.ctx.Done():
    53  				return
    54  			default:
    55  				log.Error("router packetConn error", err)
    56  				continue
    57  			}
    58  		}
    59  		address, _ := tunnel.NewAddressFromAddr("udp", addr.String())
    60  		c.packetChan <- &packetInfo{
    61  			src: &tunnel.Metadata{
    62  				Address: address,
    63  			},
    64  			payload: buf[:n],
    65  		}
    66  	}
    67  }
    68  
    69  func (c *PacketConn) Close() error {
    70  	c.cancel()
    71  	c.proxy.Close()
    72  	return c.PacketConn.Close()
    73  }
    74  
    75  func (c *PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
    76  	panic("implement me")
    77  }
    78  
    79  func (c *PacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
    80  	panic("implement me")
    81  }
    82  
    83  func (c *PacketConn) WriteWithMetadata(p []byte, m *tunnel.Metadata) (int, error) {
    84  	policy := c.Route(m.Address)
    85  	switch policy {
    86  	case Proxy:
    87  		return c.proxy.WriteWithMetadata(p, m)
    88  	case Block:
    89  		return 0, common.NewError("router blocked address (udp): " + m.Address.String())
    90  	case Bypass:
    91  		ip, err := m.Address.ResolveIP()
    92  		if err != nil {
    93  			return 0, common.NewError("router failed to resolve udp address").Base(err)
    94  		}
    95  		return c.PacketConn.WriteTo(p, &net.UDPAddr{
    96  			IP:   ip,
    97  			Port: m.Address.Port,
    98  		})
    99  	default:
   100  		panic("unknown policy")
   101  	}
   102  }
   103  
   104  func (c *PacketConn) ReadWithMetadata(p []byte) (int, *tunnel.Metadata, error) {
   105  	select {
   106  	case info := <-c.packetChan:
   107  		n := copy(p, info.payload)
   108  		return n, info.src, nil
   109  	case <-c.ctx.Done():
   110  		return 0, nil, io.EOF
   111  	}
   112  }