devplacepy/nodes/client/tunnel.go

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()
}