Added uploading

1 files changed, 126 insertions, 9 deletions

diff --git a/main.go b/main.go
index dcb7d47..74b38b5 100644
--- a/main.go
+++ b/main.go
@@ -1,31 +1,40 @@
 package main
 
 import (
-    "net/http"
-    "flag"
-    "fmt"
-    "strings"
     "crypto/md5"
+    "mime"
+    "io/ioutil"
+    "io"
     "encoding/json"
-    "os"
+    "sync"
+    "flag"
+    "fmt"
     "html/template"
     "log"
+    "strconv"
+    "net/http"
+    "os"
     "path"
+    "strings"
 
-    _"github.com/lib/pq"
     "github.com/jmoiron/sqlx"
+    _"github.com/lib/pq"
 )
 
 type Config struct {
     Listen   string `json:"listen"`
     DBStr    string `json:"db"`
     TmplPath string `json:"tmpl"`
+    DataPath string `json:"data"`
 }
 
 type Server struct {
     mux  *http.ServeMux
     conf *Config
     db   *sqlx.DB
+
+    fslock sync.RWMutex
+    filestore map[string][]string
 }
 
 type note struct {
@@ -43,6 +52,13 @@ const scheme = `
     )
 `
 
+var validMIME = map[string]bool {
+    "text/plain": true,
+    "image/jpeg": true,
+    "image/png": true,
+    "application/pdf": true,
+}
+
 func main() {
     confFlag := flag.String("c", "config.json", "config file")
     flag.Parse()
@@ -75,13 +91,27 @@ func NewServer(conf *Config) *Server {
     s := &Server{}
     s.mux = http.NewServeMux()
     s.mux.HandleFunc("/", s.httpRoot)
+    s.mux.HandleFunc("/upload", s.httpUpload)
     s.conf = conf
 
+    s.filestore = map[string][]string{}
+
     return s
 }
 
 func (s *Server) Connect() error {
-    var err error
+    files, err := ioutil.ReadDir(s.conf.DataPath)
+    if err != nil {
+        return err
+    }
+
+    for _, file := range files {
+        name := file.Name()
+
+        hash := name[0:4]
+        s.saveFile(hash, name)
+    }
+
     s.db, err = sqlx.Connect("postgres", s.conf.DBStr)
     if err != nil {
         return err
@@ -111,6 +141,21 @@ func (s *Server) renderTemplate(w http.ResponseWriter, data interface{}, pathnam
     return err
 }
 
+func (s *Server) httpLog(r *http.Request, format string, args... interface{}) {
+    args = append([]interface{}{r.URL.Path, r.RemoteAddr}, args...)
+    log.Printf("(url: %s, ip: %s) " + format, args...)
+}
+
+func (s *Server) Error(w http.ResponseWriter, r *http.Request, err string, code int) {
+    if code < 500 {
+        http.Error(w, err, code)
+    } else {
+        w.WriteHeader(code)
+    }
+
+    s.httpLog(r, "ERR: %s", err)
+}
+
 func (s *Server) httpCreateNode(w http.ResponseWriter, r *http.Request) {
     var (
         name = r.FormValue("name")
@@ -123,10 +168,12 @@ func (s *Server) httpCreateNode(w http.ResponseWriter, r *http.Request) {
         INSERT INTO notes (hash, name, location)
         VALUES ($1, $2, $3)`, hash, name, location)
     if err != nil {
-        http.Error(w, err.Error(), http.StatusInternalServerError)
+        s.Error(w, r, err.Error(), http.StatusInternalServerError)
         return
     }
 
+    s.httpLog(r, "Created node %s", hash)
+
     var ref strings.Builder
     fmt.Fprintf(&ref, "#%s", hash)
     if location != "" {
@@ -161,9 +208,79 @@ func (s *Server) httpRoot(w http.ResponseWriter, r *http.Request) {
     err := s.db.SelectContext(r.Context(), &page.Notes, `
         SELECT hash, name, location FROM notes WHERE available = True`)
     if err != nil {
-        http.Error(w, err.Error(), http.StatusInternalServerError)
+        s.Error(w, r, err.Error(), http.StatusInternalServerError)
         return
     }
 
     s.renderTemplate(w, &page, "root.template")
 }
+
+func (s *Server) httpUpload(w http.ResponseWriter, r *http.Request) {
+    if r.Method != "POST" {
+        w.WriteHeader(http.StatusMethodNotAllowed)
+        return
+    }
+
+    // Head hash
+    hash := r.FormValue("h")
+    _, err := strconv.ParseUint(hash, 16, 16)
+    if err != nil || len(hash) != 4 {
+        s.Error(w, r, "Invalid hash", http.StatusBadRequest)
+        return
+    }
+
+    // Load file
+    inf, header, err := r.FormFile("f")
+    if err != nil {
+        s.Error(w, r, err.Error(), http.StatusInternalServerError)
+        return
+    }
+    defer inf.Close()
+
+    // Check if valid type
+    mtype := header.Header.Get("Content-Type")
+    ok := validMIME[mtype]
+    if !ok {
+        s.Error(w, r, "Invalid file type", http.StatusBadRequest)
+        return
+    }
+
+    // Create file
+    fname := s.allocFile(hash, mtype)
+    f, err := os.Create(path.Join(s.conf.DataPath, fname))
+    if err != nil {
+        s.Error(w, r, err.Error(), http.StatusInternalServerError)
+        return
+    }
+    defer f.Close()
+
+    _, err = io.Copy(f, inf)
+    if err != nil {
+        s.Error(w, r, err.Error(), http.StatusInternalServerError)
+        return
+    }
+
+    s.saveFile(hash, fname)
+    s.httpLog(r, "Uploaded file %s", fname)
+}
+
+func (s *Server) allocFile(hash string, t string) string {
+    s.fslock.RLock()
+    existing := s.filestore[hash]
+    s.fslock.RUnlock()
+
+    var ext string
+    extarr, _ := mime.ExtensionsByType(t)
+    if len(extarr) > 0 {
+        ext = extarr[0]
+    }
+
+    return fmt.Sprintf("%s.%d%s", hash, len(existing), ext)
+}
+
+func (s *Server) saveFile(hash string, fname string) {
+    s.fslock.Lock()
+    s.filestore[hash] = append(s.filestore[hash], fname)
+
+    s.fslock.Unlock()
+}