Stress tester for clash tun

clash_stress.go

package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"sync"
)

// SOCKS address types as defined in RFC 1928 section 5.
const (
	AtypIPv4       = 1
	AtypDomainName = 3
	AtypIPv6       = 4
)

func main() {
	address := flag.String("a", "", "IP address")
	port := flag.Int("p", 80, "port")
	thread := flag.Int("t", 1, "threads")
	tcp := flag.Bool("tcp", false, "use tcp")
	listen := flag.Bool("l", false, "listen mode")
	socsk5 := flag.String("socks5", "", "socks5 target address")
	flag.Parse()

	if *listen { // Listen Mode
		if *tcp {

		} else {
			addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", *address, *port))
			if err != nil {
				log.Fatalln(err)
			}
			conn, err := net.ListenUDP("udp", addr)
			if err != nil {
				log.Fatalln(err)
			}

			buf := make([]byte, 2048)
			for {
				conn.ReadFromUDP(buf)
			}
		}
		return
	}

	payload := bytes.Repeat([]byte{0xc}, 512)

	for i := 0; i < *thread; i++ {
		if *tcp { // TCP
			if *socsk5 != "" {
				log.Fatalln("TCP socks5 not implemented")
			}
			addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", *address, *port))
			if err != nil {
				log.Fatalln(err)
			}
			go func() {
				for {
					conn, err := net.DialTCP("tcp", nil, addr)
					if err != nil {
						continue
					}
					conn.Write(payload)
				}
			}()

		} else {
			addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", *address, *port))
			if *socsk5 != "" {
				socks5Target, err := net.ResolveUDPAddr("udp", *socsk5)
				if err != nil {
					log.Fatalln(err)
				}
				payload = packUDP(payload, socks5Target.IP, uint16(socks5Target.Port))
			}
			if err != nil {
				log.Fatalln(err)
			}
			go func() {
				conn, err := net.DialUDP("udp", nil, addr)
				if err != nil {
					log.Fatalln(err)
				}
				for {
					conn.Write(payload)
				}
			}()
		}
	}
	var wg sync.WaitGroup
	wg.Add(1)
	wg.Wait()
}

func packUDP(b []byte, ip net.IP, port uint16) []byte {
	ipv4 := true
	lheader := 10        // ipv4
	if ip.To4() == nil { // ipv6
		ipv4 = false
		lheader += 12
	}
	// +----+------+------+----------+----------+----------+
	// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
	// +----+------+------+----------+----------+----------+
	// | 2  |  1   |  1   | Variable |    2     | Variable |
	// +----+------+------+----------+----------+----------+
	ret := make([]byte, len(b)+lheader)
	ret[0] = 0
	ret[1] = 0
	ret[2] = 0
	if ipv4 {
		ret[3] = AtypIPv4
		copy(ret[4:], ip.To4())
	} else {
		ret[3] = AtypIPv6
		copy(ret[4:], ip.To16())
	}
	binary.BigEndian.PutUint16(ret[lheader-2:], port)
	copy(ret[lheader:], b)
	return ret
}