update file list & clipboard through websocket

dhamith93 · Oct 14, 2021 · 1760f67 · 1760f67
1 parent d5e97fb
commit 1760f67
Show file tree

Hide file tree

Showing 8 changed files with 165 additions and 26 deletions.
diff --git a/src/client.go b/src/client.go
@@ -0,0 +1,57 @@
+package main
+
+import (
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	writeWait  = 10 * time.Second
+	pongWait   = 60 * time.Second
+	pingPeriod = (pongWait * 9) / 10
+)
+
+type client struct {
+	hub  *hub
+	conn *websocket.Conn
+	send chan []byte
+}
+
+func (c *client) writePump() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		c.conn.Close()
+	}()
+	for {
+		select {
+		case message, ok := <-c.send:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			w, err := c.conn.NextWriter(websocket.TextMessage)
+			if err != nil {
+				return
+			}
+			w.Write(message)
+
+			n := len(c.send)
+			for i := 0; i < n; i++ {
+				w.Write(<-c.send)
+			}
+
+			if err := w.Close(); err != nil {
+				return
+			}
+		case <-ticker.C:
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}
diff --git a/src/go.mod b/src/go.mod
@@ -5,4 +5,5 @@ go 1.16
 require (
 	github.com/gorilla/handlers v1.5.1
 	github.com/gorilla/mux v1.8.0
+	github.com/gorilla/websocket v1.4.2
 )
diff --git a/src/go.sum b/src/go.sum
@@ -4,3 +4,5 @@ github.com/gorilla/handlers v1.5.1 h1:9lRY6j8DEeeBT10CvO9hGW0gmky0BprnvDI5vfhUHH
 github.com/gorilla/handlers v1.5.1/go.mod h1:t8XrUpc4KVXb7HGyJ4/cEnwQiaxrX/hz1Zv/4g96P1Q=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
diff --git a/src/hub.go b/src/hub.go
@@ -0,0 +1,31 @@
+package main
+
+type hub struct {
+	clients    map[*client]bool
+	broadcast  chan []byte
+	register   chan *client
+	unregister chan *client
+}
+
+func newHub() *hub {
+	return &hub{
+		broadcast:  make(chan []byte),
+		register:   make(chan *client),
+		unregister: make(chan *client),
+		clients:    make(map[*client]bool),
+	}
+}
+
+func (h *hub) run() {
+	for {
+		select {
+		case client := <-h.register:
+			h.clients[client] = true
+		case client := <-h.unregister:
+			if _, ok := h.clients[client]; ok {
+				delete(h.clients, client)
+				close(client.send)
+			}
+		}
+	}
+}
diff --git a/src/main.go b/src/main.go
@@ -36,6 +36,8 @@ func start(handler handler) {
 	port := readFile("port.txt", ":5500")
 	handler.fileList = getFileList()
 	handler.clipboardItems = []clipboardItem{}
+	handler.hub = newHub()
+	go handler.hub.run()
 	handler.handleRequests(port)
 }
 

diff --git a/src/public/index.html b/src/public/index.html
@@ -30,7 +30,6 @@ <h4>Shortcut - connected to <span id="host-ip"></span></h4>
                 <p><strong>Device Name:</strong> <span id="self-device-name"></span></p>
                 <input id="file-input" type="file" name="name" style="display: none;" />
                 <a id="file-picker-btn" class="waves-effect waves-light btn"><i class="material-icons left">attach_file</i>Select files to upload</a>
-                <a id="refresh-btn" class="waves-effect waves-light btn"><i class="material-icons left">refresh</i>Reload</a>
             </div>
         </div>
         <div id="main-area">

diff --git a/src/public/scripts/main.js b/src/public/scripts/main.js
@@ -2,7 +2,6 @@ document.addEventListener('DOMContentLoaded', () => {
     const mainArea = document.getElementById('main-area');
     const hostIpSpan = document.getElementById('host-ip');
     const filePickerBtn = document.getElementById('file-picker-btn');
-    const refreshBtn = document.getElementById('refresh-btn');
     const clipboardCollection = document.getElementById('clipboard-collection');
     const clipboardSendBtn = document.getElementById('clipboard-add-btn');
     const clipboardContent = document.getElementById('clipboard-content');
@@ -12,6 +11,7 @@ document.addEventListener('DOMContentLoaded', () => {
     const modals = document.querySelectorAll('.modal');
     M.Modal.init(modals, null);
     const clipboard = new ClipboardJS('.copy');
+    let socket;
 
     if (deviceName && deviceName.length > 0) {
         document.getElementById('self-device-name').innerHTML = deviceName;
@@ -23,31 +23,21 @@ document.addEventListener('DOMContentLoaded', () => {
         const input = document.createElement('input');
         input.type = 'file';
         input.setAttribute('multiple', 'true');
-
         input.onchange = e => { 
             const files = e.target.files; 
             processFiles([...files]);
         }
-
         input.click();
     });
 
-    refreshBtn.addEventListener('click', e => {
-        getFileList();
-        getClipboardItems();
-    });
-
     clipboardSendBtn.addEventListener('click', e => {
         const content = clipboardContent.value;
         const data = {
             'DeviceName': deviceName,
             'Content': content
         };
 
-        axios.post('clipboard', data).then((response) => {
-            if (response.data) {
-                processClipboardItems(response.data);
-            }
+        axios.post('clipboard', data).then((r) => {
             clipboardContent.value = '';
         }, (error) => {
             console.error(error);
@@ -102,11 +92,7 @@ document.addEventListener('DOMContentLoaded', () => {
                     document.getElementById('progress-completed').style.width =  percentCompleted + '%';
                 }
             };
-            axios.post('upload', formData, config).then((response) => {
-                if (response.data) {
-                    populateTable(fileTable, response.data);
-                }
-            }, (error) => {
+            axios.post('upload', formData, config).then((r) => {}, (error) => {
                 console.error(error);
                 alert('Error uploading the file... ' + error);
             }).finally(() => {
@@ -120,19 +106,15 @@ document.addEventListener('DOMContentLoaded', () => {
             if (response.data) {
                 populateTable(fileTable, response.data);
             }
-        }, (error) => {
-            console.error(error);
-        });
+        }, (error) => console.error(error));
     }
 
     function getClipboardItems() {
         axios.get('clipboard').then((response) => {
             if (response.data) {
                 processClipboardItems(response.data);
             }
-        }, (error) => {
-            console.error(error);
-        });
+        }, (error) => console.error(error));
     }
 
     function processClipboardItems(data) {
@@ -149,13 +131,33 @@ document.addEventListener('DOMContentLoaded', () => {
         axios.get('meta').then((response) => {
             if (response.data) {
                 hostIpSpan.innerHTML = response.data.Url;
+                socket = new WebSocket('ws://' + response.data.Url.replace('http://', '') + '/ws');
+                socket.onmessage = (e) => {
+                    try {
+                        const data = JSON.parse(e.data);
+                        data.forEach(item => {
+                            switch (item.MsgType) {
+                                case 'clipboardItems':
+                                    processClipboardItems(item.Msg);
+                                    break;
+                                case 'fileList':
+                                    populateTable(fileTable, item.Msg);
+                                    break;
+                                default:
+                                    break;
+                            }
+                        });
+                    } catch (e) {
+                        // fall back to API if parsing JSON failed
+                        getClipboardItems();
+                        getFileList();
+                    }
+                }
             }
         }, (error) => {
             console.error(error);
         });
     }
 
     getMeta();
-    getFileList();
-    getClipboardItems();  
 });
diff --git a/src/server.go b/src/server.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/gorilla/handlers"
 	"github.com/gorilla/mux"
+	"github.com/gorilla/websocket"
 )
 
 type meta struct {
@@ -26,15 +27,26 @@ type clipboardItem struct {
 	Content    string
 }
 
+type message struct {
+	MsgType string
+	Msg     interface{}
+}
+
 type handler struct {
 	server         http.Server
+	upgrader       websocket.Upgrader
+	conn           *websocket.Conn
+	hub            *hub
 	fileList       []FileList
 	clipboardItems []clipboardItem
 }
 
 func (h *handler) handleRequests(port string) {
 	router := mux.NewRouter().StrictSlash(true)
 	router.HandleFunc("/meta", h.sendMeta)
+	router.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
+		h.handleWS(h.hub, w, r)
+	})
 	router.Path("/upload").Methods("POST").HandlerFunc(h.handleFile)
 	router.Path("/clipboard").Methods("POST").HandlerFunc(h.handleClipboardItem)
 	router.Path("/clipboard").Methods("GET").HandlerFunc(h.getClipboardItems)
@@ -94,6 +106,8 @@ func (h *handler) handleFile(w http.ResponseWriter, r *http.Request) {
 		json.NewEncoder(w).Encode(err.Error())
 		return
 	}
+
+	go h.sendUpdate(message{"fileList", h.fileList})
 	h.getFiles(w, r)
 }
 
@@ -108,6 +122,8 @@ func (h *handler) handleClipboardItem(w http.ResponseWriter, r *http.Request) {
 	}
 	clipboardItem.Content = html.EscapeString(clipboardItem.Content)
 	h.clipboardItems = append(h.clipboardItems, clipboardItem)
+
+	go h.sendUpdate(message{"clipboardItems", h.clipboardItems})
 	json.NewEncoder(w).Encode(h.clipboardItems)
 }
 
@@ -117,6 +133,35 @@ func (h *handler) getClipboardItems(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(h.clipboardItems)
 }
 
+func (h *handler) handleWS(hub *hub, w http.ResponseWriter, r *http.Request) {
+	var err error
+	h.upgrader = websocket.Upgrader{}
+	h.conn, err = h.upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		http.Error(w, "Could not open websocket connection", http.StatusBadRequest)
+	}
+
+	client := &client{hub: hub, conn: h.conn, send: make(chan []byte, 256)}
+	client.hub.register <- client
+
+	go client.writePump()
+	data, _ := json.Marshal([]message{{"fileList", h.fileList}, {"clipboardItems", h.clipboardItems}})
+	client.send <- data
+}
+
+func (h *handler) sendUpdate(v ...interface{}) {
+	data, _ := json.Marshal(v)
+
+	for c := range h.hub.clients {
+		select {
+		case c.send <- data:
+		default:
+			delete(h.hub.clients, c)
+			close(c.send)
+		}
+	}
+}
+
 func getOutboundIP() string {
 	conn, err := net.Dial("udp", "8.8.8.8:80")
 	if err != nil {