diff --git a/cmd/serve.go b/cmd/serve.go index 8b00f3c..2a23cad 100644 --- 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 index 51c47d5..ebed11b 100644 --- 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 index 55153df..b12e9f2 100644 --- 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 index aa8ff27..3c79f26 100644 --- 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 index adc83c6..242291b 100644 --- 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 index 466e679..869ba84 100644 --- 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 index 3f1426b..da729a8 100644 --- 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,7 +58,8 @@ func Run(c cache.Cache) error { } type Webserver struct { - Cache cache.Cache + Cache cache.Cache + Index404 []byte } func (h *Webserver) HandleFastHTTP(ctx *fasthttp.RequestCtx) { @@ -62,13 +67,27 @@ func (h *Webserver) HandleFastHTTP(ctx *fasthttp.RequestCtx) { 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) {