<commit_before>package backend

import (
	"net/http"
	"net/url"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/websocket"

	"github.com/rancherio/websocket-proxy/common"
)

// Implement this iterface and pass implementations into ConnectToProxy() to have messages
// routed to and from the handler.
type Handler interface {
	Handle(messageKey string, initialMessage string, incomingMessages <-chan string, response chan<- common.Message)
}

func ConnectToProxy(proxyUrl string, handlers map[string]Handler) {
	log.WithFields(log.Fields{"url": proxyUrl}).Info("Connecting to proxy.")

	dialer := &websocket.Dialer{}
	headers := http.Header{}
	ws, _, err := dialer.Dial(proxyUrl, headers)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Fatal("Failed to connect to proxy.")
	}

	connectToProxyWS(ws, handlers)
}

func connectToProxyWS(ws *websocket.Conn, handlers map[string]Handler) {
	responders := make(map[string]chan string)
	responseChannel := make(chan common.Message, 10)

	// Write messages to proxy
	go func() {
		for {
			message, ok := <-responseChannel
			if !ok {
				return
			}
			data := common.FormatMessage(message.Key, message.Type, message.Body)
			ws.WriteMessage(1, []byte(data))
		}
	}()

	// Read and route messages from proxy
	for {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Error("Received error reading from socket. Exiting.")
			close(responseChannel)
			return
		}
<commit_msg>Add client to server pings<commit_after>package backend

import (
	"net/http"
	"net/url"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/websocket"

	"github.com/rancherio/websocket-proxy/common"
)

// Implement this iterface and pass implementations into ConnectToProxy() to have messages
// routed to and from the handler.
type Handler interface {
	Handle(messageKey string, initialMessage string, incomingMessages <-chan string, response chan<- common.Message)
}

func ConnectToProxy(proxyUrl string, handlers map[string]Handler) {
	log.WithFields(log.Fields{"url": proxyUrl}).Info("Connecting to proxy.")

	dialer := &websocket.Dialer{}
	headers := http.Header{}
	ws, _, err := dialer.Dial(proxyUrl, headers)
	if err != nil {
		log.WithFields(log.Fields{
			"error": err,
		}).Fatal("Failed to connect to proxy.")
	}

	connectToProxyWS(ws, handlers)
}

func connectToProxyWS(ws *websocket.Conn, handlers map[string]Handler) {
	responders := make(map[string]chan string)
	responseChannel := make(chan common.Message, 10)

	// Write messages to proxy
	go func() {
		for {
			select {
			case message, ok := <-responseChannel:
				if !ok {
					return
				}
				data := common.FormatMessage(message.Key, message.Type, message.Body)
				ws.WriteMessage(1, []byte(data))
			case <-time.After(time.Second * 5):
				ws.WriteControl(websocket.PingMessage, []byte(""), time.Now().Add(time.Second))
			}
		}
	}()

	// Read and route messages from proxy
	for {
		_, msg, err := ws.ReadMessage()
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Error("Received error reading from socket. Exiting.")
			close(responseChannel)
			return
		}

{'fraction_non_alphanumeric': 0.13690851735015772, 'fraction_numerical': 0.0022082018927444794, 'mean_word_length': 4.198360655737705, 'pattern_counts': {'":': 6, '<': 10, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '7442027', 'n_tokens_mistral': 1139, 'n_tokens_neox': 1015, 'n_words': 259}