caddy-keydb-codebreaker/keydb_extension.go
2024-08-13 14:14:59 -07:00

96 lines
2 KiB
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)
)