Visually test your NGINX rate limiting rules

main.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sync"
	"time"
)

const CSI = "\x1B["

type Printer struct {
	w     io.Writer
	lines int
	mu    sync.Mutex
}

func NewPrinter(w io.Writer) *Printer {
	return &Printer{w: w}
}

func (p *Printer) AddLine(content string) int {
	p.mu.Lock()
	defer p.mu.Unlock()

	fmt.Fprint(p.w, fmt.Sprintf("[%d] ", p.lines+1), content, "\n")
	p.lines++
	return p.lines - 1
}

func (p *Printer) PrintLine(line int, text string) {
	p.mu.Lock()
	defer p.mu.Unlock()

	dist := p.lines - line
	fmt.Fprint(p.w,
		fmt.Sprintf("%s%dF%s2K[%d] ", CSI, dist, CSI, line+1),
		text,
		fmt.Sprintf("%s%dE", CSI, dist))
}

func statusColor(status int) string {
	switch {
	case status >= 500:
		return fmt.Sprintf("\x1B[31m%d\x1B[0m", status)
	case status >= 400:
		return fmt.Sprintf("\x1B[33m%d\x1B[0m", status)
	case status >= 300:
		return fmt.Sprintf("\x1B[36m%d\x1B[0m", status)
	case status >= 200:
		return fmt.Sprintf("\x1B[32m%d\x1B[0m", status)
	default:
		return fmt.Sprintf("%d", status)
	}
}

func makeRequest(url string, timeBase time.Time, p *Printer, wg *sync.WaitGroup) {
	defer wg.Done()
	startTime := time.Now()
	startTick := startTime.Sub(timeBase).Truncate(time.Millisecond)
	lineNo := p.AddLine(fmt.Sprintf("Requesting [%s]", startTick))

	resp, err := http.Get(url)
	endTime := time.Now()
	endTick := endTime.Sub(startTime).Truncate(time.Millisecond)
	if errors.Is(err, io.EOF) {
		resp = &http.Response{StatusCode: 444, Body: io.NopCloser(nil)}
	} else if err != nil {
		p.PrintLine(lineNo, fmt.Sprintf("Done [%s] [Error in %s: %v]", startTick, endTick, err))
		return
	}
	resp.Body.Close()
	p.PrintLine(lineNo, fmt.Sprintf("Done [%s] [%s in %s]", startTick, statusColor(resp.StatusCode), endTick))
}

func main() {
	var (
		interval time.Duration
		count    int
	)
	flag.DurationVar(&interval, "i", 100*time.Millisecond, "Interval between requests")
	flag.IntVar(&count, "c", 10, "Number of requests")
	flag.Parse()
	if flag.NArg() != 1 {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] <url>\nOptions:\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(1)
	}
	url := flag.Arg(0)

	p := NewPrinter(os.Stdout)
	wg := new(sync.WaitGroup)
	wg.Add(count)
	defer wg.Wait()
	timeBase := time.Now()
	for i := 0; i < count; i++ {
		go makeRequest(url, timeBase, p, wg)
		time.Sleep(interval)
	}
}