improve project structure

cmd/main.go

@@ -1,56 +1,42 @@
 package main
 
 import (
+	"context"
 	"flag"
-	"fmt"
 	"io"
-	"net/http"
 	"os"
-	"strconv"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/armbian/ansi-hastebin/config"
-	"github.com/armbian/ansi-hastebin/handler"
-	"github.com/armbian/ansi-hastebin/keygenerator"
-	"github.com/armbian/ansi-hastebin/storage"
-	"github.com/go-chi/chi/v5"
-	"github.com/go-chi/chi/v5/middleware"
-	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/armbian/ansi-hastebin/internal/keygenerator"
+	"github.com/armbian/ansi-hastebin/internal/server"
+	"github.com/armbian/ansi-hastebin/internal/storage"
 	"github.com/sirupsen/logrus"
 )
 
-func main() {
-	// Creater router instance
-	r := chi.NewRouter()
-
-	// Add several middlewares
-	r.Use(middleware.Logger)
-	r.Use(middleware.Recoverer)
-
-	// Check if config argument sent
-	var configLocation string
-	flag.StringVar(&configLocation, "config", "", "Pass config yaml")
-	flag.Parse()
-
-	// Parse config fields
-	cfg := config.NewConfig(configLocation)
+func handleConfig(location string) (*config.Config, storage.Storage, keygenerator.KeyGenerator) {
+	cfg := config.NewConfig(location)
+	exp := time.Duration(cfg.Expiration)
 
 	var pasteStorage storage.Storage
 	switch cfg.Storage.Type {
 	case "file":
-		pasteStorage = storage.NewFileStorage(cfg.Storage.FilePath, cfg.Expiration)
+		pasteStorage = storage.NewFileStorage(cfg.Storage.FilePath, exp)
 	case "redis":
-		pasteStorage = storage.NewRedisStorage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Expiration)
+		pasteStorage = storage.NewRedisStorage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, exp)
 	case "memcached":
 		pasteStorage = storage.NewMemcachedStorage(cfg.Storage.Host, cfg.Storage.Port, int(cfg.Expiration))
 	case "mongodb":
-		pasteStorage = storage.NewMongoDBStorage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Storage.Database, cfg.Expiration)
+		pasteStorage = storage.NewMongoDBStorage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Storage.Database, exp)
 	case "postgres":
 		pasteStorage = storage.NewPostgresStorage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Storage.Database, int(cfg.Expiration))
 	case "s3":
 		pasteStorage = storage.NewS3Storage(cfg.Storage.Host, cfg.Storage.Port, cfg.Storage.Username, cfg.Storage.Password, cfg.Storage.AWSRegion, cfg.Storage.Bucket)
 	default:
 		logrus.Fatalf("Unknown storage type: %s", cfg.Storage.Type)
-		return
+		return nil, nil, nil
 	}
 
 	// Set static documents from config
@@ -64,10 +50,11 @@ func main() {
 		if err != nil {
 			logrus.WithError(err).WithField("path", doc.Path).Fatal("Failed to read document")
 		}
-
 		file.Close()
 
-		pasteStorage.Set(doc.Key, string(content), false)
+		if err := pasteStorage.Set(doc.Key, string(content), false); err != nil {
+			logrus.WithError(err).WithField("key", doc.Key).Fatal("Failed to set document")
+		}
 	}
 
 	var keyGenerator keygenerator.KeyGenerator
@@ -79,54 +66,33 @@ func main() {
 		keyGenerator = keygenerator.NewPhoneticKeyGenerator()
 	default:
 		logrus.Fatalf("Unknown key generator: %s", cfg.KeyGenerator)
-		return
+		return nil, nil, nil
 	}
 
