add files restrictions

dal/minio/minio.go

@@ -30,7 +30,7 @@ func New(ctx context.Context, logger hlog.Logger,
 	if err != nil {
 		return nil, err
 	}
-	fmt.Println("monio", keyID,accessKey,token,region)
+	fmt.Println("monio", keyID, accessKey, token, region)
 
 	bucketExists, err := conn.BucketExists(ctx, bucket)
 	if err != nil {
@@ -78,7 +78,7 @@ func (bs *BlobStore) PutFile(
 	if err != nil {
 		return err
 	}
-	fmt.Println(info)
+	fmt.Println("info", info)
 
 	return nil
 }

service/service.go

@@ -11,6 +11,7 @@ import (
 	"heruvym/middleware"
 	"heruvym/model"
 	"heruvym/tools"
+	"heruvym/utils"
 	"net/http"
 	"strings"
 	"sync"
@@ -670,6 +671,15 @@ type PutFileResp struct {
 func (h *Heruvym) PutFile(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
+	sess := jwt_adapter.Get(r.Context())
+	if sess == nil {
+		w.WriteHeader(http.StatusBadRequest)
+		if _, err := w.Write([]byte("not authorized")); err != nil {
+			fmt.Println("CAN NOT WRITE", err)
+		}
+		return
+	}
+
 	if err := r.ParseMultipartForm(10 * MB); err != nil {
 		w.WriteHeader(http.StatusBadRequest)
 		if _, err := w.Write([]byte("can not parse multipart " + err.Error())); err != nil {
@@ -694,6 +704,46 @@ func (h *Heruvym) PutFile(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	for _, files := range r.MultipartForm.File {
+		for _, fileHeader := range files {
+			fileSize := fileHeader.Size
+			fileType := utils.GetFileType(fileHeader.Filename)
+
+			switch fileType {
+			case "image":
+				if fileSize > 5*MB {
+					w.WriteHeader(http.StatusRequestEntityTooLarge)
+					if _, err := w.Write([]byte("Image file size exceeds the limit of 5MB")); err != nil {
+						fmt.Println("CAN NOT WRITE", err)
+					}
+					return
+				}
+			case "video":
+				if fileSize > 50*MB {
+					w.WriteHeader(http.StatusRequestEntityTooLarge)
+					if _, err := w.Write([]byte("Video file size exceeds the limit of 50MB")); err != nil {
+						fmt.Println("CAN NOT WRITE", err)
+					}
+					return
+				}
+			case "document":
+				if fileSize > 10*MB {
+					w.WriteHeader(http.StatusRequestEntityTooLarge)
+					if _, err := w.Write([]byte("Document file size exceeds the limit of 10MB")); err != nil {
+						fmt.Println("CAN NOT WRITE", err)
+					}
+					return
+				}
+			default:
+				w.WriteHeader(http.StatusNotAcceptable)
+				if _, err := w.Write([]byte("Unsupported file type")); err != nil {
+					fmt.Println("CAN NOT WRITE", err)
+				}
+				return
+			}
+		}
+	}
+
 	filesCount := len(r.MultipartForm.File)
 
 	if filesCount == 0 {
@@ -704,15 +754,6 @@ func (h *Heruvym) PutFile(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	sess := jwt_adapter.Get(r.Context())
-	if sess == nil {
-		w.WriteHeader(http.StatusBadRequest)
-		if _, err := w.Write([]byte("not authorized")); err != nil {
-			fmt.Println("CAN NOT WRITE", err)
-		}
-		return
-	}
-
 	var req PutFileReq
 	req.Ticket = r.MultipartForm.Value["ticket"][0]
 

test/main_test.go

@@ -191,13 +191,62 @@ func TestTicket(t *testing.T) {
 				//assert.NoError(t, err)
 				//assert.Equal(t, "test", channed.Message)
 			})
-			bdd.Test(t, "send file", func() {
+			bdd.Test(t, "send file passed", func() {
 				body := &bytes.Buffer{}
 				writer := multipart.NewWriter(body)
 
-				partf, _ := writer.CreateFormFile("f.go", "f.go")
+				partf, _ := writer.CreateFormFile("image.jpg", "image.jpg")
-				f, err := os.OpenFile("./main_test.go", os.O_RDWR, 0644)
+				f, err := os.OpenFile("./testfiles/image.jpg", os.O_RDWR, 0644)
 				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("image2.png", "image2.png")
+				f, err = os.OpenFile("./testfiles/image2.png", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("gtaV.mp4", "gtaV.mp4")
+				f, err = os.OpenFile("./testfiles/gtaV.mp4", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.csv", "test.csv")
+				f, err = os.OpenFile("./testfiles/test.csv", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.doc", "test.doc")
+				f, err = os.OpenFile("./testfiles/test.doc", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.docx", "test.docx")
+				f, err = os.OpenFile("./testfiles/test.docx", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.pdf", "test.pdf")
+				f, err = os.OpenFile("./testfiles/test.pdf", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.txt", "test.txt")
+				f, err = os.OpenFile("./testfiles/test.txt", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
+				io.Copy(partf, f)
+
+				partf, _ = writer.CreateFormFile("test.xlsx", "test.xlsx")
+				f, err = os.OpenFile("./testfiles/test.xlsx", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				defer f.Close()
 				io.Copy(partf, f)
 
 				assert.NoError(t, writer.WriteField("ticket", ticket.Ticket))
@@ -215,6 +264,54 @@ func TestTicket(t *testing.T) {
 				str, err := ioutil.ReadAll(resp.Body)
 				fmt.Println("resp files", err, string(str))
 			})
+			bdd.Test(t, "send mp3 file", func() {
+				body := &bytes.Buffer{}
+				writer := multipart.NewWriter(body)
+
+				partf, _ := writer.CreateFormFile("audio.mp3", "audio.mp3")
+				f, err := os.OpenFile("./testfiles/audio.mp3", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				io.Copy(partf, f)
+
+				assert.NoError(t, writer.WriteField("ticket", ticket.Ticket))
+
+				writer.Close()
+				req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://localhost:1488/sendFiles", body)
+				req.Header.Set("Content-Type", "multipart/form-data; boundary="+writer.Boundary())
+				req.Header.Set("Referer", "http://localhost:1488")
+				req.Header.Add(jwt_adapter.DefaultHeaderKey, "Bearer "+token)
+
+				resp, err := http.DefaultClient.Do(req)
+				assert.NoError(t, err)
+				assert.Equal(t, http.StatusNotAcceptable, resp.StatusCode)
+
+				str, err := ioutil.ReadAll(resp.Body)
+				fmt.Println("resp files", err, string(str))
+			})
+			bdd.Test(t, "send file big size", func() {
+				body := &bytes.Buffer{}
+				writer := multipart.NewWriter(body)
+
+				partf, _ := writer.CreateFormFile("gtaV+size.mp4", "gtaV+size.mp4")
+				f, err := os.OpenFile("./testfiles/gtaV+size.mp4", os.O_RDWR, 0644)
+				assert.NoError(t, err)
+				io.Copy(partf, f)
+
+				assert.NoError(t, writer.WriteField("ticket", ticket.Ticket))
+
+				writer.Close()
+				req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://localhost:1488/sendFiles", body)
+				req.Header.Set("Content-Type", "multipart/form-data; boundary="+writer.Boundary())
+				req.Header.Set("Referer", "http://localhost:1488")
+				req.Header.Add(jwt_adapter.DefaultHeaderKey, "Bearer "+token)
+
+				resp, err := http.DefaultClient.Do(req)
+				assert.NoError(t, err)
+				assert.Equal(t, http.StatusRequestEntityTooLarge, resp.StatusCode)
+
+				str, err := ioutil.ReadAll(resp.Body)
+				fmt.Println("resp files", err, string(str))
+			})
 			//todo
 			bdd.Test(t, "/subscribe, tools.SseWrapper(h.GetList)", func() {
 				client := &http.Client{}

test/testfiles/test.csv

@@ -0,0 +1 @@
+123

test/testfiles/test.txt

@@ -0,0 +1 @@
+123

utils/utils.go

@@ -0,0 +1,20 @@
+package utils
+
+import (
+	"path/filepath"
+	"strings"
+)
+
+func GetFileType(filename string) string {
+	ext := strings.ToLower(filepath.Ext(filename))
+	switch ext {
+	case ".jpg", ".png":
+		return "image"
+	case ".mp4":
+		return "video"
+	case ".txt", ".doc", ".docx", ".csv", ".xlsx", ".pdf":
+		return "document"
+	default:
+		return "unsupported"
+	}
+}