git » gofer » main » tree

[main] / ratelimit / ratelimit.go 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
package ratelimit

import (
	"fmt"
	"net/http"
	"sort"

	"blitiri.com.ar/go/gofer/config"
	"blitiri.com.ar/go/gofer/ipratelimit"
	"blitiri.com.ar/go/gofer/trace"
	"blitiri.com.ar/go/log"
)

// Global registry for convenience.
// This is not pretty but it simplifies a lot of the handling for now.
var registry = map[string]*ipratelimit.Limiter{}

var traces = map[*ipratelimit.Limiter]*trace.Trace{}

func FromConfig(name string, conf config.RateLimit) {
	if conf.Size == 0 {
		conf.Size = 1000
	}

	rl := ipratelimit.NewLimiter(
		conf.Rate.Requests, conf.Rate.Period, conf.Size)

	// If config has custom IPv6 rates, use them.
	if conf.Rate64.Period > 0 {
		rl.SetIPv6s64Rate(conf.Rate64.Requests, conf.Rate64.Period)
	}
	if conf.Rate56.Period > 0 {
		rl.SetIPv6s56Rate(conf.Rate56.Requests, conf.Rate56.Period)
	}
	if conf.Rate48.Period > 0 {
		rl.SetIPv6s48Rate(conf.Rate48.Requests, conf.Rate48.Period)
	}

	registry[name] = rl
	traces[rl] = trace.New("ratelimit", name)
	traces[rl].SetMaxEvents(1000)

	log.Infof("ratelimit %q: %d/%s, size %d",
		name, conf.Rate.Requests, conf.Rate.Period, conf.Size)
	return
}

func FromName(name string) *ipratelimit.Limiter {
	return registry[name]
}

func Trace(rl *ipratelimit.Limiter) *trace.Trace {
	return traces[rl]
}

func DebugHandler(w http.ResponseWriter, r *http.Request) {
	names := []string{}
	for name := range registry {
		names = append(names, name)
	}
	sort.Strings(names)

	fmt.Fprintf(w, `<!DOCTYPE html>
<html>

<head>
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>ratelimit</title>
<style type="text/css">
  body {
    font-family: sans-serif;
  }
  @media (prefers-color-scheme: dark) {
    body {
      background: #121212;
	  color: #c9d1d9;
	}
	a { color: #44b4ec; }
  }
  table {
    text-align: right;
  }
  th {
    text-align: center;
  }
  td, th {
    padding: 0.15em 0.5em;
  }
  td.ip {
    min-width: 10em;
	text-align: left;
	font-family: monospace;
  }
</style>
</head>

<body>
`)

	for _, name := range names {
		fmt.Fprintf(w, "<h1>%s</h1>\n\n%s\n\n",
			name, registry[name].DebugHTML())
	}

	fmt.Fprintf(w, "</body>\n</html>\n")
}