Support redirecting 404 to index

Vilsol · Jan 16, 2022 · 5771cc7 · 5771cc7
1 parent e28bb50
commit 5771cc7
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 10 deletions.
diff --git a/cmd/serve.go b/cmd/serve.go
@@ -30,6 +30,8 @@ func init() {
 	ServeCMD.PersistentFlags().String("bot-proxy", "", "Bot proxy URL")
 	ServeCMD.PersistentFlags().String("bot-agents", botAgents, "Bot User-Agent header regex")
 
+	ServeCMD.PersistentFlags().Bool("404-index", false, "Redirect any 404 to the index file")
+
 	_ = viper.BindPFlag("paths", ServeCMD.PersistentFlags().Lookup("paths"))
 	_ = viper.BindPFlag("watch", ServeCMD.PersistentFlags().Lookup("watch"))
 
@@ -50,6 +52,8 @@ func init() {
 	_ = viper.BindPFlag("bot.proxy", ServeCMD.PersistentFlags().Lookup("bot-proxy"))
 	_ = viper.BindPFlag("bot.agents", ServeCMD.PersistentFlags().Lookup("bot-agents"))
 
+	_ = viper.BindPFlag("404-index", ServeCMD.PersistentFlags().Lookup("404-index"))
+
 	RootCMD.AddCommand(ServeCMD)
 }
 

diff --git a/docs/yeet.md b/docs/yeet.md
@@ -15,4 +15,4 @@ yeet is an in-memory indexed static file webserver
 
 * [yeet serve](yeet_serve.md)	 - Serve files with yeet
 
-###### Auto generated by spf13/cobra on 7-Jan-2022
+###### Auto generated by spf13/cobra on 16-Jan-2022
diff --git a/docs/yeet_serve.md b/docs/yeet_serve.md
@@ -5,6 +5,7 @@ Serve files with yeet
 ### Options
 
 ```
+      --404-index                  Redirect any 404 to the index file
       --bot-agents string          Bot User-Agent header regex (default "bot|crawl|spider|external|meta|scrap|archive|discourse")
       --bot-proxy string           Bot proxy URL
       --expiry                     Use cache expiry
@@ -36,4 +37,4 @@ Serve files with yeet
 * [yeet serve s3](yeet_serve_s3.md)	 - Serve an S3 bucket
 * [yeet serve s3-redis](yeet_serve_s3-redis.md)	 - Serve Redis-backed S3 buckets
 
-###### Auto generated by spf13/cobra on 7-Jan-2022
+###### Auto generated by spf13/cobra on 16-Jan-2022
diff --git a/docs/yeet_serve_local.md b/docs/yeet_serve_local.md
@@ -15,6 +15,7 @@ yeet serve local [flags]
 ### Options inherited from parent commands
 
 ```
+      --404-index                  Redirect any 404 to the index file
       --bot-agents string          Bot User-Agent header regex (default "bot|crawl|spider|external|meta|scrap|archive|discourse")
       --bot-proxy string           Bot proxy URL
       --expiry                     Use cache expiry
@@ -37,4 +38,4 @@ yeet serve local [flags]
 
 * [yeet serve](yeet_serve.md)	 - Serve files with yeet
 
-###### Auto generated by spf13/cobra on 7-Jan-2022
+###### Auto generated by spf13/cobra on 16-Jan-2022
diff --git a/docs/yeet_serve_s3-redis.md b/docs/yeet_serve_s3-redis.md
@@ -20,6 +20,7 @@ yeet serve s3-redis [flags]
 ### Options inherited from parent commands
 
 ```
+      --404-index                  Redirect any 404 to the index file
       --bot-agents string          Bot User-Agent header regex (default "bot|crawl|spider|external|meta|scrap|archive|discourse")
       --bot-proxy string           Bot proxy URL
       --expiry                     Use cache expiry
@@ -42,4 +43,4 @@ yeet serve s3-redis [flags]
 
 * [yeet serve](yeet_serve.md)	 - Serve files with yeet
 
-###### Auto generated by spf13/cobra on 7-Jan-2022
+###### Auto generated by spf13/cobra on 16-Jan-2022
diff --git a/docs/yeet_serve_s3.md b/docs/yeet_serve_s3.md
@@ -20,6 +20,7 @@ yeet serve s3 [flags]
 ### Options inherited from parent commands
 
 ```
+      --404-index                  Redirect any 404 to the index file
       --bot-agents string          Bot User-Agent header regex (default "bot|crawl|spider|external|meta|scrap|archive|discourse")
       --bot-proxy string           Bot proxy URL
       --expiry                     Use cache expiry
@@ -42,4 +43,4 @@ yeet serve s3 [flags]
 
 * [yeet serve](yeet_serve.md)	 - Serve files with yeet
 
-###### Auto generated by spf13/cobra on 7-Jan-2022
+###### Auto generated by spf13/cobra on 16-Jan-2022
diff --git a/server/webserver.go b/server/webserver.go
@@ -19,6 +19,10 @@ func Run(c cache.Cache) error {
 		Cache: c,
 	}
 
+	if viper.GetBool("404-index") {
+		ws.Index404 = []byte("/" + viper.GetString("index-file"))
+	}
+
 	address := viper.GetString("host") + ":" + strconv.Itoa(viper.GetInt("port"))
 
 	ln, err := net.Listen("tcp", address)
@@ -54,21 +58,36 @@ func Run(c cache.Cache) error {
 }
 
 type Webserver struct {
-	Cache cache.Cache
+	Cache    cache.Cache
+	Index404 []byte
 }
 
 func (h *Webserver) HandleFastHTTP(ctx *fasthttp.RequestCtx) {
 	fileType, stream, size, failed := h.Cache.Get(ctx.Path(), ctx.Host())
 	if size > 0 {
 		ctx.SetContentType(fileType)
 		ctx.SetBodyStream(stream, size)
-	} else {
-		if failed {
+		return
+	}
+
+	if failed {
+		ctx.SetStatusCode(500)
+		return
+	}
+
+	if h.Index404 != nil {
+		fileType, stream, size, failed := h.Cache.Get(h.Index404, ctx.Host())
+		if size > 0 {
+			ctx.SetContentType(fileType)
+			ctx.SetBodyStream(stream, size)
+			return
+		} else if failed {
 			ctx.SetStatusCode(500)
-		} else {
-			ctx.SetStatusCode(404)
+			return
 		}
 	}
+
+	ctx.SetStatusCode(404)
 }
 
 func (h *Webserver) HandleFastHTTPWithBotProxy(botHeaderRegex *regexp.Regexp, proxy fasthttp.RequestHandler) func(ctx *fasthttp.RequestCtx) {