Improve upload store initialization to clean up main

knadh · clementauger · Aug 25, 2020 · Aug 25, 2020 · Aug 25, 2020 · Aug 25, 2020
commit bbc04e896c1ddcb7b9b562b404546bbe9b568eb5
diff --git a/handlers.go b/handlers.go
@@ -346,7 +346,7 @@ func readJSONReq(r *http.Request, o interface{}) error {
 }
 
 // handleUpload handles file uploads.
-func handleUpload(store *upload.Store, maxUploadSize int64, rlPeriod time.Duration, rlCount float64, rlBurst int) func(w http.ResponseWriter, r *http.Request) {
+func handleUpload(store *upload.Store) func(w http.ResponseWriter, r *http.Request) {
 
 	type roomLimiter struct {
 		limiter *rate.Limiter
@@ -355,7 +355,7 @@ func handleUpload(store *upload.Store, maxUploadSize int64, rlPeriod time.Durati
 	var mu sync.Mutex
 	roomLimiters := map[string]roomLimiter{}
 	go func() {
-		t := time.NewTicker(rlPeriod + (time.Minute))
+		t := time.NewTicker(store.RlPeriod + (time.Minute))
 		defer t.Stop()
 		for range t.C {
 			now := time.Now()
@@ -370,15 +370,15 @@ func handleUpload(store *upload.Store, maxUploadSize int64, rlPeriod time.Durati
 	}()
 
 	return func(w http.ResponseWriter, r *http.Request) {
-		r.ParseMultipartForm(maxUploadSize)
+		r.ParseMultipartForm(store.MaxUploadSize)
 
 		roomID := chi.URLParam(r, "roomID")
 		mu.Lock()
 		// no defer here becasue file upload can be slow, thus lock for too long
 		x, ok := roomLimiters[roomID]
 		if !ok {
 			x = roomLimiter{
-				limiter: rate.NewLimiter(rate.Every(rlPeriod/time.Duration(rlCount)), rlBurst),
+				limiter: rate.NewLimiter(rate.Every(store.RlPeriod/time.Duration(store.RlCount)), store.RlBurst),
 				expire:  time.Now().Add(time.Minute * 10),
 			}
 			roomLimiters[roomID] = x
@@ -425,8 +425,8 @@ func handleUpload(store *upload.Store, maxUploadSize int64, rlPeriod time.Durati
 }
 
 // handleUploaded uploaded files display.
-func handleUploaded(store *upload.Store, maxAge time.Duration) func(w http.ResponseWriter, r *http.Request) {
-	maxAgeHeader := fmt.Sprintf("max-age=%v", int64(maxAge/time.Second))
+func handleUploaded(store *upload.Store) func(w http.ResponseWriter, r *http.Request) {
+	maxAgeHeader := fmt.Sprintf("max-age=%v", int64(store.MaxAge/time.Second))
 	return func(w http.ResponseWriter, r *http.Request) {
 		fileID := chi.URLParam(r, "fileID")
 		fileID = strings.Split(fileID, "_")[0]
@@ -438,7 +438,7 @@ func handleUploaded(store *upload.Store, maxAge time.Duration) func(w http.Respo
 		}
 		w.Header().Add("Content-Type", up.MimeType)
 		w.Header().Add("Content-Length", fmt.Sprint(len(up.Data)))
-		if maxAge > 0 {
+		if store.MaxAge > 0 {
 			w.Header().Add("Cache-Control", maxAgeHeader)
 		}
 		w.WriteHeader(http.StatusOK)

diff --git a/internal/upload/store.go b/internal/upload/store.go
@@ -4,8 +4,12 @@ import (
 	"crypto/sha1"
 	"errors"
 	"fmt"
+	"strconv"
 	"sync"
 	"time"
+
+	"github.com/alecthomas/units"
+	tparse "github.com/karrick/tparse/v2"
 )
 
 // Config represents the file upload options.
@@ -20,11 +24,75 @@ type Config struct {
 
 // Store file uploads in memory.
 type Store struct {
-	cfg    Config
-	maxMem int64
-	mu     sync.Mutex
-	items  map[string]File
-	size   int64
+	cfg   Config
+	mu    sync.Mutex
+	items map[string]File
+	size  int64
+
+	MaxMemory     int64
+	MaxUploadSize int64
+	MaxAge        time.Duration
+	RlPeriod      time.Duration
+	RlCount       float64
+	RlBurst       int
+}
+
+//Init the store, parsing configuration values.
+func (s *Store) Init() error {
+	s.MaxMemory = 32 << 20
+	if s.cfg.MaxMemory != "" {
+		x, err := units.ParseStrictBytes(s.cfg.MaxMemory)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.max-memory' config: %v", err)
+		}
+		s.MaxMemory = x
+	}
+
+	s.MaxUploadSize = 32 << 20
+	if s.cfg.MaxUploadSize != "" {
+		x, err := units.ParseStrictBytes(s.cfg.MaxUploadSize)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.max-upload-size' config: %v", err)
+		}
+		s.MaxUploadSize = x
+	}
+
+	s.MaxAge = time.Hour * 24 * 30 * 12
+	if s.cfg.MaxAge != "" {
+		x, err := tparse.AbsoluteDuration(time.Now(), s.cfg.MaxAge)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.max-age' config: %v", err)
+		}
+		s.MaxAge = x
+	}
+
+	s.RlPeriod = time.Minute
+	if s.cfg.RateLimitPeriod != "" {
+		x, err := tparse.AbsoluteDuration(time.Now(), s.cfg.RateLimitPeriod)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.rate-limit-period' config: %v", err)
+		}
+		s.RlPeriod = x
+	}
+
+	s.RlCount = 20.0
+	if s.cfg.RateLimitCount != "" {
+		x, err := strconv.ParseFloat(s.cfg.RateLimitCount, 64)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.rate-limit-count' config: %v", err)
+		}
+		s.RlCount = x
+	}
+
+	s.RlBurst = 1
+	if s.cfg.RateLimitBurst != "" {
+		x, err := strconv.Atoi(s.cfg.RateLimitBurst)
+		if err != nil {
+			return fmt.Errorf("error unmarshalling 'upload.rate-limit-burst' config: %v", err)
+		}
+		s.RlBurst = x
+	}
+	return nil
 }
 
 // File represents an upload.
@@ -37,11 +105,10 @@ type File struct {
 }
 
 // New returns a new file uplod store.
-func New(cfg Config, maxMemory int64) *Store {
+func New(cfg Config) *Store {
 	return &Store{
-		cfg:    cfg,
-		maxMem: maxMemory,
-		items:  make(map[string]File),
+		cfg:   cfg,
+		items: make(map[string]File),
 	}
 }
 
@@ -64,7 +131,7 @@ func (s *Store) Add(name, mimeType string, data []byte) (File, error) {
 	copy(up.Data, data)
 	s.items[id] = up
 	s.size += int64(len(data))
-	for s.size > s.maxMem {
+	for s.size > s.MaxMemory {
 		var oldest *File
 		for _, up := range s.items {
 			if oldest == nil {

diff --git a/main.go b/main.go
@@ -13,14 +13,11 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"syscall"
 	"time"
 
-	"github.com/alecthomas/units"
 	"github.com/go-chi/chi"
-	"github.com/karrick/tparse/v2"
 	"github.com/knadh/koanf"
 	"github.com/knadh/koanf/parsers/toml"
 	"github.com/knadh/koanf/providers/env"
@@ -267,60 +264,10 @@ func main() {
 	if err := ko.Unmarshal("upload", &uploadCfg); err != nil {
 		logger.Fatalf("error unmarshalling 'upload' config: %v", err)
 	}
-	var maxMemory int64 = 32 << 20
-	if uploadCfg.MaxMemory != "" {
-		x, err := units.ParseStrictBytes(uploadCfg.MaxMemory)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.max-memory' config: %v", err)
-		}
-		maxMemory = x
-	}
-
-	uploadStore := upload.New(uploadCfg, maxMemory)
-
-	var maxUploadSize int64 = 32 << 20
-	if uploadCfg.MaxUploadSize != "" {
-		x, err := units.ParseStrictBytes(uploadCfg.MaxUploadSize)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.max-upload-size' config: %v", err)
-		}
-		maxUploadSize = x
-	}
-
-	var maxAge time.Duration = time.Hour * 24 * 30 * 12
-	if uploadCfg.MaxAge != "" {
-		x, err := tparse.AbsoluteDuration(time.Now(), uploadCfg.MaxAge)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.max-age' config: %v", err)
-		}
-		maxAge = x
-	}
-
-	var rlPeriod time.Duration = time.Minute
-	if uploadCfg.RateLimitPeriod != "" {
-		x, err := tparse.AbsoluteDuration(time.Now(), uploadCfg.RateLimitPeriod)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.rate-limit-period' config: %v", err)
-		}
-		rlPeriod = x
-	}
 
-	rlCount := 20.0
-	if uploadCfg.RateLimitCount != "" {
-		x, err := strconv.ParseFloat(uploadCfg.RateLimitCount, 64)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.rate-limit-count' config: %v", err)
-		}
-		rlCount = x
-	}
-
-	rlBurst := 1
-	if uploadCfg.RateLimitBurst != "" {
-		x, err := strconv.Atoi(uploadCfg.RateLimitBurst)
-		if err != nil {
-			logger.Fatalf("error unmarshalling 'upload.rate-limit-burst' config: %v", err)
-		}
-		rlBurst = x
+	uploadStore := upload.New(uploadCfg)
+	if err := uploadStore.Init(); err != nil {
+		logger.Fatalf("error initializing upload store: %v", err)
 	}
 
 	// Register HTTP routes.
@@ -333,8 +280,8 @@ func main() {
 	r.Delete("/api/rooms/{roomID}/login", wrap(handleLogout, app, hasAuth|hasRoom))
 	r.Post("/api/rooms", wrap(handleCreateRoom, app, 0))
 
-	r.Post("/api/rooms/{roomID}/upload", handleUpload(uploadStore, maxUploadSize, rlPeriod, rlCount, rlBurst))
-	r.Get("/api/rooms/{roomID}/uploaded/{fileID}", handleUploaded(uploadStore, maxAge))
+	r.Post("/api/rooms/{roomID}/upload", handleUpload(uploadStore))
+	r.Get("/api/rooms/{roomID}/uploaded/{fileID}", handleUploaded(uploadStore))
 
 	// Views.
 	r.Get("/r/{roomID}", wrap(handleRoomPage, app, hasAuth|hasRoom))