trantor/lib/upload.go

package trantor

import (
	log "github.com/cihub/seelog"

	"bytes"
	"crypto/rand"
	"encoding/base64"
	"io/ioutil"
	"mime/multipart"

	"github.com/meskio/epubgo"
	"gitlab.com/trantor/trantor/lib/database"
	"gitlab.com/trantor/trantor/lib/parser"
	"gitlab.com/trantor/trantor/lib/storage"
)

const (
	uploadChanSize = 100
)

func InitUpload(database database.DB, store storage.Store) {
	uploadChannel = make(chan uploadRequest, uploadChanSize)
	go uploadWorker(database, store)
}

var uploadChannel chan uploadRequest

type uploadRequest struct {
	file     multipart.File
	filename string
}

func uploadWorker(database database.DB, store storage.Store) {
	for req := range uploadChannel {
		processFile(req, database, store)
	}
}

func processFile(req uploadRequest, db database.DB, store storage.Store) {
	defer req.file.Close()

	epub, err := openMultipartEpub(req.file)
	if err != nil {
		log.Warn("Not valid epub uploaded file ", req.filename, ": ", err)
		return
	}
	defer epub.Close()

	book := parser.EpubMetadata(epub)
	book.ID = genID()

	req.file.Seek(0, 0)
	size, err := store.Store(book.ID, req.file, epubFile)
	if err != nil {
		log.Error("Error storing book (", book.ID, "): ", err)
		return
	}

	book.FileSize = int(size)
	book.Cover = GetCover(epub, book.ID, store)

	err = db.AddBook(book)
	if err != nil {
		log.Error("Error storing metadata (", book.ID, "): ", err)
		return
	}
	log.Info("File uploaded: ", req.filename)
}

func uploadPostHandler(h handler) {
	if h.ro {
		h.sess.Notify("Upload failed!", "The library is in Read Only mode, no books can be uploaded", "error")
		uploadHandler(h)
		return
	}

	problem := false

	h.r.ParseMultipartForm(20000000)
	filesForm := h.r.MultipartForm.File["epub"]
	for _, f := range filesForm {
		file, err := f.Open()
		if err != nil {
			log.Error("Can not open uploaded file ", f.Filename, ": ", err)
			h.sess.Notify("Upload problem!", "There was a problem with book "+f.Filename, "error")
			problem = true
			continue
		}
		uploadChannel <- uploadRequest{file, f.Filename}
	}

	if !problem {
		if len(filesForm) > 0 {
			h.sess.Notify("Upload successful!", "Thank you for your contribution", "success")
		} else {
			h.sess.Notify("Upload problem!", "No books where uploaded.", "error")
		}
	}
	uploadHandler(h)
}

func uploadHandler(h handler) {
	var data uploadData
	data.S = GetStatus(h)
	data.S.Title = "Upload -- " + data.S.Title
	data.S.Upload = true
	h.template.load(h, "upload", data)
}

type uploadData struct {
	S Status
}

func openMultipartEpub(file multipart.File) (*epubgo.Epub, error) {
	buff, _ := ioutil.ReadAll(file)
	reader := bytes.NewReader(buff)
	return epubgo.Load(reader, int64(len(buff)))
}

func genID() string {
	b := make([]byte, 12)
	rand.Read(b)
	return base64.URLEncoding.EncodeToString(b)
}