caddy-hugo2/caddyhugo.go

package caddyhugo

import (
	"encoding/base64"
	"errors"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"git.stephensearles.com/stephen/acedoc"
	"git.stephensearles.com/stephen/caddy-hugo2/assets"

	"github.com/gorilla/websocket"
	"github.com/mholt/caddy"
	"github.com/mholt/caddy/caddyhttp/httpserver"
	"github.com/spf13/afero"
	"github.com/spf13/hugo/deps"
	"github.com/spf13/hugo/hugofs"
	"github.com/spf13/hugo/hugolib"
)

const (
	IdleWebsocketTimeout = 10 * time.Minute
	WebsocketFileTicker  = 1 * time.Second
)

func init() {
	plugin := CaddyHugo{}

	// register a "generic" plugin, like a directive or middleware
	caddy.RegisterPlugin("hugo", caddy.Plugin{
		ServerType: "http",
		Action:     plugin.Setup,
	})

	// ... there are others. See the godoc.
}

type docref struct {
	clients uint
	name    string
	doc     *acedoc.Document
	tmpdir  string
}

type CaddyHugo struct {
	ServerType string
	Site       *httpserver.SiteConfig
	HugoSites  *hugolib.HugoSites
	HugoCfg    *deps.DepsCfg

	Media *MediaSource

	docs map[string]*docref
	mtx  sync.Mutex

	authorTmpl, adminTmpl, editTmpl *template.Template

	ltime uint64
}

func (ch *CaddyHugo) ObserveLTime(ltime uint64) uint64 {
	ch.mtx.Lock()

	if ch.ltime < ltime {
		ch.ltime = ltime
	}

	ch.mtx.Unlock()
	return ch.LTime()
}

func (ch *CaddyHugo) LTime() uint64 {
	ch.mtx.Lock()
	defer ch.mtx.Unlock()
	ch.ltime++
	return ch.ltime
}

func (ch *CaddyHugo) Build() error {
	err := ch.HugoSites.Build(hugolib.BuildCfg{ResetState: true})
	if err != nil {
		return fmt.Errorf("error building hugo sites: %v", err)
	}

	return nil
}

func (ch *CaddyHugo) Setup(c *caddy.Controller) error {
	var err error

	ch.docs = make(map[string]*docref)
	ch.Site = httpserver.GetConfig(c)

	ch.HugoCfg = &deps.DepsCfg{}

	ch.HugoCfg.Cfg, err = hugolib.LoadConfig(hugofs.Os, ch.Site.Root, "")
	if err != nil {
		return fmt.Errorf("error loading hugo config: %v", err)
	}

	ch.HugoCfg.Cfg.Set("workingdir", ch.Site.Root)
	ch.HugoSites, err = hugolib.NewHugoSites(*ch.HugoCfg)
	if err != nil {
		return fmt.Errorf("error intializing hugo: %v", err)
	}

	err = ch.Build()
	if err != nil {
		return fmt.Errorf("error building initial hugo: %v", err)
	}

	ch.authorTmpl, err = template.New("").Parse(AuthorPage)
	if err != nil {
		return fmt.Errorf("author template invalid: %v", err)
	}

	ch.adminTmpl, err = template.New("").Parse(AdminPage)
	if err != nil {
		return fmt.Errorf("admin template invalid: %v", err)
	}

	ch.editTmpl, err = template.New("").Parse(EditPage)
	if err != nil {
		return fmt.Errorf("edit template invalid: %v", err)
	}

	thumbDir, err := ioutil.TempDir("", "thumbs")
	if err != nil {
		return fmt.Errorf("couldn't initialize media: %v", err)
	}

	ch.Media = &MediaSource{
		StorageDir: path.Join(ch.Site.Root, "media"),
		ThumbDir:   thumbDir,
	}

	err = os.MkdirAll(ch.Media.StorageDir, 0755)
	if err != nil {
		return fmt.Errorf("couldn't initialize media: %v", err)
	}

	// add a function that wraps listeners for the HTTP server
	// (it's more common for a directive to call this rather than a standalone plugin)
	ch.Site.AddMiddleware(ch.Middleware(c))

	return ch.Publish()
}

