package main import ( "embed" "encoding/json" "fmt" "io/fs" "io/ioutil" "log" "math" "math/rand" "net/http" "os" "runtime/debug" "strconv" "strings" "sync" "time" "github.com/go-git/go-git/v5" "github.com/go-git/go-git/v5/plumbing/object" "github.com/gorilla/mux" "github.com/gorilla/websocket" ) const ( socketBufferSize = 10 socketPingPeriod = 2 * time.Minute socketReadTimeout = 5 * time.Minute socketWriteTimeout = 20 * time.Second ) //go:embed web var webAssets embed.FS 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), true) } 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 != "fetch" && c.Command != "delete" && c.Note != nil { if c.Modified != c.Note.ModifiedAt { s.Write(socketResponse("fail", "modified")) continue } } fetchAction := c.Command == "fetch" checkAction := c.Command == "check" access := -1 if a, ok := c.Notebook.Serve[s.Author.Key]; ok { access = a } if access < authorAccessRead || (access == authorAccessRead && !fetchAction) || (access == authorAccessCheck && !checkAction && !fetchAction) { 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{"fetch": webFetch, "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 } var authorName string if authkey == "public" { authorName = "Public" } else { authorName = author.Name[0:int(math.Min(10, float64(len(author.Name))))] } log.Println() log.Printf("%-10s http://%s/#login/%s", authorName, 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() { webAssetsFS, err := fs.Sub(webAssets, "web") if err != nil { log.Fatal(err) } printServedNotebooks() r := mux.NewRouter() r.HandleFunc("/w/{author}/{modified}", webSocketHandler) r.PathPrefix("/").Handler(http.FileServer(http.FS(webAssetsFS))) 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["notebookid"] = notebookID response["noteid"] = noteID for _, notebook := range ss.Author.Notebooks { if notebook.ID != notebookID { continue } note := notebook.getNote(noteID, true) if note == nil { return } response["note"] = note } ss.Write(&response) } func sendNotesSince(ss *stickSocket, modified int64) { newModified := modified author := ss.Author var opts []string response := map[string]interface{}{} response["author"] = author.Name response["notebooks"] = map[string]*servedNotebook{} response["pinned"] = map[string][]string{} 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, false) if note == nil || note.ModifiedAt < modified { return nil } if note.ModifiedAt > newModified { newModified = note.ModifiedAt } snb.Notes[noteID] = note r, err := f.Reader() if err == nil { // Ignore err opts = optionsReader(r) for _, opt := range opts { if opt == "pin" { response["pinned"].(map[string][]string)[snb.ID] = append(response["pinned"].(map[string][]string)[snb.ID], noteID) break } } } return nil }) checkError(err) } response["notebooks"].(map[string]*servedNotebook)[snb.ID] = snb } if modified > 0 { response["delta"] = true } ss.Write(&response) } // Fetch a note func webFetch(s *stickSocket, c *stickSocketCommand) *map[string]interface{} { return &map[string]interface{}{"status": "success", "notebookid": c.Notebook.ID, "noteid": c.Note.ID, "note": c.Note} } // 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) } optSort := false optList := false opts := options(body) for _, opt := range opts { switch opt { case "sort": optSort = true case "list": optList = true } } if optSort { body = sortListItems(body, optList) } 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 sendNote(s, c.Notebook.ID, c.Note.ID) 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), true) 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{} { if stick.Config.Debug && status == "fail" { log.Println("Failure message: " + message + " - stack trace:") debug.PrintStack() } 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 _, ss := range webSockets { if ss.Author.Key == author.Key { sendNote(ss, notebookID, noteID) } } }