<commit_before>package neptulon

import (
	"io"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

	"golang.org/x/net/websocket"
)

// Conn is a client connection.
type Conn struct {
	ID         string
	Session    *cmap.CMap
	middleware []func(ctx *ReqCtx) error
	resRoutes  *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws         *websocket.Conn
	deadline   time.Duration
	closed     bool
}

// NewConn creates a new Conn object.
func NewConn() (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:        id,
		Session:   cmap.New(),
		resRoutes: cmap.New(),
		deadline:  time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
// Default value for read/write deadline is 300 seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...func(ctx *ReqCtx) error) {
	c.middleware = append(c.middleware, middleware...)
}

// Connect connects to the given WebSocket server.
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}

	c.ws = ws
	go c.startReceive()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.ws.RemoteAddr()
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *ResCtx) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *ResCtx) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes the connection.
func (c *Conn) Close() error {
	c.closed = true
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}

// UseConn reuses an established websocket.Conn.
// This function blocks and does not return until the connection is closed by another goroutine.
func (c *Conn) useConn(ws *websocket.Conn) {
	c.ws = ws
	c.startReceive()
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if peer closed the connection
			if err == io.EOF {
				log.Printf("Peer disconnected. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			// if we closed the connection
			if c.closed {
				log.Printf("Connection closed. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is a response
		if m.ID != "" && (m.Result != nil || m.Error != nil) {
			if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
				err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
				c.resRoutes.Delete(m.ID)
				if err != nil {
					log.Println("Error while handling response:", err)
					break
				}
			} else {
				log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
				break
			}
		}

		// if the message is not a JSON-RPC message
		log.Printf("Received on unknown message. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
	}
}<commit_msg>fix typo in unknown message handling<commit_after>package neptulon

import (
	"io"
	"log"
	"net"
	"time"

	"github.com/neptulon/cmap"
	"github.com/neptulon/shortid"

	"golang.org/x/net/websocket"
)

// Conn is a client connection.
type Conn struct {
	ID         string
	Session    *cmap.CMap
	middleware []func(ctx *ReqCtx) error
	resRoutes  *cmap.CMap // message ID (string) -> handler func(ctx *ResCtx) error : expected responses for requests that we've sent
	ws         *websocket.Conn
	deadline   time.Duration
	closed     bool
}

// NewConn creates a new Conn object.
func NewConn() (*Conn, error) {
	id, err := shortid.UUID()
	if err != nil {
		return nil, err
	}

	return &Conn{
		ID:        id,
		Session:   cmap.New(),
		resRoutes: cmap.New(),
		deadline:  time.Second * time.Duration(300),
	}, nil
}

// SetDeadline set the read/write deadlines for the connection, in seconds.
// Default value for read/write deadline is 300 seconds.
func (c *Conn) SetDeadline(seconds int) {
	c.deadline = time.Second * time.Duration(seconds)
}

// Middleware registers middleware to handle incoming request messages.
func (c *Conn) Middleware(middleware ...func(ctx *ReqCtx) error) {
	c.middleware = append(c.middleware, middleware...)
}

// Connect connects to the given WebSocket server.
func (c *Conn) Connect(addr string) error {
	ws, err := websocket.Dial(addr, "", "http://localhost")
	if err != nil {
		return err
	}

	c.ws = ws
	go c.startReceive()
	time.Sleep(time.Millisecond) // give receive goroutine a few cycles to start
	return nil
}

// RemoteAddr returns the remote network address.
func (c *Conn) RemoteAddr() net.Addr {
	return c.ws.RemoteAddr()
}

// SendRequest sends a JSON-RPC request through the connection with an auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequest(method string, params interface{}, resHandler func(res *ResCtx) error) (reqID string, err error) {
	id, err := shortid.UUID()
	if err != nil {
		return "", err
	}

	req := request{ID: id, Method: method, Params: params}
	if err = c.send(req); err != nil {
		return "", err
	}

	c.resRoutes.Set(req.ID, resHandler)
	return id, nil
}

// SendRequestArr sends a JSON-RPC request through the connection, with array params and auto generated request ID.
// resHandler is called when a response is returned.
func (c *Conn) SendRequestArr(method string, resHandler func(res *ResCtx) error, params ...interface{}) (reqID string, err error) {
	return c.SendRequest(method, params, resHandler)
}

// Close closes the connection.
func (c *Conn) Close() error {
	c.closed = true
	return c.ws.Close()
}

// SendResponse sends a JSON-RPC response message through the connection.
func (c *Conn) sendResponse(id string, result interface{}, err *ResError) error {
	return c.send(response{ID: id, Result: result, Error: err})
}

// Send sends the given message through the connection.
func (c *Conn) send(msg interface{}) error {
	if err := c.ws.SetWriteDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Send(c.ws, msg)
}

// Receive receives message from the connection.
func (c *Conn) receive(msg *message) error {
	if err := c.ws.SetReadDeadline(time.Now().Add(c.deadline)); err != nil {
		return err
	}

	return websocket.JSON.Receive(c.ws, &msg)
}

// UseConn reuses an established websocket.Conn.
// This function blocks and does not return until the connection is closed by another goroutine.
func (c *Conn) useConn(ws *websocket.Conn) {
	c.ws = ws
	c.startReceive()
}

// startReceive starts receiving messages. This method blocks and does not return until the connection is closed.
func (c *Conn) startReceive() {
	// append the last middleware to request stack, which will write the response to connection, if any
	c.middleware = append(c.middleware, func(ctx *ReqCtx) error {
		if ctx.Res != nil || ctx.Err != nil {
			return ctx.Conn.sendResponse(ctx.ID, ctx.Res, ctx.Err)
		}

		return nil
	})

	for {
		var m message
		err := c.receive(&m)
		if err != nil {
			// if peer closed the connection
			if err == io.EOF {
				log.Printf("Peer disconnected. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			// if we closed the connection
			if c.closed {
				log.Printf("Connection closed. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
				break
			}

			log.Println("Error while receiving message:", err)
			break
		}

		// if the message is a request
		if m.Method != "" {
			if err := newReqCtx(c, m.ID, m.Method, m.Params, c.middleware).Next(); err != nil {
				log.Println("Error while handling request:", err)
				break
			}

			continue
		}

		// if the message is not a JSON-RPC message
		if m.ID == "" || (m.Result == nil && m.Error == nil) {
			log.Printf("Received an unknown message. Conn ID: %v, Remote Addr: %v\n", c.ID, c.RemoteAddr())
		}

		// if the message is a response
		if resHandler, ok := c.resRoutes.GetOk(m.ID); ok {
			err := resHandler.(func(ctx *ResCtx) error)(newResCtx(c, m.ID, m.Result, m.Error))
			c.resRoutes.Delete(m.ID)
			if err != nil {
				log.Println("Error while handling response:", err)
				break
			}
		} else {
			log.Println("Error while handling response: got response to a request with unknown ID:", m.ID)
			break
		}
	}
}
{'fraction_non_alphanumeric': 0.12977895835309744, 'fraction_numerical': 0.0011384119153780476, 'mean_word_length': 3.388842631140716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '4562214', 'n_tokens_mistral': 3719, 'n_tokens_neox': 3465, 'n_words': 1185}