vendor/github.com/gorilla/websocket/examples/filewatch/main.go

   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 != "GET" {
 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 `