archived-rttys/client.go

/*
 * Copyright (C) 2017 Jianhui Zhao <jianhuizhao329@gmail.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package main

import (
    "log"
    "time"
    "sync"
    "errors"
    "net/http"
    "encoding/base64"
    "github.com/gorilla/websocket"
)

const (
    // Time allowed to write a message to the peer.
    writeWait = 5 * time.Second

    // Time allowed to read the next pong message from the peer.
    pongWait = 1 * time.Second

    // pings to peer with this period.
    pingPeriod = 5 * time.Second
)

var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

type wsMessage struct {
    isDev bool
    msgType int
    data []byte
}

// Representing a device or user browser
type Client struct {
    br *Bridge
    isDev bool
    // device description
    description string
    devid string
    // Registration time
    timestamp int64
    sid string
    conn *websocket.Conn
    // Buffered channel of outbound messages.
    outbound chan *wsMessage

    // Avoid repeated closes
    mutex sync.Mutex
    isClosed bool
    closeChan chan byte
}

func (c *Client) wsClose() {
    defer c.mutex.Unlock()
    c.mutex.Lock()

    if !c.isClosed {
        c.conn.Close()
        c.isClosed = true
        close(c.closeChan)
    }
}
func (c *Client) unregister() {
    c.br.unregister <- c
}

func (c *Client) wsWrite(messageType int, data []byte) error {
    select {
    case c.outbound <- &wsMessage{c.isDev, messageType, data}:
    case <- c.closeChan:
        return errors.New("websocket closed")
    }
    return nil
}

func (c *Client) readPump() {
    defer func() {
        c.unregister()
    }()

    c.conn.SetPongHandler(func(string) error {
        c.conn.SetReadDeadline(time.Now().Add(time.Hour));
        return nil
    })

    for {
        msgType, data, err := c.conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
                log.Printf("error: %v", err)
            }
            break
        }

        if msgType == websocket.BinaryMessage {
            msg := &wsMessage{c.isDev, msgType, data}

            select {
            case c.br.inbound <- msg:
            case <- c.closeChan:
                return
            }
        }
    }
}

func (c *Client) writePump() {
    ticker := time.NewTicker(pingPeriod)
    defer func() {
        ticker.Stop()
        c.unregister()
    }()

    for {
        select {
        case msg := <- c.outbound:
            if err := c.conn.WriteMessage(msg.msgType, msg.data); err != nil {
                return
            }
        case <- ticker.C:
            c.conn.SetWriteDeadline(time.Now().Add(writeWait))
            c.conn.SetReadDeadline(time.Now().Add(pongWait));
            if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return
            }
        case <- c.closeChan:
            return
        }
    }
}

/* serveWs handles websocket requests from the peer. */
func serveWs(br *Bridge, w http.ResponseWriter, r *http.Request) {
    devid := r.URL.Query().Get("devid")
    if devid == "" {
        log.Println("devid required")
        return
    }

    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println(err)
        return
    }

    client := &Client{
        br: br,
        devid: devid,
        conn: conn,
        timestamp: time.Now().Unix(),
        outbound: make(chan *wsMessage, 100),
        closeChan: make(chan byte),
        isClosed: false,
    }

    isDev := r.URL.Query().Get("device")
    if isDev == "1" {
        client.isDev = true
        
        description, _ := base64.StdEncoding.DecodeString(r.URL.Query().Get("description"))
        client.description = string(description)
    }

    client.br.register <- client

    go client.readPump()
    go client.writePump()
}