devplacepy/nodes/client/session_test.go

package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

const testSecret = "dpn_testsecret0123456789"

type harness struct {
	server  *httptest.Server
	content *httptest.Server
	session *Session
	policy  *Policy
	conns   atomic.Int32
	tickets atomic.Int32
	failure atomic.Value
}

func readFrame(t *testing.T, c *Conn) map[string]any {
	t.Helper()
	_ = c.SetReadDeadline(time.Now().Add(10 * time.Second))
	raw, err := c.ReadMessage()
	if err != nil {
		t.Errorf("fake server read: %v", err)
		return map[string]any{"type": "read_error"}
	}
	frame := map[string]any{}
	if err := json.Unmarshal(raw, &frame); err != nil {
		t.Errorf("fake server decode: %v", err)
	}
	return frame
}

func writeFrame(t *testing.T, c *Conn, frame map[string]any) {
	t.Helper()
	data, _ := json.Marshal(frame)
	if err := c.WriteText(string(data)); err != nil {
		t.Errorf("fake server write: %v", err)
	}
}

func welcome() map[string]any {
	return map[string]any{
		"type":       "welcome",
		"device_uid": "dev1",
		"limits": map[string]any{
			"max_body_bytes": 1000, "ping_seconds": 20, "reconnect_min_seconds": 1, "reconnect_max_seconds": 2,
			"max_concurrent_per_node": 2, "max_donor_mb_per_day": 20480, "job_timeout_seconds": 25,
		},
		"destinations": []any{
			map[string]any{"suffix": "reddit.com", "adapter": "reddit", "methods": []string{"GET"}, "ports": []int{80, 443}, "modes": []string{"tunnel", "fetch"}, "max_rps": 0.5},
			map[string]any{"suffix": "redd.it", "adapter": "reddit", "methods": []string{"GET"}, "ports": []int{80, 443}, "modes": []string{"tunnel", "fetch"}, "max_rps": 0.5},
		},
		"settings": map[string]any{"enabled": true, "bandwidth_mb_day": 500, "max_concurrent": 2, "active_from": "", "active_to": "", "timezone": "UTC"},
	}
}

func newHarness(t *testing.T, script func(n int, c *Conn)) *harness {
	t.Helper()
	h := &harness{}
	h.content = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Cookie") != "" || r.Header.Get("Authorization") != "" {
			h.failure.Store("identity header reached the destination")
		}
		_, _ = w.Write([]byte(strings.Repeat("r", 5000)))
	}))
	mux := http.NewServeMux()
	mux.HandleFunc("/nodes/ticket", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost || r.Header.Get("Authorization") != "Bearer "+testSecret {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		n := h.tickets.Add(1)
		_ = json.NewEncoder(w).Encode(map[string]any{"ticket": "t" + string(rune('0'+n)), "expires_in": 30, "ws_path": "/nodes/ws"})
	})
	mux.HandleFunc("/nodes/ws", func(w http.ResponseWriter, r *http.Request) {
		if !strings.HasPrefix(r.URL.Query().Get("ticket"), "t") || strings.Contains(r.URL.RawQuery, testSecret) {
			h.failure.Store("bad ticket query")
		}
		conn, err := upgrade(w, r)
		if err != nil {
			return
		}
		defer conn.Close()
		script(int(h.conns.Add(1)), conn)
	})
	h.server = httptest.NewServer(mux)
	t.Cleanup(h.server.Close)
	t.Cleanup(h.content.Close)
	cfg := DefaultConfig()
	cfg.Server = h.server.URL
	cfg.Secret = testSecret
	cfg.Timezone = "UTC"
	h.policy = testPolicy(publicAnswer)
	h.policy.dial = func(ctx context.Context, network, address string) (net.Conn, error) {
		var dialer net.Dialer
		return dialer.DialContext(ctx, "tcp", h.content.Listener.Addr().String())
	}
	limiter, err := NewLimiter(cfg, "", time.Now)
	if err != nil {
		t.Fatal(err)
	}
	h.session = NewSession(cfg, h.policy, limiter, NewFetcher(h.policy), NewStats(time.Now()), NewLogger(testWriter{t}, true))
	h.session.jitter = func() float64 { return 0 }
	return h
}

type testWriter struct{ t *testing.T }

func (w testWriter) Write(p []byte) (int, error) {
	w.t.Log(strings.TrimSpace(string(p)))
	return len(p), nil
}

func (h *harness) run(t *testing.T) error {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	err := h.session.Run(ctx)
	if failure := h.failure.Load(); failure != nil {
		t.Error(failure)
	}
	return err
}

