caddy-keydb-codebreaker/keydb_extension.go

package keydbextension

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"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/go-redis/redis/v8"
)

func init() {
	caddy.RegisterModule(KeyDBHandler{})
}

type KeyDBHandler struct {
	Address  string `json:"address"`
	Password string `json:"password"`
	DB       int    `json:"db"`
	client   *redis.Client
}

func (KeyDBHandler) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.keydb",
		New: func() caddy.Module { return new(KeyDBHandler) },
	}
}

func (h *KeyDBHandler) Provision(ctx caddy.Context) error {
	if h.DB < 0 || h.DB > 15 {
		return fmt.Errorf("invalid db value: %d", h.DB)
	}
	h.client = redis.NewClient(&redis.Options{
		Addr:     h.Address,
		Password: h.Password,
		DB:       h.DB, // Directly use h.DB
	})
	return nil
}

func (h KeyDBHandler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	hash := r.URL.Query().Get("hash")
	if hash == "" {
		http.Error(w, "Missing hash parameter", http.StatusBadRequest)
		return nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	val, err := h.client.Get(ctx, hash).Result()
	if err == redis.Nil {
		http.Error(w, "Cache miss", http.StatusNotFound)
		return nil
	} else if err != nil {
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return err
	}

	fmt.Fprintln(w, val)
	return nil
}

func (h *KeyDBHandler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
	for d.Next() {
		for d.NextBlock(0) {
			switch d.Val() {
			case "address":
				if !d.Args(&h.Address) {
					return d.ArgErr()
				}
			case "password":
				if !d.Args(&h.Password) {
					return d.ArgErr()
				}
			case "db":
				var dbString string
				if !d.Args(&dbString) {
					return d.ArgErr()
				}
				db, err := strconv.Atoi(dbString)
				if err != nil {
					return err
				}
				h.DB = db
			}
		}
	}
	return nil
}

func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	var m KeyDBHandler
	err := m.UnmarshalCaddyfile(h.Dispenser)
	return m, err
}

var (
	_ caddy.Provisioner           = (*KeyDBHandler)(nil)
	_ caddyhttp.MiddlewareHandler = (*KeyDBHandler)(nil)
	_ caddyfile.Unmarshaler       = (*KeyDBHandler)(nil)
)