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

     1  package socks_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"net"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/txthinking/socks5"
    14  	"golang.org/x/net/proxy"
    15  
    16  	"github.com/ipfans/trojan-go/common"
    17  	"github.com/ipfans/trojan-go/config"
    18  	"github.com/ipfans/trojan-go/test/util"
    19  	"github.com/ipfans/trojan-go/tunnel"
    20  	"github.com/ipfans/trojan-go/tunnel/adapter"
    21  	"github.com/ipfans/trojan-go/tunnel/socks"
    22  )
    23  
    24  func TestSocks(t *testing.T) {
    25  	port := common.PickPort("tcp", "127.0.0.1")
    26  	ctx := config.WithConfig(context.Background(), adapter.Name, &adapter.Config{
    27  		LocalHost: "127.0.0.1",
    28  		LocalPort: port,
    29  	})
    30  	ctx = config.WithConfig(ctx, socks.Name, &socks.Config{
    31  		LocalHost: "127.0.0.1",
    32  		LocalPort: port,
    33  	})
    34  	tcpServer, err := adapter.NewServer(ctx, nil)
    35  	common.Must(err)
    36  	addr := tunnel.NewAddressFromHostPort("tcp", "127.0.0.1", port)
    37  	s, err := socks.NewServer(ctx, tcpServer)
    38  	common.Must(err)
    39  	socksClient, err := proxy.SOCKS5("tcp", addr.String(), nil, proxy.Direct)
    40  	common.Must(err)
    41  	var conn1, conn2 net.Conn
    42  	wg := sync.WaitGroup{}
    43  	wg.Add(2)
    44  
    45  	time.Sleep(time.Second * 2)
    46  	go func() {
    47  		conn2, err = s.AcceptConn(nil)
    48  		common.Must(err)
    49  		wg.Done()
    50  	}()
    51  
    52  	time.Sleep(time.Second * 1)
    53  	go func() {
    54  		conn1, err = socksClient.Dial("tcp", util.EchoAddr)
    55  		common.Must(err)
    56  		wg.Done()
    57  	}()
    58  
    59  	wg.Wait()
    60  	if !util.CheckConn(conn1, conn2) {
    61  		t.Fail()
    62  	}
    63  	fmt.Println(conn2.(tunnel.Conn).Metadata())
    64  
    65  	udpConn, err := net.ListenPacket("udp", ":0")
    66  	common.Must(err)
    67  
    68  	addr = &tunnel.Address{
    69  		AddressType: tunnel.DomainName,
    70  		DomainName:  "google.com",
    71  		Port:        12345,
    72  	}
    73  
    74  	payload := util.GeneratePayload(1024)
    75  	buf := bytes.NewBuffer(make([]byte, 0, 4096))
    76  	buf.Write([]byte{0, 0, 0}) // RSV, FRAG
    77  	common.Must(addr.WriteTo(buf))
    78  	buf.Write(payload)
    79  
    80  	udpConn.WriteTo(buf.Bytes(), &net.UDPAddr{
    81  		IP:   net.ParseIP("127.0.0.1"),
    82  		Port: port,
    83  	})
    84  
    85  	packet, err := s.AcceptPacket(nil)
    86  	common.Must(err)
    87  	recvBuf := make([]byte, 4096)
    88  	n, m, err := packet.ReadWithMetadata(recvBuf)
    89  	common.Must(err)
    90  	if m.DomainName != "google.com" || m.Port != 12345 || n != 1024 || !(bytes.Equal(recvBuf[:n], payload)) {
    91  		t.Fail()
    92  	}
    93  
    94  	payload = util.GeneratePayload(1024)
    95  	_, err = packet.WriteWithMetadata(payload, &tunnel.Metadata{
    96  		Address: &tunnel.Address{
    97  			AddressType: tunnel.IPv4,
    98  			IP:          net.ParseIP("123.123.234.234"),
    99  			Port:        12345,
   100  		},
   101  	})
   102  	common.Must(err)
   103  
   104  	_, _, err = udpConn.ReadFrom(recvBuf)
   105  	common.Must(err)
   106  
   107  	r := bytes.NewReader(recvBuf)
   108  	header := [3]byte{}
   109  	r.Read(header[:])
   110  	addr = new(tunnel.Address)
   111  	common.Must(addr.ReadFrom(r))
   112  	if addr.IP.String() != "123.123.234.234" || addr.Port != 12345 {
   113  		t.Fail()
   114  	}
   115  
   116  	recvBuf, err = ioutil.ReadAll(r)
   117  	common.Must(err)
   118  
   119  	if bytes.Equal(recvBuf, payload) {
   120  		t.Fail()
   121  	}
   122  	packet.Close()
   123  	udpConn.Close()
   124  
   125  	c, _ := socks5.NewClient(fmt.Sprintf("127.0.0.1:%d", port), "", "", 0, 0)
   126  
   127  	conn, err := c.Dial("udp", util.EchoAddr)
   128  	common.Must(err)
   129  
   130  	payload = util.GeneratePayload(4096)
   131  	recvBuf = make([]byte, 4096)
   132  
   133  	conn.Write(payload)
   134  
   135  	newPacket, err := s.AcceptPacket(nil)
   136  	common.Must(err)
   137  
   138  	_, m, err = newPacket.ReadWithMetadata(recvBuf)
   139  	common.Must(err)
   140  	if m.String() != util.EchoAddr || !bytes.Equal(recvBuf, payload) {
   141  		t.Fail()
   142  	}
   143  
   144  	s.Close()
   145  }