func TestSessionServesJobsRefusesBadOnesAndStopsOnKick(t *testing.T) {
	h := newHarness(t, func(n int, c *Conn) {
		hello := readFrame(t, c)
		if hello["type"] != "hello" || hello["protocol"] != float64(1) || hello["kind"] != "native" || hello["client_version"] != Version {
			t.Errorf("hello %v", hello)
		}
		if modes, _ := hello["modes"].([]any); len(modes) == 0 {
			t.Errorf("hello modes %v", hello["modes"])
		}
		if n == 2 {
			writeFrame(t, c, welcome())
			writeFrame(t, c, map[string]any{"type": "kick", "reason": "revoked"})
			return
		}
		writeFrame(t, c, welcome())
		writeFrame(t, c, map[string]any{"type": "ping"})
		if frame := readFrame(t, c); frame["type"] != "pong" {
			t.Errorf("expected pong, got %v", frame)
		}
		writeFrame(t, c, map[string]any{"type": "job", "job_id": "ok1", "method": "GET", "url": "http://www.reddit.com/r/go.json", "timeout": 5})
		result := readFrame(t, c)
		if result["type"] != "job_result" || result["job_id"] != "ok1" || result["status"] != float64(200) || result["truncated"] != true {
			t.Errorf("result %v", result)
		} else {
			body, err := base64.StdEncoding.DecodeString(result["body_b64"].(string))
			if err != nil || len(body) != 1000 {
				t.Errorf("body %d %v", len(body), err)
			}
		}
		refused := map[string]string{
			"bad1": "http://evilreddit.com/", "bad2": "http://reddit.com:8080/", "bad3": "http://127.0.0.1/",
		}
		for id, target := range refused {
			writeFrame(t, c, map[string]any{"type": "job", "job_id": id, "method": "GET", "url": target, "timeout": 5})
			answer := readFrame(t, c)
			want := RefuseHost
			if id == "bad2" {
				want = RefusePort
			}
			if answer["type"] != "job_error" || answer["job_id"] != id || answer["code"] != want {
				t.Errorf("%s: %v", id, answer)
			}
		}
		writeFrame(t, c, map[string]any{"type": "job", "job_id": "post", "method": "POST", "url": "http://reddit.com/", "timeout": 5})
		if answer := readFrame(t, c); answer["code"] != RefuseMethod {
			t.Errorf("post: %v", answer)
		}
		writeFrame(t, c, map[string]any{"type": "tunnel_open", "conn_id": "c1", "host": "evil.com", "port": 443})
		if answer := readFrame(t, c); answer["type"] != "tunnel_opened" || answer["ok"] != false || answer["code"] != RefuseHost {
			t.Errorf("tunnel: %v", answer)
		}
		_ = c.WriteClose(CloseRetryOwner, "owner")
	})
	started := time.Now()
	err := h.run(t)
	var fatal *FatalError
	if !errors.As(err, &fatal) || fatal.Reason != "revoked" {
		t.Fatalf("run returned %v", err)
	}
	if h.conns.Load() != 2 || h.tickets.Load() != 2 {
		t.Fatalf("connections %d tickets %d", h.conns.Load(), h.tickets.Load())
	}
	if time.Since(started) > 5*time.Second {
		t.Fatal("4013 was not retried quickly")
	}
	if line := h.session.stats.Line(); !strings.Contains(line, "jobs=1") || !strings.Contains(line, "refused=") {
		t.Fatalf("stats %s", line)
	}
}

func TestSessionRelaysTunnelBothWays(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	go func() {
		conn, err := listener.Accept()
		if err != nil {
			return
		}
		defer conn.Close()
		buffer := make([]byte, 64)
		for {
			n, err := conn.Read(buffer)
			if err != nil {
				return
			}
			_, _ = conn.Write([]byte("echo:" + string(buffer[:n])))
		}
	}()
	h := newHarness(t, func(n int, c *Conn) {
		readFrame(t, c)
		writeFrame(t, c, welcome())
		if n == 2 {
			writeFrame(t, c, map[string]any{"type": "kick", "reason": "kill_switch"})
			return
		}
		writeFrame(t, c, map[string]any{"type": "tunnel_open", "conn_id": "c1", "host": "www.reddit.com", "port": 443})
		if answer := readFrame(t, c); answer["type"] != "tunnel_opened" || answer["ok"] != true {
			t.Errorf("opened: %v", answer)
			return
		}
		writeFrame(t, c, map[string]any{"type": "tunnel_data", "conn_id": "c1", "b64": base64.StdEncoding.EncodeToString([]byte("hi"))})
		data := readFrame(t, c)
		decoded, _ := base64.StdEncoding.DecodeString(data["b64"].(string))
		if data["type"] != "tunnel_data" || string(decoded) != "echo:hi" {
			t.Errorf("data: %v %q", data, decoded)
		}
		writeFrame(t, c, map[string]any{"type": "tunnel_close", "conn_id": "c1"})
		_ = c.WriteClose(CloseGoingAway, "restart")
	})
	h.policy.dial = func(ctx context.Context, network, address string) (net.Conn, error) {
		var dialer net.Dialer
		return dialer.DialContext(ctx, "tcp", listener.Addr().String())
	}
	err = h.run(t)
	var fatal *FatalError
	if !errors.As(err, &fatal) || fatal.Reason != "kill_switch" {
		t.Fatalf("run returned %v", err)
	}
}

