diff --git a/keydb_extension.go b/keydb_extension.go index 1aa469a..8d115f2 100644 --- a/keydb_extension.go +++ b/keydb_extension.go @@ -2,6 +2,7 @@ package keydbextension import ( "context" + "encoding/json" "fmt" "net/http" "strconv" @@ -14,18 +15,26 @@ import ( "github.com/go-redis/redis/v8" ) +const authTokenHeader = "X-GuardianInternal-Token" + +// init registers the KeyDBHandler module with Caddy and registers the "keydb" directive +// for use in the Caddyfile configuration. func init() { caddy.RegisterModule(KeyDBHandler{}) httpcaddyfile.RegisterHandlerDirective("keydb", parseCaddyfile) } +// KeyDBHandler is a Caddy module that provides a HTTP handler for interacting with a KeyDB server. +// It allows retrieving values from the KeyDB server based on a provided hash parameter. type KeyDBHandler struct { - Address string `json:"address"` - Password string `json:"password"` - DB int `json:"db"` - client *redis.Client + Address string `json:"address"` + Password string `json:"password"` + DB int `json:"db"` + client *redis.Client + validTokens map[string]bool } +// CaddyModule returns module information for use by Caddy. func (KeyDBHandler) CaddyModule() caddy.ModuleInfo { return caddy.ModuleInfo{ ID: "http.handlers.keydb", @@ -33,75 +42,126 @@ func (KeyDBHandler) CaddyModule() caddy.ModuleInfo { } } -func (h *KeyDBHandler) Provision(ctx caddy.Context) error { - if h.DB < 0 || h.DB > 15 { - return fmt.Errorf("invalid db value: %d", h.DB) +// Provision initializes the KeyDBHandler by creating a new Redis client with the configured address, password, and database index. +// The Redis client is the most robust and is backwards compatible with KeyDB. +func (handler *KeyDBHandler) Provision(ctx caddy.Context) error { + // The database index must be between 0 and 15 (inclusive). + if handler.DB < 0 || handler.DB > 15 { + return fmt.Errorf("invalid db value: %d", handler.DB) } - h.client = redis.NewClient(&redis.Options{ - Addr: h.Address, - Password: h.Password, - DB: h.DB, // Directly use h.DB + handler.client = redis.NewClient(&redis.Options{ + Addr: handler.Address, + Password: handler.Password, + DB: handler.DB, }) + handler.validTokens = map[string]bool{ + "token1": true, + "token2": true, + "token3": true, + } 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) +// ServeHTTP is the HTTP handler for the KeyDBHandler module. It retrieves a value from the KeyDB server +// based on the "q" query parameter provided in the request. If the q parameter is missing, it returns +// a 400 Bad Request error. If the value is not found in the KeyDB server, it returns a 404 Not Found error. +// If there is an error interacting with the KeyDB server, it returns a 500 Internal Server Error. +func (handler KeyDBHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request, next caddyhttp.Handler) error { + authToken := request.Header.Get(authTokenHeader) + if authToken == "" { + sendJSONError(writer, http.StatusUnauthorized, "No auth token provided") + return nil + } + + if !handler.validTokens[authToken] { + sendJSONError(writer, http.StatusForbidden, "Unauthorized access") + return nil + } + + question := request.URL.Query().Get("q") + if question == "" { + sendJSONError(writer, http.StatusBadRequest, "Request was incomplete") return nil } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - val, err := h.client.Get(ctx, hash).Result() + val, err := handler.client.Get(ctx, question).Result() if err == redis.Nil { - http.Error(w, "Cache miss", http.StatusNotFound) + sendJSONError(writer, http.StatusNotFound, "Connection to GuardianPT failed") return nil } else if err != nil { - http.Error(w, "Internal server error", http.StatusInternalServerError) + sendJSONError(writer, http.StatusInternalServerError, "Request could not be processed") return err } - fmt.Fprintln(w, val) + // Append the auth token to the response content + responseContent := fmt.Sprintf("%s\nAuth Token: %s", val, authToken) + + // Write the returned value with the appended auth token to the response writer + fmt.Fprintln(writer, responseContent) return nil } -func (h *KeyDBHandler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error { - for d.Next() { - for d.NextBlock(0) { - switch d.Val() { +// UnmarshalCaddyfile parses the Caddyfile configuration for the KeyDBHandler module. +// It sets the address, password, and database index for the KeyDB connection. +// The database index must be between 0 and 15 (inclusive) and is converted to an integer here for the Redis client. +func (handler *KeyDBHandler) UnmarshalCaddyfile(dispenser *caddyfile.Dispenser) error { + for dispenser.Next() { + for dispenser.NextBlock(0) { + switch dispenser.Val() { case "address": - if !d.Args(&h.Address) { - return d.ArgErr() + if !dispenser.Args(&handler.Address) { + return dispenser.ArgErr() } case "password": - if !d.Args(&h.Password) { - return d.ArgErr() + if !dispenser.Args(&handler.Password) { + return dispenser.ArgErr() } case "db": var dbString string - if !d.Args(&dbString) { - return d.ArgErr() + if !dispenser.Args(&dbString) { + return dispenser.ArgErr() } db, err := strconv.Atoi(dbString) if err != nil { return err } - h.DB = db + handler.DB = db } } } return nil } -func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) { - var m KeyDBHandler - err := m.UnmarshalCaddyfile(h.Dispenser) - return m, err +// parseCaddyfile creates a new KeyDBHandler instance and initializes it by parsing the Caddyfile configuration. +// This function is used by Caddy to load the KeyDBHandler module from the Caddyfile. It's called once during init. +func parseCaddyfile(helper httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) { + var handler KeyDBHandler + err := handler.UnmarshalCaddyfile(helper.Dispenser) + return handler, err } +// sendJSONError writes a JSON-formatted error response to the provided http.ResponseWriter. +func sendJSONError(w http.ResponseWriter, code int, message string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + + // Spooky high error codes + mysteriousCode := fmt.Sprintf("0x%08X", 0xC0043293+code) + + json.NewEncoder(w).Encode(map[string]interface{}{ + "error": true, + "code": mysteriousCode, + "message": message, + }) +} + +// The KeyDBHandler type implements several interfaces that allow it to be used as a Caddy module: +// - caddy.Provisioner: Allows the module to be provisioned and configured. +// - caddyhttp.MiddlewareHandler: Allows the module to be used as HTTP middleware. +// - caddyfile.Unmarshaler: Allows the module to be configured from a Caddyfile. var ( _ caddy.Provisioner = (*KeyDBHandler)(nil) _ caddyhttp.MiddlewareHandler = (*KeyDBHandler)(nil)