framing.go

package netpipe

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
)

// Wire format (v2 - Stage 4):
//
//   [4 bytes: payload length as big-endian uint32][1 byte: flags][N bytes: body]
//
//   length = 1 (flags) + len(body)
//
// Flags byte:
//   bit 0 (0x01): encrypted - body is [12B nonce][AES-GCM ciphertext+tag]
//   bits 1-7:     reserved for future use
//
// The developer never sees any of this. Send()/SendEncrypted() frame
// automatically, readMessage() returns the decoded payload + flags.

const (
	// headerSize is the 4-byte uint32 length prefix.
	headerSize = 4

	// flagSize is the 1-byte flags field after the length.
	flagSize = 1

	// MaxMessageSize caps a single message at 16 MB.
	MaxMessageSize = 16 * 1024 * 1024

	// flagEncrypted marks a message as AES-256-GCM encrypted.
	flagEncrypted byte = 0x01

	// flagReject is sent by the server before closing a connection it cannot accept.
	// Body contains a short UTF-8 reason string.
	flagReject byte = 0x10

	// flagHandshake is used in P2P mode for the Diffie-Hellman key exchange.
	// Body contains [32B X25519 public key][16B peer UUID].
	flagHandshake byte = 0x20

	// flagSession is used for session resumption on reconnect.
	// Server → client: body is [16B session UUID] (assigned on first connect)
	// Client → server: body is [16B session UUID] (sent on reconnect to resume)
	flagSession byte = 0x40
)

// writeMessage writes a framed message with flags to the connection.
// If conn is a *safeConn, the write is mutex-protected against interleaving.
func writeMessage(conn net.Conn, flags byte, body []byte) error {
	if len(body) == 0 && flags == 0 {
		return nil
	}
	if len(body) > MaxMessageSize {
		return fmt.Errorf("netpipe: message too large (%d bytes, max %d)", len(body), MaxMessageSize)
	}

	totalLen := flagSize + len(body) // what follows the 4-byte header
	frame := make([]byte, headerSize+totalLen)
	binary.BigEndian.PutUint32(frame[:headerSize], uint32(totalLen))
	frame[headerSize] = flags
	copy(frame[headerSize+flagSize:], body)

	// use mutex-protected write if available (prevents frame interleaving)
	if sc, ok := conn.(*safeConn); ok {
		_, err := sc.safeWrite(frame)
		return err
	}
	_, err := conn.Write(frame)
	return err
}

// readMessage reads one complete framed message and returns (flags, body, error).
func readMessage(conn net.Conn) (byte, []byte, error) {
	// 1. read 4-byte length header
	header := make([]byte, headerSize)
	if _, err := io.ReadFull(conn, header); err != nil {
		return 0, nil, err
	}

	length := binary.BigEndian.Uint32(header)
	if length == 0 {
		return 0, []byte{}, nil
	}
	if length > MaxMessageSize {
		return 0, nil, fmt.Errorf("netpipe: incoming message too large (%d bytes, max %d)", length, MaxMessageSize)
	}

	// 2. read flags + body
	payload := make([]byte, length)
	if _, err := io.ReadFull(conn, payload); err != nil {
		return 0, nil, err
	}

	flags := payload[0]
	body := payload[flagSize:]

	return flags, body, nil
}

// ---------------------------------------------------------------------------
// Convenience wrappers - keep the simple API surface clean
// ---------------------------------------------------------------------------

// writeFrame writes a plain (unencrypted) framed message.
func writeFrame(conn net.Conn, data []byte) error {
	return writeMessage(conn, 0x00, data)
}

// writeFrameEncrypted encrypts data with the given key then writes
// a framed message with the encrypted flag set.
func writeFrameEncrypted(conn net.Conn, data []byte, key string) error {
	ciphertext, err := encrypt(data, key)
	if err != nil {
		return err
	}
	return writeMessage(conn, flagEncrypted, ciphertext)
}

// readFrame reads one framed message, ignoring flags (backwards compat).
func readFrame(conn net.Conn) ([]byte, error) {
	_, body, err := readMessage(conn)
	return body, err
}

// readFrameFull reads one framed message and returns flags + body.
func readFrameFull(conn net.Conn) (byte, []byte, error) {
	return readMessage(conn)
}