From 4be7164f2041ee38f5e5b3f6aca6fb38847be406 Mon Sep 17 00:00:00 2001
From: gadotroee <55343099+gadotroee@users.noreply.github.com>
Date: Mon, 31 Jan 2022 10:49:44 +0200
Subject: [PATCH] Fix routing in frontend not working (#723)

---
 agent/main.go                       | 23 ++++++++++++++---------
 agent/pkg/routes/not_found_route.go | 18 ------------------
 2 files changed, 14 insertions(+), 27 deletions(-)
 delete mode 100644 agent/pkg/routes/not_found_route.go

diff --git a/agent/main.go b/agent/main.go
index 0c3e377c7..1149ef367 100644
--- a/agent/main.go
+++ b/agent/main.go
@@ -62,7 +62,6 @@ const (
 	socketConnectionRetries    = 30
 	socketConnectionRetryDelay = time.Second * 2
 	socketHandshakeTimeout     = time.Second * 2
-	uiIndexPath                = "./site/index.html"
 )
 
 func main() {
@@ -248,16 +247,23 @@ func hostApi(socketHarOutputChannel chan<- *tapApi.OutputChannelItem) {
 
 	app.Use(DisableRootStaticCache())
 
-	if err := setUIFlags(); err != nil {
-		logger.Log.Errorf("Error setting ui mode, err: %v", err)
+	var staticFolder string
+	if config.Config.StandaloneMode {
+		staticFolder = "./site-standalone"
+	} else {
+		staticFolder = "./site"
 	}
 
-	if config.Config.StandaloneMode {
-		app.Use(static.ServeRoot("/", "./site-standalone"))
-	} else {
-		app.Use(static.ServeRoot("/", "./site"))
+	indexStaticFile := staticFolder + "/index.html"
+	if err := setUIFlags(indexStaticFile); err != nil {
+		logger.Log.Errorf("Error setting ui flags, err: %v", err)
 	}
 
+	app.Use(static.ServeRoot("/", staticFolder))
+	app.NoRoute(func(c *gin.Context) {
+		c.File(indexStaticFile)
+	})
+
 	app.Use(middlewares.CORSMiddleware()) // This has to be called after the static middleware, does not work if its called before
 
 	api.WebSocketRoutes(app, &eventHandlers, startTime)
@@ -278,7 +284,6 @@ func hostApi(socketHarOutputChannel chan<- *tapApi.OutputChannelItem) {
 	routes.EntriesRoutes(app)
 	routes.MetadataRoutes(app)
 	routes.StatusRoutes(app)
-	routes.NotFoundRoute(app)
 	utils.StartServer(app)
 }
 
@@ -293,7 +298,7 @@ func DisableRootStaticCache() gin.HandlerFunc {
 	}
 }
 
-func setUIFlags() error {
+func setUIFlags(uiIndexPath string) error {
 	read, err := ioutil.ReadFile(uiIndexPath)
 	if err != nil {
 		return err
diff --git a/agent/pkg/routes/not_found_route.go b/agent/pkg/routes/not_found_route.go
deleted file mode 100644
index bde062ff0..000000000
--- a/agent/pkg/routes/not_found_route.go
+++ /dev/null
@@ -1,18 +0,0 @@
-package routes
-
-import (
-	"github.com/gin-gonic/gin"
-	"net/http"
-)
-
-// NotFoundRoute defines the 404 Error route.
-func NotFoundRoute(app *gin.Engine) {
-	app.Use(
-		func(c *gin.Context) {
-			c.JSON(http.StatusNotFound, map[string]interface{}{
-				"error": true,
-				"msg":   "sorry, endpoint is not found",
-			})
-		},
-	)
-}