From 8c0f8d4151ce4085969cf7d157826d2431e9a759 Mon Sep 17 00:00:00 2001
From: Piotr Biernat <keedosn@gmail.com>
Date: Sun, 25 Jul 2021 20:13:03 +0200
Subject: [PATCH] [fix] Added support for headers cache

---
 pkg/cache/cache.go            |  3 +--
 pkg/cache/memory_datastore.go |  2 +-
 pkg/cache/response.go         |  8 +++++---
 pkg/server/server.go          | 33 +++++++++++++++++++++++----------
 4 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go
index 9296fb3..526319b 100644
--- a/pkg/cache/cache.go
+++ b/pkg/cache/cache.go
@@ -28,10 +28,9 @@ func GetCacheDatastore(config config.Cache) *CacheDatastore {
 		datastore = NewMemoryDatastore()
 	}
 
-	//fail-safe switch to memory datasource
+	// fail-safe switch to memory datasource
 	if !datastore.IsConnected() {
 		log.Println("Cache server is not responding, switching to memory cache.")
-
 		datastore = NewMemoryDatastore()
 	}
 
diff --git a/pkg/cache/memory_datastore.go b/pkg/cache/memory_datastore.go
index 9b5e065..411c453 100644
--- a/pkg/cache/memory_datastore.go
+++ b/pkg/cache/memory_datastore.go
@@ -16,7 +16,7 @@ import (
 
 type TtlItem struct {
 	ts  int // timestamp
-	ttl int
+	ttl int // ttl in seconds
 }
 
 func NewMemoryDatastore() *MemoryDatastore {
diff --git a/pkg/cache/response.go b/pkg/cache/response.go
index 2ea9d91..bd036e4 100644
--- a/pkg/cache/response.go
+++ b/pkg/cache/response.go
@@ -14,10 +14,12 @@ import (
 	"log"
 )
 
+type Headers map[string]string
+
 type ResponseCache struct {
-	URL  string
-	Body string
-	// Headers map[string]string
+	URL     string
+	Body    string
+	Headers Headers
 }
 
 type ResponseCacheManager struct {
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 826434e..af33c57 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -23,8 +23,8 @@ import (
 )
 
 const (
-	Version = "0.1-dev"
-	Name    = "Vegvisir/" + Version
+	ServerVersion = "0.1-dev"
+	ServerName    = "Vegvisir/" + ServerVersion
 )
 
 type Server struct {
@@ -79,11 +79,11 @@ func (s *Server) Shutdown(ctx context.Context) { // TODO: wait for all connectio
 }
 
 func (s *Server) mainHandler(ctx *fasthttp.RequestCtx) {
-	ctx.Response.Header.Add(fasthttp.HeaderServer, Name)
+	// http := client.NewHttpClient(ctx)
 
 	// move all below logic to concrete handler or sth....
 	reqUrl, sReqUrl, sReqMethod := ctx.RequestURI(), string(ctx.RequestURI()), string(ctx.Method())
-	log.Println("Incomming request:", sReqMethod, sReqUrl)
+	log.Println("Incoming request:", sReqMethod, sReqUrl)
 
 	found, route := s.router.FindByRequestURL(reqUrl)
 	if !found {
@@ -96,7 +96,12 @@ func (s *Server) mainHandler(ctx *fasthttp.RequestCtx) {
 	if ok, data := s.respCM.Load(sReqUrl); ok {
 		log.Println("Read resp from cache: ", route.TargetUrl)
 
-		ctx.SetBody([]byte(data.Body)) // FIXME missing headers etc...
+		// copy headers and body from cache
+		ctx.Response.Header.DisableNormalizing()
+		for key, value := range data.Headers {
+			ctx.Response.Header.Set(key, value)
+		}
+		ctx.SetBody([]byte(data.Body))
 	} else {
 		log.Println("Send req to backend url: ", route.TargetUrl)
 
@@ -116,15 +121,23 @@ func (s *Server) mainHandler(ctx *fasthttp.RequestCtx) {
 			return
 		}
 
-		ctx.Response.Header.SetBytesV(fasthttp.HeaderContentType, bckResp.Header.ContentType())
-		ctx.SetStatusCode(bckResp.StatusCode())
+		headers := make(cache.Headers)
+		// rewrite headers from backend to gateway response
+		bckResp.Header.Set(fasthttp.HeaderServer, ServerName)
+		bckResp.Header.Del(fasthttp.HeaderXPoweredBy)
+		ctx.Response.Header.DisableNormalizing()
+		bckResp.Header.VisitAll(func(key, value []byte) {
+			headers[string(key)] = string(value)
+			ctx.Response.Header.SetBytesKV(key, value)
+		})
+		// ctx.SetStatusCode(bckResp.StatusCode())
 		ctx.SetBody(bckResp.Body())
 
 		// save response to cache
 		respCache := cache.ResponseCache{
-			URL:  sReqUrl,
-			Body: string(bckResp.Body()),
-			// Headers: []
+			URL:     sReqUrl,
+			Body:    string(bckResp.Body()),
+			Headers: headers,
 		} // FIXME: prepare resp cache struct in respCM.Save method or other service...
 		s.respCM.Save(sReqUrl, respCache)
 	}