stream.go

package netpipe

import (
	"encoding/binary"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/google/uuid"
)

// Stream wire format (per chunk, inside the framing body):
//
//   [16 bytes: stream UUID][4 bytes: chunk index][4 bytes: total chunks][N bytes: chunk data]
//
// The flagStream bit is set on every chunk frame. The receiver buffers
// chunks keyed by stream UUID and fires the OnStream callback once all
// chunks for a stream have arrived.

const (
	// flagStream marks a frame as part of a chunked stream.
	flagStream byte = 0x02

	// streamHeaderSize = 16 (UUID) + 4 (chunk index) + 4 (total chunks)
	streamHeaderSize = 24

	// DefaultChunkSize is the default chunk body size (64 KB).
	DefaultChunkSize = 64 * 1024
)

// ---------------------------------------------------------------------------
// Stream reassembly buffer - used by both server and client
// ---------------------------------------------------------------------------

// streamBuffer collects chunks for in-progress streams.
type streamBuffer struct {
	mu              sync.Mutex
	streams         map[string]*streamState // key = stream UUID hex
	perClient       map[string]int          // key = client/peer ID → active stream count
	maxPerClient    int                     // 0 = unlimited
}

// streamState tracks one in-progress stream.
type streamState struct {
	total     uint32
	received  uint32
	chunks    map[uint32][]byte // index → chunk data
	createdAt time.Time
	clientID  string // who started this stream
}

// maxPendingStreams limits how many incomplete streams can exist at once.
const maxPendingStreams = 256

// streamTimeout is how long an incomplete stream can sit before being cleaned up.
const streamTimeout = 60 * time.Second

func newStreamBuffer() *streamBuffer {
	return &streamBuffer{
		streams:   make(map[string]*streamState),
		perClient: make(map[string]int),
	}
}

// addChunk stores a chunk and returns the fully reassembled data if this
// was the last missing chunk. Returns (nil, false) if still waiting.
// clientID is used for per-client stream limits.
func (sb *streamBuffer) addChunk(streamID string, index, total uint32, data []byte) ([]byte, bool) {
	return sb.addChunkFrom("", streamID, index, total, data)
}

// addChunkFrom is like addChunk but tracks which client sent it.
func (sb *streamBuffer) addChunkFrom(clientID, streamID string, index, total uint32, data []byte) ([]byte, bool) {
	if total == 0 || index >= total {
		return nil, false
	}

	sb.mu.Lock()
	defer sb.mu.Unlock()

	// clean up expired streams
	now := time.Now()
	for id, s := range sb.streams {
		if now.Sub(s.createdAt) > streamTimeout {
			if s.clientID != "" {
				sb.perClient[s.clientID]--
				if sb.perClient[s.clientID] <= 0 {
					delete(sb.perClient, s.clientID)
				}
			}
			delete(sb.streams, id)
		}
	}

	state, ok := sb.streams[streamID]
	if !ok {
		// reject if global limit hit
		if len(sb.streams) >= maxPendingStreams {
			return nil, false
		}

		// reject if per-client limit hit
		if sb.maxPerClient > 0 && clientID != "" && sb.perClient[clientID] >= sb.maxPerClient {
			return nil, false
		}

		state = &streamState{
			total:     total,
			chunks:    make(map[uint32][]byte, total),
			createdAt: now,
			clientID:  clientID,
		}
		sb.streams[streamID] = state
		if clientID != "" {
			sb.perClient[clientID]++
		}
	}

	if state.total != total {
		return nil, false
	}

	if _, exists := state.chunks[index]; exists {
		return nil, false
	}

	chunkCopy := make([]byte, len(data))
	copy(chunkCopy, data)
	state.chunks[index] = chunkCopy
	state.received++

	if state.received < state.total {
		return nil, false
	}

	// reassemble in order
	totalSize := 0
	for _, c := range state.chunks {
		totalSize += len(c)
	}

	assembled := make([]byte, 0, totalSize)
	for i := uint32(0); i < state.total; i++ {
		assembled = append(assembled, state.chunks[i]...)
	}

	// clean up
	if state.clientID != "" {
		sb.perClient[state.clientID]--
		if sb.perClient[state.clientID] <= 0 {
			delete(sb.perClient, state.clientID)
		}
	}
	delete(sb.streams, streamID)

	return assembled, true
}

