<commit_before>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero 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 Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package client

import (
	"code.google.com/p/go.net/websocket"
	"errors"
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/pct"
	"sync"
	"time"
)

const (
	SEND_BUFFER_SIZE = 10
	RECV_BUFFER_SIZE = 10
)

type WebsocketClient struct {
	logger *pct.Logger
	api    pct.APIConnector
	link   string
	// --
	conn        *websocket.Conn
	recvChan    chan *proto.Cmd
	sendChan    chan *proto.Reply
	connectChan chan bool
	errChan     chan error
	backoff     *pct.Backoff
	started     bool
	sendSync    *pct.SyncChan
	recvSync    *pct.SyncChan
	mux         *sync.Mutex
	name        string
	status      *pct.Status
	connected   bool
}

func NewWebsocketClient(logger *pct.Logger, api pct.APIConnector, link string) (*WebsocketClient, error) {
	name := logger.Service()
	c := &WebsocketClient{
		logger: logger,
		api:    api,
		link:   link,
		// --
		conn:        nil,
		recvChan:    make(chan *proto.Cmd, RECV_BUFFER_SIZE),
		sendChan:    make(chan *proto.Reply, SEND_BUFFER_SIZE),
		connectChan: make(chan bool, 1),
		errChan:     make(chan error, 2),
		backoff:     pct.NewBackoff(5 * time.Minute),
		sendSync:    pct.NewSyncChan(),
		recvSync:    pct.NewSyncChan(),
		mux:         new(sync.Mutex),
		name:        name,
		status:      pct.NewStatus([]string{name, name + "-link"}),
	}
	return c, nil
}

func (c *WebsocketClient) Start() {
	// Start send() and recv() goroutines, but they wait for successful Connect().
	if !c.started {
		go c.send()
		go c.recv()
		c.started = true
	}
}

func (c *WebsocketClient) Stop() {
	if c.started {
		c.sendSync.Stop()
		c.recvSync.Stop()
		c.sendSync.Wait()
		c.recvSync.Wait()
		c.started = false
	}
}

func (c *WebsocketClient) Connect() {
	c.logger.Debug("Connect:call")
	defer c.logger.Debug("Connect:return")

	for {
		// Wait before attempt to avoid DDoS'ing the API
		// (there are many other agents in the world).
		c.logger.Debug("Connect:backoff.Wait")
		c.status.Update(c.name, "Connect wait")
		time.Sleep(c.backoff.Wait())

		if err := c.ConnectOnce(); err != nil {
			c.logger.Warn(err)
			continue
		}

		// Start/resume send() and recv() goroutines if Start() was called.
		if c.started {
			c.recvSync.Start()
			c.sendSync.Start()
		}

		c.backoff.Success()
		c.notifyConnect(true)
		return // success
	}
}

func (c *WebsocketClient) ConnectOnce() error {
	c.logger.Debug("ConnectOnce:call")
	defer c.logger.Debug("ConnectOnce:return")

	// Make websocket connection.  If this fails, either API is down or the ws
	// address is wrong.
	link := c.api.AgentLink(c.link)
	c.logger.Debug("ConnectOnce:link:" + link)
	config, err := websocket.NewConfig(link, c.api.Origin())
	if err != nil {
		return err
	}
	config.Header.Add("X-Percona-API-Key", c.api.ApiKey())

	c.logger.Debug("ConnectOnce:websocket.DialConfig")
	c.status.Update(c.name, "Connecting "+link)
	conn, err := websocket.DialConfig(config)
	if err != nil {
		return err
	}

	c.mux.Lock()
	defer c.mux.Unlock()
	c.connected = true
	c.conn = conn
	c.status.Update(c.name, "Connected "+link)

	return nil
}

