[MI-3487] Fixed the issue of checking file size before downloading it…

… from Teams (#293)

* [MI-3487] Fixed the issue of checking file size before downloading it from Teams

* [MI-3487] Fixed the failing unit tests and modified an error message

server/handlers/attachments.go

@@ -55,7 +55,8 @@ func (ah *ActivityHandler) handleDownloadFile(weburl string, client msteams.Clie
 		return client.GetHostedFileContent(activityIDs)
 	}
 
-	data, err := client.GetFileContent(weburl)
+	fileSizeAllowed := *ah.plugin.GetAPI().GetConfig().FileSettings.MaxFileSize
+	data, err := client.GetFileContent(weburl, fileSizeAllowed)
 	if err != nil {
 		return nil, err
 	}
@@ -109,13 +110,6 @@ func (ah *ActivityHandler) handleAttachments(channelID, text string, msg *msteam
 			continue
 		}
 
-		fileSizeAllowed := *ah.plugin.GetAPI().GetConfig().FileSettings.MaxFileSize
-		if len(attachmentData) > int(fileSizeAllowed) {
-			ah.plugin.GetAPI().LogError("cannot upload file to Mattermost as its size is greater than allowed size", "filename", a.Name)
-			errorFound = true
-			continue
-		}
-
 		contentType := http.DetectContentType(attachmentData)
 		if strings.HasPrefix(contentType, "image") && contentType != "image/svg+xml" {
 			w, h, imageErr := imaging.GetDimensions(bytes.NewReader(attachmentData))

server/handlers/attachments_test.go

@@ -20,26 +20,37 @@ import (
 func TestHandleDownloadFile(t *testing.T) {
 	ah := ActivityHandler{}
 	client := mocksClient.NewClient(t)
+	mockAPI := &plugintest.API{}
+	mockAPI.On("GetConfig").Return(&model.Config{
+		FileSettings: model.FileSettings{
+			MaxFileSize: model.NewInt64(5),
+		},
+	})
 
 	for _, testCase := range []struct {
 		description   string
 		userID        string
 		weburl        string
 		expectedError string
+		setupPlugin   func(plugin *mocksPlugin.PluginIface)
 		setupClient   func()
 	}{
 		{
 			description: "Successfully file downloaded",
 			userID:      testutils.GetUserID(),
 			weburl:      "https://example.com/file1.txt",
+			setupPlugin: func(p *mocksPlugin.PluginIface) {
+				p.On("GetAPI").Return(mockAPI)
+			},
 			setupClient: func() {
-				client.On("GetFileContent", "https://example.com/file1.txt").Return([]byte("data"), nil)
+				client.On("GetFileContent", "https://example.com/file1.txt", int64(5)).Return([]byte("data"), nil)
 			},
 		},
 		{
 			description: "Successfully downloaded hosted content file",
 			userID:      testutils.GetUserID(),
 			weburl:      "https://graph.microsoft.com/beta/teams/mock-teamID/channels/mock-channelID/messages/mock-messageID/hostedContents/mock-hostedContentsID/$value",
+			setupPlugin: func(p *mocksPlugin.PluginIface) {},
 			setupClient: func() {
 				client.On("GetHostedFileContent", mock.AnythingOfType("*msteams.ActivityIds")).Return([]byte("data"), nil)
 			},
@@ -48,13 +59,17 @@ func TestHandleDownloadFile(t *testing.T) {
 			description:   "Unable to get file content",
 			userID:        testutils.GetUserID(),
 			expectedError: "Error while getting file content",
+			setupPlugin: func(p *mocksPlugin.PluginIface) {
+				p.On("GetAPI").Return(mockAPI)
+			},
 			setupClient: func() {
-				client.On("GetFileContent", "").Return(nil, errors.New("Error while getting file content"))
+				client.On("GetFileContent", "", int64(5)).Return(nil, errors.New("Error while getting file content"))
 			},
 		},
 	} {
 		t.Run(testCase.description, func(t *testing.T) {
 			p := mocksPlugin.NewPluginIface(t)
+			testCase.setupPlugin(p)
 			testCase.setupClient()
 			ah.plugin = p
 
@@ -288,7 +303,7 @@ func TestHandleAttachments(t *testing.T) {
 				}, nil)
 			},
 			setupClient: func(client *mocksClient.Client) {
-				client.On("GetFileContent", "").Return([]byte{}, nil)
+				client.On("GetFileContent", "", int64(5)).Return([]byte{}, nil)
 			},
 			attachments: []msteams.Attachment{
 				{
@@ -314,31 +329,6 @@ func TestHandleAttachments(t *testing.T) {
 				},
 			},
 		},
-		{
-			description: "File size is greater than the max allowed size",
-			setupPlugin: func(p *mocksPlugin.PluginIface, mockAPI *plugintest.API, client *mocksClient.Client, store *mocksStore.Store) {
-				p.On("GetClientForApp").Return(client)
-				p.On("GetAPI").Return(mockAPI)
-			},
-			setupAPI: func(mockAPI *plugintest.API) {
-				mockAPI.On("GetConfig").Return(&model.Config{
-					FileSettings: model.FileSettings{
-						MaxFileSize: model.NewInt64(-1),
-					},
-				})
-				mockAPI.On("LogError", "cannot upload file to Mattermost as its size is greater than allowed size", "filename", "mock-name").Return()
-			},
-			setupClient: func(client *mocksClient.Client) {
-				client.On("GetFileContent", "").Return([]byte{}, nil)
-			},
-			attachments: []msteams.Attachment{
-				{
-					Name: "mock-name",
-				},
-			},
-			expectedText:  "mock-text",
-			expectedError: true,
-		},
 		{
 			description: "Error uploading the file",
 			setupPlugin: func(p *mocksPlugin.PluginIface, mockAPI *plugintest.API, client *mocksClient.Client, store *mocksStore.Store) {
@@ -355,7 +345,7 @@ func TestHandleAttachments(t *testing.T) {
 				mockAPI.On("LogError", "upload file to Mattermost failed", "filename", "mock-name", "error", "error uploading the file").Return()
 			},
 			setupClient: func(client *mocksClient.Client) {
-				client.On("GetFileContent", "").Return([]byte{}, nil)
+				client.On("GetFileContent", "", int64(5)).Return([]byte{}, nil)
 			},
 			attachments: []msteams.Attachment{
 				{
@@ -380,7 +370,7 @@ func TestHandleAttachments(t *testing.T) {
 				mockAPI.On("LogDebug", "discarding the rest of the attachments as Mattermost supports only 10 attachments per post").Return()
 			},
 			setupClient: func(client *mocksClient.Client) {
-				client.On("GetFileContent", "").Return([]byte{}, nil)
+				client.On("GetFileContent", "", int64(5)).Return([]byte{}, nil)
 			},
 			attachments: []msteams.Attachment{
 				{}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {},

server/msteams/client.go

@@ -1162,7 +1162,7 @@ func (tc *ClientImpl) GetUser(userID string) (*User, error) {
 	return &user, nil
 }
 
-func (tc *ClientImpl) GetFileContent(weburl string) ([]byte, error) {
+func (tc *ClientImpl) GetFileContent(weburl string, fileSizeAllowed int64) ([]byte, error) {
 	u, err := url.Parse(weburl)
 	if err != nil {
 		return nil, err
@@ -1218,6 +1218,11 @@ func (tc *ClientImpl) GetFileContent(weburl string) ([]byte, error) {
 		return nil, errors.New("downloadUrl not found")
 	}
 
+	fileSize := item.GetSize()
+	if fileSize != nil && *fileSize > fileSizeAllowed {
+		return nil, fmt.Errorf("skipping file download from MS Teams because the file size is greater than the allowed size")
+	}
+
 	data, err := drives.NewItemItemsItemContentRequestBuilder(*(downloadURL.(*string)), tc.client.RequestAdapter).Get(tc.ctx, nil)
 	if err != nil {
 		return nil, NormalizeGraphAPIError(err)

server/msteams/interface.go

@@ -40,7 +40,7 @@ type Client interface {
 	GetUser(userID string) (*User, error)
 	GetMyID() (string, error)
 	GetMe() (*User, error)
-	GetFileContent(weburl string) ([]byte, error)
+	GetFileContent(weburl string, fileSizeAllowed int64) ([]byte, error)
 	GetHostedFileContent(activityIDs *ActivityIds) ([]byte, error)
 	GetCodeSnippet(url string) (string, error)
 	ListUsers() ([]User, error)