Caddy module for counting visits per host using Redis Cluster

count_visits.go

package zaddy

import (
	"context"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/redis/go-redis/v9"
)

func init() {
	caddy.RegisterModule(CountVisitsMiddleware{})
	httpcaddyfile.RegisterHandlerDirective("count_visits", parseCountVisitsCaddyfile)
}

// CountVisitsMiddleware implements an HTTP handler that count visits in Redis
type CountVisitsMiddleware struct {
	rdb *redis.ClusterClient
}

// CaddyModule returns the Caddy module information.
func (CountVisitsMiddleware) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.count_visits",
		New: func() caddy.Module { return new(CountVisitsMiddleware) },
	}
}

// Provision implements caddy.Provisioner.
func (m *CountVisitsMiddleware) Provision(ctx caddy.Context) error {
	host := os.Getenv("REDIS_HOST")

	if !strings.Contains(host, ":") {
		host += ":6379"
	}
	m.rdb = redis.NewClusterClient(&redis.ClusterOptions{
		Addrs: []string{host},
	})
	return nil
}

// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (m CountVisitsMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	key := time.Now().Format(time.DateOnly) + ":" + r.Host
	m.rdb.IncrBy(context.Background(), "c:"+key, 1)
	m.rdb.PFAdd(context.Background(), "ucnt:"+key, GetUserIP(r)) // Unique Count
	return next.ServeHTTP(w, r)
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (m *CountVisitsMiddleware) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	return nil
}

func parseCountVisitsCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	var m CountVisitsMiddleware
	return m, nil
}

// Interface guards
var (
	_ caddy.Module                = (*CountVisitsMiddleware)(nil)
	_ caddy.Provisioner           = (*CountVisitsMiddleware)(nil)
	_ caddyhttp.MiddlewareHandler = (*CountVisitsMiddleware)(nil)
)