vendor/github.com/gorilla/websocket/examples/autobahn/server.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 // Command server is a test server for the Autobahn WebSockets Test Suite.
   6 package main
   7
   8 import (
   9         "errors"
  10         "flag"
  11         "io"
  12         "log"
  13         "net/http"
  14         "time"
  15         "unicode/utf8"
  16
  17         "github.com/gorilla/websocket"
  18 )
  19
  20 var upgrader = websocket.Upgrader{
  21         ReadBufferSize:    4096,
  22         WriteBufferSize:   4096,
  23         EnableCompression: true,
  24         CheckOrigin: func(r *http.Request) bool {
  25                 return true
  26         },
  27 }
  28
  29 // echoCopy echoes messages from the client using io.Copy.
  30 func echoCopy(w http.ResponseWriter, r *http.Request, writerOnly bool) {
  31         conn, err := upgrader.Upgrade(w, r, nil)
  32         if err != nil {
  33                 log.Println("Upgrade:", err)
  34                 return
  35         }
  36         defer conn.Close()
  37         for {
  38                 mt, r, err := conn.NextReader()
  39                 if err != nil {
  40                         if err != io.EOF {
  41                                 log.Println("NextReader:", err)
  42                         }
  43                         return
  44                 }
  45                 if mt == websocket.TextMessage {
  46                         r = &validator{r: r}
  47                 }
  48                 w, err := conn.NextWriter(mt)
  49                 if err != nil {
  50                         log.Println("NextWriter:", err)
  51                         return
  52                 }
  53                 if mt == websocket.TextMessage {
  54                         r = &validator{r: r}
  55                 }
  56                 if writerOnly {
  57                         _, err = io.Copy(struct{ io.Writer }{w}, r)
  58                 } else {
  59                         _, err = io.Copy(w, r)
  60                 }
  61                 if err != nil {
  62                         if err == errInvalidUTF8 {
  63                                 conn.WriteControl(websocket.CloseMessage,
  64                                         websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""),
  65                                         time.Time{})
  66                         }
  67                         log.Println("Copy:", err)
  68                         return
  69                 }
  70                 err = w.Close()
  71                 if err != nil {
  72                         log.Println("Close:", err)
  73                         return
  74                 }
  75         }
  76 }
  77
  78 func echoCopyWriterOnly(w http.ResponseWriter, r *http.Request) {
  79         echoCopy(w, r, true)
  80 }
  81
  82 func echoCopyFull(w http.ResponseWriter, r *http.Request) {
  83         echoCopy(w, r, false)
  84 }
  85
  86 // echoReadAll echoes messages from the client by reading the entire message
  87 // with ioutil.ReadAll.
  88 func echoReadAll(w http.ResponseWriter, r *http.Request, writeMessage, writePrepared bool) {
  89         conn, err := upgrader.Upgrade(w, r, nil)
  90         if err != nil {
  91                 log.Println("Upgrade:", err)
  92                 return
  93         }
  94         defer conn.Close()
  95         for {
  96                 mt, b, err := conn.ReadMessage()
  97                 if err != nil {
  98                         if err != io.EOF {
  99                                 log.Println("NextReader:", err)
 100                         }
 101                         return
 102                 }
 103                 if mt == websocket.TextMessage {
 104                         if !utf8.Valid(b) {
 105                                 conn.WriteControl(websocket.CloseMessage,
 106                                         websocket.FormatCloseMessage(websocket.CloseInvalidFramePayloadData, ""),
 107                                         time.Time{})
 108                                 log.Println("ReadAll: invalid utf8")
 109                         }
 110                 }
 111                 if writeMessage {
 112                         if !writePrepared {
 113                                 err = conn.WriteMessage(mt, b)
 114                                 if err != nil {
 115                                         log.Println("WriteMessage:", err)
 116                                 }
 117                         } else {
 118                                 pm, err := websocket.NewPreparedMessage(mt, b)
 119                                 if err != nil {
 120                                         log.Println("NewPreparedMessage:", err)
 121                                         return
 122                                 }
 123                                 err = conn.WritePreparedMessage(pm)
 124                                 if err != nil {
 125                                         log.Println("WritePreparedMessage:", err)
 126                                 }
 127                         }
 128                 } else {
 129                         w, err := conn.NextWriter(mt)
 130                         if err != nil {
 131                                 log.Println("NextWriter:", err)
 132                                 return
 133                         }
 134                         if _, err := w.Write(b); err != nil {
 135                                 log.Println("Writer:", err)
 136                                 return
 137                         }
 138                         if err := w.Close(); err != nil {
 139                                 log.Println("Close:", err)
 140                                 return
 141                         }
 142                 }
 143         }
 144 }
 145
 146 func echoReadAllWriter(w http.ResponseWriter, r *http.Request) {
 147         echoReadAll(w, r, false, false)
 148 }
 149
 150 func echoReadAllWriteMessage(w http.ResponseWriter, r *http.Request) {
 151         echoReadAll(w, r, true, false)
 152 }
 153
 154 func echoReadAllWritePreparedMessage(w http.ResponseWriter, r *http.Request) {
 155         echoReadAll(w, r, true, true)
 156 }
 157
 158 func serveHome(w http.ResponseWriter, r *http.Request) {
 159         if r.URL.Path != "/" {
 160                 http.Error(w, "Not found.", http.StatusNotFound)
 161                 return
 162         }
 163         if r.Method != "GET" {
 164                 http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
 165                 return
 166         }
 167         w.Header().Set("Content-Type", "text/html; charset=utf-8")
 168         io.WriteString(w, "<html><body>Echo Server</body></html>")
 169 }
 170
 171 var addr = flag.String("addr", ":9000", "http service address")
 172
 173 func main() {
 174         flag.Parse()
 175         http.HandleFunc("/", serveHome)
 176         http.HandleFunc("/c", echoCopyWriterOnly)
 177         http.HandleFunc("/f", echoCopyFull)
 178         http.HandleFunc("/r", echoReadAllWriter)
 179         http.HandleFunc("/m", echoReadAllWriteMessage)
 180         http.HandleFunc("/p", echoReadAllWritePreparedMessage)
 181         err := http.ListenAndServe(*addr, nil)
 182         if err != nil {
 183                 log.Fatal("ListenAndServe: ", err)
 184         }
 185 }
 186
 187 type validator struct {
 188         state int
 189         x     rune
 190         r     io.Reader
 191 }
 192
 193 var errInvalidUTF8 = errors.New("invalid utf8")
 194
 195 func (r *validator) Read(p []byte) (int, error) {
 196         n, err := r.r.Read(p)
 197         state := r.state
 198         x := r.x
 199         for _, b := range p[:n] {
 200                 state, x = decode(state, x, b)
 201                 if state == utf8Reject {
 202                         break
 203                 }
 204         }
 205         r.state = state
 206         r.x = x
 207         if state == utf8Reject || (err == io.EOF && state != utf8Accept) {
 208                 return n, errInvalidUTF8
 209         }
 210         return n, err
 211 }
 212
 213 // UTF-8 decoder from http://bjoern.hoehrmann.de/utf-8/decoder/dfa/
 214 //
 215 // Copyright (c) 2008-2009 Bjoern Hoehrmann <bjoern@hoehrmann.de>
 216 //
 217 // Permission is hereby granted, free of charge, to any person obtaining a copy
 218 // of this software and associated documentation files (the "Software"), to
 219 // deal in the Software without restriction, including without limitation the
 220 // rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 221 // sell copies of the Software, and to permit persons to whom the Software is
 222 // furnished to do so, subject to the following conditions:
 223 //
 224 // The above copyright notice and this permission notice shall be included in
 225 // all copies or substantial portions of the Software.
 226 //
 227 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 228 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 229 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 230 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 231 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 232 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 233 // IN THE SOFTWARE.
 234 var utf8d = [...]byte{
 235         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 00..1f
 236         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 20..3f
 237         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 40..5f
 238         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 60..7f
 239         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, // 80..9f
 240         7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, // a0..bf
 241         8, 8, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, // c0..df
 242         0xa, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x3, 0x4, 0x3, 0x3, // e0..ef
 243         0xb, 0x6, 0x6, 0x6, 0x5, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, 0x8, // f0..ff
 244         0x0, 0x1, 0x2, 0x3, 0x5, 0x8, 0x7, 0x1, 0x1, 0x1, 0x4, 0x6, 0x1, 0x1, 0x1, 0x1, // s0..s0
 245         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, // s1..s2
 246         1, 2, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, // s3..s4
 247         1, 2, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 3, 1, 3, 1, 1, 1, 1, 1, 1, // s5..s6
 248         1, 3, 1, 1, 1, 1, 1, 3, 1, 3, 1, 1, 1, 1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, // s7..s8
 249 }
 250
 251 const (
 252         utf8Accept = 0
 253         utf8Reject = 1
 254 )
 255
 256 func decode(state int, x rune, b byte) (int, rune) {
 257         t := utf8d[b]
 258         if state != utf8Accept {
 259                 x = rune(b&0x3f) | (x << 6)
 260         } else {
 261                 x = rune((0xff >> t) & b)
 262         }
 263         state = int(utf8d[256+state*16+int(t)])
 264         return state, x
 265 }