From b5ee197ac749cb4a4ae001b7767a78838349f297 Mon Sep 17 00:00:00 2001 From: Muhammed Efe Cetin Date: Sun, 16 Mar 2025 13:32:49 +0300 Subject: [PATCH] improve project structure --- cmd/main.go | 122 +++++++----------- config.yaml | 2 +- config/config.go | 3 +- handler/document.go | 19 ++- .../keygenerator}/key_generator.go | 0 .../keygenerator}/phonetic.go | 0 .../keygenerator}/phonetic_test.go | 0 .../keygenerator}/random.go | 0 .../keygenerator}/random_test.go | 0 internal/server/server.go | 106 +++++++++++++++ {storage => internal/storage}/file.go | 6 + {storage => internal/storage}/file_test.go | 4 + {storage => internal/storage}/memcached.go | 4 + .../storage}/memcached_test.go | 4 + {storage => internal/storage}/mongodb.go | 4 + {storage => internal/storage}/mongodb_test.go | 4 + {storage => internal/storage}/postgres.go | 5 + .../storage}/postgres_test.go | 2 + {storage => internal/storage}/redis.go | 4 + {storage => internal/storage}/redis_test.go | 4 + {storage => internal/storage}/s3.go | 4 + {storage => internal/storage}/s3_test.go | 2 + {storage => internal/storage}/storage.go | 1 + 23 files changed, 216 insertions(+), 84 deletions(-) rename {keygenerator => internal/keygenerator}/key_generator.go (100%) rename {keygenerator => internal/keygenerator}/phonetic.go (100%) rename {keygenerator => internal/keygenerator}/phonetic_test.go (100%) rename {keygenerator => internal/keygenerator}/random.go (100%) rename {keygenerator => internal/keygenerator}/random_test.go (100%) create mode 100644 internal/server/server.go rename {storage => internal/storage}/file.go (92%) rename {storage => internal/storage}/file_test.go (94%) rename {storage => internal/storage}/memcached.go (94%) rename {storage => internal/storage}/memcached_test.go (97%) rename {storage => internal/storage}/mongodb.go (97%) rename {storage => internal/storage}/mongodb_test.go (97%) rename {storage => internal/storage}/postgres.go (97%) rename {storage => internal/storage}/postgres_test.go (98%) rename {storage => internal/storage}/redis.go (95%) rename {storage => internal/storage}/redis_test.go (97%) rename {storage => internal/storage}/s3.go (98%) rename {storage => internal/storage}/s3_test.go (97%) rename {storage => internal/storage}/storage.go (91%) diff --git a/cmd/main.go b/cmd/main.go index 4782184..8c3b38b 100644 --- a/cmd/main.go +++ b/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) } diff --git a/config.yaml b/config.yaml index f3bc996..c849f4f 100644 --- a/config.yaml +++ b/config.yaml @@ -12,7 +12,7 @@ storage: file_path: "./test" documents: - - key: "about.md" + - key: "about" path: "./about.md" logging: diff --git a/config/config.go b/config/config.go index 73bda69..9438424 100644 --- a/config/config.go +++ b/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"` diff --git a/handler/document.go b/handler/document.go index 92363cf..5db1c3d 100644 --- a/handler/document.go +++ b/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] diff --git a/keygenerator/key_generator.go b/internal/keygenerator/key_generator.go similarity index 100% rename from keygenerator/key_generator.go rename to internal/keygenerator/key_generator.go diff --git a/keygenerator/phonetic.go b/internal/keygenerator/phonetic.go similarity index 100% rename from keygenerator/phonetic.go rename to internal/keygenerator/phonetic.go diff --git a/keygenerator/phonetic_test.go b/internal/keygenerator/phonetic_test.go similarity index 100% rename from keygenerator/phonetic_test.go rename to internal/keygenerator/phonetic_test.go diff --git a/keygenerator/random.go b/internal/keygenerator/random.go similarity index 100% rename from keygenerator/random.go rename to internal/keygenerator/random.go diff --git a/keygenerator/random_test.go b/internal/keygenerator/random_test.go similarity index 100% rename from keygenerator/random_test.go rename to internal/keygenerator/random_test.go diff --git a/internal/server/server.go b/internal/server/server.go new file mode 100644 index 0000000..cf0980a --- /dev/null +++ b/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") + } +} diff --git a/storage/file.go b/internal/storage/file.go similarity index 92% rename from storage/file.go rename to internal/storage/file.go index 6f4ec27..30c6404 100644 --- a/storage/file.go +++ b/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 +} diff --git a/storage/file_test.go b/internal/storage/file_test.go similarity index 94% rename from storage/file_test.go rename to internal/storage/file_test.go index 8a27e59..6b6c41e 100644 --- a/storage/file_test.go +++ b/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()) } diff --git a/storage/memcached.go b/internal/storage/memcached.go similarity index 94% rename from storage/memcached.go rename to internal/storage/memcached.go index eee6a73..83c33b5 100644 --- a/storage/memcached.go +++ b/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() +} diff --git a/storage/memcached_test.go b/internal/storage/memcached_test.go similarity index 97% rename from storage/memcached_test.go rename to internal/storage/memcached_test.go index ce16233..bcc5a5f 100644 --- a/storage/memcached_test.go +++ b/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()) } diff --git a/storage/mongodb.go b/internal/storage/mongodb.go similarity index 97% rename from storage/mongodb.go rename to internal/storage/mongodb.go index 9868363..8044262 100644 --- a/storage/mongodb.go +++ b/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()) +} diff --git a/storage/mongodb_test.go b/internal/storage/mongodb_test.go similarity index 97% rename from storage/mongodb_test.go rename to internal/storage/mongodb_test.go index 3da8155..268f450 100644 --- a/storage/mongodb_test.go +++ b/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()) } diff --git a/storage/postgres.go b/internal/storage/postgres.go similarity index 97% rename from storage/postgres.go rename to internal/storage/postgres.go index 12de8f4..5f5ed4c 100644 --- a/storage/postgres.go +++ b/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 +} diff --git a/storage/postgres_test.go b/internal/storage/postgres_test.go similarity index 98% rename from storage/postgres_test.go rename to internal/storage/postgres_test.go index 5fc4ca9..22d5a2a 100644 --- a/storage/postgres_test.go +++ b/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()) } diff --git a/storage/redis.go b/internal/storage/redis.go similarity index 95% rename from storage/redis.go rename to internal/storage/redis.go index bfb79a1..0e97a55 100644 --- a/storage/redis.go +++ b/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() +} diff --git a/storage/redis_test.go b/internal/storage/redis_test.go similarity index 97% rename from storage/redis_test.go rename to internal/storage/redis_test.go index 6cae015..a4aefc8 100644 --- a/storage/redis_test.go +++ b/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()) } diff --git a/storage/s3.go b/internal/storage/s3.go similarity index 98% rename from storage/s3.go rename to internal/storage/s3.go index 6ac519c..ef31cda 100644 --- a/storage/s3.go +++ b/internal/storage/s3.go @@ -94,3 +94,7 @@ func (s *S3Storage) Get(key string, skip_expiration bool) (string, error) { return string(buf.Bytes()), err } + +func (s *S3Storage) Close() error { + return nil +} diff --git a/storage/s3_test.go b/internal/storage/s3_test.go similarity index 97% rename from storage/s3_test.go rename to internal/storage/s3_test.go index 3d241e8..a99309b 100644 --- a/storage/s3_test.go +++ b/internal/storage/s3_test.go @@ -61,4 +61,6 @@ func TestS3Storage(t *testing.T) { val, err = store.Get("nonExistingKey", false) require.ErrorIs(t, ErrNotFound, err) require.Equal(t, "", val) + + require.NoError(t, store.Close()) } diff --git a/storage/storage.go b/internal/storage/storage.go similarity index 91% rename from storage/storage.go rename to internal/storage/storage.go index 237f184..160f9d9 100644 --- a/storage/storage.go +++ b/internal/storage/storage.go @@ -3,4 +3,5 @@ package storage type Storage interface { Set(key string, value string, skip_expiration bool) error Get(key string, skip_expiration bool) (string, error) + Close() error }