Added websocket hub for handling websocket clients, with a wsserve command
which runs a plain TCP server for receiving websocket hub clients, with server config, CSRF and rate limit middleware. - internal/websocket: Hub pushing messages to WebSocket clients (ServeHTTP, example.com/ws upgrade with optional CSRF check) and plain TCP clients (ServeTCP, newline-delimited JSON envelopes); ping/pong keepalive, per-client message rate limit, slow consumer handling, graceful Shutdown, debug logging of messages sent and received - internal/websocket/tcpserver: standalone TCP server for a hub, with accept back-off, a per-IP connection limit (server.websocket.connection_limit) and graceful shutdown - internal/cli/wsserve: `wsserve` command running the TCP server, with --listen/--host/--port overrides - internal/middleware: keyed token bucket rate limiter, real client IP from trusted proxies - internal/httpx: RFC 9457 problem responses, request IDs in context - internal/config: server section (host, port, shutdown timeout, rate limiter idle TTL, websocket); util.ByteSize settings (e.g. "32KiB"); lists of sections as single keys; empty lists stored as nil; isExplicitlySet for printLoadedConfig; JobSearchConfigSource renamed to JSCSourcesSource; TestSections checks every section's methods - internal/util: ByteSize, parsed and formatted with go-humanize
This commit is contained in:
bwilliams committed
1 parent
803a8929e6
commit
c9f42b62ad
42 files changed
+3619
-57
No files matched your search
@@ -0,0 +1,144 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"hash/maphash"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/time/rate"
|
||||
|
||||
"git.avsdev.uk/Go/job-search/internal/httpx"
|
||||
)
|
||||
|
||||
// Limiter is the rate limiting hook. Implementations decide whether a request
|
||||
// identified by key may proceed. Swap in a Redis/memcached-backed limiter for
|
||||
// multi-instance deployments by implementing this interface.
|
||||
type Limiter interface {
|
||||
Allow(key string) (allowed bool, retryAfter time.Duration)
|
||||
}
|
||||
|
||||
// LimiterFunc adapts a function to the Limiter interface.
|
||||
type LimiterFunc func(key string) (bool, time.Duration)
|
||||
|
||||
func (f LimiterFunc) Allow(key string) (bool, time.Duration) { return f(key) }
|
||||
|
||||
// KeyFunc derives the rate limit key from a request.
|
||||
type KeyFunc func(*http.Request) string
|
||||
|
||||
// KeyByIP keys on the real client IP (after proxy processing).
|
||||
func KeyByIP(r *http.Request) string { return ClientIP(r) }
|
||||
|
||||
// RateLimit turns a Limiter into middleware. Rejected requests get 429 with
|
||||
// a Retry-After header. A nil limiter yields a pass-through middleware.
|
||||
func RateLimit(l Limiter, key KeyFunc) func(http.Handler) http.Handler {
|
||||
if l == nil {
|
||||
return func(next http.Handler) http.Handler { return next }
|
||||
}
|
||||
if key == nil {
|
||||
key = KeyByIP
|
||||
}
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
ok, retry := l.Allow(key(r))
|
||||
if !ok {
|
||||
secs := int(math.Ceil(retry.Seconds()))
|
||||
if secs < 1 {
|
||||
secs = 1
|
||||
}
|
||||
w.Header().Set("Retry-After", strconv.Itoa(secs))
|
||||
httpx.WriteProblem(w, r, http.StatusTooManyRequests, "rate limit exceeded")
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// KeyedLimiter is an in-memory token bucket per key, sharded to reduce lock
|
||||
// contention, with idle-entry eviction.
|
||||
type KeyedLimiter struct {
|
||||
limit rate.Limit
|
||||
burst int
|
||||
ttl time.Duration
|
||||
seed maphash.Seed
|
||||
shards [32]limiterShard
|
||||
}
|
||||
|
||||
type limiterShard struct {
|
||||
mu sync.Mutex
|
||||
entries map[string]*limiterEntry
|
||||
}
|
||||
|
||||
type limiterEntry struct {
|
||||
lim *rate.Limiter
|
||||
last time.Time
|
||||
}
|
||||
|
||||
// NewKeyedLimiter creates a limiter allowing rps sustained requests per
|
||||
// second with the given burst. A janitor goroutine evicts entries idle for
|
||||
// longer than ttl until ctx is cancelled.
|
||||
func NewKeyedLimiter(ctx context.Context, rps float64, burst int, ttl time.Duration) *KeyedLimiter {
|
||||
k := &KeyedLimiter{limit: rate.Limit(rps), burst: burst, ttl: ttl, seed: maphash.MakeSeed()}
|
||||
for i := range k.shards {
|
||||
k.shards[i].entries = make(map[string]*limiterEntry)
|
||||
}
|
||||
if ttl > 0 {
|
||||
go k.janitor(ctx)
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
func (k *KeyedLimiter) shard(key string) *limiterShard {
|
||||
return &k.shards[maphash.String(k.seed, key)%uint64(len(k.shards))]
|
||||
}
|
||||
|
||||
// Allow implements Limiter.
|
||||
func (k *KeyedLimiter) Allow(key string) (bool, time.Duration) {
|
||||
now := time.Now()
|
||||
s := k.shard(key)
|
||||
s.mu.Lock()
|
||||
e := s.entries[key]
|
||||
if e == nil {
|
||||
e = &limiterEntry{lim: rate.NewLimiter(k.limit, k.burst)}
|
||||
s.entries[key] = e
|
||||
}
|
||||
e.last = now
|
||||
lim := e.lim
|
||||
s.mu.Unlock()
|
||||
|
||||
res := lim.ReserveN(now, 1)
|
||||
if !res.OK() {
|
||||
return false, time.Second
|
||||
}
|
||||
if d := res.DelayFrom(now); d > 0 {
|
||||
res.CancelAt(now)
|
||||
return false, d
|
||||
}
|
||||
return true, 0
|
||||
}
|
||||
|
||||
func (k *KeyedLimiter) janitor(ctx context.Context) {
|
||||
t := time.NewTicker(k.ttl / 2)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case now := <-t.C:
|
||||
for i := range k.shards {
|
||||
s := &k.shards[i]
|
||||
s.mu.Lock()
|
||||
for key, e := range s.entries {
|
||||
if now.Sub(e.last) > k.ttl {
|
||||
delete(s.entries, key)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRateLimit(t *testing.T) {
|
||||
l := NewKeyedLimiter(t.Context(), 1, 2, 0)
|
||||
h := RateLimit(l, nil)(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
|
||||
var last int
|
||||
for i := 0; i < 4; i++ {
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil))
|
||||
last = rec.Code
|
||||
if last == http.StatusTooManyRequests && rec.Header().Get("Retry-After") == "" {
|
||||
t.Fatal("missing Retry-After")
|
||||
}
|
||||
}
|
||||
if last != http.StatusTooManyRequests {
|
||||
t.Fatalf("want 429, got %d", last)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ========== RealIP - ClientInfo ==========
|
||||
|
||||
// ClientInfo describes the client as seen after proxy processing.
|
||||
type ClientInfo struct {
|
||||
IP netip.Addr // best-effort real client IP
|
||||
Peer netip.Addr // immediate TCP peer (after PROXY protocol, if used)
|
||||
Scheme string // "http" or "https", honouring trusted X-Forwarded-Proto
|
||||
ViaProxy bool // true when IP came from a trusted forwarding header
|
||||
}
|
||||
|
||||
// Client returns the ClientInfo stored by RealIP. Without RealIP it falls back
|
||||
// to the connection itself (see peerInfo) and ignores forwarding headers.
|
||||
func Client(r *http.Request) ClientInfo {
|
||||
if ci, ok := r.Context().Value(clientInfoKey).(ClientInfo); ok {
|
||||
return ci
|
||||
}
|
||||
return peerInfo(r)
|
||||
}
|
||||
|
||||
// ClientIP returns the client IP as a string ("" if unknown).
|
||||
func ClientIP(r *http.Request) string {
|
||||
ip := Client(r).IP
|
||||
if !ip.IsValid() {
|
||||
return ""
|
||||
}
|
||||
return ip.String()
|
||||
}
|
||||
|
||||
// IsSecure reports whether the request reached the edge over HTTPS. Without
|
||||
// RealIP it only reports whether the connection to this server is TLS.
|
||||
func IsSecure(r *http.Request) bool {
|
||||
return Client(r).Scheme == "https"
|
||||
}
|
||||
|
||||
// peerInfo builds a ClientInfo from the connection alone: the TCP peer (after
|
||||
// PROXY protocol, if used) and whether this server terminated TLS.
|
||||
func peerInfo(r *http.Request) ClientInfo {
|
||||
peer := addrFromRemote(r.RemoteAddr)
|
||||
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
|
||||
return ClientInfo{
|
||||
IP: peer,
|
||||
Peer: peer,
|
||||
Scheme: scheme,
|
||||
}
|
||||
}
|
||||
|
||||
// ========== RealIP ==========
|
||||
|
||||
// RealIPOptions configures RealIP.
|
||||
type RealIPOptions struct {
|
||||
Trusted []netip.Prefix
|
||||
// Headers is processed in increasing priority order - the last header in the
|
||||
// list that is found is the header used to determine a client IP address.
|
||||
Headers []string // e.g. X-Forwarded-For, Forwarded, X-Real-IP
|
||||
StripUntrusted bool
|
||||
}
|
||||
|
||||
type ctxKey int
|
||||
|
||||
const clientInfoKey ctxKey = iota
|
||||
|
||||
var forwardingHeaders = []string{
|
||||
"X-Forwarded-For",
|
||||
"X-Forwarded-Proto",
|
||||
"X-Forwarded-Host",
|
||||
"X-Forwarded-Port",
|
||||
"X-Forwarded-Ssl",
|
||||
"Forwarded",
|
||||
"X-Real-IP",
|
||||
}
|
||||
|
||||
// RealIP determines the client IP and scheme. Forwarding headers are only
|
||||
// honoured when the TCP peer is a trusted proxy; the X-Forwarded-For and
|
||||
// Forwarded chains are walked right-to-left skipping trusted hops, so a
|
||||
// client cannot spoof its address by prepending entries.
|
||||
//
|
||||
// PROXY protocol is handled at the listener (see server package), so by the
|
||||
// time this runs r.RemoteAddr already holds the address from the PROXY header.
|
||||
func RealIP(opts RealIPOptions) func(http.Handler) http.Handler {
|
||||
trusted := func(a netip.Addr) bool {
|
||||
for _, p := range opts.Trusted {
|
||||
if p.Contains(a) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
info := peerInfo(r)
|
||||
|
||||
if info.Peer.IsValid() && trusted(info.Peer) {
|
||||
for _, h := range opts.Headers {
|
||||
if ip, ok := clientFromHeader(r.Header, h, trusted); ok {
|
||||
info.IP = ip
|
||||
info.ViaProxy = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if p := forwardedProto(r.Header); p == "http" || p == "https" {
|
||||
info.Scheme = p
|
||||
}
|
||||
} else if opts.StripUntrusted {
|
||||
for _, h := range forwardingHeaders {
|
||||
r.Header.Del(h)
|
||||
}
|
||||
}
|
||||
|
||||
ctx := context.WithValue(r.Context(), clientInfoKey, info)
|
||||
next.ServeHTTP(w, r.WithContext(ctx))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func clientFromHeader(
|
||||
h http.Header, name string, trusted func(netip.Addr) bool,
|
||||
) (netip.Addr, bool) {
|
||||
switch http.CanonicalHeaderKey(name) {
|
||||
case "X-Forwarded-For":
|
||||
return walkChain(splitList(h.Values("X-Forwarded-For")), trusted)
|
||||
|
||||
case "Forwarded":
|
||||
var chain []string
|
||||
for _, el := range splitList(h.Values("Forwarded")) {
|
||||
if v, ok := forwardedParam(el, "for"); ok {
|
||||
chain = append(chain, v)
|
||||
} else {
|
||||
chain = append(chain, "") // element without for= breaks the chain
|
||||
}
|
||||
}
|
||||
return walkChain(chain, trusted)
|
||||
|
||||
case "X-Real-Ip":
|
||||
if v := strings.TrimSpace(h.Get("X-Real-IP")); v != "" {
|
||||
return parseIPToken(v)
|
||||
}
|
||||
}
|
||||
|
||||
return netip.Addr{}, false
|
||||
}
|
||||
|
||||
// walkChain returns the right-most address that is not a trusted proxy. If
|
||||
// every hop is trusted, the left-most is returned. Any unparsable hop
|
||||
// ("unknown", obfuscated identifiers, garbage) stops the walk.
|
||||
func walkChain(
|
||||
chain []string, trusted func(netip.Addr) bool,
|
||||
) (netip.Addr, bool) {
|
||||
var last netip.Addr
|
||||
for i := len(chain) - 1; i >= 0; i-- {
|
||||
ip, ok := parseIPToken(chain[i])
|
||||
if !ok {
|
||||
return netip.Addr{}, false
|
||||
}
|
||||
if !trusted(ip) {
|
||||
return ip, true
|
||||
}
|
||||
last = ip
|
||||
}
|
||||
return last, last.IsValid()
|
||||
}
|
||||
|
||||
func forwardedProto(h http.Header) string {
|
||||
if v := h.Get("X-Forwarded-Proto"); v != "" {
|
||||
first, _, _ := strings.Cut(v, ",")
|
||||
return strings.ToLower(strings.TrimSpace(first))
|
||||
}
|
||||
|
||||
if els := splitList(h.Values("Forwarded")); len(els) > 0 {
|
||||
if v, ok := forwardedParam(els[0], "proto"); ok {
|
||||
return strings.ToLower(v)
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// forwardedParam extracts a parameter from one Forwarded HTTP Extension
|
||||
// element (RFC 7239).
|
||||
func forwardedParam(element, key string) (string, bool) {
|
||||
for _, pair := range strings.Split(element, ";") {
|
||||
k, v, ok := strings.Cut(strings.TrimSpace(pair), "=")
|
||||
if ok && strings.EqualFold(k, key) {
|
||||
return strings.Trim(strings.TrimSpace(v), `"`), true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func splitList(values []string) []string {
|
||||
var out []string
|
||||
for _, v := range values {
|
||||
for _, part := range strings.Split(v, ",") {
|
||||
if part = strings.TrimSpace(part); part != "" {
|
||||
out = append(out, part)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// parseIPToken accepts "1.2.3.4", "1.2.3.4:80", "::1", "[::1]" and "[::1]:80".
|
||||
func parseIPToken(s string) (netip.Addr, bool) {
|
||||
s = strings.Trim(strings.TrimSpace(s), `"`)
|
||||
if ap, err := netip.ParseAddrPort(s); err == nil {
|
||||
return ap.Addr().Unmap(), true
|
||||
}
|
||||
|
||||
s = strings.TrimSuffix(strings.TrimPrefix(s, "["), "]")
|
||||
if a, err := netip.ParseAddr(s); err == nil {
|
||||
return a.Unmap().WithZone(""), true
|
||||
}
|
||||
|
||||
return netip.Addr{}, false
|
||||
}
|
||||
|
||||
func addrFromRemote(remote string) netip.Addr {
|
||||
if a, ok := parseIPToken(remote); ok {
|
||||
return a
|
||||
}
|
||||
return netip.Addr{}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// serveRealIP runs a request from peer with the given headers through RealIP
|
||||
// and returns the ClientInfo the next handler sees.
|
||||
func serveRealIP(t *testing.T, opts RealIPOptions, peer string, h map[string]string) ClientInfo {
|
||||
t.Helper()
|
||||
var got ClientInfo
|
||||
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { got = Client(r) })
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
req.RemoteAddr = peer
|
||||
for k, v := range h {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
RealIP(opts)(next).ServeHTTP(httptest.NewRecorder(), req)
|
||||
return got
|
||||
}
|
||||
|
||||
var defaultHeaders = []string{"X-Forwarded-For", "Forwarded", "X-Real-IP"}
|
||||
|
||||
func TestRealIP(t *testing.T) {
|
||||
opts := RealIPOptions{
|
||||
Trusted: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")},
|
||||
Headers: defaultHeaders,
|
||||
}
|
||||
// Spoofed left-most entry is ignored; right-most untrusted wins.
|
||||
c := serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{
|
||||
"X-Forwarded-For": "6.6.6.6, 203.0.113.9, 127.0.0.1", "X-Forwarded-Proto": "https",
|
||||
})
|
||||
if c.IP.String() != "203.0.113.9" || c.Scheme != "https" || !c.ViaProxy {
|
||||
t.Fatalf("xff: %+v", c)
|
||||
}
|
||||
c = serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{
|
||||
"Forwarded": `for="[2001:db8::7]:4711";proto=https`,
|
||||
})
|
||||
if c.IP.String() != "2001:db8::7" {
|
||||
t.Fatalf("forwarded: %+v", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUntrustedPeerHeadersIgnored(t *testing.T) {
|
||||
opts := RealIPOptions{
|
||||
Trusted: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
|
||||
Headers: defaultHeaders,
|
||||
}
|
||||
c := serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{"X-Forwarded-For": "203.0.113.9"})
|
||||
if c.IP.String() != "127.0.0.1" || c.ViaProxy {
|
||||
t.Fatalf("untrusted XFF honoured: %+v", c)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user