-	// Add document handler
-	document_handler := handler.NewDocumentHandler(cfg.KeyLength, cfg.MaxLength, pasteStorage, keyGenerator)
-
-	// Add prometheus metrics
-	r.Get("/metrics", promhttp.Handler().ServeHTTP)
-
-	// Add document routes
-	r.Get("/raw/{id}", document_handler.HandleRawGet)
-	r.Head("/raw/{id}", document_handler.HandleRawGet)
-
-	r.Post("/log", document_handler.HandlePutLog)
-	r.Put("/log", document_handler.HandlePutLog)
-
-	r.Post("/documents", document_handler.HandlePost)
-
-	r.Get("/documents/{id}", document_handler.HandleGet)
-	r.Head("/documents/{id}", document_handler.HandleGet)
-
-	static := os.DirFS("static")
-	r.Get("/{id}", func(w http.ResponseWriter, r *http.Request) {
-		id := chi.URLParam(r, "id")
-		if file, err := static.Open(id); err == nil {
-			defer file.Close()
-			io.Copy(w, file)
-			return
-		}
-
-		index, err := static.Open("index.html")
-		if err != nil {
-			http.Error(w, "Not found", http.StatusNotFound)
-			return
-		}
-
-		defer index.Close()
-
-		io.Copy(w, index)
-	})
-
-	fileServer := http.StripPrefix("/", http.FileServer(http.FS(static)))
-	r.Get("/*", func(w http.ResponseWriter, r *http.Request) {
-		fileServer.ServeHTTP(w, r)
-	})
-
-	if err := http.ListenAndServe(cfg.Host+":"+strconv.Itoa(cfg.Port), r); err != nil {
-		fmt.Printf("Failed to start server: %v\n", err)
-		return
-	}
+	return cfg, pasteStorage, keyGenerator
+}
+
+func main() {
+	// Parse command line arguments
+	var configFile string
+	flag.StringVar(&configFile, "config", "config.yaml", "Configuration file")
+	flag.Parse()
+
+	srv := server.NewServer(handleConfig(configFile))
+	srv.RegisterRoutes()
+
+	// Start the server in a separate goroutine
+	go func() {
+		srv.Start()
+	}()
+
+	// Wait for signal to stop the server
+	stopCh := make(chan os.Signal, 1)
+	signal.Notify(stopCh, syscall.SIGTERM, syscall.SIGINT)
+	<-stopCh
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	srv.Shutdown(ctx)
 }

config.yaml

@@ -12,7 +12,7 @@ storage:
   file_path: "./test"
 
 documents:
-  - key: "about.md"
+  - key: "about"
     path: "./about.md"
 
 logging:

config/config.go