// ---------------------------------------------------------------------------
// Send-side: split data into chunks and write
// ---------------------------------------------------------------------------

// writeStream splits data into chunks and sends each as a flagStream frame.
// TCP only - uses writeMessage for framing.
func writeStream(conn net.Conn, data []byte, chunkSize int) error {
	if chunkSize <= 0 {
		chunkSize = DefaultChunkSize
	}

	streamID := uuid.New()
	idBytes, err := streamID.MarshalBinary() // 16 bytes
	if err != nil {
		return fmt.Errorf("netpipe: stream uuid: %w", err)
	}

	totalChunks := (len(data) + chunkSize - 1) / chunkSize
	if totalChunks == 0 {
		totalChunks = 1 // send at least one chunk for empty data
	}

	for i := 0; i < totalChunks; i++ {
		start := i * chunkSize
		end := start + chunkSize
		if end > len(data) {
			end = len(data)
		}
		chunk := data[start:end]

		// build stream header: [16B UUID][4B index][4B total]
		header := make([]byte, streamHeaderSize)
		copy(header[:16], idBytes)
		binary.BigEndian.PutUint32(header[16:20], uint32(i))
		binary.BigEndian.PutUint32(header[20:24], uint32(totalChunks))

		// body = stream header + chunk data
		body := make([]byte, streamHeaderSize+len(chunk))
		copy(body[:streamHeaderSize], header)
		copy(body[streamHeaderSize:], chunk)

		if err := writeMessage(conn, flagStream, body); err != nil {
			return fmt.Errorf("netpipe: stream chunk %d/%d: %w", i+1, totalChunks, err)
		}
	}

	return nil
}

// writeStreamUDP splits data into chunks and sends each as a UDP packet
// with the flagStream flag.
func writeStreamUDPClient(conn net.Conn, data []byte, chunkSize int) error {
	if chunkSize <= 0 {
		chunkSize = DefaultChunkSize
	}

	streamID := uuid.New()
	idBytes, err := streamID.MarshalBinary()
	if err != nil {
		return fmt.Errorf("netpipe: stream uuid: %w", err)
	}

	totalChunks := (len(data) + chunkSize - 1) / chunkSize
	if totalChunks == 0 {
		totalChunks = 1
	}

	for i := 0; i < totalChunks; i++ {
		start := i * chunkSize
		end := start + chunkSize
		if end > len(data) {
			end = len(data)
		}
		chunk := data[start:end]

		header := make([]byte, streamHeaderSize)
		copy(header[:16], idBytes)
		binary.BigEndian.PutUint32(header[16:20], uint32(i))
		binary.BigEndian.PutUint32(header[20:24], uint32(totalChunks))

		// UDP packet: [1B flags][24B stream header][NB chunk]
		packet := make([]byte, 1+streamHeaderSize+len(chunk))
		packet[0] = flagStream
		copy(packet[1:1+streamHeaderSize], header)
		copy(packet[1+streamHeaderSize:], chunk)

		if _, err := conn.Write(packet); err != nil {
			return fmt.Errorf("netpipe: stream udp chunk %d/%d: %w", i+1, totalChunks, err)
		}
	}

	return nil
}

// ---------------------------------------------------------------------------
// Receive-side: parse stream header from body
// ---------------------------------------------------------------------------

// parseStreamBody extracts the stream UUID, chunk index, total, and chunk data
// from a flagStream frame body.
func parseStreamBody(body []byte) (streamID string, index, total uint32, chunk []byte, err error) {
	if len(body) < streamHeaderSize {
		return "", 0, 0, nil, fmt.Errorf("netpipe: stream body too short (%d bytes)", len(body))
	}

	var uid uuid.UUID
	if err := uid.UnmarshalBinary(body[:16]); err != nil {
		return "", 0, 0, nil, fmt.Errorf("netpipe: stream uuid parse: %w", err)
	}

	index = binary.BigEndian.Uint32(body[16:20])
	total = binary.BigEndian.Uint32(body[20:24])
	chunk = body[streamHeaderSize:]

	return uid.String(), index, total, chunk, nil
}