initial commit
This commit is contained in:
iceBear67
@@ -0,0 +1,324 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/iceBear67/redapricot/client/wire"
|
||||
)
|
||||
|
||||
// WorkerPool manages up to maxConn worker connections and allocates streams
|
||||
// using the least-loaded strategy (PROTOCOL.md §7.1).
|
||||
type WorkerPool struct {
|
||||
client *Client
|
||||
maxConn int
|
||||
|
||||
mu sync.Mutex
|
||||
conns []*WorkerConn
|
||||
}
|
||||
|
||||
func newWorkerPool(c *Client, maxConn int) *WorkerPool {
|
||||
return &WorkerPool{client: c, maxConn: maxConn}
|
||||
}
|
||||
|
||||
// Allocate returns a worker conn and a fresh stream id to place a new stream on.
|
||||
func (p *WorkerPool) Allocate() (*WorkerConn, int, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
var best *WorkerConn
|
||||
bestCount := 0
|
||||
for _, wc := range p.conns {
|
||||
n := wc.streamCount()
|
||||
if best == nil || n < bestCount {
|
||||
best = wc
|
||||
bestCount = n
|
||||
}
|
||||
}
|
||||
|
||||
needNew := best == nil || (bestCount > SaturationThreshold && len(p.conns) < p.maxConn)
|
||||
if needNew {
|
||||
wc, err := p.dialWorker()
|
||||
if err != nil {
|
||||
if best == nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
log.Printf("worker dial failed, reusing existing conn: %v", err)
|
||||
} else {
|
||||
p.conns = append(p.conns, wc)
|
||||
best = wc
|
||||
}
|
||||
}
|
||||
return best, best.newSid(), nil
|
||||
}
|
||||
|
||||
func (p *WorkerPool) dialWorker() (*WorkerConn, error) {
|
||||
fc, err := p.client.dialSession(MagicWorker)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
wc := &WorkerConn{
|
||||
pool: p,
|
||||
fc: fc,
|
||||
streams: make(map[int]*Stream),
|
||||
nextSid: 1,
|
||||
}
|
||||
go wc.readLoop()
|
||||
log.Printf("opened worker conn (#%d in pool)", len(p.conns)+1)
|
||||
return wc, nil
|
||||
}
|
||||
|
||||
func (p *WorkerPool) count() int {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return len(p.conns)
|
||||
}
|
||||
|
||||
func (p *WorkerPool) remove(wc *WorkerConn) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
for i, c := range p.conns {
|
||||
if c == wc {
|
||||
p.conns = append(p.conns[:i], p.conns[i+1:]...)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *WorkerPool) closeAll() {
|
||||
p.mu.Lock()
|
||||
conns := append([]*WorkerConn(nil), p.conns...)
|
||||
p.mu.Unlock()
|
||||
for _, wc := range conns {
|
||||
_ = wc.fc.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// WorkerConn is one multiplexed worker connection to the hub.
|
||||
type WorkerConn struct {
|
||||
pool *WorkerPool
|
||||
fc *wire.FramedConn
|
||||
|
||||
mu sync.Mutex
|
||||
streams map[int]*Stream
|
||||
nextSid int
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) streamCount() int {
|
||||
wc.mu.Lock()
|
||||
defer wc.mu.Unlock()
|
||||
return len(wc.streams)
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) newSid() int {
|
||||
wc.mu.Lock()
|
||||
defer wc.mu.Unlock()
|
||||
sid := wc.nextSid
|
||||
wc.nextSid++
|
||||
return sid
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) registerStream(sid int, st *Stream) {
|
||||
wc.mu.Lock()
|
||||
wc.streams[sid] = st
|
||||
wc.mu.Unlock()
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) getStream(sid int) *Stream {
|
||||
wc.mu.Lock()
|
||||
defer wc.mu.Unlock()
|
||||
return wc.streams[sid]
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) removeStream(sid int) *Stream {
|
||||
wc.mu.Lock()
|
||||
defer wc.mu.Unlock()
|
||||
st := wc.streams[sid]
|
||||
delete(wc.streams, sid)
|
||||
return st
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) readLoop() {
|
||||
for {
|
||||
payload, err := wc.fc.ReadFrame()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
r := wire.NewReader(payload)
|
||||
ftype, err := r.U8()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
sid, err := r.VarInt()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
switch ftype {
|
||||
case MuxData:
|
||||
if st := wc.getStream(sid); st != nil {
|
||||
st.deliverFromHub(r.Remaining())
|
||||
}
|
||||
case MuxFin, MuxRst:
|
||||
if st := wc.removeStream(sid); st != nil {
|
||||
st.shutdown(false)
|
||||
}
|
||||
default:
|
||||
log.Printf("worker: unknown mux type %d", ftype)
|
||||
}
|
||||
}
|
||||
// Connection lost: tear down all streams and drop from pool.
|
||||
wc.pool.remove(wc)
|
||||
wc.mu.Lock()
|
||||
streams := make([]*Stream, 0, len(wc.streams))
|
||||
for _, st := range wc.streams {
|
||||
streams = append(streams, st)
|
||||
}
|
||||
wc.streams = make(map[int]*Stream)
|
||||
wc.mu.Unlock()
|
||||
for _, st := range streams {
|
||||
st.shutdown(false)
|
||||
}
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) sendSyn(sid int, cid []byte) {
|
||||
_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxSyn).VarInt(sid).Bytes(cid).Out())
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) sendData(sid int, data []byte) error {
|
||||
return wc.fc.WriteFrame(wire.NewWriter().U8(MuxData).VarInt(sid).Bytes(data).Out())
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) sendFin(sid int) {
|
||||
_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxFin).VarInt(sid).Out())
|
||||
}
|
||||
|
||||
func (wc *WorkerConn) sendRst(sid int) {
|
||||
_ = wc.fc.WriteFrame(wire.NewWriter().U8(MuxRst).VarInt(sid).Out())
|
||||
}
|
||||
|
||||
// Stream bridges one player (via the hub) to one destination connection.
|
||||
type Stream struct {
|
||||
wc *WorkerConn
|
||||
sid int
|
||||
cid []byte
|
||||
mapping Mapping
|
||||
srcIP string
|
||||
srcPort int
|
||||
|
||||
mu sync.Mutex
|
||||
dest net.Conn
|
||||
connected bool
|
||||
preBuf []byte
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newStream(wc *WorkerConn, sid int, cid []byte, m Mapping, ip string, port int) *Stream {
|
||||
return &Stream{wc: wc, sid: sid, cid: cid, mapping: m, srcIP: ip, srcPort: port}
|
||||
}
|
||||
|
||||
// run dials the destination, optionally writes the PROXY v2 header, flushes any
|
||||
// buffered hub bytes, then pumps destination -> hub.
|
||||
func (s *Stream) run() {
|
||||
dest, err := net.DialTimeout("tcp", s.mapping.Destination, 10*time.Second)
|
||||
if err != nil {
|
||||
log.Printf("stream %d: dial %s failed: %v", s.sid, s.mapping.Destination, err)
|
||||
s.wc.removeStream(s.sid)
|
||||
s.wc.sendRst(s.sid)
|
||||
return
|
||||
}
|
||||
if tcp, ok := dest.(*net.TCPConn); ok {
|
||||
_ = tcp.SetNoDelay(true)
|
||||
}
|
||||
|
||||
if s.mapping.ProxyProtocol {
|
||||
if hdr := s.buildProxyHeader(dest); hdr != nil {
|
||||
if _, err := dest.Write(hdr); err != nil {
|
||||
log.Printf("stream %d: proxy header write: %v", s.sid, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Atomically flush pre-connect buffer and enable direct writes.
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
_ = dest.Close()
|
||||
return
|
||||
}
|
||||
s.dest = dest
|
||||
if len(s.preBuf) > 0 {
|
||||
_, _ = dest.Write(s.preBuf)
|
||||
s.preBuf = nil
|
||||
}
|
||||
s.connected = true
|
||||
s.mu.Unlock()
|
||||
|
||||
// destination -> hub
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
n, err := dest.Read(buf)
|
||||
if n > 0 {
|
||||
if werr := s.wc.sendData(s.sid, buf[:n]); werr != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
s.shutdown(true)
|
||||
}
|
||||
|
||||
func (s *Stream) buildProxyHeader(dest net.Conn) []byte {
|
||||
srcIP := net.ParseIP(s.srcIP)
|
||||
if srcIP == nil {
|
||||
return nil
|
||||
}
|
||||
dstTCP, ok := dest.RemoteAddr().(*net.TCPAddr)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return BuildProxyV2(srcIP, s.srcPort, dstTCP.IP, dstTCP.Port)
|
||||
}
|
||||
|
||||
// deliverFromHub writes bytes coming from the hub to the destination, buffering
|
||||
// until the destination connection is established.
|
||||
func (s *Stream) deliverFromHub(data []byte) {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if !s.connected {
|
||||
s.preBuf = append(s.preBuf, data...)
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
dest := s.dest
|
||||
s.mu.Unlock()
|
||||
if _, err := dest.Write(data); err != nil {
|
||||
s.shutdown(true)
|
||||
}
|
||||
}
|
||||
|
||||
// shutdown closes the stream; notifyHub sends a FIN to the hub when true.
|
||||
func (s *Stream) shutdown(notifyHub bool) {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
dest := s.dest
|
||||
s.mu.Unlock()
|
||||
|
||||
if dest != nil {
|
||||
_ = dest.Close()
|
||||
}
|
||||
s.wc.removeStream(s.sid)
|
||||
if notifyHub {
|
||||
s.wc.sendFin(s.sid)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user