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