From dc16506932f35b6ed34ba7ef748f1978c8f2358a Mon Sep 17 00:00:00 2001
From: Lukas Schulte Pelkum <lockas.sp@proton.me>
Date: Thu, 8 Jun 2023 19:44:32 +0200
Subject: [PATCH] integrate the web server into the startup logic

---
 .gitignore                          |  5 +----
 cmd/pasty/main.go                   | 29 +++++++++++++++++++++++++++++
 internal/storage/postgres/driver.go |  1 -
 internal/web/server.go              | 18 +++++++++++++++++-
 4 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index 776188e..497aa10 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,3 @@
-
 # Created by https://www.toptal.com/developers/gitignore/api/jetbrains+all,go
 # Edit at https://www.toptal.com/developers/gitignore?templates=jetbrains+all,go
 
@@ -114,6 +113,4 @@ modules.xml
 
 # End of https://www.toptal.com/developers/gitignore/api/jetbrains+all,go
 
-web/*.gz
-data/
-.env
\ No newline at end of file
+.env
diff --git a/cmd/pasty/main.go b/cmd/pasty/main.go
index 4e8eff9..bbb3a9e 100644
--- a/cmd/pasty/main.go
+++ b/cmd/pasty/main.go
@@ -2,12 +2,15 @@ package main
 
 import (
 	"context"
+	"errors"
 	"github.com/lus/pasty/internal/config"
 	"github.com/lus/pasty/internal/meta"
 	"github.com/lus/pasty/internal/storage"
 	"github.com/lus/pasty/internal/storage/postgres"
+	"github.com/lus/pasty/internal/web"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+	"net/http"
 	"os"
 	"os/signal"
 	"strings"
@@ -65,6 +68,32 @@ func main() {
 		}
 	}()
 
+	// Start the web server
+	log.Info().Str("address", cfg.WebAddress).Msg("Starting the web server...")
+	webServer := &web.Server{
+		Address:                   cfg.WebAddress,
+		Storage:                   driver,
+		HastebinSupport:           cfg.HastebinSupport,
+		PasteIDLength:             cfg.IDLength,
+		PasteIDCharset:            cfg.IDCharacters,
+		PasteLengthCap:            cfg.LengthCap,
+		ModificationTokensEnabled: cfg.ModificationTokens,
+		ModificationTokenLength:   cfg.ModificationTokenLength,
+		ModificationTokenCharset:  cfg.ModificationTokenCharacters,
+		AdminTokens:               []string{cfg.ModificationTokenMaster},
+	}
+	go func() {
+		if err := webServer.Start(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatal().Err(err).Msg("Could not start the web server.")
+		}
+	}()
+	defer func() {
+		log.Info().Msg("Shutting down the web server...")
+		if err := webServer.Shutdown(context.Background()); err != nil {
+			log.Err(err).Msg("Could not shut down the web server.")
+		}
+	}()
+
 	// Wait for an interrupt signal
 	log.Info().Msg("The application has been started. Use Ctrl+C to shut it down.")
 	shutdownChan := make(chan os.Signal, 1)
diff --git a/internal/storage/postgres/driver.go b/internal/storage/postgres/driver.go
index e401044..9ab10fc 100644
--- a/internal/storage/postgres/driver.go
+++ b/internal/storage/postgres/driver.go
@@ -54,7 +54,6 @@ func (driver *Driver) Initialize(ctx context.Context) error {
 		pool.Close()
 		return err
 	}
-	log.Info().Msg("Successfully performed PostgreSQL database migrations.")
 
 	driver.connPool = pool
 	driver.pastes = &pasteRepository{
diff --git a/internal/web/server.go b/internal/web/server.go
index 4ee2e49..21684b9 100644
--- a/internal/web/server.go
+++ b/internal/web/server.go
@@ -1,6 +1,7 @@
 package web
 
 import (
+	"context"
 	"github.com/go-chi/chi/v5"
 	"github.com/lus/pasty/internal/pastes"
 	"github.com/lus/pasty/internal/storage"
@@ -35,6 +36,8 @@ type Server struct {
 
 	// The administration tokens.
 	AdminTokens []string
+
+	httpServer *http.Server
 }
 
 func (server *Server) Start() error {
@@ -60,5 +63,18 @@ func (server *Server) Start() error {
 	router.With(server.v2MiddlewareInjectPaste, server.v2MiddlewareAuthorize).Patch("/api/v2/pastes/{paste_id}", server.v2EndpointModifyPaste)
 	router.Delete("/api/v2/pastes/{paste_id}", server.v2EndpointDeletePaste)
 
-	return http.ListenAndServe(server.Address, router)
+	// Start the HTTP server
+	server.httpServer = &http.Server{
+		Addr:    server.Address,
+		Handler: router,
+	}
+	return server.httpServer.ListenAndServe()
+}
+
+func (server *Server) Shutdown(ctx context.Context) error {
+	if err := server.httpServer.Shutdown(ctx); err != nil {
+		return err
+	}
+	server.httpServer = nil
+	return nil
 }