working pre-unique

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)