github.com/ipfans/trojan-go@v0.11.0/test/util/target.go (about)

     1  package util
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"net"
     9  	"net/http"
    10  	"sync"
    11  	"time"
    12  
    13  	"golang.org/x/net/websocket"
    14  
    15  	"github.com/ipfans/trojan-go/common"
    16  	"github.com/ipfans/trojan-go/log"
    17  )
    18  
    19  var (
    20  	HTTPAddr string
    21  	HTTPPort string
    22  )
    23  
    24  func runHelloHTTPServer() {
    25  	httpHello := func(w http.ResponseWriter, req *http.Request) {
    26  		w.Write([]byte("HelloWorld"))
    27  	}
    28  
    29  	wsConfig, err := websocket.NewConfig("wss://127.0.0.1/websocket", "https://127.0.0.1")
    30  	common.Must(err)
    31  	wsServer := websocket.Server{
    32  		Config: *wsConfig,
    33  		Handler: func(conn *websocket.Conn) {
    34  			conn.Write([]byte("HelloWorld"))
    35  		},
    36  		Handshake: func(wsConfig *websocket.Config, httpRequest *http.Request) error {
    37  			log.Debug("websocket url", httpRequest.URL, "origin", httpRequest.Header.Get("Origin"))
    38  			return nil
    39  		},
    40  	}
    41  	mux := &http.ServeMux{}
    42  	mux.HandleFunc("/", httpHello)
    43  	mux.HandleFunc("/websocket", wsServer.ServeHTTP)
    44  	HTTPAddr = GetTestAddr()
    45  	_, HTTPPort, _ = net.SplitHostPort(HTTPAddr)
    46  	server := http.Server{
    47  		Addr:    HTTPAddr,
    48  		Handler: mux,
    49  	}
    50  	go server.ListenAndServe()
    51  	time.Sleep(time.Second * 1) // wait for http server
    52  	fmt.Println("http test server listening on", HTTPAddr)
    53  	wg.Done()
    54  }
    55  
    56  var (
    57  	EchoAddr string
    58  	EchoPort int
    59  )
    60  
    61  func runTCPEchoServer() {
    62  	listener, err := net.Listen("tcp", EchoAddr)
    63  	common.Must(err)
    64  	wg.Done()
    65  	go func() {
    66  		defer listener.Close()
    67  		for {
    68  			conn, err := listener.Accept()
    69  			if err != nil {
    70  				return
    71  			}
    72  			go func(conn net.Conn) {
    73  				defer conn.Close()
    74  				for {
    75  					buf := make([]byte, 2048)
    76  					conn.SetDeadline(time.Now().Add(time.Second * 5))
    77  					n, err := conn.Read(buf)
    78  					conn.SetDeadline(time.Time{})
    79  					if err != nil {
    80  						return
    81  					}
    82  					_, err = conn.Write(buf[0:n])
    83  					if err != nil {
    84  						return
    85  					}
    86  				}
    87  			}(conn)
    88  		}
    89  	}()
    90  }
    91  
    92  func runUDPEchoServer() {
    93  	conn, err := net.ListenPacket("udp", EchoAddr)
    94  	common.Must(err)
    95  	wg.Done()
    96  	go func() {
    97  		for {
    98  			buf := make([]byte, 1024*8)
    99  			n, addr, err := conn.ReadFrom(buf)
   100  			if err != nil {
   101  				return
   102  			}
   103  			log.Info("Echo from", addr)
   104  			conn.WriteTo(buf[0:n], addr)
   105  		}
   106  	}()
   107  }
   108  
   109  func GeneratePayload(length int) []byte {
   110  	buf := make([]byte, length)
   111  	io.ReadFull(rand.Reader, buf)
   112  	return buf
   113  }
   114  
   115  var (
   116  	BlackHoleAddr string
   117  	BlackHolePort int
   118  )
   119  
   120  func runTCPBlackHoleServer() {
   121  	listener, err := net.Listen("tcp", BlackHoleAddr)
   122  	common.Must(err)
   123  	wg.Done()
   124  	go func() {
   125  		defer listener.Close()
   126  		for {
   127  			conn, err := listener.Accept()
   128  			if err != nil {
   129  				return
   130  			}
   131  			go func(conn net.Conn) {
   132  				io.Copy(ioutil.Discard, conn)
   133  				conn.Close()
   134  			}(conn)
   135  		}
   136  	}()
   137  }
   138  
   139  func runUDPBlackHoleServer() {
   140  	conn, err := net.ListenPacket("udp", BlackHoleAddr)
   141  	common.Must(err)
   142  	wg.Done()
   143  	go func() {
   144  		defer conn.Close()
   145  		buf := make([]byte, 1024*8)
   146  		for {
   147  			_, _, err := conn.ReadFrom(buf)
   148  			if err != nil {
   149  				return
   150  			}
   151  		}
   152  	}()
   153  }
   154  
   155  var wg = sync.WaitGroup{}
   156  
   157  func init() {
   158  	wg.Add(5)
   159  	runHelloHTTPServer()
   160  
   161  	EchoPort = common.PickPort("tcp", "127.0.0.1")
   162  	EchoAddr = fmt.Sprintf("127.0.0.1:%d", EchoPort)
   163  
   164  	BlackHolePort = common.PickPort("tcp", "127.0.0.1")
   165  	BlackHoleAddr = fmt.Sprintf("127.0.0.1:%d", BlackHolePort)
   166  
   167  	runTCPEchoServer()
   168  	runUDPEchoServer()
   169  
   170  	runTCPBlackHoleServer()
   171  	runUDPBlackHoleServer()
   172  
   173  	wg.Wait()
   174  }