Adding FileWillBeUploaded plugin hook (#9169)

* Adding file upload hook.

* Adding hook test for FileWillBeUploaded

* Some debugging fixes.

* Fix typo.

* Fixing double close

* Fix capitalization on docs.

app/file.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/mattermost/mattermost-server/mlog"
 	"github.com/mattermost/mattermost-server/model"
+	"github.com/mattermost/mattermost-server/plugin"
 	"github.com/mattermost/mattermost-server/utils"
 )
 
@@ -366,7 +367,7 @@ func (a *App) UploadFiles(teamId string, channelId string, userId string, files
 		io.Copy(buf, file)
 		data := buf.Bytes()
 
-		info, err := a.DoUploadFile(time.Now(), teamId, channelId, userId, filenames[i], data)
+		info, data, err := a.DoUploadFileExpectModification(time.Now(), teamId, channelId, userId, filenames[i], data)
 		if err != nil {
 			return nil, err
 		}
@@ -390,6 +391,11 @@ func (a *App) UploadFiles(teamId string, channelId string, userId string, files
 }
 
 func (a *App) DoUploadFile(now time.Time, rawTeamId string, rawChannelId string, rawUserId string, rawFilename string, data []byte) (*model.FileInfo, *model.AppError) {
+	info, _, err := a.DoUploadFileExpectModification(now, rawTeamId, rawChannelId, rawUserId, rawFilename, data)
+	return info, err
+}
+
+func (a *App) DoUploadFileExpectModification(now time.Time, rawTeamId string, rawChannelId string, rawUserId string, rawFilename string, data []byte) (*model.FileInfo, []byte, *model.AppError) {
 	filename := filepath.Base(rawFilename)
 	teamId := filepath.Base(rawTeamId)
 	channelId := filepath.Base(rawChannelId)
@@ -398,7 +404,7 @@ func (a *App) DoUploadFile(now time.Time, rawTeamId string, rawChannelId string,
 	info, err := model.GetInfoForBytes(filename, data)
 	if err != nil {
 		err.StatusCode = http.StatusBadRequest
-		return nil, err
+		return nil, data, err
 	}
 
 	if orientation, err := getImageOrientation(bytes.NewReader(data)); err == nil &&
@@ -419,7 +425,7 @@ func (a *App) DoUploadFile(now time.Time, rawTeamId string, rawChannelId string,
 		// Check dimensions before loading the whole thing into memory later on
 		if info.Width*info.Height > MaxImageSize {
 			err := model.NewAppError("uploadFile", "api.file.upload_file.large_image.app_error", map[string]interface{}{"Filename": filename}, "", http.StatusBadRequest)
-			return nil, err
+			return nil, data, err
 		}
 
 		nameWithoutExtension := filename[:strings.LastIndex(filename, ".")]
@@ -427,15 +433,33 @@ func (a *App) DoUploadFile(now time.Time, rawTeamId string, rawChannelId string,
 		info.ThumbnailPath = pathPrefix + nameWithoutExtension + "_thumb.jpg"
 	}
 
+	if a.PluginsReady() {
+		pluginContext := &plugin.Context{}
+		var rejectionReason string
+		a.Plugins.RunMultiPluginHook(func(hooks plugin.Hooks) bool {
+			var newBytes bytes.Buffer
+			info, rejectionReason = hooks.FileWillBeUploaded(pluginContext, info, bytes.NewReader(data), &newBytes)
+			rejected := info == nil
+			if !rejected && newBytes.Len() != 0 {
+				data = newBytes.Bytes()
+				info.Size = int64(len(data))
+			}
+			return !rejected
+		}, plugin.FileWillBeUploadedId)
+		if info == nil {
+			return nil, data, model.NewAppError("DoUploadFile", "File rejected by plugin. "+rejectionReason, nil, "", http.StatusBadRequest)
+		}
+	}
+
 	if _, err := a.WriteFile(bytes.NewReader(data), info.Path); err != nil {
-		return nil, err
+		return nil, data, err
 	}
 
 	if result := <-a.Srv.Store.FileInfo().Save(info); result.Err != nil {
-		return nil, result.Err
+		return nil, data, result.Err
 	}
 
-	return info, nil
+	return info, data, nil
 }
 
 func (a *App) HandleImages(previewPathList []string, thumbnailPathList []string, fileData [][]byte) {

app/plugin_api.go

@@ -45,9 +45,14 @@ func (api *PluginAPI) LoadPluginConfiguration(dest interface{}) error {
 	}
 
 	if pluginSettingsJsonBytes, err := json.Marshal(finalConfig); err != nil {
-		return err
+		api.logger.Error("Error marshaling config for plugin", mlog.Err(err))
+		return nil
 	} else {
-		return json.Unmarshal(pluginSettingsJsonBytes, dest)
+		err := json.Unmarshal(pluginSettingsJsonBytes, dest)
+		if err != nil {
+			api.logger.Error("Error unmarshaling config for plugin", mlog.Err(err))
+		}
+		return nil
 	}
 }
 

app/plugin_hooks_test.go

@@ -4,6 +4,8 @@
 package app
 
 import (
+	"bytes"
+	"io"
 	"io/ioutil"
 	"os"
 	"os/exec"
@@ -302,3 +304,69 @@ func TestHookMessageHasBeenUpdated(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestHookFileWillBeUploaded(t *testing.T) {
+	th := Setup().InitBasic()
+	defer th.TearDown()
+
+	var mockAPI plugintest.API
+	mockAPI.On("LoadPluginConfiguration", mock.Anything).Return(nil)
+	mockAPI.On("DeleteUser", "testhook.txt").Return(nil)
+	mockAPI.On("DeleteTeam", "inputfile").Return(nil)
+	SetAppEnvironmentWithPlugins(t,
+		[]string{
+			`
+		package main
+
+		import (
+			"io"
+			"bytes"
+			"github.com/mattermost/mattermost-server/plugin"
+			"github.com/mattermost/mattermost-server/model"
+		)
+
+		type MyPlugin struct {
+			plugin.MattermostPlugin
+		}
+
+		func (p *MyPlugin) FileWillBeUploaded(c *plugin.Context, info *model.FileInfo, file io.Reader, output io.Writer) (*model.FileInfo, string) {
+			p.API.DeleteUser(info.Name)
+			var buf bytes.Buffer
+			buf.ReadFrom(file)
+			p.API.DeleteTeam(buf.String())
+
+			outbuf := bytes.NewBufferString("changedtext")
+			io.Copy(output, outbuf)
+			info.Name = "modifiedinfo"
+			return info, ""
+		}
+
+		func main() {
+			plugin.ClientMain(&MyPlugin{})
+		}
+	`}, th.App, func(*model.Manifest) plugin.API { return &mockAPI })
+
+	response, err := th.App.UploadFiles(
+		"noteam",
+		th.BasicChannel.Id,
+		th.BasicUser.Id,
+		[]io.ReadCloser{ioutil.NopCloser(bytes.NewBufferString("inputfile"))},
+		[]string{"testhook.txt"},
+		[]string{},
+	)
+	assert.Nil(t, err)
+	assert.NotNil(t, response)
+	assert.Equal(t, 1, len(response.FileInfos))
+	fileId := response.FileInfos[0].Id
+
+	fileInfo, err := th.App.GetFileInfo(fileId)
+	assert.Nil(t, err)
+	assert.NotNil(t, fileInfo)
+	assert.Equal(t, "modifiedinfo", fileInfo.Name)
+
+	fileReader, err := th.App.FileReader(fileInfo.Path)
+	assert.Nil(t, err)
+	var resultBuf bytes.Buffer
+	io.Copy(&resultBuf, fileReader)
+	assert.Equal(t, "changedtext", resultBuf.String())
+}

plugin/client_rpc.go

@@ -10,6 +10,7 @@ import (
 	"encoding/gob"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"net/http"
@@ -199,7 +200,10 @@ func (g *apiRPCClient) LoadPluginConfiguration(dest interface{}) error {
 	if err := g.client.Call("Plugin.LoadPluginConfiguration", _args, _returns); err != nil {
 		g.log.Error("RPC call to LoadPluginConfiguration API failed.", mlog.Err(err))
 	}
-	return json.Unmarshal(_returns.A, dest)
+	if err := json.Unmarshal(_returns.A, dest); err != nil {
+		g.log.Error("LoadPluginConfiguration API failed to unmarshal.", mlog.Err(err))
+	}
+	return nil
 }
 
 func (s *apiRPCServer) LoadPluginConfiguration(args *Z_LoadPluginConfigurationArgsArgs, returns *Z_LoadPluginConfigurationArgsReturns) error {
@@ -326,3 +330,86 @@ func (s *hooksRPCServer) ServeHTTP(args *Z_ServeHTTPArgs, returns *struct{}) err
 
 	return nil
 }
+
+func init() {
+	hookNameToId["FileWillBeUploaded"] = FileWillBeUploadedId
+}
+
+type Z_FileWillBeUploadedArgs struct {
+	A                     *Context
+	B                     *model.FileInfo
+	UploadedFileStream    uint32
+	ReplacementFileStream uint32
+}
+
+type Z_FileWillBeUploadedReturns struct {
+	A *model.FileInfo
+	B string
+}
+
+func (g *hooksRPCClient) FileWillBeUploaded(c *Context, info *model.FileInfo, file io.Reader, output io.Writer) (*model.FileInfo, string) {
+	if !g.implemented[FileWillBeUploadedId] {
+		return info, ""
+	}
+
+	uploadedFileStreamId := g.muxBroker.NextId()
+	go func() {
+		uploadedFileConnection, err := g.muxBroker.Accept(uploadedFileStreamId)
+		if err != nil {
+			g.log.Error("Plugin failed to serve upload file stream. MuxBroker could not Accept connection", mlog.Err(err))
+			return
+		}
+		defer uploadedFileConnection.Close()
+		serveIOReader(file, uploadedFileConnection)
+	}()
+
+	replacementFileStreamId := g.muxBroker.NextId()
+	go func() {
+		replacementFileConnection, err := g.muxBroker.Accept(replacementFileStreamId)
+		if err != nil {
+			g.log.Error("Plugin failed to serve replacement file stream. MuxBroker could not Accept connection", mlog.Err(err))
+			return
+		}
+		defer replacementFileConnection.Close()
+		if _, err := io.Copy(output, replacementFileConnection); err != nil && err != io.EOF {
+			g.log.Error("Error reading replacement file.", mlog.Err(err))
+		}
+	}()
+
+	_args := &Z_FileWillBeUploadedArgs{c, info, uploadedFileStreamId, replacementFileStreamId}
+	_returns := &Z_FileWillBeUploadedReturns{}
+	if g.implemented[FileWillBeUploadedId] {
+		if err := g.client.Call("Plugin.FileWillBeUploaded", _args, _returns); err != nil {
+			g.log.Error("RPC call FileWillBeUploaded to plugin failed.", mlog.Err(err))
+		}
+	}
+	return _returns.A, _returns.B
+}
+
+func (s *hooksRPCServer) FileWillBeUploaded(args *Z_FileWillBeUploadedArgs, returns *Z_FileWillBeUploadedReturns) error {
+	uploadFileConnection, err := s.muxBroker.Dial(args.UploadedFileStream)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "[ERROR] Can't connect to remote upload file stream, error: %v", err.Error())
+		return err
+	}
+	defer uploadFileConnection.Close()
+	fileReader := connectIOReader(uploadFileConnection)
+	defer fileReader.Close()
+
+	replacementFileConnection, err := s.muxBroker.Dial(args.ReplacementFileStream)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "[ERROR] Can't connect to remote replacement file stream, error: %v", err.Error())
+		return err
+	}
+	defer replacementFileConnection.Close()
+	returnFileWriter := replacementFileConnection
+
+	if hook, ok := s.impl.(interface {
+		FileWillBeUploaded(c *Context, info *model.FileInfo, file io.Reader, output io.Writer) (*model.FileInfo, string)
+	}); ok {
+		returns.A, returns.B = hook.FileWillBeUploaded(args.A, args.B, fileReader, returnFileWriter)
+	} else {
+		return fmt.Errorf("Hook FileWillBeUploaded called but not implemented.")
+	}
+	return nil
+}

plugin/hooks.go

@@ -4,6 +4,7 @@
 package plugin
 
 import (
+	"io"
 	"net/http"
 
 	"github.com/mattermost/mattermost-server/model"
@@ -28,6 +29,7 @@ const (
 	UserHasJoinedTeamId     = 11
 	UserHasLeftTeamId       = 12
 	ChannelHasBeenCreatedId = 13
+	FileWillBeUploadedId    = 14
 	TotalHooksId            = iota
 )
 
@@ -113,4 +115,11 @@ type Hooks interface {
 	// UserHasLeftTeam is invoked after the membership has been removed from the database.
 	// If actor is not nil, the user was removed from the team by the actor.
 	UserHasLeftTeam(c *Context, teamMember *model.TeamMember, actor *model.User)
+
+	// FileWillBeUploaded is invoked when a file is uploaded, but before it is committed to backing store.
+	// Read from file to retrieve the body of the uploaded file. You may modify the body of the file by writing to output.
+	// Returned FileInfo will be used instead of input FileInfo. Return nil to reject the file upload and include a text reason as the second argument.
+	// Note that this method will be called for files uploaded by plugins, including the plugin that uploaded the post.
+	// FileInfo.Size will be automatically set properly if you modify the file.
+	FileWillBeUploaded(c *Context, info *model.FileInfo, file io.Reader, output io.Writer) (*model.FileInfo, string)
 }

plugin/interface_generator/main.go

@@ -355,6 +355,7 @@ func removeExcluded(info *PluginInterfaceInfo) *PluginInterfaceInfo {
 			"Implemented",
 			"LoadPluginConfiguration",
 			"ServeHTTP",
+			"FileWillBeUploaded",
 		}
 		for _, exclusion := range excluded {
 			if exclusion == item {