caddy-keydb-codebreaker/guardian_extension.go

package guardianextension

import (
	"fmt"
	"html"
	"net/http"
	"sync"

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

const authTokenHeader = "X-GuardianInternal-Token"

// init registers the GuardianCacheHandler module with Caddy and registers the "guardiancache" directive
// for use in the Caddyfile configuration.
func init() {
	caddy.RegisterModule((*GuardianCacheHandler)(nil))
	httpcaddyfile.RegisterHandlerDirective("guardiancache", parseCaddyfile)
}

// GuardianCacheHandler is a Caddy module that provides a cache handler for the Guardian application.
// It manages a cache of data stored in a JSON file, and validates access tokens.
// The mutex is necessary because maps are not thread-safe for concurrent read (and/or write) operations in Go.
type GuardianCacheHandler struct {
	DataFile    string `json:"data_file"`
	validTokens map[string]bool
	data        map[string]string
	mutex       sync.RWMutex
}

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

// Provision initializes the GuardianCacheHandler by loading data from a JSON file and setting up valid tokens.
func (handler *GuardianCacheHandler) Provision(ctx caddy.Context) error {
	handler.validTokens = map[string]bool{
		"token1": true,
		"token2": true,
		"token3": true,
	}

	errorResponsesOnce.Do(initErrorResponses)

	// Load data from file
	if err := handler.LoadData(); err != nil {
		return fmt.Errorf("failed to load data: %v", err)
	}

	return nil
}

// validateAuthToken checks if the auth token is valid. If the token is valid, it returns the escaped token.
// If the token is invalid, it sends an error response and returns an empty string which returns nil in the calling ServeHTTP method.
func (handler *GuardianCacheHandler) validateAuthToken(authToken string) (string, error) {
	if authToken == "" {
		return "", ErrNoAuthToken
	}

	if len(authToken) > 256 {
		return "", ErrUnauthorized
	}

	escapedAuthToken := html.EscapeString(authToken)

	// TODO: Temporary - this isn't how validity will be checked
	if !handler.validTokens[escapedAuthToken] {
		return "", ErrUnauthorized
	}

	return escapedAuthToken, nil
}

// ServeHTTP is the HTTP handler for the GuardianCacheHandler module. It retrieves a value from the loaded JSON data
// 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 loaded JSON data, it returns a 404 Not Found error.
// If there is an error interacting with the loaded JSON data, it returns a 500 Internal Server Error.
func (handler *GuardianCacheHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request, next caddyhttp.Handler) error {
	// Step 1: Get authentication token from the request header and validate it
	_, err := handler.validateAuthToken(request.Header.Get(authTokenHeader))

	if err != nil {
		sendJSONError(writer, err)
		return nil
	}

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

	if len(question) > 256 {
		sendJSONError(writer, ErrConnFailed)
		return nil
	}

	hashedQuestion := questionHash(question)

	handler.mutex.RLock()
	val, exists := handler.data[hashedQuestion]
	handler.mutex.RUnlock()

	if !exists {
		sendJSONError(writer, ErrConnFailed)
		return nil
	}

	fmt.Fprintln(writer, val)
	return nil
}

// UnmarshalCaddyfile parses the Caddyfile configuration for the GuardianCacheHandler module.
// It sets the path to the JSON data file for the GuardianCacheHandler.
// The parameter "data_file" specifies the location of the JSON file to be loaded.
func (handler *GuardianCacheHandler) UnmarshalCaddyfile(dispenser *caddyfile.Dispenser) error {
	for dispenser.Next() {
		for dispenser.NextBlock(0) {
			switch dispenser.Val() {
			case "filepath":
				if !dispenser.Args(&handler.DataFile) {
					return dispenser.ArgErr()
				}
			default:
				return dispenser.Errf("unrecognized parameter: %s", dispenser.Val())
			}
		}
	}
	return nil
}

// The GuardianCacheHandler 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           = (*GuardianCacheHandler)(nil)
	_ caddyhttp.MiddlewareHandler = (*GuardianCacheHandler)(nil)
	_ caddyfile.Unmarshaler       = (*GuardianCacheHandler)(nil)
)