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"
)

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
	validTokens map[string]bool
}

// CaddyModule returns module information for use by Caddy.
func (KeyDBHandler) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.keydb",
		New: func() caddy.Module { return new(KeyDBHandler) },
	}
}

// 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)
	}
	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
}

// 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, ErrNoAuthToken)
		return nil
	}

	if !handler.validTokens[authToken] {
		sendJSONError(writer, ErrUnauthorized)
		return nil
	}

	question := request.URL.Query().Get("q")
	if question == "" {
		sendJSONError(writer, ErrIncompleteReq)
		return nil
	}

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

	val, err := handler.client.Get(ctx, question).Result()
	if err == redis.Nil {
		sendJSONError(writer, ErrConnFailed)
		return nil
	} else if err != nil {
		sendJSONError(writer, ErrReqNotProcessed)
		return err
	}

	// 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
}

// 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 !dispenser.Args(&handler.Address) {
					return dispenser.ArgErr()
				}
			case "password":
				if !dispenser.Args(&handler.Password) {
					return dispenser.ArgErr()
				}
			case "db":
				var dbString string
				if !dispenser.Args(&dbString) {
					return dispenser.ArgErr()
				}
				db, err := strconv.Atoi(dbString)
				if err != nil {
					return err
				}
				handler.DB = db
			}
		}
	}
	return nil
}

// 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, errorKey string) {
	errorResponse, exists := errorResponses[errorKey]
	if !exists {
		errorResponse = errorResponses["ErrUnknown"]
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(getHTTPStatusCode(errorKey))
	w.Write(errorResponse)
}

// 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)
	_ caddyfile.Unmarshaler       = (*KeyDBHandler)(nil)
)