func main() {
flag.Parse()
log.SetFlags(0)
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, os.Interrupt)
u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
log.Printf("connecting to %s", u.String())
c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
if err != nil {
log.Fatal("dial:", err)
}
defer c.Close()
done := make(chan struct{})
go func() {
defer c.Close()
defer close(done)
for {
_, message, err := c.ReadMessage()
if err != nil {
log.Println("read:", err)
return
}
log.Printf("recv: %s", message)
}
}()
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
for {
select {
case t := <-ticker.C:
err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
if err != nil {
log.Println("write:", err)
return
}
case <-interrupt:
log.Println("interrupt")
// To cleanly close a connection, a client should send a close
// frame and wait for the server to close the connection.
err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil {
log.Println("write close:", err)
return
}
select {
case <-done:
case <-time.After(time.Second):
}
c.Close()
return
}
}
}
func closeWithErr(log lager.Logger, conn *websocket.Conn, code int, reason string) {
err := conn.WriteControl(
websocket.CloseMessage,
websocket.FormatCloseMessage(code, reason),
time.Time{},
)
if err != nil {
log.Error("failed-to-close-websocket-connection", err)
}
}
// Sends dispose signals to all peers and deletes the room record.
func (room *Room) dispose(status int) {
if status == 0 {
status = websocket.CloseNormalClosure
}
// Close all Websocket connections.
for peer := range room.peers {
peer.ws.WriteControl(
websocket.CloseMessage,
websocket.FormatCloseMessage(status, "Room disposed"),
time.Time{})
delete(room.peers, peer)
}
// Close all the room channels.
close(room.broadcastQueue)
close(room.register)
close(room.unregister)
delete(rooms, room.Id)
// Delete the room record from the DB.
db := dbPool.Get()
defer db.Close()
db.Delete(config.CachePrefixRoom + room.Id)
}
// Start connects to chat and starts bot.
func (c *Client) Start() {
conn, _, err := websocket.DefaultDialer.Dial(c.WebSocketURL, nil)
if err != nil {
log.Fatal(err)
}
defer conn.Close()
c.conn = conn
done := make(chan struct{})
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, os.Interrupt)
go c.processMessages(done)
for {
select {
case <-interrupt:
err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil {
log.Fatal(err)
}
select {
case <-done:
case <-time.After(time.Second):
}
return
}
}
}
func (c *conn) close(closeCode int, closeErr error) error {
var err error
var errMsg string
if closeErr != nil {
errMsg = closeErr.Error()
c.server.logError(closeErr)
// Attempt to set close code from error message
errMsgLen := len(errMsg)
if errMsgLen >= 21 && errMsg[:17] == "websocket: close " {
closeCode, err = strconv.Atoi(errMsg[17:21])
if errMsgLen > 21 {
errMsg = errMsg[22:]
}
}
}
// Default close code
if closeCode == 0 {
closeCode = websocket.CloseNoStatusReceived
}
// Send close message
closeMessage := websocket.FormatCloseMessage(closeCode, errMsg)
deadline := time.Now().Add(time.Second)
err = c.conn.WriteControl(websocket.CloseMessage, closeMessage, deadline)
// Kill and remove connection
c.closeChan <- closeSignal{}
c.server.connSet.remove(c)
return err
}
func (w *WebsocketServer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
var handler func(string, *gorilla.Conn)
switch r.URL.Path {
case TAIL_LOGS_PATH:
handler = w.streamLogs
case RECENT_LOGS_PATH:
handler = w.recentLogs
case DUMP_LOGS_PATH:
handler = w.recentLogs
default:
http.Error(rw, "invalid path "+r.URL.Path, 400)
return
}
appId, err := w.validate(r)
if err != nil {
http.Error(rw, err.Error(), 400)
return
}
ws, err := gorilla.Upgrade(rw, r, nil, 1024, 1024)
if err != nil {
http.Error(rw, err.Error(), 400)
return
}
defer ws.Close()
defer ws.WriteControl(gorilla.CloseMessage, gorilla.FormatCloseMessage(gorilla.CloseNormalClosure, ""), time.Time{})
handler(appId, ws)
}
// StartMockServer starts a mock websocket server.
func StartMockServer(t *testing.T, closeWS <-chan bool) (*httptest.Server, chan<- string, <-chan string, <-chan error, error) {
serverChan := make(chan string)
requestsChan := make(chan string)
errChan := make(chan error)
upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ws, err := upgrader.Upgrade(w, r, nil)
go func() {
<-closeWS
ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
}()
if err != nil {
errChan <- err
}
go func() {
_, msg, err := ws.ReadMessage()
if err != nil {
errChan <- err
} else {
requestsChan <- string(msg)
}
}()
for str := range serverChan {
err := ws.WriteMessage(websocket.TextMessage, []byte(str))
if err != nil {
errChan <- err
}
}
})
server := httptest.NewTLSServer(handler)
return server, serverChan, requestsChan, errChan, nil
}
// Close the Base connection. Closes the send Handler and all channels used
// Since all channels are either internal or channels this middleware is sending on.
func (c *Connection) Close(closeCode int) error {
c.disconnectSend <- true
//TODO look for a better way to unblock the reader
c.ws.SetReadDeadline(time.Now())
// Send close message to the client
c.log("Sending close message to client", LEVEL_DEBUG)
c.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(closeCode, ""), time.Now().Add(c.WriteWait))
// If the connection can not be closed, return the error
c.log("Closing websocket connection", LEVEL_DEBUG)
if err := c.ws.Close(); err != nil {
c.log("Connection could not be closed: %s", LEVEL_ERROR, err.Error())
return err
}
// Send disconnect message to the next handler
c.log("Sending disconnect to handler", LEVEL_DEBUG)
c.Done <- true
// Close disconnect and error channels this connection was sending on
close(c.Done)
close(c.Error)
return nil
}
func (c *migrationFields) disconnect() {
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
c.controlLock.Lock()
if c.controlConn != nil {
c.controlConn.WriteMessage(websocket.CloseMessage, closeMsg)
c.controlConn = nil /* don't close twice */
}
c.controlLock.Unlock()
/* Below we just Close(), which doesn't actually write to the
* websocket, it just closes the underlying connection. If e.g. there
* is still a filesystem transfer going on, but the other side has run
* out of disk space, writing an actual CloseMessage here will cause
* gorilla websocket to panic. Instead, we just force close this
* connection, since we report the error over the control channel
* anyway.
*/
if c.fsConn != nil {
c.fsConn.Close()
}
if c.criuConn != nil {
c.criuConn.Close()
}
}
func (w *WebsocketServer) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
w.logger.Debug("WebsocketServer.ServeHTTP: starting")
var handler wsHandler
var err error
paths := strings.Split(request.URL.Path, "/")
endpointName := paths[1]
if endpointName == "firehose" {
handler, err = w.firehoseHandler(paths, writer, request)
} else {
handler, err = w.appHandler(paths, writer, request)
}
if err != nil {
w.logger.Errorf("WebsocketServer.ServeHTTP: %s", err.Error())
return
}
ws, err := gorilla.Upgrade(writer, request, nil, 1024, 1024)
if err != nil {
w.logger.Errorf("WebsocketServer.ServeHTTP: Upgrade error (returning 400): %s", err.Error())
http.Error(writer, err.Error(), 400)
return
}
defer func() {
ws.WriteControl(gorilla.CloseMessage, gorilla.FormatCloseMessage(gorilla.CloseNormalClosure, ""), time.Time{})
ws.Close()
}()
handler(ws)
}
func (f *fakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if r.Method == "HEAD" {
w.WriteHeader(http.StatusOK)
return
}
if r.Method != "GET" {
http.Error(w, "Method not allowed", 405)
return
}
ws, err := websocket.Upgrade(w, r, nil, 0, 0)
if _, ok := err.(websocket.HandshakeError); ok {
http.Error(w, "Not a websocket handshake", 400)
return
} else if err != nil {
log.Println(err)
return
}
defer ws.Close()
defer ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
f.Lock()
f.lastWSConn = ws
f.Unlock()
go ws.ReadMessage()
for msg := range f.messages {
if err := ws.WriteMessage(websocket.BinaryMessage, msg); err != nil {
return
}
}
}
func fakeFirehoseHandler(rw http.ResponseWriter, r *http.Request) {
defer GinkgoRecover()
authorization := r.Header.Get("Authorization")
if authorization != "bearer good-token" {
log.Printf("Bad token passed to firehose: %s", authorization)
rw.WriteHeader(403)
r.Body.Close()
return
}
upgrader := websocket.Upgrader{
CheckOrigin: func(*http.Request) bool { return true },
}
ws, _ := upgrader.Upgrade(rw, r, nil)
defer ws.Close()
defer ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
for envelope := range fakeFirehoseInputChan {
buffer, err := proto.Marshal(envelope)
Expect(err).NotTo(HaveOccurred())
err = ws.WriteMessage(websocket.BinaryMessage, buffer)
Expect(err).NotTo(HaveOccurred())
}
}
func WebsocketSendStream(conn *websocket.Conn, r io.Reader) chan bool {
ch := make(chan bool)
go func(conn *websocket.Conn, r io.Reader) {
in := ReaderToChannel(r)
for {
buf, ok := <-in
if !ok {
break
}
w, err := conn.NextWriter(websocket.BinaryMessage)
if err != nil {
Debugf("got error getting next writer %s", err)
break
}
_, err = w.Write(buf)
w.Close()
if err != nil {
Debugf("got err writing %s", err)
break
}
}
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
conn.WriteMessage(websocket.CloseMessage, closeMsg)
ch <- true
}(conn, r)
return ch
}
func shouldCloseConnection(actual interface{}, _ ...interface{}) string {
conn := actual.(*websocket.Conn)
err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
if err != nil {
return "Websocket connection didn't close properly: " + err.Error()
}
return ""
}
func (c *Connection) SendClose(closeCode int, text string) {
// XXX: we're allowed to send control frames from any thread, so it
// might be easier to write the message directly to the web socket.
c.send <- message{
websocket.CloseMessage,
websocket.FormatCloseMessage(closeCode, text),
}
}
func NewFakeFirehoseInAppMode(validToken, appName string) *FakeFirehose {
return &FakeFirehose{
AppMode: true,
AppName: appName,
validToken: validToken,
closeMessage: websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
}
}
func (ws *wsconn) Close() error {
ws.connClosing = true
ws.wlock.Lock()
ws.wsc.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add((time.Second * 5)))
ws.wlock.Unlock()
err := ws.wsc.Close()
return err
}
func (pipe *Pipe) Close(msg string) {
pipe.logger.Printf("pipe: closing (msg=%q)", msg)
pipe.conn.WriteMessage(
websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, msg),
)
pipe.conn.Close()
}
func (ep *websocketPeer) Close() error {
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye")
err := ep.conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(5*time.Second))
if err != nil {
//log.Println("error sending close message:", err)
}
ep.closed = true
return ep.conn.Close()
}
// Close terminates the websocket connection to traffic controller.
func (cnsmr *Consumer) Close() error {
cnsmr.Lock()
defer cnsmr.Unlock()
if cnsmr.ws == nil {
return errors.New("connection does not exist")
}
cnsmr.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
return cnsmr.ws.Close()
}
func (p *pingingWebsocket) Close() error {
p.writeLock.Lock()
defer p.writeLock.Unlock()
p.pinger.Stop()
if err := p.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "ok"), time.Now().Add(writeWait)); err != nil {
p.conn.Close()
return err
}
return p.conn.Close()
}
func socket_handler(w http.ResponseWriter, r *http.Request) {
if redirect_to_host(w, r) {
return
}
if r.Method != "GET" {
http.Error(w, "Method not allowed", 405)
return
}
ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
if err != nil {
Log.Println(err)
return
}
user := NewUser()
Log.Printf("New user %v (%v) - (%v)\n", user.UserId, ws.RemoteAddr(), r.UserAgent())
close_msg := ""
defer func() {
ws.WriteControl(
websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, close_msg),
time.Now().Add(WRITE_WAIT),
)
ws.Close()
Log.Printf("User %v left (%v)\n", user.UserId, close_msg)
}()
// Retrieve the site the user wants to draw over:
location_url, err := url.QueryUnescape(r.RequestURI[6:]) // skip "/ws?u="
if err != nil {
close_msg = fmt.Sprintf("Invalid query: %v", err)
return
}
location := GetLocation(location_url)
request := &JoinRequest{user, make(chan bool)}
location.Join <- request
user_joined := <-request.resultChan
if !user_joined {
close_msg = fmt.Sprintf(
"Too much users at %v, try adding #something at the end of the URL.",
location_url,
)
return
}
go user.Sender(ws)
go user.Receiver(location, ws)
close_msg = <-user.kick
location.Quit <- user
}
// reader is the guts of this package. It takes the stdin and stdout pipes
// of the cmd we created in ServeWS and pipes them between the client and server
// over websockets.
func reader(conn *websocket.Conn, stdout io.ReadCloser, stdin io.WriteCloser) {
// Setup our connection's websocket ping/pong handlers from our const values.
conn.SetReadLimit(maxMessageSize)
conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
tickerChan := make(chan bool)
defer close(tickerChan) // make sure to close the ticker when we are done.
go ticker(conn, tickerChan)
for {
msgType, r, err := conn.NextReader()
if err != nil {
if msgType == -1 {
return // we got a disconnect from the client. We are good to close.
}
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
w, err := conn.NextWriter(msgType)
if err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
if _, err := io.Copy(stdin, r); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
go func() {
if _, err := io.Copy(w, stdout); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
if err := w.Close(); err != nil {
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, ""), time.Time{})
return
}
}()
}
}
func (router *EventRouter) Stop() (err error) {
if router.eventStream != nil {
router.mu.Lock()
defer router.mu.Unlock()
if router.eventStream != nil {
router.eventStream.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
router.eventStream = nil
}
}
return nil
}
// close connection with an error string.
func closeMessage(ws *websocket.Conn, errStr string) {
// close code 1008 is used for a generic "policy violation" message.
msg := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, errStr)
log.Println("[DEBUG] Writing a close message of", string(msg))
err := ws.WriteMessage(websocket.CloseMessage, msg)
if err != nil {
log.Println("[ERROR] Could not write message back to user", err)
}
ws.Close()
return
}
func (t *WebSocketTransport) Close() error {
if t._connection != nil {
r := websocket.FormatCloseMessage(websocket.CloseAbnormalClosure, "server closed")
t._connection.WriteMessage(websocket.TextMessage, r)
t._connection.Close()
t._connection = nil
}
atomic.StoreInt32(&t._stop, 1)
return nil
}
func (c *connection) close() error {
c.lock.Lock()
defer c.lock.Unlock()
c.isClosed = true
if c.ws == nil {
return nil
}
err := c.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
if err != nil {
return err
}
return c.ws.Close()
}
func (l *websocketListener) Start(url string, appId string, outputChan OutputChannel, stopChan StopChannel) error {
conn, _, err := websocket.DefaultDialer.Dial(url, nil)
if err != nil {
return err
}
go func() {
<-stopChan
conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
conn.Close()
}()
return l.listenWithTimeout(l.timeout, url, appId, conn, outputChan)
}
func (c *migrationFields) disconnect() {
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
if c.controlConn != nil {
c.controlConn.WriteMessage(websocket.CloseMessage, closeMsg)
}
if c.fsConn != nil {
c.fsConn.WriteMessage(websocket.CloseMessage, closeMsg)
}
if c.criuConn != nil {
c.criuConn.WriteMessage(websocket.CloseMessage, closeMsg)
}
}
// Who the hell do we call close first on? App or connection?
// Either way one or the other may have to check on the other, which is no good
func (ep *WebsocketConnection) Close(reason string) error {
core.Info("Closing connection with reason: %s", reason)
closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "goodbye")
err := ep.conn.WriteControl(websocket.CloseMessage, closeMsg, time.Now().Add(5*time.Second))
if err != nil {
log.Println("error sending close message:", err)
}
ep.lock = nil
ep.closed = true
return ep.conn.Close()
}