diff --git a/pkg/config/config.go b/pkg/config/config.go index 4c50bde..629a372 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -22,11 +22,15 @@ type Config struct { NameLanguage string `toml:"name_language"` Receive struct { Directory string `toml:"directory"` + MaxFileSize int `toml:"max_file_size"` SaveUserID int `toml:"save_user_id"` SaveGroupID int `toml:"save_group_id"` Clipboard bool `toml:"clipboard"` ExitAfterFileCount int `toml:"exit_after_file_count"` } `toml:"receive"` + Send struct { + Directory string `toml:"directory"` + } `toml:"send"` Functions struct { HttpFileServer bool `toml:"http_file_server"` LocalSendServer bool `toml:"local_send_server"` diff --git a/pkg/config/config.toml b/pkg/config/config.toml index 029c4c3..def1e9a 100644 --- a/pkg/config/config.toml +++ b/pkg/config/config.toml @@ -3,10 +3,14 @@ name_language = "en" [receive] directory = "uploads" +max_file_size = 1000000000 save_user_id = 0 save_group_id = 0 clipboard = false +[send] +directory = "uploads" + [functions] http_file_server = true local_send_server = true diff --git a/pkg/handlers/download.go b/pkg/handlers/download.go index 3c66a9f..43db95e 100644 --- a/pkg/handlers/download.go +++ b/pkg/handlers/download.go @@ -6,6 +6,8 @@ import ( "net/http" "os" "path/filepath" + + "github.com/ilius/localsend-go/pkg/config" ) // DownloadHandler handles file download requests @@ -16,8 +18,8 @@ func DownloadHandler(w http.ResponseWriter, r *http.Request) { return } - // Assuming the files are stored in the "uploads" directory - filePath := filepath.Join("uploads", fileName) + // Assuming the files are stored in the configured directory + filePath := filepath.Join(config.ConfigData.Send.Directory, fileName) file, err := os.Open(filePath) if err != nil { http.Error(w, fmt.Sprintf("Could not open file: %v", err), http.StatusNotFound) diff --git a/pkg/handlers/upload_api.go b/pkg/handlers/upload_api.go index 2fe551d..5a5cdb5 100644 --- a/pkg/handlers/upload_api.go +++ b/pkg/handlers/upload_api.go @@ -95,6 +95,11 @@ func UploadAPIHandler(w http.ResponseWriter, r *http.Request) { slog.Error("Error creating directory", "err", err) return } + + if config.ConfigData.Receive.ExitAfterFileCount > 0 { + defer checkExitAfterFileCount() + } + // Create a file file, err := os.Create(filePath) if err != nil { @@ -105,6 +110,8 @@ func UploadAPIHandler(w http.ResponseWriter, r *http.Request) { defer file.Close() buffer := make([]byte, 2*1024*1024) // 2MB buffer + size := 0 + maxSize := config.ConfigData.Receive.MaxFileSize for { n, err := r.Body.Read(buffer) if err != nil && err != io.EOF { @@ -115,28 +122,33 @@ func UploadAPIHandler(w http.ResponseWriter, r *http.Request) { if n == 0 { break } + size += n + if size > maxSize { + http.Error(w, "Failed to write file", http.StatusInternalServerError) + slog.Error("Max file size reached", "size", size, "maxSize", maxSize) + } _, err = file.Write(buffer[:n]) if err != nil { http.Error(w, "Failed to write file", http.StatusInternalServerError) slog.Error("Error writing file", "err", err) return - } } changeFileOwnerGroup(filePath) slog.Info("Saved file", "filePath", filePath) w.WriteHeader(http.StatusOK) +} - if config.ConfigData.Receive.ExitAfterFileCount > 0 { - count := int(uploadCount.Add(1)) - if count >= config.ConfigData.Receive.ExitAfterFileCount { - slog.Info("Exiting due to max recieved file count reached") - go func() { - time.Sleep(100 * time.Millisecond) - os.Exit(0) - }() - } +func checkExitAfterFileCount() { + count := int(uploadCount.Add(1)) + if count < config.ConfigData.Receive.ExitAfterFileCount { + return } + slog.Info("Exiting due to max recieved file count reached") + go func() { + time.Sleep(100 * time.Millisecond) + os.Exit(0) + }() }