contented/Server.go

package contented

import (
	"crypto/sha512"
	"encoding/hex"
	"encoding/json"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/boltdb/bolt"
)

var SERVER_HEADER string = `contented/0.1`

type ServerOptions struct {
	DataDirectory  string
	DBPath         string
	AppTitle       string
	MaxUploadBytes int64
}

type Server struct {
	opts           ServerOptions
	db             *bolt.DB
	metadataBucket []byte
}

func NewServer(opts *ServerOptions) (*Server, error) {
	s := &Server{
		opts:           *opts,
		metadataBucket: []byte(`METADATA`),
	}

	b, err := bolt.Open(opts.DBPath, 0644, bolt.DefaultOptions)
	if err != nil {
		return nil, err
	}

	err = b.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(s.metadataBucket)
		return err
	})
	if err != nil {
		return nil, err
	}

	s.db = b

	return s, nil
}

type Metadata struct {
	FileSize   int64
	UploadTime time.Time
	UploadIP   string
	Filename   string
}

func (this *Server) handleUpload(src multipart.File, hdr *multipart.FileHeader, UploadIP string) (string, error) {

	// Get file length
	srcLen, err := src.Seek(0, io.SeekEnd)
	if err != nil {
		return "", err
	}

	_, err = src.Seek(0, io.SeekStart)
	if err != nil {
		return "", err
	}

	// Get file hash
	hasher := sha512.New512_256()
	_, err = io.CopyN(hasher, src, int64(srcLen))
	if err != nil {
		return "", err
	}

	_, err = src.Seek(0, io.SeekStart)
	if err != nil {
		return "", err
	}

	// Save file to disk
	fileID := hex.EncodeToString(hasher.Sum(nil))
	dest, err := os.OpenFile(filepath.Join(this.opts.DataDirectory, fileID), os.O_CREATE|os.O_WRONLY, 0600)
	if err != nil {
		if os.IsExist(err) {
			return fileID, nil // hash matches existing upload
		}

		return "", err
	}
	defer dest.Close()

	_, err = io.CopyN(dest, src, int64(srcLen))
	if err != nil {
		return "", err
	}

	// Persist metadata to DB
	m := Metadata{
		Filename:   hdr.Filename,
		UploadTime: time.Now(),
		UploadIP:   UploadIP,
		FileSize:   srcLen,
	}
	err = this.SetMetadata(fileID, m)
	if err != nil {
		return "", err
	}

	// Done
	return fileID, nil
}

func (this *Server) Metadata(fileID string) (*Metadata, error) {
	var m Metadata
	err := this.db.View(func(tx *bolt.Tx) error {
		content := tx.Bucket(this.metadataBucket).Get([]byte(fileID))
		if len(content) == 0 {
			return os.ErrNotExist
		}

		return json.Unmarshal(content, &m)
	})
	if err != nil {
		return nil, err
	}

	return &m, nil
}

func (this *Server) SetMetadata(fileID string, m Metadata) error {
	jb, err := json.Marshal(m)
	if err != nil {
		return err
	}

	return this.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket(this.metadataBucket).Put([]byte(fileID), jb)
	})
}

func (this *Server) handleView(w http.ResponseWriter, r *http.Request, fileID string) error {
	// Load file
	f, err := os.Open(filepath.Join(this.opts.DataDirectory, fileID))
	if err != nil {
		return err
	}

	defer f.Close()

	// Load metadata
	m, err := this.Metadata(fileID)
	if err != nil {
		return err
	}

	http.ServeContent(w, r, m.Filename, m.UploadTime, f)
	return nil
}

func (this *Server) handleInformation(w http.ResponseWriter, fileID string) {
	m, err := this.Metadata(fileID)
	if err != nil {
		if os.IsNotExist(err) {
			http.Error(w, "Not found", 404)
			return
		}

		log.Println(err.Error())
		http.Error(w, "Internal error", 500)
		return
	}

	jb, err := json.Marshal(m)
	if err != nil {
		log.Println(err.Error())
		http.Error(w, "Internal error", 500)
		return
	}

	w.Header().Set(`Content-Type`, `application/json`)
	w.WriteHeader(200)
	w.Write(jb)
}

func (this *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	r.Header.Set(`Server`, SERVER_HEADER)
	if this.opts.MaxUploadBytes > 0 {
		r.Body = http.MaxBytesReader(w, r.Body, this.opts.MaxUploadBytes)
	}

	if r.Method == "GET" && strings.HasPrefix(r.URL.Path, `/view/`) {
		err := this.handleView(w, r, r.URL.Path[6:])
		if err != nil {
			log.Printf("%s View failed: %s\n", r.RemoteAddr, err.Error())
			if os.IsNotExist(err) {
				http.Error(w, "File not found", 404)
			} else {
				http.Error(w, "Couldn't provide content", 500)
			}
		}

	} else if r.Method == "GET" && strings.HasPrefix(r.URL.Path, `/info/`) {
		this.handleInformation(w, r.URL.Path[6:])

	} else if r.Method == "POST" && r.URL.Path == `/upload` {
		mp, mph, err := r.FormFile("f")
		if err != nil {
			log.Printf("%s Invalid request: %s\n", r.RemoteAddr, err.Error())
			http.Error(w, "Invalid request", 400)
		}
		path, err := this.handleUpload(mp, mph, strings.TrimRight(r.RemoteAddr, "0123456789:")) // strip remote-port
		if err != nil {
			log.Printf("%s Upload failed: %s\n", r.RemoteAddr, err.Error())
			http.Error(w, "Upload failed", 500)
		} else {
			http.Redirect(w, r, `/view/`+path, http.StatusFound)
		}

	} else if r.Method == "GET" && r.URL.Path == `/` {
		this.handleHomepage(w)

	} else {
		http.Error(w, "Not found", 404)

	}
}