Add packet reactor

Author: nekohasekai

common/bufio/channel_demux.go

@@ -0,0 +1,143 @@
+package bufio
+
+import (
+	"context"
+	"reflect"
+	"sync"
+	"sync/atomic"
+
+	N "github.com/sagernet/sing/common/network"
+)
+
+type channelDemuxEntry struct {
+	channel <-chan *N.PacketBuffer
+	stream  *reactorStream
+}
+
+type ChannelDemultiplexer struct {
+	ctx        context.Context
+	cancel     context.CancelFunc
+	mutex      sync.Mutex
+	entries    map[<-chan *N.PacketBuffer]*channelDemuxEntry
+	updateChan chan struct{}
+	running    bool
+	closed     atomic.Bool
+	wg         sync.WaitGroup
+}
+
+func NewChannelDemultiplexer(ctx context.Context) *ChannelDemultiplexer {
+	ctx, cancel := context.WithCancel(ctx)
+	demux := &ChannelDemultiplexer{
+		ctx:        ctx,
+		cancel:     cancel,
+		entries:    make(map[<-chan *N.PacketBuffer]*channelDemuxEntry),
+		updateChan: make(chan struct{}, 1),
+	}
+	return demux
+}
+
+func (p *ChannelDemultiplexer) Add(stream *reactorStream, channel <-chan *N.PacketBuffer) {
+	p.mutex.Lock()
+
+	if p.closed.Load() {
+		p.mutex.Unlock()
+		return
+	}
+
+	entry := &channelDemuxEntry{
+		channel: channel,
+		stream:  stream,
+	}
+	p.entries[channel] = entry
+	if !p.running {
+		p.running = true
+		p.wg.Add(1)
+		go p.run()
+	}
+	p.mutex.Unlock()
+	p.signalUpdate()
+}
+
+func (p *ChannelDemultiplexer) Remove(channel <-chan *N.PacketBuffer) {
+	p.mutex.Lock()
+	delete(p.entries, channel)
+	p.mutex.Unlock()
+	p.signalUpdate()
+}
+
+func (p *ChannelDemultiplexer) signalUpdate() {
+	select {
+	case p.updateChan <- struct{}{}:
+	default:
+	}
+}
+
+func (p *ChannelDemultiplexer) Close() error {
+	p.mutex.Lock()
+	p.closed.Store(true)
+	p.mutex.Unlock()
+
+	p.cancel()
+	p.signalUpdate()
+	p.wg.Wait()
+	return nil
+}
+
+func (p *ChannelDemultiplexer) run() {
+	defer p.wg.Done()
+
+	for {
+		p.mutex.Lock()
+		if len(p.entries) == 0 {
+			p.running = false
+			p.mutex.Unlock()
+			return
+		}
+
+		cases := make([]reflect.SelectCase, 0, len(p.entries)+2)
+
+		cases = append(cases, reflect.SelectCase{
+			Dir:  reflect.SelectRecv,
+			Chan: reflect.ValueOf(p.ctx.Done()),
+		})
+
+		cases = append(cases, reflect.SelectCase{
+			Dir:  reflect.SelectRecv,
+			Chan: reflect.ValueOf(p.updateChan),
+		})
+
+		entryList := make([]*channelDemuxEntry, 0, len(p.entries))
+		for _, entry := range p.entries {
+			cases = append(cases, reflect.SelectCase{
+				Dir:  reflect.SelectRecv,
+				Chan: reflect.ValueOf(entry.channel),
+			})
+			entryList = append(entryList, entry)
+		}
+		p.mutex.Unlock()
+
+		chosen, recv, recvOK := reflect.Select(cases)
+
+		switch chosen {
+		case 0:
+			p.mutex.Lock()
+			p.running = false
+			p.mutex.Unlock()
+			return
+		case 1:
+			continue
+		default:
+			entry := entryList[chosen-2]
+			p.mutex.Lock()
+			delete(p.entries, entry.channel)
+			p.mutex.Unlock()
+
+			if recvOK {
+				packet := recv.Interface().(*N.PacketBuffer)
+				go entry.stream.runActiveLoop(packet)
+			} else {
+				go entry.stream.closeWithError(nil)
+			}
+		}
+	}
+}

common/bufio/fd_demux_darwin.go