@@ -4,7 +4,6 @@ import (
 	"os"
 	"strconv"
 	"strings"
-	"time"
 
 	"github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v3"
@@ -79,7 +78,7 @@ type Config struct {
 	// Expiration is the maximum lifetime of paste entry
 	// 0 means there will be no expiration.
 	// "file" and "s3" storages don't support expiration control.
-	Expiration time.Duration `yaml:"expiration"`
+	Expiration int `yaml:"expiration"`
 
 	// RecompressStaticAssets is a flag to recompress static assets by default
 	RecompressStaticAssets bool `yaml:"recompress_static_assets"`

handler/document.go

@@ -7,8 +7,8 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/armbian/ansi-hastebin/keygenerator"
-	"github.com/armbian/ansi-hastebin/storage"
+	"github.com/armbian/ansi-hastebin/internal/keygenerator"
+	"github.com/armbian/ansi-hastebin/internal/storage"
 	"github.com/go-chi/chi/v5"
 	"github.com/sirupsen/logrus"
 )
@@ -22,7 +22,6 @@ type DocumentHandler struct {
 }
 
 func NewDocumentHandler(keyLength, maxLength int, store storage.Storage, keyGenerator keygenerator.KeyGenerator) *DocumentHandler {
-	keyLength = 10
 	return &DocumentHandler{
 		KeyLength:    keyLength,
 		MaxLength:    maxLength,
@@ -31,6 +30,20 @@ func NewDocumentHandler(keyLength, maxLength int, store storage.Storage, keyGene
 	}
 }
 
+// RegisterRoutes registers document routes
+func (h *DocumentHandler) RegisterRoutes(r chi.Router) {
+	r.Get("/raw/{id}", h.HandleRawGet)
+	r.Head("/raw/{id}", h.HandleRawGet)
+
+	r.Post("/log", h.HandlePutLog)
+	r.Put("/log", h.HandlePutLog)
+
+	r.Post("/documents", h.HandlePost)
+
+	r.Get("/documents/{id}", h.HandleGet)
+	r.Head("/documents/{id}", h.HandleGet)
+}
+
 // Handle retrieving a document
 func (h *DocumentHandler) HandleGet(w http.ResponseWriter, r *http.Request) {
 	key := strings.Split(chi.URLParam(r, "id"), ".")[0]

internal/server/server.go

@@ -0,0 +1,106 @@
+package server
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"os"
+	"strconv"
+
+	"github.com/armbian/ansi-hastebin/config"
+	"github.com/armbian/ansi-hastebin/handler"
+	"github.com/armbian/ansi-hastebin/internal/keygenerator"
+	"github.com/armbian/ansi-hastebin/internal/storage"
+	"github.com/go-chi/chi/v5"
+	"github.com/go-chi/chi/v5/middleware"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/sirupsen/logrus"
+)
+
+type Server struct {
+	config       *config.Config
+	storage      storage.Storage
+	keyGenerator keygenerator.KeyGenerator
+	server       *http.Server
+	mux          *chi.Mux
+}
+
+func NewServer(config *config.Config, storage storage.Storage, keyGenerator keygenerator.KeyGenerator) *Server {
+	mux := chi.NewRouter()
+	httpServer := &http.Server{
+		Addr:    config.Host + ":" + strconv.Itoa(config.Port),
+		Handler: mux,
+	}
+
+	return &Server{
+		config:       config,
+		storage:      storage,
+		keyGenerator: keyGenerator,
+		server:       httpServer,
+		mux:          mux,
+	}
+}
+
+func (s *Server) RegisterRoutes() {
+	// Register middlewares
+	s.mux.Use(middleware.Logger)
+	s.mux.Use(middleware.Recoverer)
+
+	// Register promhttp middleware
+	s.mux.Get("/metrics", promhttp.Handler().ServeHTTP)
+
+	// Register document handler
+	documentHandler := handler.NewDocumentHandler(s.config.KeyLength, s.config.MaxLength, s.storage, s.keyGenerator)
+	documentHandler.RegisterRoutes(s.mux)
+
+	// Register health check
+	s.mux.Get("/health", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	// Register static files
+	static := os.DirFS("static")
+	s.mux.Get("/{id}", func(w http.ResponseWriter, r *http.Request) {
+		id := chi.URLParam(r, "id")
+		if file, err := static.Open(id); err == nil {
+			defer file.Close()
+			io.Copy(w, file)
+			return
+		}
+
+		index, err := static.Open("index.html")
+		if err != nil {
+			http.Error(w, "Not found", http.StatusNotFound)
+			return
+		}
+
+		defer index.Close()
+
+		io.Copy(w, index)
+	})
+
+	fileServer := http.StripPrefix("/", http.FileServer(http.FS(static)))
+	s.mux.Get("/*", func(w http.ResponseWriter, r *http.Request) {
+		fileServer.ServeHTTP(w, r)
+	})
+}
+
+func (s *Server) Start() {
+	logrus.Infof("Starting server on %s", s.server.Addr)
+
+	if err := s.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		logrus.WithError(err).Fatal("Failed to start server")
+	}
+}
+
+func (s *Server) Shutdown(ctx context.Context) {
+	logrus.Info("Gracefully shutting down server")
+
+	if err := s.storage.Close(); err != nil {
+		logrus.WithError(err).Error("Failed to close storage")
+	}
+
+	if err := s.server.Shutdown(ctx); err != nil {
+		logrus.WithError(err).Error("Failed to shutdown server")
+	}
+}

internal/storage/file.go

@@ -37,6 +37,8 @@ func (fs *FileStorage) Set(key string, value string, skip_expiration bool) error
 		return err
 	}
 
+	defer file.Close()
+
 	_, err = file.WriteString(value)
 	return err
 }
@@ -50,3 +52,7 @@ func (fs *FileStorage) Get(key string, skip_expiration bool) (string, error) {
 
 	return string(file), nil
 }
+
+func (fs *FileStorage) Close() error {
+	return nil
+}

internal/storage/file_test.go