func (c *WebsocketClient) Disconnect() error {
	c.logger.DebugOffline("Disconnect:call")
	defer c.logger.DebugOffline("Disconnect:return")

	/**
	 * Must guard c.conn here to prevent duplicate notifyConnect() because Close()
	 * causes recv() to error which calls Disconnect(), and normally we want this:
	 * to call Disconnect() on recv error so that notifyConnect(false) is called
	 * to let user know that remote end hung up.  However, when user hangs up
	 * the Disconnect() call from recv() is duplicate and not needed.
	 */
	c.mux.Lock()<commit_msg>timeout added<commit_after>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero 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 Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package client

import (
	"code.google.com/p/go.net/websocket"
	"errors"
	"fmt"
	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/pct"
	"sync"
	"time"
	"net"
	"crypto/tls"
	"strings"
)

const (
	SEND_BUFFER_SIZE = 10
	RECV_BUFFER_SIZE = 10
)

type WebsocketClient struct {
	logger *pct.Logger
	api    pct.APIConnector
	link   string
	// --
	conn        *websocket.Conn
	recvChan    chan *proto.Cmd
	sendChan    chan *proto.Reply
	connectChan chan bool
	errChan     chan error
	backoff     *pct.Backoff
	started     bool
	sendSync    *pct.SyncChan
	recvSync    *pct.SyncChan
	mux         *sync.Mutex
	name        string
	status      *pct.Status
	connected   bool
}

func NewWebsocketClient(logger *pct.Logger, api pct.APIConnector, link string) (*WebsocketClient, error) {
	name := logger.Service()
	c := &WebsocketClient{
		logger: logger,
		api:    api,
		link:   link,
		// --
		conn:        nil,
		recvChan:    make(chan *proto.Cmd, RECV_BUFFER_SIZE),
		sendChan:    make(chan *proto.Reply, SEND_BUFFER_SIZE),
		connectChan: make(chan bool, 1),
		errChan:     make(chan error, 2),
		backoff:     pct.NewBackoff(5 * time.Minute),
		sendSync:    pct.NewSyncChan(),
		recvSync:    pct.NewSyncChan(),
		mux:         new(sync.Mutex),
		name:        name,
		status:      pct.NewStatus([]string{name, name + "-link"}),
	}
	return c, nil
}

func (c *WebsocketClient) Start() {
	// Start send() and recv() goroutines, but they wait for successful Connect().
	if !c.started {
		go c.send()
		go c.recv()
		c.started = true
	}
}

func (c *WebsocketClient) Stop() {
	if c.started {
		c.sendSync.Stop()
		c.recvSync.Stop()
		c.sendSync.Wait()
		c.recvSync.Wait()
		c.started = false
	}
}

func (c *WebsocketClient) Connect() {
	c.logger.Debug("Connect:call")
	defer c.logger.Debug("Connect:return")

	for {
		// Wait before attempt to avoid DDoS'ing the API
		// (there are many other agents in the world).
		c.logger.Debug("Connect:backoff.Wait")
		c.status.Update(c.name, "Connect wait")
		time.Sleep(c.backoff.Wait())

		if err := c.ConnectOnce(); err != nil {
			c.logger.Warn(err)
			continue
		}

		// Start/resume send() and recv() goroutines if Start() was called.
		if c.started {
			c.recvSync.Start()
			c.sendSync.Start()
		}

		c.backoff.Success()
		c.notifyConnect(true)
		return // success
	}
}

func TlsDial(network, addr string, config *tls.Config, timeout uint) (*tls.Conn, error) {
	raddr := addr
	c, err := net.DialTimeout(network, raddr, time.Duration(timeout) * time.Second)
	if err != nil {
		return nil, err
	}

	colonPos := strings.LastIndex(raddr, ":")
	if colonPos == -1 {
		colonPos = len(raddr)
	}
	hostname := raddr[:colonPos]

	if config == nil {
		return nil, fmt.Errorf("No config passed to TlsDial")
	}
	// If no ServerName is set, infer the ServerName
	// from the hostname we're connecting to.
	if config.ServerName == "" {
		// Make a copy to avoid polluting argument or default.
		c := *config
		c.ServerName = hostname
		config = &c
	}
	conn := tls.Client(c, config)
	if err = conn.Handshake(); err != nil {
		c.Close()
		return nil, err
	}
	return conn, nil
}

// We need this function, because websocket.DialConfig is limited to connection without timeout.
// DialConfigTimeout opens a new client connection to a WebSocket with a config and timeout.
func DialConfigTimeout(config *websocket.Config, timeout uint) (ws *websocket.Conn, err error) {
	var client net.Conn
	if config.Location == nil {
			return nil, &websocket.DialError{config, websocket.ErrBadWebSocketLocation}
	}
	if config.Origin == nil {
			return nil, &websocket.DialError{config, websocket.ErrBadWebSocketOrigin}
	}
	switch config.Location.Scheme {
	case "ws":
			client, err = net.DialTimeout("tcp", config.Location.Host, time.Duration(timeout) * time.Second)

	case "wss":
			client, err = TlsDial("tcp", config.Location.Host, config.TlsConfig, timeout)

	default:
			err = websocket.ErrBadScheme
	}
	if err != nil {
			goto Error
	}

	ws, err = websocket.NewClient(config, client)
	if err != nil {
			goto Error
	}
	return

Error:
	return nil, &websocket.DialError{config, err}
}

func (c *WebsocketClient) ConnectOnce() error {
	c.logger.Debug("ConnectOnce:call")
	defer c.logger.Debug("ConnectOnce:return")

	// Make websocket connection.  If this fails, either API is down or the ws
	// address is wrong.
	link := c.api.AgentLink(c.link)
	c.logger.Debug("ConnectOnce:link:" + link)
	config, err := websocket.NewConfig(link, c.api.Origin())
	if err != nil {
		return err
	}
	config.Header.Add("X-Percona-API-Key", c.api.ApiKey())

	c.logger.Debug("ConnectOnce:websocket.DialConfig")
	c.status.Update(c.name, "Connecting "+link)
	conn, err := DialConfigTimeout(config, 10) // 10 seconds timeout
	if err != nil {
		return err
	}

	c.mux.Lock()
	defer c.mux.Unlock()
	c.connected = true
	c.conn = conn
	c.status.Update(c.name, "Connected "+link)

	return nil
}

func (c *WebsocketClient) Disconnect() error {
	c.logger.DebugOffline("Disconnect:call")
	defer c.logger.DebugOffline("Disconnect:return")

	/**
	 * Must guard c.conn here to prevent duplicate notifyConnect() because Close()
	 * causes recv() to error which calls Disconnect(), and normally we want this:
	 * to call Disconnect() on recv error so that notifyConnect(false) is called
	 * to let user know that remote end hung up.  However, when user hangs up
	 * the Disconnect() call from recv() is duplicate and not needed.
	 */
	c.mux.Lock()
{'fraction_non_alphanumeric': 0.12073689182805858, 'fraction_numerical': 0.0027397260273972603, 'mean_word_length': 3.4126719466444353, 'pattern_counts': {'":': 3, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3160770', 'n_tokens_mistral': 3809, 'n_tokens_neox': 3588, 'n_words': 1069}