client.go

package netpipe

import (
	"fmt"
	"net"
)

// Client connects to a NetPipe server and exchanges raw data.
// Create one with NewClient(), register callbacks, Connect(), then go Listen().
type Client struct {
	config        Config
	conn          net.Conn // *safeConn for TCP, *net.UDPConn for UDP
	encryptionKey string   // set by SetEncryptionKey
	sessionID     []byte   // 16-byte session UUID received from server (for resumption)

	// Stage 6 - stream reassembly
	streams *streamBuffer

	onData       func(data []byte)
	onDisconnect func()
	onStream     func(data []byte)
	onReject     func(reason string)
}

// NewClient creates a client with the given config.
// Pass Config{} for all defaults (localhost:5000, TCP).
func NewClient(cfg Config) *Client {
	cfg.applyDefaults()
	return &Client{
		config:  cfg,
		streams: newStreamBuffer(),
	}
}

// ---------------------------------------------------------------------------
// Event registration
// ---------------------------------------------------------------------------

// OnData registers a callback fired when data arrives from the server.
// Data is raw bytes - the developer decides what they mean.
func (c *Client) OnData(fn func(data []byte)) {
	c.onData = fn
}

// OnDisconnect registers a callback fired when the server disconnects.
func (c *Client) OnDisconnect(fn func()) {
	c.onDisconnect = fn
}

// OnStream registers a callback fired when a chunked stream transfer completes.
// Data is the fully reassembled payload.
func (c *Client) OnStream(fn func(data []byte)) {
	c.onStream = fn
}

// OnReject registers a callback fired when the server rejects the connection.
// Reason is a short string like "server full" or "too many connections from your IP".
func (c *Client) OnReject(fn func(reason string)) {
	c.onReject = fn
}

// ---------------------------------------------------------------------------
// Lifecycle
// ---------------------------------------------------------------------------

// Connect establishes the connection to the server.
// Works for both TCP and UDP - the protocol is set in Config.
func (c *Client) Connect() error {
	if c.config.Protocol == "udp" {
		return c.connectUDP()
	}
	addr := fmt.Sprintf("%s:%d", c.config.Host, c.config.Port)
	conn, err := net.Dial(c.config.Protocol, addr)
	if err != nil {
		return fmt.Errorf("netpipe: connect failed: %w", err)
	}
	// wrap in safeConn for thread-safe writes (pong from read goroutine + send from main)
	c.conn = wrapConn(conn)
	return nil
}

// SetEncryptionKey sets the key used to auto-decrypt incoming encrypted messages.
// Call before Listen().
func (c *Client) SetEncryptionKey(key string) {
	c.encryptionKey = key
}

// Disconnect closes the connection cleanly.
func (c *Client) Disconnect() error {
	if c.conn != nil {
		err := c.conn.Close()
		c.conn = nil
		return err
	}
	return nil
}

// IsConnected returns true if the client has an active connection.
func (c *Client) IsConnected() bool {
	return c.conn != nil
}

// ---------------------------------------------------------------------------
// Data
// ---------------------------------------------------------------------------

// Send writes raw bytes to the server. Returns an error if not connected.
// The message is framed automatically so it arrives as a complete unit.
func (c *Client) Send(data []byte) error {
	if c.conn == nil {
		return fmt.Errorf("netpipe: not connected")
	}
	if c.config.Protocol == "udp" {
		return c.sendUDP(0x00, data)
	}
	return writeFrame(c.conn, data)
}

// SendEncrypted writes AES-256-GCM encrypted bytes to the server.
// The key can be any string - it is hashed to 32 bytes internally.
func (c *Client) SendEncrypted(data []byte, key string) error {
	if c.conn == nil {
		return fmt.Errorf("netpipe: not connected")
	}
	ciphertext, err := encrypt(data, key)
	if err != nil {
		return err
	}
	if c.config.Protocol == "udp" {
		return c.sendUDP(flagEncrypted, ciphertext)
	}
	return writeMessage(c.conn, flagEncrypted, ciphertext)
}

// SendStream sends large data as a chunked stream.
// The server receives it as a single complete payload via OnStream.
func (c *Client) SendStream(data []byte) error {
	if c.conn == nil {
		return fmt.Errorf("netpipe: not connected")
	}
	if c.config.Protocol == "udp" {
		return writeStreamUDPClient(c.conn, data, c.config.ChunkSize)
	}
	return writeStream(c.conn, data, c.config.ChunkSize)
}

// ---------------------------------------------------------------------------
// Listen
// ---------------------------------------------------------------------------

// Listen runs a blocking read loop that fires OnData with complete messages.
// Designed to be run in a goroutine: go client.Listen()
// When the connection drops, fires OnDisconnect and returns (or auto-reconnects).
func (c *Client) Listen() {
	if c.conn == nil {
		return
	}

	if c.config.Protocol == "udp" {
		c.listenUDP()
		return
	}

	// TCP: run the read loop
	c.readLoop()

	// connection dropped - readLoop sets c.conn = nil
	if c.config.AutoReconnect {
		c.reconnectLoop()
		return
	}

	if c.onDisconnect != nil {
		c.onDisconnect()
	}
}

// readLoop is the TCP read loop extracted so reconnectLoop can reuse it
// without recursive Listen() calls. Returns when the connection drops.
func (c *Client) readLoop() {
	// if we have a session ID from a previous connection, send it to resume
	if len(c.sessionID) == sessionUUIDSize && c.conn != nil {
		writeMessage(c.conn, flagSession, c.sessionID)
	}

	for {
		flags, body, err := readFrameFull(c.conn)
		if err != nil {
			break
		}

		// session assignment from server - store for reconnection
		if flags&flagSession != 0 && len(body) >= sessionUUIDSize {
			c.sessionID = make([]byte, sessionUUIDSize)
			copy(c.sessionID, body[:sessionUUIDSize])
			continue
		}

		if flags&flagPing != 0 {
			c.sendPong()
			continue
		}
		if flags&flagPong != 0 {
			continue
		}
		if flags&flagReject != 0 {
			if c.onReject != nil {
				c.onReject(string(body))
			}
			break
		}

		if flags&flagStream != 0 {
			streamID, index, total, chunk, err := parseStreamBody(body)
			if err != nil {
				continue
			}
			assembled, complete := c.streams.addChunk(streamID, index, total, chunk)
			if complete && c.onStream != nil {
				c.onStream(assembled)
			}
			continue
		}

		data := body
		if flags&flagEncrypted != 0 && c.encryptionKey != "" {
			plaintext, err := decrypt(body, c.encryptionKey)
			if err != nil {
				continue
			}
			data = plaintext
		}

		if c.onData != nil {
			c.onData(data)
		}
	}

	c.conn = nil
}