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

     1  package dokodemo
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  
     7  	"github.com/ipfans/trojan-go/common"
     8  	"github.com/ipfans/trojan-go/tunnel"
     9  )
    10  
    11  const MaxPacketSize = 1024 * 8
    12  
    13  type Conn struct {
    14  	net.Conn
    15  	src            *tunnel.Address
    16  	targetMetadata *tunnel.Metadata
    17  }
    18  
    19  func (c *Conn) Metadata() *tunnel.Metadata {
    20  	return c.targetMetadata
    21  }
    22  
    23  // PacketConn receive packet info from the packet dispatcher
    24  type PacketConn struct {
    25  	net.PacketConn
    26  	metadata *tunnel.Metadata
    27  	input    chan []byte
    28  	output   chan []byte
    29  	src      net.Addr
    30  	ctx      context.Context
    31  	cancel   context.CancelFunc
    32  }
    33  
    34  func (c *PacketConn) Close() error {
    35  	c.cancel()
    36  	// don't close the underlying udp socket
    37  	return nil
    38  }
    39  
    40  func (c *PacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
    41  	return c.ReadWithMetadata(p)
    42  }
    43  
    44  func (c *PacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
    45  	address, err := tunnel.NewAddressFromAddr("udp", addr.String())
    46  	if err != nil {
    47  		return 0, err
    48  	}
    49  	return c.WriteWithMetadata(p, &tunnel.Metadata{
    50  		Address: address,
    51  	})
    52  }
    53  
    54  func (c *PacketConn) ReadWithMetadata(p []byte) (int, *tunnel.Metadata, error) {
    55  	select {
    56  	case payload := <-c.input:
    57  		n := copy(p, payload)
    58  		return n, c.metadata, nil
    59  	case <-c.ctx.Done():
    60  		return 0, nil, common.NewError("dokodemo packet conn closed")
    61  	}
    62  }
    63  
    64  func (c *PacketConn) WriteWithMetadata(p []byte, m *tunnel.Metadata) (int, error) {
    65  	select {
    66  	case c.output <- p:
    67  		return len(p), nil
    68  	case <-c.ctx.Done():
    69  		return 0, common.NewError("dokodemo packet conn failed to write")
    70  	}
    71  }