stick/web.go

683 lines
14 KiB
Go

package main
import (
"encoding/json"
"fmt"
"io/ioutil"
"log"
"math"
"math/rand"
"net/http"
"os"
"strconv"
"strings"
"sync"
"time"
rice "github.com/GeertJohan/go.rice"
"github.com/gorilla/mux"
"github.com/gorilla/websocket"
"gopkg.in/src-d/go-git.v4"
"gopkg.in/src-d/go-git.v4/plumbing/object"
)
const (
SocketBufferSize = 10
SocketPingPeriod = 2 * time.Minute
SocketReadTimeout = 5 * time.Minute
SocketWriteTimeout = 20 * time.Second
)
type StickSocket struct {
ID int
Status int
Conn *websocket.Conn
Author *Author
writebuffer chan *map[string]interface{}
writebufferwg sync.WaitGroup
}
func (s *StickSocket) Write(m *map[string]interface{}) {
if s.Status == 0 {
return
}
s.writebufferwg.Add(1)
s.writebuffer <- m
}
func (s *StickSocket) handleRead() {
var (
message []byte
c *StickSocketCommand
err error
)
for {
_, message, err = s.Conn.ReadMessage()
if err != nil {
s.Close(fmt.Sprintf("failed to read from client: %+v", err))
return
}
if stick.Config.Debug {
log.Printf("WS read %d %s", s.ID, message)
}
c = &StickSocketCommand{}
err = json.Unmarshal(message, &c.Data)
CheckError(err)
if command, ok := c.Data["command"]; ok {
c.Command = command.(string)
} else {
s.Write(socketResponse("fail", "invalid-command"))
continue
}
if notebookID, ok := c.Data["notebook"]; ok {
for _, notebook := range s.Author.Notebooks {
if notebook.ID != notebookID {
continue
}
c.Notebook = notebook
if noteID, ok := c.Data["note"]; ok {
c.Note = notebook.getNote(noteID.(string))
}
break
}
if m, ok := c.Data["modified"]; ok {
c.Modified = int64(m.(float64))
}
}
if c.Notebook == nil {
s.Write(socketResponse("fail", "invalid-notebook"))
continue
} else if c.Command != "edit" {
if c.Note == nil {
s.Write(socketResponse("fail", "invalid-note"))
continue
}
}
if c.Command != "delete" && c.Note != nil {
if c.Modified != c.Note.ModifiedAt {
s.Write(socketResponse("fail", "modified"))
continue
}
}
access := -1
if a, ok := c.Notebook.Serve[s.Author.Key]; ok {
access = a
}
if access <= AuthorAccessRead || (access == AuthorAccessCheck && c.Command != "check") {
s.Write(socketResponse("fail", "unauthorized"))
continue
}
if f, ok := SocketCommands[c.Command]; ok {
SocketCommandLock.Lock()
s.Write(f(s, c))
SocketCommandLock.Unlock()
} else {
s.Write(socketResponse("fail", "invalid-command"))
}
}
}
func (s *StickSocket) handleWrite() {
var (
m *map[string]interface{}
out []byte
err error
)
for m = range s.writebuffer {
if s.Status == 0 {
s.writebufferwg.Done()
continue
}
out, err = json.Marshal(m)
CheckError(err)
err = s.Conn.SetWriteDeadline(time.Now().Add(SocketWriteTimeout))
if err != nil {
go s.Close("failed to set write timeout")
continue
}
err = s.Conn.WriteMessage(websocket.TextMessage, out)
if err != nil {
go s.Close("failed to write message")
}
if stick.Config.Debug {
log.Printf("WS write %d %s", s.ID, out)
}
s.writebufferwg.Done()
}
}
func (s *StickSocket) handlePing() {
err := s.Conn.SetReadDeadline(time.Now().Add(SocketReadTimeout))
if err != nil {
s.Close("failed to set read timeout")
return
}
s.Conn.SetPongHandler(func(string) error { return s.Conn.SetReadDeadline(time.Now().Add(SocketReadTimeout)) })
ticker := time.NewTicker(SocketPingPeriod)
defer ticker.Stop()
for {
<-ticker.C
if s.Status == 0 {
return
}
err := s.Conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(SocketWriteTimeout))
if err != nil {
s.Close("failed to ping client")
return
}
}
}
func (s *StickSocket) Close(reason string) {
WebSocketsLock.Lock()
defer WebSocketsLock.Unlock()
if s.Status == 0 {
return
}
s.Status = 0
_ = s.Conn.Close()
if stick.Config.Debug {
log.Printf("WS close %d: %s", s.ID, reason)
}
s.writebufferwg.Wait()
close(s.writebuffer)
if s.ID > 0 {
delete(WebSockets, s.ID)
}
}
type StickSocketCommand struct {
Command string
Modified int64
Notebook *Notebook
Note *Note
Data map[string]interface{} // Command JSON from client
}
type StickSocketHandler func(*StickSocket, *StickSocketCommand) *map[string]interface{}
var (
SocketCommands = map[string]StickSocketHandler{"edit": webEdit, "check": webCheck, "delete": webDelete}
SocketCommandLock sync.Mutex
WebSockets = make(map[int]*StickSocket)
WebSocketsLock sync.RWMutex
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool {
return true
},
EnableCompression: true,
}
func printServedNotebooks() {
stick.AuthorsLock.RLock()
defer stick.AuthorsLock.RUnlock()
for authkey, author := range stick.Authors {
if len(author.Notebooks) == 0 {
continue
}
log.Println()
log.Printf("%-10s http://%s/#login/%s", author.Name[0:int(math.Min(10, float64(len(author.Name))))], stick.Config.Serve, authkey)
for _, notebook := range author.Notebooks {
log.Printf("- %s", notebook.Label[0:int(math.Min(20, float64(len(notebook.Label))))])
}
}
}
func webSocketHandler(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
vars := mux.Vars(r)
var (
auth string
syncmodified int64
)
if a, ok := vars["author"]; ok {
auth = a
}
var author *Author
stick.AuthorsLock.RLock()
for stickauthkey, stickauth := range stick.Authors {
if stickauthkey == auth {
author = stickauth
break
}
}
stick.AuthorsLock.RUnlock()
s := &StickSocket{Status: 1, Conn: conn, Author: author, writebuffer: make(chan *map[string]interface{}, SocketBufferSize)}
go s.handleWrite()
go s.handlePing()
defer s.Close("client disconnected")
if author == nil {
s.Write(socketResponse("fail", "unauthorized"))
time.Sleep(2 * time.Second)
s.Close("unauthorized")
return
}
var (
socketID int
)
WebSocketsLock.Lock()
for {
socketID = rand.Intn(math.MaxInt32-1) + 1
if _, ok := WebSockets[socketID]; !ok {
s.ID = socketID
WebSockets[socketID] = s
break
}
}
WebSocketsLock.Unlock()
if m, ok := vars["modified"]; ok {
syncmodified, err = strconv.ParseInt(m, 10, 64)
if err != nil {
syncmodified = 0
}
}
if stick.Config.Debug {
log.Printf("WS new %d", s.ID)
}
sendNotesSince(s, syncmodified)
s.handleRead()
}
func serveWeb() {
printServedNotebooks()
r := mux.NewRouter()
r.HandleFunc("/w/{author}/{modified}", webSocketHandler)
r.PathPrefix("/").Handler(http.FileServer(rice.MustFindBox("web").HTTPBox()))
if err := http.ListenAndServe(stick.Config.Serve, r); err != nil {
log.Fatal("Web server error: ", err)
}
}
func sendNote(ss *StickSocket, notebookID string, noteID string) {
response := map[string]interface{}{}
response["author"] = ss.Author.Name
response["notebooks"] = map[string]*ServedNotebook{}
for _, notebook := range ss.Author.Notebooks {
if notebook.ID != notebookID {
continue
}
notes := make(map[string]*Note)
if noteID != "" {
note := notebook.getNote(noteID)
if note == nil {
return
}
notes[note.ID] = note
} else {
notes = notebook.allNotes()
}
snb := &ServedNotebook{Notebook: notebook, Notes: notes}
response["notebooks"].(map[string]*ServedNotebook)[snb.ID] = snb
break
}
response["delta"] = true
ss.Write(&response)
}
func sendNotesSince(ss *StickSocket, modified int64) {
newModified := modified
author := ss.Author
response := map[string]interface{}{}
response["author"] = author.Name
response["notebooks"] = map[string]*ServedNotebook{}
for _, notebook := range author.Notebooks {
snb := &ServedNotebook{Notebook: notebook, Notes: make(map[string]*Note)}
ref, err := notebook.Repository.Head()
if err == nil { // Non-empty repository
commit, err := notebook.Repository.CommitObject(ref.Hash())
CheckError(err)
tree, err := commit.Tree()
CheckError(err)
err = tree.Files().ForEach(func(f *object.File) error {
noteID := hash(f.Name)
note := notebook.getNote(noteID)
if note == nil || note.ModifiedAt < modified {
return nil
}
if note.ModifiedAt > newModified {
newModified = note.ModifiedAt
}
snb.Notes[noteID] = note
return nil
})
CheckError(err)
}
response["notebooks"].(map[string]*ServedNotebook)[snb.ID] = snb
}
if modified > 0 {
response["delta"] = true
}
ss.Write(&response)
}
// Edit a note
func webEdit(s *StickSocket, c *StickSocketCommand) *map[string]interface{} {
if stick.Config.Debug {
defer elapsed("webEdit")()
}
var (
title string
body string
err error
)
if b, ok := c.Data["body"]; ok {
body = b.(string)
}
if t, ok := c.Data["title"]; ok {
title = t.(string)
}
wt, err := c.Notebook.Repository.Worktree()
CheckError(err)
fileName := ""
addNew := false
if c.Note == nil {
addNew = true
fileName = strings.Replace(strings.ToLower(title), " ", "-", -1)
if !strings.HasSuffix(fileName, ".md") {
fileName += ".md"
}
_, err = wt.Filesystem.Stat(fileName)
if addNew && os.IsExist(err) {
return socketResponse("fail", "note-already-exists")
} else if !os.IsExist(err) && !os.IsNotExist(err) {
CheckError(err)
}
} else {
file := c.Notebook.File(c.Note.ID)
if file == nil {
return socketResponse("fail", "failed-to-read-note")
}
fileName = file.Name
data, err := file.Contents()
CheckError(err)
if data == body {
// TODO: Send unmodified message
sendNotesSince(s, c.Note.ModifiedAt)
return socketResponse("success", "")
}
}
if fileName == "" {
return socketResponse("fail", "invalid-note-name")
}
// Write
status, err := wt.Status()
CheckError(err)
if len(status) > 0 {
return socketResponse("fail", "modified-externally")
}
err = ioutil.WriteFile(wt.Filesystem.Join(wt.Filesystem.Root(), fileName), []byte(body), 0644)
CheckError(err)
// Commit
_, err = wt.Add(fileName)
CheckError(err)
commitMessage := "add"
if !addNew {
commitMessage = "update"
}
_, err = wt.Commit(commitMessage+" "+fileName, &git.CommitOptions{
Author: &object.Signature{
Name: s.Author.Name,
Email: s.Author.Email,
When: time.Now(),
},
})
CheckError(err)
if c.Note == nil {
c.Note = c.Notebook.getNote(hash(fileName))
if c.Note == nil {
return socketResponse("fail", "modified")
}
}
updateWebSockets(s.Author, c.Notebook.ID, c.Note.ID)
return socketResponse("success", "")
}
// Mark/un-mark a checkbox via its item number (index)
func webCheck(s *StickSocket, c *StickSocketCommand) *map[string]interface{} {
if stick.Config.Debug {
defer elapsed("webCheck")()
}
var (
checkItem int64
)
if ci, ok := c.Data["item"]; ok {
checkItem = int64(ci.(float64))
}
file := c.Notebook.File(c.Note.ID)
if file == nil {
return socketResponse("fail", "failed-to-read-note")
}
data, err := file.Contents()
CheckError(err)
// Update
var currentCheckbox int64
checkSearch := 0
for {
uncheckedCheckboxLocation := strings.Index(data[checkSearch:], "[ ]")
checkedCheckboxLocation := strings.Index(strings.ToLower(data[checkSearch:]), "[x]")
if uncheckedCheckboxLocation == -1 && checkedCheckboxLocation == -1 {
return socketResponse("fail", "invalid-check-item")
}
var checkboxLocation int
var replace string
if uncheckedCheckboxLocation > -1 && (checkedCheckboxLocation == -1 || uncheckedCheckboxLocation < checkedCheckboxLocation) {
checkboxLocation = uncheckedCheckboxLocation
replace = "[x]"
} else {
checkboxLocation = checkedCheckboxLocation
replace = "[ ]"
}
if currentCheckbox == checkItem {
data = data[0:checkSearch+checkboxLocation] + replace + data[checkSearch+checkboxLocation+3:]
break
}
checkSearch += checkboxLocation + 3
currentCheckbox++
}
wt, err := c.Notebook.Repository.Worktree()
CheckError(err)
status, err := wt.Status()
CheckError(err)
if len(status) > 0 {
return socketResponse("fail", "modified-externally")
}
// Write
err = ioutil.WriteFile(wt.Filesystem.Join(wt.Filesystem.Root(), file.Name), []byte(data), 0644)
CheckError(err)
// Commit
_, err = wt.Add(file.Name)
CheckError(err)
_, err = wt.Commit("update "+file.Name, &git.CommitOptions{
Author: &object.Signature{
Name: s.Author.Name,
Email: s.Author.Email,
When: time.Now(),
},
})
CheckError(err)
updateWebSockets(s.Author, c.Notebook.ID, c.Note.ID)
return socketResponse("success", "")
}
// Delete a note
func webDelete(s *StickSocket, c *StickSocketCommand) *map[string]interface{} {
if stick.Config.Debug {
defer elapsed("webDelete")()
}
var (
err error
)
file := c.Notebook.File(c.Note.ID)
if file == nil || file.Name == "" {
return socketResponse("fail", "failed-to-read-note")
}
// Delete
wt, err := c.Notebook.Repository.Worktree()
CheckError(err)
status, err := wt.Status()
CheckError(err)
if len(status) > 0 {
return socketResponse("fail", "modified-externally")
}
err = wt.Filesystem.Remove(file.Name)
CheckError(err)
// Commit
_, err = wt.Add(file.Name)
CheckError(err)
_, err = wt.Commit("delete "+file.Name, &git.CommitOptions{
Author: &object.Signature{
Name: s.Author.Name,
Email: s.Author.Email,
When: time.Now(),
},
})
CheckError(err)
updateWebSockets(s.Author, c.Notebook.ID, "")
return socketResponse("success", "")
}
func socketResponse(status string, message string) *map[string]interface{} {
return &map[string]interface{}{"status": status, "message": message}
}
func updateWebSockets(author *Author, notebookID string, noteID string) {
stick.NotebooksLock.RLock()
defer stick.NotebooksLock.RUnlock()
WebSocketsLock.RLock()
defer WebSocketsLock.RUnlock()
if noteID == "" {
for _, ss := range WebSockets {
if ss.Author.Key == author.Key {
sendNotesSince(ss, 0)
}
}
return
}
for _, notebook := range author.Notebooks {
if notebook.ID != notebookID {
continue
}
note := notebook.getNote(noteID)
if note == nil {
continue
}
for _, ss := range WebSockets {
if ss.Author.Key == author.Key {
sendNote(ss, notebookID, noteID)
}
}
return
}
}