main.go

Documentation: github.com/gorilla/websocket/examples/filewatch

     1  // Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"flag"
     9  	"html/template"
    10  	"io/ioutil"
    11  	"log"
    12  	"net/http"
    13  	"os"
    14  	"strconv"
    15  	"time"
    16  
    17  	"github.com/gorilla/websocket"
    18  )
    19  
    20  const (
    21  	// Time allowed to write the file to the client.
    22  	writeWait = 10 * time.Second
    23  
    24  	// Time allowed to read the next pong message from the client.
    25  	pongWait = 60 * time.Second
    26  
    27  	// Send pings to client with this period. Must be less than pongWait.
    28  	pingPeriod = (pongWait * 9) / 10
    29  
    30  	// Poll file for changes with this period.
    31  	filePeriod = 10 * time.Second
    32  )
    33  
    34  var (
    35  	addr      = flag.String("addr", ":8080", "http service address")
    36  	homeTempl = template.Must(template.New("").Parse(homeHTML))
    37  	filename  string
    38  	upgrader  = websocket.Upgrader{
    39  		ReadBufferSize:  1024,
    40  		WriteBufferSize: 1024,
    41  	}
    42  )
    43  
    44  func readFileIfModified(lastMod time.Time) ([]byte, time.Time, error) {
    45  	fi, err := os.Stat(filename)
    46  	if err != nil {
    47  		return nil, lastMod, err
    48  	}
    49  	if !fi.ModTime().After(lastMod) {
    50  		return nil, lastMod, nil
    51  	}
    52  	p, err := ioutil.ReadFile(filename)
    53  	if err != nil {
    54  		return nil, fi.ModTime(), err
    55  	}
    56  	return p, fi.ModTime(), nil
    57  }
    58  
    59  func reader(ws *websocket.Conn) {
    60  	defer ws.Close()
    61  	ws.SetReadLimit(512)
    62  	ws.SetReadDeadline(time.Now().Add(pongWait))
    63  	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
    64  	for {
    65  		_, _, err := ws.ReadMessage()
    66  		if err != nil {
    67  			break
    68  		}
    69  	}
    70  }
    71  
    72  func writer(ws *websocket.Conn, lastMod time.Time) {
    73  	lastError := ""
    74  	pingTicker := time.NewTicker(pingPeriod)
    75  	fileTicker := time.NewTicker(filePeriod)
    76  	defer func() {
    77  		pingTicker.Stop()
    78  		fileTicker.Stop()
    79  		ws.Close()
    80  	}()
    81  	for {
    82  		select {
    83  		case <-fileTicker.C:
    84  			var p []byte
    85  			var err error
    86  
    87  			p, lastMod, err = readFileIfModified(lastMod)
    88  
    89  			if err != nil {
    90  				if s := err.Error(); s != lastError {
    91  					lastError = s
    92  					p = []byte(lastError)
    93  				}
    94  			} else {
    95  				lastError = ""
    96  			}
    97  
    98  			if p != nil {
    99  				ws.SetWriteDeadline(time.Now().Add(writeWait))
   100  				if err := ws.WriteMessage(websocket.TextMessage, p); err != nil {
   101  					return
   102  				}
   103  			}
   104  		case <-pingTicker.C:
   105  			ws.SetWriteDeadline(time.Now().Add(writeWait))
   106  			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
   107  				return
   108  			}
   109  		}
   110  	}
   111  }
   112  
   113  func serveWs(w http.ResponseWriter, r *http.Request) {
   114  	ws, err := upgrader.Upgrade(w, r, nil)
   115  	if err != nil {
   116  		if _, ok := err.(websocket.HandshakeError); !ok {
   117  			log.Println(err)
   118  		}
   119  		return
   120  	}
   121  
   122  	var lastMod time.Time
   123  	if n, err := strconv.ParseInt(r.FormValue("lastMod"), 16, 64); err == nil {
   124  		lastMod = time.Unix(0, n)
   125  	}
   126  
   127  	go writer(ws, lastMod)
   128  	reader(ws)
   129  }
   130  
   131  func serveHome(w http.ResponseWriter, r *http.Request) {
   132  	if r.URL.Path != "/" {
   133  		http.Error(w, "Not found", http.StatusNotFound)
   134  		return
   135  	}
   136  	if r.Method != http.MethodGet {
   137  		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
   138  		return
   139  	}
   140  	w.Header().Set("Content-Type", "text/html; charset=utf-8")
   141  	p, lastMod, err := readFileIfModified(time.Time{})
   142  	if err != nil {
   143  		p = []byte(err.Error())
   144  		lastMod = time.Unix(0, 0)
   145  	}
   146  	var v = struct {
   147  		Host    string
   148  		Data    string
   149  		LastMod string
   150  	}{
   151  		r.Host,
   152  		string(p),
   153  		strconv.FormatInt(lastMod.UnixNano(), 16),
   154  	}
   155  	homeTempl.Execute(w, &v)
   156  }
   157  
   158  func main() {
   159  	flag.Parse()
   160  	if flag.NArg() != 1 {
   161  		log.Fatal("filename not specified")
   162  	}
   163  	filename = flag.Args()[0]
   164  	http.HandleFunc("/", serveHome)
   165  	http.HandleFunc("/ws", serveWs)
   166  	if err := http.ListenAndServe(*addr, nil); err != nil {
   167  		log.Fatal(err)
   168  	}
   169  }
   170  
   171  const homeHTML = `<!DOCTYPE html>
   172  <html lang="en">
   173      <head>
   174          <title>WebSocket Example</title>
   175      </head>
   176      <body>
   177          <pre id="fileData">{{.Data}}</pre>
   178          <script type="text/javascript">
   179              (function() {
   180                  var data = document.getElementById("fileData");
   181                  var conn = new WebSocket("ws://{{.Host}}/ws?lastMod={{.LastMod}}");
   182                  conn.onclose = function(evt) {
   183                      data.textContent = 'Connection closed';
   184                  }
   185                  conn.onmessage = function(evt) {
   186                      console.log('file updated');
   187                      data.textContent = evt.data;
   188                  }
   189              })();
   190          </script>
   191      </body>
   192  </html>
   193  `
   194
View as plain text