@@ -35,6 +35,8 @@ func TestFileStorage(t *testing.T) {
 
 	_, err = store.Get("testKey2", false)
 	require.True(t, os.IsNotExist(err))
+
+	require.NoError(t, store.Close())
 }
 
 func TestFileStorageSkipExpiration(t *testing.T) {
@@ -52,4 +54,6 @@ func TestFileStorageSkipExpiration(t *testing.T) {
 	val, err := store.Get("persistentKey", true)
 	require.NoError(t, err)
 	require.Equal(t, "persistentValue", val)
+
+	require.NoError(t, store.Close())
 }

internal/storage/memcached.go

@@ -53,3 +53,7 @@ func (s *MemcachedStorage) Get(key string, skip_expiration bool) (string, error)
 
 	return string(item.Value), nil
 }
+
+func (s *MemcachedStorage) Close() error {
+	return s.client.Close()
+}

internal/storage/memcached_test.go

@@ -65,6 +65,8 @@ func TestMemcachedStorage(t *testing.T) {
 
 	_, err = store.Get("testKey2", false)
 	require.ErrorIs(t, memcache.ErrCacheMiss, err) // Should not exist
+
+	require.NoError(t, store.Close())
 }
 
 func TestMemcachedStorageSkipExpiration(t *testing.T) {
@@ -88,4 +90,6 @@ func TestMemcachedStorageSkipExpiration(t *testing.T) {
 	val, err = store.Get("persistentKey", true)
 	require.NoError(t, err)
 	require.Equal(t, "persistentValue", val)
+
+	require.NoError(t, store.Close())
 }

internal/storage/mongodb.go

@@ -142,3 +142,7 @@ func (s *MongoDBStorage) Get(key string, skip_expiration bool) (string, error) {
 
 	return string(i.Value), nil
 }
+
+func (s *MongoDBStorage) Close() error {
+	return s.db.Client().Disconnect(context.Background())
+}

internal/storage/mongodb_test.go

@@ -74,6 +74,8 @@ func TestMongoDBStorage(t *testing.T) {
 	val, err = store.Get("testKey2", false)
 	require.Equal(t, "", val)
 	require.ErrorIs(t, mongo.ErrNoDocuments, err)
+
+	require.NoError(t, store.Close())
 }
 
 func TestMongoDBStorageSkipExpiration(t *testing.T) {
@@ -97,4 +99,6 @@ func TestMongoDBStorageSkipExpiration(t *testing.T) {
 	val, err = store.Get("persistentKey", true)
 	require.NoError(t, err)
 	require.Equal(t, "persistentValue", val)
+
+	require.NoError(t, store.Close())
 }

internal/storage/postgres.go

@@ -86,3 +86,8 @@ func (s *PostgresStorage) Get(key string, skip_expiration bool) (string, error)
 
 	return value, nil
 }
+
+func (s *PostgresStorage) Close() error {
+	s.pool.Close()
+	return nil
+}

internal/storage/postgres_test.go

@@ -67,4 +67,6 @@ func TestPostgresStorage(t *testing.T) {
 	val, err = store.Get("key1", false)
 	require.NoError(t, err)
 	require.Empty(t, val)
+
+	require.NoError(t, store.Close())
 }

internal/storage/redis.go

@@ -60,3 +60,7 @@ func (s *RedisStorage) Get(key string, skip_expiration bool) (string, error) {
 
 	return res, nil
 }
+
+func (s *RedisStorage) Close() error {
+	return s.client.Close()
+}

internal/storage/redis_test.go

@@ -85,6 +85,8 @@ func TestRedisStorage_SetAndGet(t *testing.T) {
 	ttlAfterGetNoExpire, err := storage.client.TTL(context.Background(), keyNoExpire).Result()
 	require.NoError(t, err)
 	require.Equal(t, -1*time.Nanosecond, ttlAfterGetNoExpire) // -1ns means no expiration
+
+	require.NoError(t, storage.Close())
 }
 
 func TestRedisStorage_GetNonExistentKey(t *testing.T) {
@@ -97,4 +99,6 @@ func TestRedisStorage_GetNonExistentKey(t *testing.T) {
 	_, err := storage.Get("nonExistentKey", false)
 	require.Error(t, err)
 	require.Equal(t, redis.Nil, err)
+
+	require.NoError(t, storage.Close())
 }