handler.go

package httpcache

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"time"

	"net/http"

	"github.com/caddyserver/caddy/v2"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
	"github.com/hashicorp/consul/api"
	"github.com/adammakowskidev/caddy-cache-engine/backends"
	"github.com/adammakowskidev/caddy-cache-engine/extends/distributed"
	"github.com/adammakowskidev/caddy-cache-engine/pkg/helper"
	"go.uber.org/zap"
)

const (
	cacheHit    = "hit"
	cacheMiss   = "miss"
	cacheSkip   = "skip"
	cacheBypass = "bypass"
)

var (
	contextKeysToPreserve = [...]caddy.CtxKey{
		caddy.ReplacerCtxKey,
		"path_prefix",
		"mitm",
	}

	cache  *HTTPCache
	config *Config
)

func init() {
	caddy.RegisterModule(Handler{})
}

// getHandlerCache is a singleton of HTTPCache
func getHandlerCache() *HTTPCache {
	l.RLock()
	defer l.RUnlock()
	return cache
}

// Handler is a http handler as a middleware to cache the response
type Handler struct {
	Config   *Config    `json:"config,omitempty"`
	Cache    *HTTPCache `json:"-"`
	URLLocks *URLLock   `json:"-"`

	DistributedRaw json.RawMessage            `json:"distributed,omitempty" caddy:"namespace=distributed inline_key=distributed"`
	Distributed    *distributed.ConsulService `json:"-"`

	logger *zap.Logger
}

func (h *Handler) addStatusHeaderIfConfigured(w http.ResponseWriter, status string) {
	if h.Config.StatusHeader != "" {
		w.Header().Set(h.Config.StatusHeader, status)
	}
}

func (h *Handler) respond(w http.ResponseWriter, entry *Entry, cacheStatus string) error {
	h.addStatusHeaderIfConfigured(w, cacheStatus)
	copyHeaders(entry.Response.snapHeader, w.Header())

	// when the request method is head, we don't need ot perform write body
	if entry.Request.Method == "HEAD" {
		w.WriteHeader(entry.Response.Code)
		return nil
	}

	err := entry.WriteBodyTo(w)
	return err
}

func popOrNil(h *Handler, errChan chan error) (err error) {
	select {
	case err := <-errChan:
		if err != nil {
			h.logger.Error(fmt.Sprintf("popOrNil: %s", err.Error()))
		}
		return err
	default:
		return nil
	}

}

func (h *Handler) fetchUpstream(req *http.Request, next caddyhttp.Handler, key string) (*Entry, error) {
	// Create a new empty response
	response := NewResponse()

	errChan := make(chan error, 1)

	// Do the upstream fetching in background
	go func(req *http.Request, response *Response) {

		upstreamError := next.ServeHTTP(response, req)
		errChan <- upstreamError
		response.Close()

	}(req, response)

	// Wait headers to be sent
	response.WaitHeaders()

	// Create a new CacheEntry
	return NewEntry(key, req, response, h.Config), popOrNil(h, errChan)
}

// CaddyModule returns the Caddy module information
func (Handler) CaddyModule() caddy.ModuleInfo {
	return caddy.ModuleInfo{
		ID:  "http.handlers.http_cache",
		New: func() caddy.Module { return new(Handler) },
	}
}

func (h *Handler) provisionRuleMatchers() error {

	for _, raw := range h.Config.RuleMatchersRaws {

		switch raw.Type {
		case MatcherTypePath:

			var content *PathRuleMatcher
			err := json.Unmarshal(raw.Data, &content)
			if err != nil {
				return err
			}

			h.Config.RuleMatchers = append(h.Config.RuleMatchers, content)

		case MatcherTypeHeader:

			var content *HeaderRuleMatcher
			err := json.Unmarshal(raw.Data, &content)
			if err != nil {
				return err
			}

			h.Config.RuleMatchers = append(h.Config.RuleMatchers, content)
		}

	}

	return nil
}

func handleDeleteKey(data interface{}) error {
	// the type is api.KVPairs,
	// https://learn.hashicorp.com/tutorials/consul/distributed-semaphore?in=consul/developer-configuration
	// implement a distributed lock to handle del
	kvs, ok := data.(api.KVPairs)
	if !ok {
		return fmt.Errorf("non expected data type: %s", reflect.TypeOf(data))
	}

	for _, kv := range kvs {
		myIP, _ := helper.IPAddr()
		ip := string(kv.Value)

		// and to exclude self-trigger event
		if kv.Session != "" && ip != myIP.String() {

			caddy.Log().Named("distributed cache").
				Debug(fmt.Sprintf("perform cache delete: ip: %s=%s , content: %+v\n", ip, myIP.String(), kv))

			key := kv.Key

			// handle the dispatching delete key event, need to trim the prefix
			// to get the original key
			if strings.HasPrefix(kv.Key, distributed.Keyprefix) {
				key = strings.TrimLeft(kv.Key, distributed.Keyprefix+"/")
			}

			cache := getHandlerCache()
			cache.Del(key)
		}
	}

	return nil
}

func (h *Handler) provisionDistributed(ctx caddy.Context) error {
	if h.DistributedRaw != nil {

		// Register the watch handlers before the distributed module is provisioned
		wh := distributed.WatchHandler{
			Pg:       distributed.GetKeyPrefixParams(distributed.Keyprefix),
			Callback: handleDeleteKey,
		}

		distributed.RegisterWatchHandler(wh)

		val, err := ctx.LoadModule(h, "DistributedRaw") // this will call provision
		if err != nil {
			return fmt.Errorf("loading distributed module: %s", err.Error())
		}
		h.Distributed = val.(*distributed.ConsulService)
	}

	return nil
}

