Zo doen wijdat in oost

nodes/client/tunnel.go

@@ -0,0 +1,201 @@
+package main
+
+import (
+	"context"
+	"encoding/base64"
+	"net"
+	"strconv"
+	"sync"
+	"time"
+)
+
+type Sender func(frameType string, fields map[string]any) error
+
+type tunnel struct {
+	cancel context.CancelFunc
+	out    chan []byte
+	remote bool
+}
+
+type Tunnels struct {
+	policy  *Policy
+	limiter *Limiter
+	stats   *Stats
+	send    Sender
+	max     int
+	mu      sync.Mutex
+	items   map[string]*tunnel
+	wg      sync.WaitGroup
+}
+
+func NewTunnels(policy *Policy, limiter *Limiter, stats *Stats, send Sender, max int) *Tunnels {
+	return &Tunnels{policy: policy, limiter: limiter, stats: stats, send: send, max: max, items: map[string]*tunnel{}}
+}
+
+func (t *Tunnels) deny(id, code string) {
+	_ = t.send("tunnel_opened", map[string]any{"conn_id": id, "ok": false, "code": code})
+}
+
+func (t *Tunnels) Open(ctx context.Context, frame Frame) {
+	id := frame.ConnID
+	if id == "" || len(id) > MaxIDBytes {
+		return
+	}
+	if _, err := t.policy.CheckHostPort(frame.Host, frame.Port, "tunnel"); err != nil {
+		t.deny(id, ErrorCode(err))
+		t.stats.Refused()
+		return
+	}
+	t.mu.Lock()
+	_, exists := t.items[id]
+	full := len(t.items) >= t.max
+	t.mu.Unlock()
+	if exists || full {
+		t.deny(id, RefuseConcurrency)
+		t.stats.Refused()
+		return
+	}
+	if code := t.limiter.Acquire(); code != "" {
+		t.deny(id, code)
+		t.stats.Refused()
+		return
+	}
+	runCtx, cancel := context.WithTimeout(ctx, time.Duration(TunnelLifetimeSeconds)*time.Second)
+	item := &tunnel{cancel: cancel, out: make(chan []byte, TunnelQueueDepth)}
+	t.mu.Lock()
+	if _, taken := t.items[id]; taken || len(t.items) >= t.max {
+		t.mu.Unlock()
+		cancel()
+		t.limiter.Release()
+		t.deny(id, RefuseConcurrency)
+		t.stats.Refused()
+		return
+	}
+	t.items[id] = item
+	t.mu.Unlock()
+	t.wg.Add(1)
+	go t.run(runCtx, id, frame.Host, frame.Port, item)
+}
+
+func (t *Tunnels) finish(id string, item *tunnel, announce bool) {
+	t.mu.Lock()
+	if t.items[id] == item {
+		delete(t.items, id)
+	}
+	t.mu.Unlock()
+	item.cancel()
+	t.limiter.Release()
+	if announce {
+		_ = t.send("tunnel_close", map[string]any{"conn_id": id})
+	}
+}
+
+func (t *Tunnels) run(ctx context.Context, id, host string, port int, item *tunnel) {
+	defer t.wg.Done()
+	conn, err := t.policy.DialContext(ctx, "tcp", net.JoinHostPort(host, strconv.Itoa(port)))
+	if err != nil {
+		t.finish(id, item, false)
+		t.deny(id, ErrorCode(err))
+		t.stats.Failed()
+		return
+	}
+	t.stats.Tunnel()
+	if err := t.send("tunnel_opened", map[string]any{"conn_id": id, "ok": true}); err != nil {
+		conn.Close()
+		t.finish(id, item, false)
+		return
+	}
+	go func() {
+		<-ctx.Done()
+		conn.Close()
+	}()
+	var writers sync.WaitGroup
+	writers.Add(1)
+	go func() {
+		defer writers.Done()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case chunk := <-item.out:
+				_ = conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+				if _, err := conn.Write(chunk); err != nil {
+					item.cancel()
+					return
+				}
+				t.stats.Bytes(int64(len(chunk)), 0)
+				t.limiter.AddBytes(int64(len(chunk)))
+			}
+		}
+	}()
+	buffer := make([]byte, TunnelChunkBytes)
+	for ctx.Err() == nil {
+		_ = conn.SetReadDeadline(time.Now().Add(time.Duration(TunnelIdleSeconds) * time.Second))
+		n, err := conn.Read(buffer)
+		if n > 0 {
+			t.stats.Bytes(0, int64(n))
+			t.limiter.AddBytes(int64(n))
+			encoded := base64.StdEncoding.EncodeToString(buffer[:n])
+			if sendErr := t.send("tunnel_data", map[string]any{"conn_id": id, "b64": encoded}); sendErr != nil {
+				break
+			}
+			if t.limiter.Remaining() == 0 {
+				break
+			}
+		}
+		if err != nil {
+			break
+		}
+	}
+	conn.Close()
+	item.cancel()
+	writers.Wait()
+	t.mu.Lock()
+	announce := !item.remote
+	t.mu.Unlock()
+	t.finish(id, item, announce)
+}
+
+func (t *Tunnels) Data(frame Frame) {
+	t.mu.Lock()
+	item := t.items[frame.ConnID]
+	t.mu.Unlock()
+	if item == nil {
+		return
+	}
+	if len(frame.B64) > MaxInboundMessage {
+		item.cancel()
+		return
+	}
+	chunk, err := base64.StdEncoding.Strict().DecodeString(frame.B64)
+	if err != nil {
+		item.cancel()
+		return
+	}
+	select {
+	case item.out <- chunk:
+	default:
+		item.cancel()
+	}
+}
+
+func (t *Tunnels) Close(id string) {
+	t.mu.Lock()
+	item := t.items[id]
+	if item != nil {
+		item.remote = true
+	}
+	t.mu.Unlock()
+	if item != nil {
+		item.cancel()
+	}
+}
+
+func (t *Tunnels) CloseAll() {
+	t.mu.Lock()
+	for _, item := range t.items {
+		item.cancel()
+	}
+	t.mu.Unlock()
+	t.wg.Wait()
+}