func (ch *CaddyHugo) ServeHTTPWithNext(next httpserver.Handler, c *caddy.Controller, w http.ResponseWriter, r *http.Request) (int, error) {
	if !ch.Match(r) {
		p := path.Join(ch.Site.Root, "public", r.URL.Path)
		http.ServeFile(w, r, p)
		return 200, nil
	}

	if !ch.Auth(r) {
		return http.StatusUnauthorized, errors.New("not authorized")
	}

	if strings.HasPrefix(r.URL.Path, "/hugo/publish") {
		err := ch.Publish()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return http.StatusInternalServerError, nil
		}

		http.Redirect(w, r, "/", http.StatusFound)
		return http.StatusFound, nil
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/simplemde.css") {
		w.Write(assets.MustAsset("simplemde/dist/simplemde.min.css"))
		return http.StatusOK, nil
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/simplemde.js") {
		w.Write(assets.MustAsset("simplemde/debug/simplemde.js"))
		return http.StatusOK, nil
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/admin") {
		return ch.Admin().ServeHTTP(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/author") {
		return ch.AuthorHome().ServeHTTP(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/edit/") {
		return ch.Edit(c).ServeHTTP(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/draft/") {
		return ch.serveDraft(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/media") {
		return ch.serveMediaPage(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/hugo/upload") {
		return ch.uploadMedia(w, r)
	}
	if strings.HasPrefix(r.URL.Path, "/media/") {
		return ch.serveMedia(w, r)
	}

	http.NotFound(w, r)
	return 404, nil
}

func (ch CaddyHugo) Middleware(c *caddy.Controller) httpserver.Middleware {
	return func(next httpserver.Handler) httpserver.Handler {
		return httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
			return ch.ServeHTTPWithNext(next, c, w, r)
		})
	}
}

func (ch CaddyHugo) Auth(r *http.Request) bool {
	return true
}

func (ch CaddyHugo) Match(r *http.Request) bool {
	if strings.HasPrefix(r.URL.Path, "/media/") {
		return true
	}

	if r.URL.Path == "/hugo" {
		return true
	}

	return strings.HasPrefix(r.URL.Path, "/hugo/")
}

func (ch CaddyHugo) BasePath() string {
	return "/hugo"
}

func (ch CaddyHugo) Admin() httpserver.Handler {
	return httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
		err := ch.adminTmpl.Execute(w, ch.TmplData(r, nil))
		if err != nil {
			fmt.Println(err)
			return http.StatusInternalServerError, err
		}

		return http.StatusOK, nil

	})
}

func (ch CaddyHugo) AuthorHome() httpserver.Handler {
	return httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
		td := ch.TmplData(r, nil)
		err := ch.authorTmpl.Execute(w, td)
		if err != nil {
			fmt.Println(err)
			return http.StatusInternalServerError, err
		}

		return http.StatusOK, nil
	})
}

func (ch *CaddyHugo) Edit(c *caddy.Controller) httpserver.Handler {
	return httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
		if r.URL.Path == "/hugo/edit/new" {
			return ch.NewContent(w, r)
		}

		if r.Header.Get("Upgrade") == "websocket" {
			return ch.DeltaWebsocket(w, r)
		}

		doc, err := ch.doc(r)
		if err != nil {
			fmt.Println(err)
			return http.StatusNotFound, err
		}

		err = ch.editTmpl.Execute(w, ch.TmplData(r, doc))
		if err != nil {

			fmt.Println(err)
			return http.StatusInternalServerError, err
		}

		return http.StatusOK, nil
	})
}

func (ch *CaddyHugo) serveDraft(w http.ResponseWriter, r *http.Request) (int, error) {
	pathSegments := strings.SplitN(r.URL.Path, "/", 5)
	if len(pathSegments) < 4 {

		return http.StatusNotFound, nil
	}

	encoded := pathSegments[3]

	nameBytes, err := base64.RawURLEncoding.DecodeString(encoded)
	if err != nil {
		return http.StatusNotFound, err
	}

	name := string(nameBytes)

	ch.mtx.Lock()
	defer ch.mtx.Unlock()

	docref, ok := ch.docs[name]
	if !ok {
		return http.StatusNotFound, fmt.Errorf("draft not found")
	}

	http.StripPrefix(
		"/hugo/draft/"+encoded,
		http.FileServer(http.Dir(docref.tmpdir))).ServeHTTP(w, r)
	return 200, nil
}