// Provision setups the configs
func (h *Handler) Provision(ctx caddy.Context) error {

	h.logger = ctx.Logger(h)

	if h.Config == nil {
		h.Config = getDefaultConfig()
	}

	err := h.provisionRuleMatchers()
	if err != nil {
		return err
	}

	// NOTE: A dirty work to assign the config and cache to global vars
	// There will be the corresponding functions to get each of them.
	// Therefore, we can call its Del to purge the cache via the admin interface
	distributedOn := h.DistributedRaw != nil
	cache = NewHTTPCache(h.Config, distributedOn)
	h.Cache = cache
	h.URLLocks = NewURLLock(h.Config)

	// Some type of the backends need extra initialization.
	switch h.Config.Type {
	case inMemory:
		if err := backends.InitGroupCacheRes(h.Config.CacheMaxMemorySize); err != nil {
			return err
		}
	}
	// load the guest module distributed
	err = h.provisionDistributed(ctx)
	if err != nil {
		return err
	}

	config = h.Config
	return nil
}

// Validate validates httpcache's configuration.
func (h *Handler) Validate() error {
	return nil
}

// Cleanup release the resources
func (h *Handler) Cleanup() error {
	var err error

	if h.Config.Type == inMemory {
		err = backends.ReleaseGroupCacheRes()
	}

	return err
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
	// add a log here to record the elapsed time (from receiving the request to send the response)
	start := time.Now()
	upstreamDuration := time.Duration(0)

	// TODO: think a proper way to log these info
	// research why can not write the log into files.
	defer func(h *Handler, t time.Time) {

		if upstreamDuration == 0 {
			duration := time.Since(t)
			h.logger.Debug("cache handler",
				zap.String("host", r.Host), // find a way to get upstream
				zap.String("method", r.Method),
				zap.String("uri", r.RequestURI),
				zap.Duration("request time", duration))
		} else {
			duration := time.Since(t)
			h.logger.Debug("cache handler",
				zap.String("host", r.Host), // find a way to get upstream
				zap.String("method", r.Method),
				zap.String("uri", r.RequestURI),
				zap.Duration("request time", duration),
				zap.Duration("upstream request time", upstreamDuration))
		}

	}(h, start)

	if !shouldUseCache(r, h.Config) {
		h.addStatusHeaderIfConfigured(w, cacheBypass)
		return next.ServeHTTP(w, r)
	}

	key := getKey(h.Config.CacheKeyTemplate, r)
	lock := h.URLLocks.Acquire(key)
	defer lock.Unlock()

	previousEntry, exists := h.Cache.Get(key, r, false)

	// First case: CACHE HIT
	// The response exists in cache and is public
	// It should be served as saved
	if exists && previousEntry.isPublic {
		if err := h.respond(w, previousEntry, cacheHit); err == nil {
			return nil
		} else if _, ok := err.(backends.NoPreCollectError); ok {
			// if the err is No pre collect, just return nil
			w.WriteHeader(previousEntry.Response.Code)
			return nil
		}
	}

	// Check whether the key exists in the groupcahce when the
	// distributed cache is enabled.
	// Currently, only support the memory backends
	if h.Distributed != nil {
		// new an entry without fetching the upstream
		response := NewResponse()
		entry := NewEntry(key, r, response, h.Config)
		err := entry.setBackend(r.Context(), h.Config)
		if err != nil {
			return caddyhttp.Error(http.StatusInternalServerError, err)
		}

		h.Cache.Put(r, entry, h.Config)
		response.Close()

		// NOTE: should set the content-length to the header manually when distributed
		// cache is enabled because we get the content from the other peer.
		// In this case, the snapHeader will not contain the content-length info
		if err = h.respond(w, entry, cacheHit); err == nil {
			return nil
		}

		// when error is NoPreCollectError, we need to fetch the resource from the
		// upstream.
		if e, ok := err.(backends.NoPreCollectError); !ok {
			return caddyhttp.Error(entry.Response.Code, e)
		}
	}

	// Second case: CACHE SKIP
	// The response is in cache but it is not public
	// It should NOT be served from cache
	// It should be fetched from upstream and check the new headers
	// To check if the new response changes to public

	// Third case: CACHE MISS
	// The response is not in cache
	// It should be fetched from upstream and save it in cache

	t := time.Now()
	entry, err := h.fetchUpstream(r, next, key)
	upstreamDuration = time.Since(t)

	if entry.Response.Code >= 500 {
		// using stale entry when available
		previousEntry, exists := h.Cache.Get(key, r, true)

		if exists && previousEntry.isPublic {
			if err := h.respond(w, previousEntry, cacheHit); err == nil {
				return nil
			} else if _, ok := err.(backends.NoPreCollectError); ok {
				// if the err is No pre collect, just return nil
				w.WriteHeader(previousEntry.Response.Code)
				return nil
			}
		}
	}

	if err != nil {
		return caddyhttp.Error(entry.Response.Code, err)
	}

	// Case when response was private but now is public
	if entry.isPublic {
		err := entry.setBackend(r.Context(), h.Config)
		if err != nil {
			return caddyhttp.Error(http.StatusInternalServerError, err)
		}

		h.Cache.Put(r, entry, h.Config)
		err = h.respond(w, entry, cacheMiss)
		if err != nil {
			h.logger.Error("cache handler", zap.Error(err))
			return caddyhttp.Error(entry.Response.Code, err)
		}

		return nil
	}

	err = h.respond(w, entry, cacheSkip)
	if err != nil {
		h.logger.Error("cache handler", zap.Error(err))
		return caddyhttp.Error(entry.Response.Code, err)
	}

	return nil
}

var (
	_ caddyhttp.MiddlewareHandler = (*Handler)(nil)
	_ caddy.Provisioner           = (*Handler)(nil)
	_ caddy.Validator             = (*Handler)(nil)
)