From fc229b1f06c6b8db110a596681809cb469f4e20c Mon Sep 17 00:00:00 2001
From: Yiyang Kang <kkyy@kkyy.me>
Date: Sun, 12 Mar 2023 17:51:37 +0800
Subject: [PATCH] fix: OOM on large file

Remove the etag middleware since it requires the entire file to be read
into memory.
---
 README.md |  2 +-
 main.go   | 81 ++++++++++++++++++++++++++++++++-----------------------
 2 files changed, 48 insertions(+), 35 deletions(-)

diff --git a/README.md b/README.md
index 0021cb3..8dadf9f 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,7 @@
 
 A **very simple** HTTP server that might come in handy from time to time.
 
-The only thing it does is serving static files in the current directory.
+The only thing it does is serve static files in the current directory.
 
 ## Usage
 
diff --git a/main.go b/main.go
index cf01b81..fc79c93 100644
--- a/main.go
+++ b/main.go
@@ -6,12 +6,10 @@ import (
 	"os"
 	"os/signal"
 	"strconv"
-	"sync"
 	"syscall"
 	"time"
 
 	"github.com/gofiber/fiber/v2"
-	"github.com/gofiber/fiber/v2/middleware/etag"
 	"github.com/kanna5/qrtxt"
 )
 
@@ -26,27 +24,8 @@ func getLocalAddr() (string, error) {
 	return host, err
 }
 
-func main() {
-	listenPort := 8080
-
-	var err error
-	if len(os.Args) >= 2 {
-		portStr := os.Args[1]
-		listenPort, err = strconv.Atoi(portStr)
-	}
-	if err != nil {
-		panic(fmt.Errorf("cannot parse port number: %w", err))
-	}
-	if listenPort < 1 || listenPort > 65535 {
-		panic(fmt.Errorf("invalid port number: %d", listenPort))
-	}
-
-	localAddr, err := getLocalAddr()
-	if err != nil {
-		panic(fmt.Errorf("cannot get local address: %w", err))
-	}
-
-	fullAddress := fmt.Sprintf("http://%s:%d/", localAddr, listenPort)
+func serve(port int) (*fiber.App, <-chan error) {
+	ch := make(chan error)
 
 	app := fiber.New(fiber.Config{
 		AppName:               "StaticFileServer",
@@ -56,21 +35,51 @@ func main() {
 		RequestMethods:        []string{"GET", "HEAD"},
 		ServerHeader:          "StaticFileServer",
 	})
-	app.Use(etag.New())
 	app.Static("/", "./", fiber.Static{Browse: true})
 
+	go func() {
+		defer close(ch)
+		ch <- app.Listen(":" + strconv.Itoa(port))
+	}()
+
+	return app, ch
+}
+
+func main() {
+	listenPort := 8080
+
+	var err error
+	if len(os.Args) >= 2 {
+		portStr := os.Args[1]
+		listenPort, err = strconv.Atoi(portStr)
+	}
+	if err != nil {
+		panic(fmt.Errorf("Cannot parse port number: %w", err))
+	}
+	if listenPort < 1 || listenPort > 65535 {
+		panic(fmt.Errorf("Invalid port number: %d", listenPort))
+	}
+
+	localAddr, err := getLocalAddr()
+	if err != nil {
+		panic(fmt.Errorf("Cannot get local address: %w", err))
+	}
+
+	fullAddress := fmt.Sprintf("http://%s:%d/", localAddr, listenPort)
+
+	app, errCh := serve(listenPort)
+	// Early check to see if the server is started successfully. This might not be reliable
+	select {
+	case err := <-errCh:
+		panic(fmt.Errorf("Cannot start server: %w", err))
+	case <-time.After(1 * time.Millisecond):
+	}
+
 	if qrText, err := qrtxt.Encode(fullAddress, qrtxt.Low); err == nil {
 		fmt.Println(qrText)
 	}
 	fmt.Printf("\nListening on %s\n", fullAddress)
 
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		err = app.Listen(":" + strconv.Itoa(listenPort))
-	}()
-
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
 	go func() {
@@ -78,9 +87,13 @@ func main() {
 		_ = app.Shutdown()
 	}()
 
-	wg.Wait()
-	if err != nil {
-		panic(fmt.Errorf("cannot start server: %w", err))
+	err, ok := <-errCh
+	if !ok {
+		panic("Server closed unexpectedly")
 	}
+	if err != nil {
+		panic(fmt.Errorf("Failed to run server: %w", err))
+	}
+
 	fmt.Println("Bye!")
 }