middleware_redis_cache.go

package main

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"github.com/gorilla/context"
	"io"
	"net/http"
	"strconv"
	"strings"
)

const (
	UPSTREAM_CACHE_HEADER_NAME     = "x-tyk-cache-action-set"
	UPSTREAM_CACHE_TTL_HEADER_NAME = "x-tyk-cache-action-set-ttl"
)

// RedisCacheMiddleware is a caching middleware that will pull data from Redis instead of the upstream proxy
type RedisCacheMiddleware struct {
	TykMiddleware
	CacheStore StorageHandler
}

type RedisCacheMiddlewareConfig struct{}

// New lets you do any initialisations for the object can be done here
func (m *RedisCacheMiddleware) New() {}

// GetConfig retrieves the configuration from the API config - we user mapstructure for this for simplicity
func (m *RedisCacheMiddleware) GetConfig() (interface{}, error) {
	var thisModuleConfig RedisCacheMiddlewareConfig
	return thisModuleConfig, nil
}

func (m RedisCacheMiddleware) CreateCheckSum(req *http.Request, keyName string) string {
	h := md5.New()
	toEncode := strings.Join([]string{req.Method, req.URL.RawQuery}, "-")
	io.WriteString(h, toEncode)
	reqChecksum := hex.EncodeToString(h.Sum(nil))

	cacheKey := m.Spec.APIDefinition.APIID + keyName + reqChecksum

	return cacheKey
}

// ProcessRequest will run any checks on the request on the way through the system, return an error to have the chain fail
func (m *RedisCacheMiddleware) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {

	// Allow global cache disabe
	if !m.Spec.APIDefinition.CacheOptions.EnableCache {
		return nil, 200
	}

	var stat RequestStatus
	// Only allow idempotent (safe) methods
	if r.Method == "GET" || r.Method == "OPTIONS" || r.Method == "HEAD" {
		// Lets see if we can throw a sledgehammer at this
		if m.Spec.APIDefinition.CacheOptions.CacheAllSafeRequests {
			stat = StatusCached
		} else {
			// New request checker, more targetted, less likely to fail
			_, versionPaths, _, _ := m.TykMiddleware.Spec.GetVersionData(r)
			found, _ := m.TykMiddleware.Spec.CheckSpecMatchesStatus(r.URL.Path, r.Method, versionPaths, Cached)
			if found {
				stat = StatusCached
			}
		}

		// Cached route matched, let go
		if stat == StatusCached {
			authHeaderValue := context.Get(r, AuthHeaderValue).(string)
			thisKey := m.CreateCheckSum(r, authHeaderValue)
			retBlob, found := m.CacheStore.GetKey(thisKey)
			if found != nil {
				log.Debug("Cache enabled, but record not found")
				// Pass through to proxy AND CACHE RESULT
				sNP := SuccessHandler{m.TykMiddleware}

				// This passes through and will write the value to the writer, but spit out a copy for the cache
				reqVal := sNP.ServeHTTP(w, r)

				cacheThisRequest := true
				cacheTTL := m.Spec.APIDefinition.CacheOptions.CacheTimeout
				// Are we using upstream cache control?
				if m.Spec.APIDefinition.CacheOptions.EnableUpstreamCacheControl {
					log.Debug("Upstream control enabled")
					// Do we cache?
					if reqVal.Header.Get(UPSTREAM_CACHE_HEADER_NAME) == "" {
						log.Warning("Upstream cache action not found, not caching")
						cacheThisRequest = false
					}
					// Do we override TTL?
					ttl := reqVal.Header.Get(UPSTREAM_CACHE_TTL_HEADER_NAME)
					if ttl != "" {
						log.Debug("TTL Set upstream")
						cacheAsInt, valErr := strconv.Atoi(ttl)
						if valErr != nil {
							log.Error("Failed to decode TTL cache value: ", valErr)
							cacheTTL = m.Spec.APIDefinition.CacheOptions.CacheTimeout
						}
						cacheTTL = int64(cacheAsInt)
					}
				}

				if cacheThisRequest {
					log.Debug("Caching request to redis")
					var wireFormatReq bytes.Buffer
					reqVal.Write(&wireFormatReq)
					log.Debug("Cache TTL is:", cacheTTL)
					go m.CacheStore.SetKey(thisKey, wireFormatReq.String(), cacheTTL)

				}
				return nil, 666

			}

			retObj := bytes.NewReader([]byte(retBlob))

			asBufioReader := bufio.NewReader(retObj)
			newRes, resErr := http.ReadResponse(asBufioReader, r)
			if resErr != nil {
				log.Error("Could not create response object: ", resErr)
			}

			defer newRes.Body.Close()
			for _, h := range hopHeaders {
				newRes.Header.Del(h)
			}

			copyHeader(w.Header(), newRes.Header)
			thisSessionState := context.Get(r, SessionData).(SessionState)
			w.Header().Add("x-tyk-cached-response", "1")
			w.Header().Set("X-RateLimit-Limit", strconv.Itoa(int(thisSessionState.QuotaMax)))
			w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(int(thisSessionState.QuotaRemaining)))
			w.Header().Set("X-RateLimit-Reset", strconv.Itoa(int(thisSessionState.QuotaRenews)))

			w.WriteHeader(newRes.StatusCode)
			m.Proxy.copyResponse(w, newRes.Body)

			// Record analytics
			sNP := SuccessHandler{m.TykMiddleware}
			go sNP.RecordHit(w, r, 0)

			// Stop any further execution
			return nil, 666
		}
	}

	return nil, 200
}