diff --git a/go.mod b/go.mod index 1129052..c668d2c 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.22.5 require ( github.com/go-redis/redis/v8 v8.11.5 + github.com/google/uuid v1.6.0 github.com/gorilla/mux v1.8.1 github.com/gorilla/websocket v1.5.3 github.com/joho/godotenv v1.5.1 diff --git a/go.sum b/go.sum index 61f4be9..dfd8e3c 100644 --- a/go.sum +++ b/go.sum @@ -8,6 +8,8 @@ github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWo github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI= github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= diff --git a/internal/core/middleware/file_uploader.go b/internal/core/middleware/file_uploader.go index 9735834..eb75088 100644 --- a/internal/core/middleware/file_uploader.go +++ b/internal/core/middleware/file_uploader.go @@ -1,50 +1,80 @@ package middleware import ( + "fmt" "io" + "log" "net/http" "os" "path/filepath" - "time" + "strings" + + "github.com/google/uuid" ) -// FileUploadMiddleware handles file uploads and saves them to the specified directory type FileUploadMiddleware struct { - uploadDir string + uploadDir string + maxFileSize int64 // Maximum file size in bytes + allowedExts []string // Allowed file extensions } // NewFileUploadMiddleware creates a new instance of FileUploadMiddleware -func NewFileUploadMiddleware(uploadDir string) *FileUploadMiddleware { - return &FileUploadMiddleware{uploadDir: uploadDir} +func NewFileUploadMiddleware(uploadDir string, maxFileSize int64, allowedExts []string) *FileUploadMiddleware { + // Ensure the upload directory exists + if err := os.MkdirAll(uploadDir, os.ModePerm); err != nil { + log.Fatalf("Failed to create upload directory: %v", err) + } + + if len(allowedExts) == 0 { + allowedExts = []string{".jpg"} // Default allowed extension if none provided + } + + return &FileUploadMiddleware{ + uploadDir: uploadDir, + maxFileSize: maxFileSize, + allowedExts: allowedExts, + } } // Handle is the middleware function that processes file uploads func (f *FileUploadMiddleware) Handle(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if err := r.ParseMultipartForm(32 << 20); err != nil { // 32 MB limit - http.Error(w, "Unable to parse form", http.StatusBadRequest) + if err := r.ParseMultipartForm(f.maxFileSize); err != nil { + http.Error(w, "File too large or unable to parse form", http.StatusBadRequest) + log.Printf("Error parsing form: %v", err) return } - // Get file from form data file, fileHeader, err := r.FormFile("file") if err != nil { http.Error(w, "Unable to get file from form data", http.StatusBadRequest) + log.Printf("Error retrieving file: %v", err) return } defer file.Close() - // Extract file extension - ext := filepath.Ext(fileHeader.Filename) - if ext == "" { - ext = ".bin" // Default extension if none provided + // Validate file size + if fileHeader.Size > f.maxFileSize { + http.Error(w, "File size exceeds limit", http.StatusRequestEntityTooLarge) + return + } + + // Validate file extension + ext := strings.ToLower(filepath.Ext(fileHeader.Filename)) + if !f.isAllowedExt(ext) { + http.Error(w, "File type not allowed", http.StatusUnsupportedMediaType) + return } - // Create file on server with timestamp and original extension - filePath := filepath.Join(f.uploadDir, generateFileName()+ext) + // Generate a unique file name + fileName := generateFileName() + ext + filePath := filepath.Join(f.uploadDir, fileName) + + // Create the file destFile, err := os.Create(filePath) if err != nil { http.Error(w, "Unable to save file", http.StatusInternalServerError) + log.Printf("Error creating file: %v", err) return } defer destFile.Close() @@ -52,15 +82,30 @@ func (f *FileUploadMiddleware) Handle(next http.Handler) http.Handler { // Copy file content if _, err := io.Copy(destFile, file); err != nil { http.Error(w, "Unable to copy file content", http.StatusInternalServerError) + log.Printf("Error copying file content: %v", err) return } + // Optionally, you can add a response to inform about the successful upload + w.WriteHeader(http.StatusCreated) + fmt.Fprintf(w, "File uploaded successfully: %s", fileName) + // Call the next handler next.ServeHTTP(w, r) }) } -// generateFileName generates a unique file name with timestamp +// isAllowedExt checks if the file extension is allowed +func (f *FileUploadMiddleware) isAllowedExt(ext string) bool { + for _, allowedExt := range f.allowedExts { + if ext == allowedExt { + return true + } + } + return false +} + +// generateFileName generates a unique file name using UUID func generateFileName() string { - return time.Now().Format("20060102150405") + return uuid.New().String() } diff --git a/internal/core/router/router.go b/internal/core/router/router.go index baf9367..a25c953 100644 --- a/internal/core/router/router.go +++ b/internal/core/router/router.go @@ -257,9 +257,9 @@ func WithCookieParser() Option { // Example usage: // // r := router.NewRouter(router.WithFileUpload("/uploads")) -func WithFileUpload(uploadDir string) Option { +func WithFileUpload(uploadDir string, maxFileSize int64, allowedExts []string) Option { return func(r *Router) { - fileUploadMiddleware := middleware.NewFileUploadMiddleware(uploadDir) + fileUploadMiddleware := middleware.NewFileUploadMiddleware(uploadDir, maxFileSize, allowedExts) r.Use(fileUploadMiddleware) } } diff --git a/pkg/lessgo/less.go b/pkg/lessgo/less.go index abfe253..c58af6c 100644 --- a/pkg/lessgo/less.go +++ b/pkg/lessgo/less.go @@ -209,8 +209,8 @@ func WithCookieParser() router.Option { // Example usage: // // r := router.NewRouter(router.WithFileUpload("/uploads")) -func WithFileUpload(uploadDir string) router.Option { - return router.WithFileUpload(uploadDir) +func WithFileUpload(uploadDir string, maxFileSize int64, allowedExts []string) router.Option { + return router.WithFileUpload(uploadDir, maxFileSize, allowedExts) } // WithCaching is an option function that enables caching for the router using Redis.