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.

3 files changed, 98 insertions, 1 deletions

diff --git a/plugin/client_rpc.go b/plugin/client_rpc.go
index ed76dc6e8..f5831445e 100644
--- a/plugin/client_rpc.go
+++ b/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
+}
diff --git a/plugin/hooks.go b/plugin/hooks.go
index 5200291f2..944909077 100644
--- a/plugin/hooks.go
+++ b/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)
 }
diff --git a/plugin/interface_generator/main.go b/plugin/interface_generator/main.go
index 4b8b6786f..b804c5d4f 100644
--- a/plugin/interface_generator/main.go
+++ b/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 {