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/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 {
	db, err := strconv.Atoi(h.DB)
	if err != nil {
		return fmt.Errorf("invalid db value: %v", err)
	}
	h.client = redis.NewClient(&redis.Options{
		Addr:     h.Address,
		Password: h.Password,
		DB:       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":
				if !d.Args(&h.DB) {
					return d.ArgErr()
				}
			}
		}
	}
	return nil
}

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