From 6a01e13858ca792713d1ede2a93babf07b7c1e28 Mon Sep 17 00:00:00 2001
From: Denis Arh <denis.arh@kendu.si>
Date: Sat, 27 Apr 2019 15:38:27 +0200
Subject: [PATCH] Cleanup attachment servig logic for CRM

---
 crm/rest/attachment.go                 | 43 +++++++++++-------
 crm/rest/handlers/attachment_custom.go | 61 --------------------------
 crm/rest/router.go                     |  2 +-
 3 files changed, 29 insertions(+), 77 deletions(-)
 delete mode 100644 crm/rest/handlers/attachment_custom.go

diff --git a/crm/rest/attachment.go b/crm/rest/attachment.go
index 3f1bcef24..5b4bb1b37 100644
--- a/crm/rest/attachment.go
+++ b/crm/rest/attachment.go
@@ -4,11 +4,11 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/http"
 	"net/url"
 	"time"
 
 	"github.com/crusttech/crust/crm/internal/service"
-	"github.com/crusttech/crust/crm/rest/handlers"
 	"github.com/crusttech/crust/crm/rest/request"
 	"github.com/crusttech/crust/crm/types"
 	"github.com/crusttech/crust/internal/auth"
@@ -84,15 +84,15 @@ func (ctrl Attachment) Original(ctx context.Context, r *request.AttachmentOrigin
 		return nil, err
 	}
 
-	return ctrl.get(ctx, r.AttachmentID, false, r.Download)
+	return ctrl.serve(ctx, r.AttachmentID, false, r.Download)
 }
 
-func (ctrl Attachment) Preview(ctx context.Context, r *request.AttachmentPreview) (interface{}, error) {
+func (ctrl *Attachment) Preview(ctx context.Context, r *request.AttachmentPreview) (interface{}, error) {
 	if err := ctrl.isAccessible(r.AttachmentID, r.UserID, r.Sign); err != nil {
 		return nil, err
 	}
 
-	return ctrl.get(ctx, r.AttachmentID, true, false)
+	return ctrl.serve(ctx, r.AttachmentID, true, false)
 }
 
 func (ctrl Attachment) isAccessible(attachmentID, userID uint64, signature string) error {
@@ -111,25 +111,38 @@ func (ctrl Attachment) isAccessible(attachmentID, userID uint64, signature strin
 	return nil
 }
 
-func (ctrl Attachment) get(ctx context.Context, ID uint64, preview, download bool) (handlers.Downloadable, error) {
-	rval := &file{download: download}
+func (ctrl Attachment) serve(ctx context.Context, ID uint64, preview, download bool) (interface{}, error) {
+	return func(w http.ResponseWriter, req *http.Request) {
+		att, err := ctrl.attachment.With(ctx).FindByID(ID)
+		if err != nil {
+			// Simplify error handling for now
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+
+		var fh io.ReadSeeker
 
-	if att, err := ctrl.attachment.FindByID(ID); err != nil {
-		return nil, err
-	} else {
-		rval.Attachment = att
 		if preview {
-			rval.content, err = ctrl.attachment.OpenPreview(att)
+			fh, err = ctrl.attachment.OpenPreview(att)
 		} else {
-			rval.content, err = ctrl.attachment.OpenOriginal(att)
+			fh, err = ctrl.attachment.OpenOriginal(att)
 		}
 
 		if err != nil {
-			return nil, err
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
 		}
-	}
 
-	return rval, nil
+		name := url.QueryEscape(att.Name)
+
+		if download {
+			w.Header().Add("Content-Disposition", "attachment; filename="+name)
+		} else {
+			w.Header().Add("Content-Disposition", "inline; filename="+name)
+		}
+
+		http.ServeContent(w, req, name, att.CreatedAt, fh)
+	}, nil
 }
 
 func makeAttachmentPayload(a *types.Attachment, userID uint64) *attachmentPayload {
diff --git a/crm/rest/handlers/attachment_custom.go b/crm/rest/handlers/attachment_custom.go
deleted file mode 100644
index 2e6c49e39..000000000
--- a/crm/rest/handlers/attachment_custom.go
+++ /dev/null
@@ -1,61 +0,0 @@
-package handlers
-
-import (
-	"io"
-	"net/http"
-	"net/url"
-	"time"
-
-	"github.com/crusttech/crust/crm/rest/request"
-)
-
-type Downloadable interface {
-	Name() string
-	Download() bool
-	ModTime() time.Time
-	Content() io.ReadSeeker
-	Valid() bool
-}
-
-func NewAttachmentDownloadable(ctrl AttachmentAPI) *Attachment {
-	h := NewAttachment(ctrl)
-	h.Original = func(w http.ResponseWriter, r *http.Request) {
-		defer r.Body.Close()
-		params := request.NewAttachmentOriginal()
-		params.Fill(r)
-
-		f, err := ctrl.Original(r.Context(), params)
-		serveFile(f, err, w, r)
-	}
-
-	h.Preview = func(w http.ResponseWriter, r *http.Request) {
-		defer r.Body.Close()
-		params := request.NewAttachmentPreview()
-		params.Fill(r)
-
-		f, err := ctrl.Preview(r.Context(), params)
-		serveFile(f, err, w, r)
-	}
-
-	return h
-}
-
-func serveFile(f interface{}, err error, w http.ResponseWriter, r *http.Request) {
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-	} else if dl, ok := f.(Downloadable); ok {
-		if !dl.Valid() {
-			w.WriteHeader(http.StatusNotFound)
-		} else {
-			if dl.Download() {
-				w.Header().Add("Content-Disposition", "attachment; filename="+url.QueryEscape(dl.Name()))
-			} else {
-				w.Header().Add("Content-Disposition", "inline; filename="+url.QueryEscape(dl.Name()))
-			}
-
-			http.ServeContent(w, r, dl.Name(), dl.ModTime(), dl.Content())
-		}
-	} else {
-		http.Error(w, "Got incompatible type from controller", http.StatusInternalServerError)
-	}
-}
diff --git a/crm/rest/router.go b/crm/rest/router.go
index 110119433..b8dde6b61 100644
--- a/crm/rest/router.go
+++ b/crm/rest/router.go
@@ -41,6 +41,6 @@ func MountRoutes() func(chi.Router) {
 		})
 
 		// Use alternative handlers that support file serving
-		handlers.NewAttachmentDownloadable(attachment).MountRoutes(r)
+		handlers.NewAttachment(attachment).MountRoutes(r)
 	}
 }