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

     1  package trojan
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"net"
     7  	"sync"
     8  	"sync/atomic"
     9  	"time"
    10  
    11  	"github.com/ipfans/trojan-go/api"
    12  	"github.com/ipfans/trojan-go/common"
    13  	"github.com/ipfans/trojan-go/config"
    14  	"github.com/ipfans/trojan-go/log"
    15  	"github.com/ipfans/trojan-go/statistic"
    16  	"github.com/ipfans/trojan-go/statistic/memory"
    17  	"github.com/ipfans/trojan-go/tunnel"
    18  	"github.com/ipfans/trojan-go/tunnel/mux"
    19  )
    20  
    21  const (
    22  	MaxPacketSize = 1024 * 8
    23  )
    24  
    25  const (
    26  	Connect   tunnel.Command = 1
    27  	Associate tunnel.Command = 3
    28  	Mux       tunnel.Command = 0x7f
    29  )
    30  
    31  type OutboundConn struct {
    32  	// WARNING: do not change the order of these fields.
    33  	// 64-bit fields that use `sync/atomic` package functions
    34  	// must be 64-bit aligned on 32-bit systems.
    35  	// Reference: https://github.com/golang/go/issues/599
    36  	// Solution: https://github.com/golang/go/issues/11891#issuecomment-433623786
    37  	sent uint64
    38  	recv uint64
    39  
    40  	metadata          *tunnel.Metadata
    41  	user              statistic.User
    42  	headerWrittenOnce sync.Once
    43  	net.Conn
    44  }
    45  
    46  func (c *OutboundConn) Metadata() *tunnel.Metadata {
    47  	return c.metadata
    48  }
    49  
    50  func (c *OutboundConn) WriteHeader(payload []byte) (bool, error) {
    51  	var err error
    52  	written := false
    53  	c.headerWrittenOnce.Do(func() {
    54  		hash := c.user.Hash()
    55  		buf := bytes.NewBuffer(make([]byte, 0, MaxPacketSize))
    56  		crlf := []byte{0x0d, 0x0a}
    57  		buf.Write([]byte(hash))
    58  		buf.Write(crlf)
    59  		c.metadata.WriteTo(buf)
    60  		buf.Write(crlf)
    61  		if payload != nil {
    62  			buf.Write(payload)
    63  		}
    64  		_, err = c.Conn.Write(buf.Bytes())
    65  		if err == nil {
    66  			written = true
    67  		}
    68  	})
    69  	return written, err
    70  }
    71  
    72  func (c *OutboundConn) Write(p []byte) (int, error) {
    73  	written, err := c.WriteHeader(p)
    74  	if err != nil {
    75  		return 0, common.NewError("trojan failed to flush header with payload").Base(err)
    76  	}
    77  	if written {
    78  		return len(p), nil
    79  	}
    80  	n, err := c.Conn.Write(p)
    81  	c.user.AddTraffic(n, 0)
    82  	atomic.AddUint64(&c.sent, uint64(n))
    83  	return n, err
    84  }
    85  
    86  func (c *OutboundConn) Read(p []byte) (int, error) {
    87  	n, err := c.Conn.Read(p)
    88  	c.user.AddTraffic(0, n)
    89  	atomic.AddUint64(&c.recv, uint64(n))
    90  	return n, err
    91  }
    92  
    93  func (c *OutboundConn) Close() error {
    94  	log.Info("connection to", c.metadata, "closed", "sent:", common.HumanFriendlyTraffic(atomic.LoadUint64(&c.sent)), "recv:", common.HumanFriendlyTraffic(atomic.LoadUint64(&c.recv)))
    95  	return c.Conn.Close()
    96  }
    97  
    98  type Client struct {
    99  	underlay tunnel.Client
   100  	user     statistic.User
   101  	ctx      context.Context
   102  	cancel   context.CancelFunc
   103  }
   104  
   105  func (c *Client) Close() error {
   106  	c.cancel()
   107  	return c.underlay.Close()
   108  }
   109  
   110  func (c *Client) DialConn(addr *tunnel.Address, overlay tunnel.Tunnel) (tunnel.Conn, error) {
   111  	conn, err := c.underlay.DialConn(addr, &Tunnel{})
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	newConn := &OutboundConn{
   116  		Conn: conn,
   117  		user: c.user,
   118  		metadata: &tunnel.Metadata{
   119  			Command: Connect,
   120  			Address: addr,
   121  		},
   122  	}
   123  	if _, ok := overlay.(*mux.Tunnel); ok {
   124  		newConn.metadata.Command = Mux
   125  	}
   126  
   127  	go func(newConn *OutboundConn) {
   128  		// if the trojan header is still buffered after 100 ms, the client may expect data from the server
   129  		// so we flush the trojan header
   130  		time.Sleep(time.Millisecond * 100)
   131  		newConn.WriteHeader(nil)
   132  	}(newConn)
   133  	return newConn, nil
   134  }
   135  
   136  func (c *Client) DialPacket(tunnel.Tunnel) (tunnel.PacketConn, error) {
   137  	fakeAddr := &tunnel.Address{
   138  		DomainName:  "UDP_CONN",
   139  		AddressType: tunnel.DomainName,
   140  	}
   141  	conn, err := c.underlay.DialConn(fakeAddr, &Tunnel{})
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  	return &PacketConn{
   146  		Conn: &OutboundConn{
   147  			Conn: conn,
   148  			user: c.user,
   149  			metadata: &tunnel.Metadata{
   150  				Command: Associate,
   151  				Address: fakeAddr,
   152  			},
   153  		},
   154  	}, nil
   155  }
   156  
   157  func NewClient(ctx context.Context, client tunnel.Client) (*Client, error) {
   158  	ctx, cancel := context.WithCancel(ctx)
   159  	auth, err := statistic.NewAuthenticator(ctx, memory.Name)
   160  	if err != nil {
   161  		cancel()
   162  		return nil, err
   163  	}
   164  
   165  	cfg := config.FromContext(ctx, Name).(*Config)
   166  	if cfg.API.Enabled {
   167  		go api.RunService(ctx, Name+"_CLIENT", auth)
   168  	}
   169  
   170  	var user statistic.User
   171  	for _, u := range auth.ListUsers() {
   172  		user = u
   173  		break
   174  	}
   175  	if user == nil {
   176  		cancel()
   177  		return nil, common.NewError("no valid user found")
   178  	}
   179  
   180  	log.Debug("trojan client created")
   181  	return &Client{
   182  		underlay: client,
   183  		ctx:      ctx,
   184  		user:     user,
   185  		cancel:   cancel,
   186  	}, nil
   187  }