Add some timeouts.

This may or may not be a good idea.

1 files changed, 15 insertions, 7 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index f4783c5..c95bb5e 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"time"
 )
 
 var (
@@ -46,15 +47,22 @@ func main() {
 	s := &server{
 		jellyfinProxy: httputil.NewSingleHostReverseProxy(jellyfinURL),
 	}
-	http.HandleFunc("/", s.index)
-	http.HandleFunc("/cinema", s.cinema)
-	http.HandleFunc("/cinema/", s.cinema)
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", s.index)
+	mux.HandleFunc("/cinema", s.cinema)
+	mux.HandleFunc("/cinema/", s.cinema)
 
-	addr := fmt.Sprintf(":%d", *port)
-	log.Printf("Listening on %q", addr)
+	httpServer := http.Server{
+		Addr:              fmt.Sprintf(":%d", *port),
+		Handler:           mux,
+		ReadHeaderTimeout: 5 * time.Second,
+		WriteTimeout:      5 * time.Second,
+	}
+
+	log.Printf("Listening on %q", httpServer.Addr)
 	if *https {
-		log.Fatal(http.ListenAndServeTLS(addr, "/var/lib/acme/roseh.moe/cert.pem", "/var/lib/acme/roseh.moe/key.pem", nil))
+		log.Fatal(httpServer.ListenAndServeTLS("/var/lib/acme/roseh.moe/cert.pem", "/var/lib/acme/roseh.moe/key.pem"))
 	} else {
-		log.Fatal(http.ListenAndServe(addr, nil))
+		log.Fatal(httpServer.ListenAndServe())
 	}
 }