@@ -0,0 +1,213 @@
+//go:build darwin
+
+package bufio
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+
+	"golang.org/x/sys/unix"
+)
+
+type fdDemuxEntry struct {
+	fd     int
+	stream *reactorStream
+}
+
+type FDDemultiplexer struct {
+	ctx     context.Context
+	cancel  context.CancelFunc
+	kqFD    int
+	mutex   sync.Mutex
+	entries map[int]*fdDemuxEntry
+	running bool
+	closed  atomic.Bool
+	wg      sync.WaitGroup
+	pipeFDs [2]int
+}
+
+func NewFDDemultiplexer(ctx context.Context) (*FDDemultiplexer, error) {
+	kqFD, err := unix.Kqueue()
+	if err != nil {
+		return nil, err
+	}
+
+	var pipeFDs [2]int
+	err = unix.Pipe(pipeFDs[:])
+	if err != nil {
+		unix.Close(kqFD)
+		return nil, err
+	}
+
+	unix.SetNonblock(pipeFDs[0], true)
+	unix.SetNonblock(pipeFDs[1], true)
+
+	_, err = unix.Kevent(kqFD, []unix.Kevent_t{{
+		Ident:  uint64(pipeFDs[0]),
+		Filter: unix.EVFILT_READ,
+		Flags:  unix.EV_ADD,
+	}}, nil, nil)
+	if err != nil {
+		unix.Close(pipeFDs[0])
+		unix.Close(pipeFDs[1])
+		unix.Close(kqFD)
+		return nil, err
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	demux := &FDDemultiplexer{
+		ctx:     ctx,
+		cancel:  cancel,
+		kqFD:    kqFD,
+		entries: make(map[int]*fdDemuxEntry),
+		pipeFDs: pipeFDs,
+	}
+	return demux, nil
+}
+
+func (p *FDDemultiplexer) Add(stream *reactorStream, fd int) error {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	if p.closed.Load() {
+		return unix.EINVAL
+	}
+
+	_, err := unix.Kevent(p.kqFD, []unix.Kevent_t{{
+		Ident:  uint64(fd),
+		Filter: unix.EVFILT_READ,
+		Flags:  unix.EV_ADD,
+	}}, nil, nil)
+	if err != nil {
+		return err
+	}
+
+	entry := &fdDemuxEntry{
+		fd:     fd,
+		stream: stream,
+	}
+	p.entries[fd] = entry
+
+	if !p.running {
+		p.running = true
+		p.wg.Add(1)
+		go p.run()
+	}
+
+	return nil
+}
+
+func (p *FDDemultiplexer) Remove(fd int) {
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	_, ok := p.entries[fd]
+	if !ok {
+		return
+	}
+
+	unix.Kevent(p.kqFD, []unix.Kevent_t{{
+		Ident:  uint64(fd),
+		Filter: unix.EVFILT_READ,
+		Flags:  unix.EV_DELETE,
+	}}, nil, nil)
+	delete(p.entries, fd)
+}
+
+func (p *FDDemultiplexer) wakeup() {
+	unix.Write(p.pipeFDs[1], []byte{0})
+}
+
+func (p *FDDemultiplexer) Close() error {
+	p.mutex.Lock()
+	p.closed.Store(true)
+	p.mutex.Unlock()
+
+	p.cancel()
+	p.wakeup()
+	p.wg.Wait()
+
+	p.mutex.Lock()
+	defer p.mutex.Unlock()
+
+	if p.kqFD != -1 {
+		unix.Close(p.kqFD)
+		p.kqFD = -1
+	}
+	if p.pipeFDs[0] != -1 {
+		unix.Close(p.pipeFDs[0])
+		unix.Close(p.pipeFDs[1])
+		p.pipeFDs[0] = -1
+		p.pipeFDs[1] = -1
+	}
+	return nil
+}
+
+func (p *FDDemultiplexer) run() {
+	defer p.wg.Done()
+
+	events := make([]unix.Kevent_t, 64)
+	var buf [1]byte
+
+	for {
+		select {
+		case <-p.ctx.Done():
+			p.mutex.Lock()
+			p.running = false
+			p.mutex.Unlock()
+			return
+		default:
+		}
+
+		n, err := unix.Kevent(p.kqFD, nil, events, nil)
+		if err != nil {
+			if err == unix.EINTR {
+				continue
+			}
+			p.mutex.Lock()
+			p.running = false
+			p.mutex.Unlock()
+			return
+		}
+
+		for i := 0; i < n; i++ {
+			ev := events[i]
+			fd := int(ev.Ident)
+
+			if fd == p.pipeFDs[0] {
+				unix.Read(p.pipeFDs[0], buf[:])
+				continue
+			}
+
+			if ev.Flags&unix.EV_ERROR != 0 {
+				continue
+			}
+
+			p.mutex.Lock()
+			entry, ok := p.entries[fd]
+			if !ok {
+				p.mutex.Unlock()
+				continue
+			}
+
+			unix.Kevent(p.kqFD, []unix.Kevent_t{{
+				Ident:  uint64(fd),
+				Filter: unix.EVFILT_READ,
+				Flags:  unix.EV_DELETE,
+			}}, nil, nil)
+			delete(p.entries, fd)
+			p.mutex.Unlock()
+
+			go entry.stream.runActiveLoop(nil)
+		}
+
+		p.mutex.Lock()
+		if len(p.entries) == 0 {
+			p.running = false
+			p.mutex.Unlock()
+			return
+		}
+		p.mutex.Unlock()
+	}
+}