From e16f8c38bb9a487254ce879029d0e42ce8b20ac2 Mon Sep 17 00:00:00 2001
From: Jacob Gunther <jacobgunther03+github@gmail.com>
Date: Sun, 21 Aug 2022 19:24:39 -0500
Subject: [PATCH] Add X-Cache-Time-Remaining header

---
 config.example.yml |  4 ++--
 src/redis.go       | 32 ++++++++++++++++++++++++++++----
 src/routes.go      | 23 +++++++++++------------
 src/status.go      | 37 +++++++++++++++++++++++++++----------
 4 files changed, 68 insertions(+), 28 deletions(-)

diff --git a/config.example.yml b/config.example.yml
index c5dd470..ae103c9 100644
--- a/config.example.yml
+++ b/config.example.yml
@@ -3,6 +3,6 @@ port: 3001
 redis: redis://127.0.0.1:6379/0
 cache:
   enable: false
-  java_cache_duration: 5m
-  bedrock_cache_duration: 5m
+  java_cache_duration: 1m
+  bedrock_cache_duration: 1m
   icon_cache_duration: 15m
\ No newline at end of file
diff --git a/src/redis.go b/src/redis.go
index a6c3157..970b4a0 100644
--- a/src/redis.go
+++ b/src/redis.go
@@ -57,9 +57,27 @@ func (r *Redis) Exists(key string) (bool, error) {
 	return val == 1, err
 }
 
-func (r *Redis) GetString(key string) (string, error) {
+func (r *Redis) TTL(key string) (time.Duration, error) {
 	if !r.Enabled {
-		return "", nil
+		return 0, nil
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+
+	defer cancel()
+
+	res := r.Client.TTL(ctx, key)
+
+	if err := res.Err(); err != nil {
+		return 0, err
+	}
+
+	return res.Result()
+}
+
+func (r *Redis) GetJSON(key string, value interface{}) error {
+	if !r.Enabled {
+		return nil
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
@@ -69,10 +87,16 @@ func (r *Redis) GetString(key string) (string, error) {
 	res := r.Client.Get(ctx, key)
 
 	if err := res.Err(); err != nil {
-		return "", err
+		return err
 	}
 
-	return res.Result()
+	data, err := res.Bytes()
+
+	if err != nil {
+		return err
+	}
+
+	return json.Unmarshal(data, value)
 }
 
 func (r *Redis) GetBytes(key string) ([]byte, error) {
diff --git a/src/routes.go b/src/routes.go
index e7287b9..829b4a9 100644
--- a/src/routes.go
+++ b/src/routes.go
@@ -3,6 +3,7 @@ package main
 import (
 	"log"
 	"net/http"
+	"strconv"
 
 	"github.com/gofiber/fiber/v2"
 )
@@ -27,18 +28,17 @@ func JavaStatusHandler(ctx *fiber.Ctx) error {
 		return ctx.Status(http.StatusBadRequest).SendString("Invalid address value")
 	}
 
-	response, err := GetJavaStatus(host, port)
+	response, expiresAt, err := GetJavaStatus(host, port)
 
 	if err != nil {
 		return err
 	}
 
-	switch v := response.(type) {
-	case string:
-		return ctx.Type("json").SendString(v)
-	default:
-		return ctx.JSON(response)
+	if expiresAt != nil {
+		ctx.Set("X-Cache-Time-Remaining", strconv.Itoa(int(expiresAt.Seconds())))
 	}
+
+	return ctx.JSON(response)
 }
 
 func BedrockStatusHandler(ctx *fiber.Ctx) error {
@@ -48,18 +48,17 @@ func BedrockStatusHandler(ctx *fiber.Ctx) error {
 		return ctx.Status(http.StatusBadRequest).SendString("Invalid address value")
 	}
 
-	response, err := GetBedrockStatus(host, port)
+	response, expiresAt, err := GetBedrockStatus(host, port)
 
 	if err != nil {
 		return err
 	}
 
-	switch v := response.(type) {
-	case string:
-		return ctx.Type("json").SendString(v)
-	default:
-		return ctx.JSON(response)
+	if expiresAt != nil {
+		ctx.Set("X-Cache-Time-Remaining", strconv.Itoa(int(expiresAt.Seconds())))
 	}
+
+	return ctx.JSON(response)
 }
 
 func IconHandler(ctx *fiber.Ctx) error {
diff --git a/src/status.go b/src/status.go
index f91b002..732bdbd 100644
--- a/src/status.go
+++ b/src/status.go
@@ -4,6 +4,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/mcstatus-io/mcutil"
 )
@@ -75,48 +76,64 @@ type Mod struct {
 	Version string `json:"version"`
 }
 
-func GetJavaStatus(host string, port uint16) (interface{}, error) {
+func GetJavaStatus(host string, port uint16) (interface{}, *time.Duration, error) {
 	cacheKey := fmt.Sprintf("java:%s-%d", host, port)
 
 	cacheExists, err := r.Exists(cacheKey)
 
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	if cacheExists {
-		return r.GetString(cacheKey)
+		var response JavaStatusResponse
+
+		if err = r.GetJSON(cacheKey, &response); err != nil {
+			return nil, nil, err
+		}
+
+		expiresAt, err := r.TTL(cacheKey)
+
+		return response, &expiresAt, err
 	}
 
 	response := FetchJavaStatus(host, port)
 
 	if err := r.SetJSON(cacheKey, response, config.Cache.JavaCacheDuration); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return response, nil
+	return response, nil, nil
 }
 
-func GetBedrockStatus(host string, port uint16) (interface{}, error) {
+func GetBedrockStatus(host string, port uint16) (interface{}, *time.Duration, error) {
 	cacheKey := fmt.Sprintf("bedrock:%s-%d", host, port)
 
 	cacheExists, err := r.Exists(cacheKey)
 
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	if cacheExists {
-		return r.GetString(cacheKey)
+		var response BedrockStatusResponse
+
+		if err = r.GetJSON(cacheKey, &response); err != nil {
+			return nil, nil, err
+		}
+
+		expiresAt, err := r.TTL(cacheKey)
+
+		return response, &expiresAt, err
 	}
 
 	response := FetchBedrockStatus(host, port)
 
 	if err := r.SetJSON(cacheKey, response, config.Cache.BedrockCacheDuration); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return response, nil
+	return response, nil, nil
 }
 
 func GetServerIcon(host string, port uint16) ([]byte, error) {