<commit_before>package networkwebsockets

import (
	"encoding/json"
	"time"

	"github.com/richtr/websocket"
)

type PeerConnection struct {
	// Unique identifier for this peer connection
	id string

	// The Network Web Socket channel to which this peer connection belongs
	channel *NetworkWebSocket

	// WebSocket connection object
	conn *websocket.Conn
}

func NewPeerConnection(channel *NetworkWebSocket, conn *websocket.Conn) *PeerConnection {
	peerConn := &PeerConnection{
		id:      GenerateId(),
		channel: channel,
		conn:    conn,
	}

	// Start websocket read/write pumps
	peerConn.Start()

	return peerConn
}

func (peer *PeerConnection) Start() {
	// Start connection read/write pumps
	go peer.writeConnectionPump()
	go peer.readConnectionPump()

	// Add reference to this peer connection to channel
	peer.addConnection()
}

// Send a message to the target websocket connection
func (peer *PeerConnection) send(action string, source string, target string, payload string) {
	// Construct proxy wire message
	m := NetworkWebSocketWireMessage{
		Action:  action,
		Source:  source,
		Target:  target,
		Payload: payload,
	}
	messagePayload, err := json.Marshal(m)
	if err != nil {
		return
	}

	peer.conn.SetWriteDeadline(time.Now().Add(writeWait))
	peer.conn.WriteMessage(websocket.TextMessage, messagePayload)
}

// readConnectionPump pumps messages from an individual websocket connection to the dispatcher
func (peer *PeerConnection) readConnectionPump() {
	defer func() {
		peer.Stop()
	}()
	peer.conn.SetReadLimit(maxMessageSize)
	peer.conn.SetReadDeadline(time.Now().Add(pongWait))
	peer.conn.SetPongHandler(func(string) error { peer.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		opCode, buf, err := peer.conn.ReadMessage()
		if err != nil || opCode != websocket.TextMessage {
			break
		}

		var message NetworkWebSocketWireMessage
		if err := json.Unmarshal(buf, &message); err != nil {
			continue // ignore unrecognized message format
		}

		switch message.Action {
		// 'connect' and 'disconnect' events are write-only so will not be handled here

		case "status":

			// Echo peer id back to callee
			peer.send("status", peer.id, peer.id, "")

		case "broadcast":

			wsBroadcast := &NetworkWebSocketWireMessage{
				Action:    "broadcast",
				Source:    peer.id,
				Target:    "", // target all connections
				Payload:   string(message.Payload),
				fromProxy: false,
			}
			peer.channel.broadcastBuffer <- wsBroadcast

		case "message":

			messageSent := false

			// Relay message to peer channel that matches target
			for _, _peer := range peer.channel.peers {
				if _peer.id == message.Target {
					_peer.send("message", peer.id, message.Target, message.Payload)
					messageSent = true
					break
				}
			}

			if !messageSent {
				// Hunt for proxy that owns target peer id in known proxies
				for _, proxy := range peer.channel.proxies {
					if proxy.peerIds[message.Target] {
						proxy.send("message", peer.id, message.Target, message.Payload)
						messageSent = true
						break
					}
				}
			}
		}

	}
}

// writeConnectionPump keeps an individual websocket connection alive
func (peer *PeerConnection) writeConnectionPump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
	}()
	for {
		select {
		case <-ticker.C:
			peer.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := peer.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}<commit_msg>Expose broadcast and direct message capabilities from peer connections<commit_after>package networkwebsockets

import (
	"encoding/json"
	"time"

	"github.com/richtr/websocket"
)

type PeerConnection struct {
	// Unique identifier for this peer connection
	id string

	// The Network Web Socket channel to which this peer connection belongs
	channel *NetworkWebSocket

	// WebSocket connection object
	conn *websocket.Conn
}

func NewPeerConnection(channel *NetworkWebSocket, conn *websocket.Conn) *PeerConnection {
	peerConn := &PeerConnection{
		id:      GenerateId(),
		channel: channel,
		conn:    conn,
	}

	// Start websocket read/write pumps
	peerConn.Start()

	return peerConn
}

func (peer *PeerConnection) Start() {
	// Start connection read/write pumps
	go peer.writeConnectionPump()
	go peer.readConnectionPump()

	// Add reference to this peer connection to channel
	peer.addConnection()
}

func (peer *PeerConnection) SendBroadcast(data string) {
	wsBroadcast := &NetworkWebSocketWireMessage{
		Action:    "broadcast",
		Source:    peer.id,
		Target:    "", // target all connections
		Payload:   data,
		fromProxy: false,
	}
	peer.channel.broadcastBuffer <- wsBroadcast
}

func (peer *PeerConnection) SendMessage(data string, targetId string) {
	if targetId == "" {
		return
	}

	// Relay message to peer channel that matches target
	for _, _peer := range peer.channel.peers {
		if _peer.id == targetId {
			_peer.send("message", peer.id, targetId, data)
			return
		}
	}

	// If we have not delivered the message yet then hunt for a
	// proxy that owns target peer id in known proxies
	for _, proxy := range peer.channel.proxies {
		if proxy.peerIds[targetId] {
			proxy.send("message", peer.id, targetId, data)
			return
		}
	}
}

// Send a message to the target websocket connection
func (peer *PeerConnection) send(action string, source string, target string, payload string) {
	// Construct proxy wire message
	m := NetworkWebSocketWireMessage{
		Action:  action,
		Source:  source,
		Target:  target,
		Payload: payload,
	}
	messagePayload, err := json.Marshal(m)
	if err != nil {
		return
	}

	peer.conn.SetWriteDeadline(time.Now().Add(writeWait))
	peer.conn.WriteMessage(websocket.TextMessage, messagePayload)
}

// readConnectionPump pumps messages from an individual websocket connection to the dispatcher
func (peer *PeerConnection) readConnectionPump() {
	defer func() {
		peer.Stop()
	}()
	peer.conn.SetReadLimit(maxMessageSize)
	peer.conn.SetReadDeadline(time.Now().Add(pongWait))
	peer.conn.SetPongHandler(func(string) error { peer.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		opCode, buf, err := peer.conn.ReadMessage()
		if err != nil || opCode != websocket.TextMessage {
			break
		}

		var message NetworkWebSocketWireMessage
		if err := json.Unmarshal(buf, &message); err != nil {
			continue // ignore unrecognized message format
		}

		switch message.Action {
		// 'connect' and 'disconnect' events are write-only so will not be handled here

		case "status":

			// Echo peer id back to callee
			peer.send("status", peer.id, peer.id, "")

		case "broadcast":

			peer.SendBroadcast(string(message.Payload))

		case "message":

			peer.SendMessage(string(message.Payload), message.Target)

		}

	}
}

// writeConnectionPump keeps an individual websocket connection alive
func (peer *PeerConnection) writeConnectionPump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
	}()
	for {
		select {
		case <-ticker.C:
			peer.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := peer.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
{'fraction_non_alphanumeric': 0.1115777684446311, 'fraction_numerical': 0.0, 'mean_word_length': 3.8010752688172045, 'pattern_counts': {'":': 6, '<': 7, '<?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': '4830637', 'n_tokens_mistral': 2447, 'n_tokens_neox': 2179, 'n_words': 614}