func TestSessionLimitsAnswerWithDonorCodes(t *testing.T) {
	h := newHarness(t, func(n int, c *Conn) {
		readFrame(t, c)
		writeFrame(t, c, welcome())
		if n == 1 {
			writeFrame(t, c, map[string]any{"type": "job", "job_id": "a", "method": "GET", "url": "http://reddit.com/", "timeout": 5})
			if answer := readFrame(t, c); answer["code"] != RefuseBandwidth {
				t.Errorf("bandwidth: %v", answer)
			}
			_ = c.WriteClose(CloseGoingAway, "bye")
			return
		}
		writeFrame(t, c, map[string]any{"type": "kick", "reason": "revoked"})
	})
	h.session.limiter.AddBytes(1 << 40)
	err := h.run(t)
	var fatal *FatalError
	if !errors.As(err, &fatal) || fatal.Reason != "revoked" {
		t.Fatalf("run returned %v", err)
	}
}

func TestFatalTicketStatuses(t *testing.T) {
	cases := []struct {
		status int
		body   string
		fatal  string
	}{
		{401, "", "unauthorized"},
		{403, `{"error":{"status":403,"message":"x","reason":"no_consent"}}`, "no_consent"},
		{403, `{"error":{"status":403,"message":"x","reason":"revoked"}}`, "revoked"},
		{403, `{"error":{"status":403,"message":"x","reason":"kill_switch"}}`, "kill_switch"},
	}
	for _, item := range cases {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(item.status)
			_, _ = w.Write([]byte(item.body))
		}))
		cfg := DefaultConfig()
		cfg.Server, cfg.Secret = server.URL, testSecret
		limiter, _ := NewLimiter(cfg, "", time.Now)
		policy := testPolicy(publicAnswer)
		session := NewSession(cfg, policy, limiter, NewFetcher(policy), NewStats(time.Now()), NewLogger(testWriter{t}, true))
		err := session.Run(context.Background())
		var fatal *FatalError
		if !errors.As(err, &fatal) || !strings.Contains(fatal.Reason, item.fatal) {
			t.Errorf("%d: %v", item.status, err)
		}
		server.Close()
	}
}

func TestTransientTicketFailureRetries(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if calls.Add(1) < 3 {
			w.WriteHeader(http.StatusBadGateway)
			return
		}
		w.WriteHeader(http.StatusUnauthorized)
	}))
	defer server.Close()
	cfg := DefaultConfig()
	cfg.Server, cfg.Secret = server.URL, testSecret
	limiter, _ := NewLimiter(cfg, "", time.Now)
	policy := testPolicy(publicAnswer)
	session := NewSession(cfg, policy, limiter, NewFetcher(policy), NewStats(time.Now()), NewLogger(testWriter{t}, true))
	session.jitter = func() float64 { return 0 }
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	var fatal *FatalError
	if err := session.Run(ctx); !errors.As(err, &fatal) || calls.Load() != 3 {
		t.Fatalf("err %v calls %d", err, calls.Load())
	}
}

func TestBackoffStaysInsideBounds(t *testing.T) {
	session := &Session{limits: ServerLimits{ReconnectMinSeconds: 1, ReconnectMaxSeconds: 60}}
	for _, jitter := range []float64{0, 0.5, 0.999999} {
		session.jitter = func() float64 { return jitter }
		for attempt := 0; attempt < 40; attempt++ {
			d := session.Backoff(attempt)
			if d < time.Second || d > 60*time.Second {
				t.Fatalf("attempt %d jitter %v: %v", attempt, jitter, d)
			}
		}
	}
	session.jitter = func() float64 { return 0.999999 }
	if first, later := session.Backoff(0), session.Backoff(5); first > time.Second+time.Millisecond || later <= first {
		t.Fatalf("growth: %v %v", first, later)
	}
}

func TestFatalReasonsMatchPlan(t *testing.T) {
	for _, reason := range []string{"revoked", "kill_switch", "no_consent"} {
		if !FatalReasons[reason] {
			t.Errorf("%s must stop reconnecting", reason)
		}
	}
	for _, reason := range []string{"heartbeat", "ticket", "capacity", "suspended"} {
		if FatalReasons[reason] {
			t.Errorf("%s must reconnect", reason)
		}
	}
}

