vk-turn-proxy/internal/dcmux/mux.go

package dcmux

import (
	"encoding/binary"
	"sync"
)

type Stream struct {
	ID         uint16
	ClientID   uint32
	recvBuf    []byte
	closed     bool
	nextSeq    uint32
	outOfOrder map[uint32][]byte
}

type Multiplexer struct {
	streams       map[uint16]*Stream
	nextID        uint16
	clientID      uint32
	onSend        func([]byte) error
	mu            sync.RWMutex
	maxStreams    int
	maxBufferSize int
	activityCh    chan struct{}
	dataReady     map[uint16]chan struct{}
	dataReadyMu   sync.Mutex
	sendSeq       map[uint16]uint32
	sendSeqMu     sync.Mutex
}

func New(clientID uint32, onSend func([]byte) error) *Multiplexer {
	return &Multiplexer{
		streams:       make(map[uint16]*Stream),
		nextID:        1,
		clientID:      clientID,
		onSend:        onSend,
		maxStreams:    10000,
		maxBufferSize: 16 * 1024 * 1024,
		activityCh:    make(chan struct{}, 1),
		dataReady:     make(map[uint16]chan struct{}),
		sendSeq:       make(map[uint16]uint32),
	}
}

func (m *Multiplexer) OpenStream() uint16 {
	m.mu.Lock()
	defer m.mu.Unlock()

	for {
		sid := m.nextID
		m.nextID++
		if m.nextID == 0 {
			m.nextID = 1
		}

		if _, exists := m.streams[sid]; !exists {
			m.streams[sid] = &Stream{
				ID:         sid,
				recvBuf:    make([]byte, 0),
				nextSeq:    0,
				outOfOrder: make(map[uint32][]byte),
			}
			return sid
		}
	}
}

func (m *Multiplexer) SendData(sid uint16, data []byte) error {
	m.mu.RLock()
	stream, exists := m.streams[sid]
	m.mu.RUnlock()

	if !exists || stream.closed {
		return nil
	}

	const chunkSize = 7168
	for i := 0; i < len(data); i += chunkSize {
		end := i + chunkSize
		if end > len(data) {
			end = len(data)
		}

		chunk := data[i:end]

		m.sendSeqMu.Lock()
		seq := m.sendSeq[sid]
		m.sendSeq[sid]++
		m.sendSeqMu.Unlock()

		frame := make([]byte, 12+len(chunk))
		binary.BigEndian.PutUint32(frame[0:4], m.clientID)
		binary.BigEndian.PutUint16(frame[4:6], sid)
		binary.BigEndian.PutUint16(frame[6:8], uint16(len(chunk)))
		binary.BigEndian.PutUint32(frame[8:12], seq)
		copy(frame[12:], chunk)

		if err := m.onSend(frame); err != nil {
			return err
		}
	}

	return nil
}

func (m *Multiplexer) CloseStream(sid uint16) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if stream, exists := m.streams[sid]; exists {
		stream.closed = true
	}

	m.sendSeqMu.Lock()
	delete(m.sendSeq, sid)
	m.sendSeqMu.Unlock()

	m.signalDataReady(sid)

	frame := make([]byte, 12)
	binary.BigEndian.PutUint32(frame[0:4], m.clientID)
	binary.BigEndian.PutUint16(frame[4:6], sid)
	binary.BigEndian.PutUint16(frame[6:8], 0)
	binary.BigEndian.PutUint32(frame[8:12], 0)

	return m.onSend(frame)
}

