M main.go → main.go

@@ -1,12 +1,14 @@ 
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	bolt "go.etcd.io/bbolt"
 	"math/rand"
 	"net/http"
 	"os"
+	"os/signal"
 	"strings"
 	"time"
 )
@@ -186,7 +188,7 @@ err = r.ParseForm()
 	if hec(w, err, 500) {
 		return
 	}
-	if r.Form.Get("api") != "" {
+	if r.Form.Get("mode") == "api" {
 		api = true
 	}
 	url.url = r.Form.Get("url")
@@ -224,6 +226,7 @@ createHandler(w, r)
 			return
 		}
 		http.ServeFile(w, r, "html/index.html")
+		return
 	}
 	path := strings.TrimSuffix(r.URL.Path[1:], "/")
 	var err error
@@ -256,6 +259,38 @@ _ = writeURL(url)
 }
 
 func main() {
+	var (
+		port   string = ":8088"
+		err    error
+		server http.Server
+	)
+
+	if len(os.Args) == 2 {
+		port = ":" + os.Args[1]
+	}
+
+	println("Starting server on port " + port)
+
 	openFiles()
 
+	http.HandleFunc("/", httpHandler)
+
+	server = http.Server{
+		Addr:    port,
+		Handler: nil,
+	}
+
+	go func() {
+		err = server.ListenAndServe()
+		if err != http.ErrServerClosed && err != nil {
+			urlDB.Close()
+			panic(err)
+		}
+	}()
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, os.Interrupt)
+	<-quit
+
+	urlDB.Close()
+	server.Shutdown(context.Background())
 }