Repository: trustwallet/blockatlas
Pull request #1016 (2020-04-07 20:58:30)

[Gincache] Add Cache-Control headers for CacheMiddleware


Patch:
diff --git a/Makefile b/Makefile
index e4dcf75..18f743f 100644
--- a/Makefile
+++ b/Makefile
@@ -275,7 +275,7 @@ go-lint-install:
 
 go-lint:
 	@echo "  >  Running golint"
-	bin/golangci-lint run
+	bin/golangci-lint run --timeout=2m
 
 .PHONY: help
 all: help
diff --git a/api/endpoint/transaction.go b/api/endpoint/transaction.go
index 52ba419..6769021 100644
--- a/api/endpoint/transaction.go
+++ b/api/endpoint/transaction.go
@@ -78,6 +78,11 @@ func GetTransactionsHistory(c *gin.Context, txAPI blockatlas.TxAPI, tokenTxAPI b
 
 	page := make(blockatlas.TxPage, 0)
 	for _, tx := range txs {
+		if tx.Direction != "" {
+			page = append(page, tx)
+			continue
+		}
+
 		tx.Direction = tx.GetTransactionDirection(address)
 		page = append(page, tx)
 	}
diff --git a/api/middleware/cache.go b/api/middleware/cache.go
index efba3bb..8d87aa0 100644
--- a/api/middleware/cache.go
+++ b/api/middleware/cache.go
@@ -120,21 +120,21 @@ func (mc *memCache) setCache(k string, x interface{}, d time.Duration) {
 	memoryCache.cache.Set(k, b, d)
 }
 
-func (mc *memCache) getCache(key string) (cacheResponse, error) {
+func (mc *memCache) getCache(key string) (cacheResponse, time.Time, error) {
 	var result cacheResponse
-	c, ok := mc.cache.Get(key)
+	c, exp, ok := mc.cache.GetWithExpiration(key)
 	if !ok {
-		return result, fmt.Errorf("gin-cache: invalid cache key %s", key)
+		return result, time.Time{}, fmt.Errorf("gin-cache: invalid cache key %s", key)
 	}
 	r, ok := c.([]byte)
 	if !ok {
-		return result, errors.E("validator cache: failed to cast cache to bytes")
+		return result, time.Time{}, errors.E("validator cache: failed to cast cache to bytes")
 	}
 	err := json.Unmarshal(r, &result)
 	if err != nil {
-		return result, errors.E(err, "not found")
+		return result, time.Time{}, errors.E(err, "not found")
 	}
-	return result, nil
+	return result, exp, nil
 }
 
 func generateKey(c *gin.Context) string {
@@ -154,12 +154,15 @@ func CacheMiddleware(expiration time.Duration, handle gin.HandlerFunc) gin.Handl
 	return func(c *gin.Context) {
 		defer c.Next()
 		key := generateKey(c)
-		mc, err := memoryCache.getCache(key)
+		mc, exp, err := memoryCache.getCache(key)
 		if err != nil || mc.Data == nil {
 			writer := newCachedWriter(expiration, c.Writer, key)
+
+			expForSet := uint(expiration.Seconds())
+			writer.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", expForSet))
+
 			c.Writer = writer
 			handle(c)
-
 			if c.IsAborted() {
 				memoryCache.deleteCache(key)
 			}
@@ -172,6 +175,15 @@ func CacheMiddleware(expiration time.Duration, handle gin.HandlerFunc) gin.Handl
 				c.Writer.Header().Set(k, v)
 			}
 		}
+
+		expCacheTime := (exp.UnixNano() - time.Now().UnixNano()) / int64(time.Second)
+
+		if expCacheTime < 0 {
+			logger.Error("Invalid expiration time")
+		} else {
+			c.Writer.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", expCacheTime))
+		}
+
 		_, err = c.Writer.Write(mc.Data)
 		if err != nil {
 			memoryCache.deleteCache(key)


Test patch:
diff --git a/api/middleware/cache_test.go b/api/middleware/cache_test.go
index c2cda6a..72d0aaf 100644
--- a/api/middleware/cache_test.go
+++ b/api/middleware/cache_test.go
@@ -64,3 +64,24 @@ func TestCachePageExpire(t *testing.T) {
 	assert.Equal(t, http.StatusOK, w2.Code)
 	assert.NotEqual(t, w1.Body.String(), w2.Body.String())
 }
+
+func TestCacheControl(t *testing.T) {
+	router := gin.New()
+	router.GET("/cache_ping_control", CacheMiddleware(time.Second*30, func(c *gin.Context) {
+		c.JSON(http.StatusOK, "pong "+fmt.Sprint(time.Now().UnixNano()))
+	}))
+
+	w1 := performRequest("GET", "/cache_ping_control", router)
+	w1CacheControl := w1.Header().Get("Cache-Control")
+	assert.NotEqual(t, "no-cache", w1CacheControl)
+	time.Sleep(time.Second * 1)
+	w2 := performRequest("GET", "/cache_ping_control", router)
+	w2CacheControl := w2.Header().Get("Cache-Control")
+
+	assert.NotEqual(t, w1CacheControl, w2CacheControl)
+	assert.Equal(t, w1.Body.String(), w2.Body.String())
+
+	assert.Equal(t, http.StatusOK, w1.Code)
+	assert.Equal(t, http.StatusOK, w2.Code)
+
+}