func (m *Multiplexer) HandleFrame(frame []byte) {
	if len(frame) < 12 {
		return
	}

	clientID := binary.BigEndian.Uint32(frame[0:4])
	sid := binary.BigEndian.Uint16(frame[4:6])
	length := binary.BigEndian.Uint16(frame[6:8])
	seq := binary.BigEndian.Uint32(frame[8:12])

	if length == 0 {
		m.mu.Lock()
		if stream, exists := m.streams[sid]; exists && stream.ClientID == clientID {
			stream.closed = true
		}
		m.mu.Unlock()
		m.signalDataReady(sid)
		return
	}

	if len(frame) < 12+int(length) {
		return
	}

	data := frame[12 : 12+length]

	m.mu.Lock()
	defer m.mu.Unlock()

	stream, exists := m.streams[sid]
	if !exists {
		if len(m.streams) >= m.maxStreams {
			return
		}
		stream = &Stream{
			ID:         sid,
			ClientID:   clientID,
			recvBuf:    make([]byte, 0),
			nextSeq:    0,
			outOfOrder: make(map[uint32][]byte),
		}
		m.streams[sid] = stream
	} else if stream.ClientID != clientID {
		stream.ClientID = clientID
		stream.recvBuf = make([]byte, 0)
		stream.closed = false
		stream.nextSeq = 0
		stream.outOfOrder = make(map[uint32][]byte)
	}

	if seq == stream.nextSeq {
		if len(stream.recvBuf)+len(data) > m.maxBufferSize {
			stream.closed = true
			m.signalDataReady(sid)
			return
		}
		stream.recvBuf = append(stream.recvBuf, data...)
		stream.nextSeq++

		for {
			nextData, ok := stream.outOfOrder[stream.nextSeq]
			if !ok {
				break
			}
			if len(stream.recvBuf)+len(nextData) > m.maxBufferSize {
				stream.closed = true
				m.signalDataReady(sid)
				return
			}
			stream.recvBuf = append(stream.recvBuf, nextData...)
			delete(stream.outOfOrder, stream.nextSeq)
			stream.nextSeq++
		}

		m.signalDataReady(sid)
		return
	}

	if seq > stream.nextSeq && len(stream.outOfOrder) < 100 {
		stream.outOfOrder[seq] = append([]byte(nil), data...)
	}
}

func (m *Multiplexer) ReadStream(sid uint16) []byte {
	m.mu.Lock()
	defer m.mu.Unlock()

	stream, exists := m.streams[sid]
	if !exists || len(stream.recvBuf) == 0 {
		return nil
	}

	data := stream.recvBuf
	stream.recvBuf = make([]byte, 0)
	return data
}

func (m *Multiplexer) StreamClosed(sid uint16) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	stream, exists := m.streams[sid]
	return !exists || stream.closed
}

func (m *Multiplexer) GetStreams() []uint16 {
	m.mu.RLock()
	defer m.mu.RUnlock()

	sids := make([]uint16, 0, len(m.streams))
	for sid := range m.streams {
		sids = append(sids, sid)
	}
	return sids
}

func (m *Multiplexer) WaitForData(sid uint16) <-chan struct{} {
	m.dataReadyMu.Lock()
	defer m.dataReadyMu.Unlock()

	if _, ok := m.dataReady[sid]; !ok {
		m.dataReady[sid] = make(chan struct{}, 1)
	}
	return m.dataReady[sid]
}

func (m *Multiplexer) WaitForActivity() <-chan struct{} {
	return m.activityCh
}

func (m *Multiplexer) CleanupStream(sid uint16) {
	m.mu.Lock()
	delete(m.streams, sid)
	m.mu.Unlock()

	m.sendSeqMu.Lock()
	delete(m.sendSeq, sid)
	m.sendSeqMu.Unlock()

	m.dataReadyMu.Lock()
	defer m.dataReadyMu.Unlock()

	if ch, ok := m.dataReady[sid]; ok {
		close(ch)
		delete(m.dataReady, sid)
	}
}

func (m *Multiplexer) Reset() {
	m.mu.Lock()
	for _, stream := range m.streams {
		stream.closed = true
	}
	m.streams = make(map[uint16]*Stream)
	m.nextID = 1
	m.mu.Unlock()

	m.sendSeqMu.Lock()
	m.sendSeq = make(map[uint16]uint32)
	m.sendSeqMu.Unlock()

	m.dataReadyMu.Lock()
	for sid, ch := range m.dataReady {
		close(ch)
		delete(m.dataReady, sid)
	}
	m.dataReadyMu.Unlock()

	m.signalActivity()
}

func (m *Multiplexer) signalDataReady(sid uint16) {
	m.signalActivity()

	m.dataReadyMu.Lock()
	defer m.dataReadyMu.Unlock()

	ch, ok := m.dataReady[sid]
	if !ok {
		return
	}

	select {
	case ch <- struct{}{}:
	default:
	}
}

func (m *Multiplexer) signalActivity() {
	select {
	case m.activityCh <- struct{}{}:
	default:
	}
}