type sleepLog struct{ waits []time.Duration }

func (l *sleepLog) sleep(_ context.Context, d time.Duration) { l.waits = append(l.waits, d) }

func fixedNow(text string) func() time.Time {
	moment, _ := time.Parse(time.RFC3339, text)
	return func() time.Time { return moment }
}

func TestSuspendedTicketWaitsUntilTheEndThenRetries(t *testing.T) {
	cases := []struct {
		name  string
		until string
		want  time.Duration
	}{
		{"future", "2026-01-01T01:00:00Z", time.Hour},
		{"empty", "", 300 * time.Second},
		{"garbage", "tomorrow", 300 * time.Second},
		{"expired", "2025-12-31T00:00:00Z", 60 * time.Second},
		{"far", "2030-01-01T00:00:00Z", 24 * time.Hour},
	}
	for _, item := range cases {
		t.Run(item.name, func(t *testing.T) {
			var calls atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if calls.Add(1) == 1 {
					w.WriteHeader(http.StatusForbidden)
					body, _ := json.Marshal(map[string]any{"error": map[string]any{"status": 403, "message": "x", "reason": "suspended", "until": item.until}})
					_, _ = w.Write(body)
					return
				}
				w.WriteHeader(http.StatusUnauthorized)
			}))
			defer server.Close()
			cfg := DefaultConfig()
			cfg.Server, cfg.Secret = server.URL, testSecret
			limiter, _ := NewLimiter(cfg, "", time.Now)
			policy := testPolicy(publicAnswer)
			session := NewSession(cfg, policy, limiter, NewFetcher(policy), NewStats(time.Now()), NewLogger(testWriter{t}, true))
			session.jitter = func() float64 { return 0 }
			session.now = fixedNow("2026-01-01T00:00:00Z")
			slept := &sleepLog{}
			session.sleep = slept.sleep
			var fatal *FatalError
			if err := session.Run(context.Background()); !errors.As(err, &fatal) || calls.Load() != 2 {
				t.Fatalf("err %v calls %d", err, calls.Load())
			}
			if len(slept.waits) != 1 || slept.waits[0] != item.want {
				t.Fatalf("waits %v want %v", slept.waits, item.want)
			}
		})
	}
}

func TestSuspendedKickHoldsThenReconnects(t *testing.T) {
	h := newHarness(t, func(n int, c *Conn) {
		readFrame(t, c)
		writeFrame(t, c, welcome())
		if n == 1 {
			writeFrame(t, c, map[string]any{"type": "kick", "reason": "suspended", "until": "2026-01-01T00:10:00Z"})
			return
		}
		writeFrame(t, c, map[string]any{"type": "kick", "reason": "revoked"})
	})
	h.session.now = fixedNow("2026-01-01T00:00:00Z")
	slept := &sleepLog{}
	h.session.sleep = slept.sleep
	var fatal *FatalError
	if err := h.run(t); !errors.As(err, &fatal) || fatal.Reason != "revoked" {
		t.Fatalf("run returned %v", err)
	}
	if h.conns.Load() != 2 || len(slept.waits) != 1 || slept.waits[0] != 10*time.Minute {
		t.Fatalf("conns %d waits %v", h.conns.Load(), slept.waits)
	}
}

func TestSuspendedWaitHonoursCancel(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusForbidden)
		_, _ = w.Write([]byte(`{"error":{"status":403,"message":"x","reason":"suspended","until":""}}`))
	}))
	defer server.Close()
	cfg := DefaultConfig()
	cfg.Server, cfg.Secret = server.URL, testSecret
	limiter, _ := NewLimiter(cfg, "", time.Now)
	policy := testPolicy(publicAnswer)
	session := NewSession(cfg, policy, limiter, NewFetcher(policy), NewStats(time.Now()), NewLogger(testWriter{t}, true))
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	started := time.Now()
	if err := session.Run(ctx); err != nil || time.Since(started) > 5*time.Second {
		t.Fatalf("err %v after %v", err, time.Since(started))
	}
}

func TestWelcomeDecodesDeviceLabel(t *testing.T) {
	frame, err := decodeFrame([]byte(`{"type":"welcome","device_uid":"d","device_label":"Laptop","settings":{"pause_on_battery":true}}`))
	if err != nil || frame.DeviceLabel != "Laptop" || !frame.Settings.PauseOnBattery {
		t.Fatalf("%+v %v", frame, err)
	}
	kick, _ := decodeFrame([]byte(`{"type":"kick","reason":"suspended","until":"2026-01-01T00:00:00Z"}`))
	if kick.Until == "" {
		t.Fatal("until not decoded")
	}
}