func (ch *CaddyHugo) doc(r *http.Request) (*docref, error) {
	ch.mtx.Lock()
	defer ch.mtx.Unlock()

	name := r.URL.Path[len("/hugo/edit/"):]
	name = filepath.Join(ch.Site.Root, name)

	_, ok := ch.docs[name]
	if !ok {
		fmt.Println("opening", name)
		contents, err := ioutil.ReadFile(name)
		if err != nil {
			return nil, err
		}

		draftPrefix := fmt.Sprintf("draft-%s", base64.RawURLEncoding.EncodeToString([]byte(name)))
		tmpdir := path.Join(os.TempDir(), draftPrefix)

		ref := &docref{
			name:   name,
			doc:    acedoc.NewString(string(contents)),
			tmpdir: tmpdir,
		}

		err = ref.doc.LogToFile(path.Join(ch.Site.Root, "logs", r.URL.Path[len("/hugo/edit/"):]))
		if err != nil {
			fmt.Println(err)
			return nil, err
		}

		ch.docs[name] = ref

		hugoCmd := exec.Command("hugo", "--watch", "-D", "-d", ref.tmpdir)
		hugoCmd.Dir = ch.Site.Root
		err = hugoCmd.Start()
		if err != nil {
			return nil, fmt.Errorf("error starting hugo: %v", err)
		}

		go func() {
			ticker := time.NewTicker(WebsocketFileTicker)
			idleTicks := 0

			defer func() {
				err := hugoCmd.Process.Signal(os.Interrupt)
				if err != nil {
					fmt.Println("error signaling to hugo:", err)
				}
				err = hugoCmd.Wait()
				if err != nil {
					fmt.Println("error waiting for hugo:", err)
				}
			}()

			for {
				<-ticker.C
				ch.mtx.Lock()

				err := ioutil.WriteFile(name, []byte(ref.doc.Contents()), 0644)
				if err != nil {
					fmt.Println("error saving document contents:", err)
				}

				if ref.clients == 0 {
					idleTicks++
					idleTime := time.Duration(idleTicks) * WebsocketFileTicker
					if idleTime >= IdleWebsocketTimeout {
						err := ch.Publish()
						fmt.Printf("idle for %v, quitting\n", idleTime)
						if err != nil {
							fmt.Printf(", error publishing: %v\n", err)
						}

						ref.doc.Close()
						os.RemoveAll(tmpdir)
						delete(ch.docs, name)
						ch.mtx.Unlock()
						return
					}
				} else {
					idleTicks = 0
				}
				ch.mtx.Unlock()
			}
		}()
	}

	return ch.docs[name], nil
}

func (ch *CaddyHugo) Publish() error {
	cmd := exec.Command("hugo")
	cmd.Dir = ch.Site.Root
	_, err := cmd.CombinedOutput()
	if err != nil {
		return err
	}

	return nil
}

func (ch *CaddyHugo) DeltaWebsocket(w http.ResponseWriter, r *http.Request) (int, error) {
	var upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println(err)
		return http.StatusBadRequest, err
	}

	doc, err := ch.doc(r)
	if err != nil {
		fmt.Println(err)
		return http.StatusBadRequest, err
	}

	const idlePing = 15 * time.Second
	const idlePingShort = 1 * time.Millisecond
	var timer *time.Timer

	timer = time.AfterFunc(idlePing, func() {
		conn.WriteJSON(Message{
			Deltas: []acedoc.Delta{},
			LTime:  ch.LTime(),
		})
		timer.Reset(idlePing)
	})

	client := doc.doc.Client(acedoc.DeltaHandlerFunc(func(ds []acedoc.Delta) error {
		timer.Reset(idlePing)
		err := conn.WriteJSON(Message{
			Deltas: ds,
			LTime:  ch.LTime(),
		})
		return err
	}))

	ch.mtx.Lock()
	doc.clients++
	ch.mtx.Unlock()

	defer func() {
		ch.mtx.Lock()
		client.Close()
		doc.clients--
		ch.mtx.Unlock()
	}()

	for {
		var message Message
		err := conn.ReadJSON(&message)
		if err != nil {
			return http.StatusBadRequest, err
		}

		ch.ObserveLTime(message.LTime)
		timer.Reset(idlePingShort)

		err = client.PushDeltas(message.Deltas...)
		if err != nil {
			return http.StatusBadRequest, err
		}

	}
}

type Message struct {
	Deltas []acedoc.Delta `json:"deltas"`
	LTime  uint64         `json:"ltime"`
}

func (ch CaddyHugo) NewContent(w http.ResponseWriter, r *http.Request) (int, error) {
	name := r.FormValue("name")
	ctype := r.FormValue("type")

	if filepath.Ext(name) != ".md" {
		name += ".md"
	}

	filename := path.Join(ctype, name)
	if ctype == "default" {
		filename = name
	}

	cmd := exec.Command("hugo", "new", filename)
	cmd.Dir = ch.Site.Root
	out, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Println("error running hugo:\n", string(out))
		return http.StatusInternalServerError, err
	}

	// serve redirect
	http.Redirect(w, r, filepath.Join("/hugo/edit/", "content", filename), http.StatusFound)
	return http.StatusFound, nil
}

func (ch CaddyHugo) TmplData(r *http.Request, docref *docref) interface{} {
	var doc *acedoc.Document
	if docref != nil {
		doc = docref.doc
	}
	if ch.HugoSites != nil && ch.HugoSites.Fs != nil {
		ch.HugoSites.Fs.Destination = afero.NewMemMapFs()
		ch.Build()
	}
	return tmplData{ch.Site, r, ch, doc, docref}
}