corteza/compose/automation/attachment_handler.go

package automation

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/cortezaproject/corteza-server/compose/types"
)

type (
	attachmentService interface {
		FindByID(ctx context.Context, namespaceID, attachmentID uint64) (*types.Attachment, error)
		CreateRecordAttachment(ctx context.Context, namespaceID uint64, name string, size int64, fh io.ReadSeeker, moduleID, recordID uint64, fieldName string) (att *types.Attachment, err error)
		DeleteByID(ctx context.Context, namespaceID uint64, attachmentID uint64) error
		OpenOriginal(att *types.Attachment) (io.ReadSeekCloser, error)
		OpenPreview(att *types.Attachment) (io.ReadSeekCloser, error)
	}

	attachmentHandler struct {
		reg modulesHandlerRegistry
		svc attachmentService
	}

	attachmentLookup interface {
		GetAttachment() (bool, uint64, *types.Attachment)
	}
)

func AttachmentHandler(reg modulesHandlerRegistry, svc attachmentService) *attachmentHandler {
	h := &attachmentHandler{
		reg: reg,
		svc: svc,
	}

	h.register()
	return h
}

func (h attachmentHandler) lookup(ctx context.Context, args *attachmentLookupArgs) (results *attachmentLookupResults, err error) {
	results = &attachmentLookupResults{}
	results.Attachment, err = h.svc.FindByID(ctx, 0, args.Attachment)
	return
}

func (h attachmentHandler) delete(ctx context.Context, args *attachmentDeleteArgs) error {
	return h.svc.DeleteByID(ctx, 0, args.Attachment)
}

func (h attachmentHandler) openOriginal(ctx context.Context, args *attachmentOpenOriginalArgs) (*attachmentOpenOriginalResults, error) {
	att, err := lookupAttachment(ctx, h.svc, args)
	if err != nil {
		return nil, err
	}

	r := &attachmentOpenOriginalResults{}
	r.Content, err = h.svc.OpenOriginal(att)
	if err != nil {
		return nil, err
	}

	// @todo we need to call Close() when file is read (or at the end of the workflow)
	//       some kind of workflow-cleanup facility is needed

	return r, nil
}

func (h attachmentHandler) openPreview(ctx context.Context, args *attachmentOpenPreviewArgs) (*attachmentOpenPreviewResults, error) {
	att, err := lookupAttachment(ctx, h.svc, args)
	if err != nil {
		return nil, err
	}

	r := &attachmentOpenPreviewResults{}
	r.Content, err = h.svc.OpenOriginal(att)
	if err != nil {
		return nil, err
	}

	// @todo we need to call Close() when file is read (or at the end of the workflow)
	//       some kind of workflow-cleanup facility is needed

	return r, nil
}

func (h attachmentHandler) create(ctx context.Context, args *attachmentCreateArgs) (*attachmentCreateResults, error) {
	var (
		err error
		att *types.Attachment

		fh   io.ReadSeeker
		size int64
	)

	switch {
	case len(args.contentBytes) > 0:
		size = int64(len(args.contentBytes))
		fh = bytes.NewReader(args.contentBytes)

	case args.contentStream != nil:
		if rs, is := args.contentStream.(io.ReadSeeker); is {
			_, err = rs.Seek(0, 0)
			if err != nil {
				return nil, err
			}

			size, err = getReaderSize(rs)
			if err != nil {
				return nil, err
			}

			_, err = rs.Seek(0, 0)
			if err != nil {
				return nil, err
			}

			fh = rs
		} else {
			// In case we only got a reader...
			//
			// For future proofing, for handling larger attachment, we create a temp.
			// file which we then use as a reader

			// Preparations
			tmpf, err := ioutil.TempFile("", "reader")
			if err != nil {
				return nil, err
			}
			defer tmpf.Close()
			defer os.Remove(tmpf.Name())

			// Writing content to file
			w := bufio.NewWriter(tmpf)
			r := args.contentStream
			buf := make([]byte, 1024)
			for {
				// read
				n, err := r.Read(buf)
				if err != nil && err != io.EOF {
					return nil, err
				}
				if n == 0 {
					break
				}

				// on-the-fly size calculation
				size += int64(n)

				// write
				if _, err := w.Write(buf[:n]); err != nil {
					return nil, err
				}
			}
			if err = w.Flush(); err != nil {
				return nil, err
			}

			_, err = tmpf.Seek(0, 0)
			if err != nil {
				return nil, err
			}
			fh = tmpf
		}

	default:
		fh = strings.NewReader(args.contentString)
		size = int64(len(args.contentString))
	}

	switch {
	case args.Resource != nil:
		att, err = h.svc.CreateRecordAttachment(
			ctx,
			args.Resource.NamespaceID,
			args.Name,
			size,
			fh,
			args.Resource.ModuleID,
			args.Resource.ID,
			args.FieldName,
		)
	default:
		return nil, fmt.Errorf("unknown resource")
	}

	if err != nil {
		return nil, err
	}

	return &attachmentCreateResults{att}, nil
}

func lookupAttachment(ctx context.Context, svc attachmentService, args attachmentLookup) (*types.Attachment, error) {
	_, ID, attachment := args.GetAttachment()

	switch {
	case attachment != nil:
		return attachment, nil
	case ID > 0:
		return svc.FindByID(ctx, 0, ID)
	}

	return nil, fmt.Errorf("empty attachment lookup params")
}

func getReaderSize(r io.Reader) (size int64, err error) {
	buf := make([]byte, 1024)
	var n int
	for {
		// read a chunk
		n, err = r.Read(buf)
		if err != nil && err != io.EOF {
			return 0, err
		}
		if n == 0 {
			break
		}
		size += int64(n)
	}

	return size, nil
}