Xray-core/transport/internet/xdrive/xdrive_test.go

package xdrive

import (
	"bytes"
	"context"
	"crypto/rand"
	gotls "crypto/tls"
	"io"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/xtls/xray-core/common/net"
	"github.com/xtls/xray-core/transport/internet"
	"github.com/xtls/xray-core/transport/internet/splithttp"
	"github.com/xtls/xray-core/transport/internet/stat"
	"github.com/xtls/xray-core/transport/internet/tls"
)

const testPatience = 30 * time.Second

func settings(folder string) *internet.MemoryStreamConfig {
	return &internet.MemoryStreamConfig{
		ProtocolName: protocolName,
		ProtocolSettings: &Config{
			RemoteFolder:      folder,
			Service:           "local",
			FlushIntervalMs:   5,
			PollIntervalMs:    5,
			MaxPollIntervalMs: 20,
			SessionTtlSeconds: 5,
		},
	}
}

func pair(t *testing.T) (client, server stat.Connection, cleanup func()) {
	t.Helper()
	return pairWith(t, settings(t.TempDir()))
}

func pairWith(t *testing.T, streamSettings *internet.MemoryStreamConfig) (client, server stat.Connection, cleanup func()) {
	t.Helper()

	accepted := make(chan stat.Connection, 1)

	listener, err := Serve(context.Background(), net.LocalHostIP, net.Port(0), streamSettings, func(conn stat.Connection) {
		accepted <- conn
	})
	if err != nil {
		t.Fatalf("Serve: %v", err)
	}

	client, err = Dial(context.Background(), net.Destination{}, streamSettings)
	if err != nil {
		listener.Close()
		t.Fatalf("Dial: %v", err)
	}

	select {
	case server = <-accepted:
	case <-time.After(testPatience):
		client.Close()
		listener.Close()
		t.Fatal("listener did not accept the session")
	}

	return client, server, func() {
		client.Close()
		server.Close()
		listener.Close()
	}
}

func expectRead(t *testing.T, conn stat.Connection, want string) {
	t.Helper()

	if err := conn.SetReadDeadline(time.Now().Add(testPatience)); err != nil {
		t.Fatalf("SetReadDeadline: %v", err)
	}
	buf := make([]byte, len(want))
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatalf("ReadFull: %v", err)
	}
	if string(buf) != want {
		t.Fatalf("read %q, want %q", buf, want)
	}
}

func TestRoundTrip(t *testing.T) {
	client, server, cleanup := pair(t)
	defer cleanup()

	if _, err := client.Write([]byte("ping")); err != nil {
		t.Fatalf("client write: %v", err)
	}
	expectRead(t, server, "ping")

	if _, err := server.Write([]byte("pong")); err != nil {
		t.Fatalf("server write: %v", err)
	}
	expectRead(t, client, "pong")
}

func TestInterleaved(t *testing.T) {
	client, server, cleanup := pair(t)
	defer cleanup()

	for i := 0; i < 20; i++ {
		if _, err := client.Write([]byte("up")); err != nil {
			t.Fatalf("client write %d: %v", i, err)
		}
		expectRead(t, server, "up")

		if _, err := server.Write([]byte("down")); err != nil {
			t.Fatalf("server write %d: %v", i, err)
		}
		expectRead(t, client, "down")
	}
}

func TestMultiSegmentTransfer(t *testing.T) {
	client, server, cleanup := pair(t)
	defer cleanup()

	payload := make([]byte, 3*defaultSegmentBytes+1234)
	if _, err := rand.Read(payload); err != nil {
		t.Fatalf("rand: %v", err)
	}

	go func() {
		client.Write(payload)
	}()

	if err := server.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil {
		t.Fatalf("SetReadDeadline: %v", err)
	}
	got := make([]byte, len(payload))
	if _, err := io.ReadFull(server, got); err != nil {
		t.Fatalf("ReadFull: %v", err)
	}
	if !bytes.Equal(got, payload) {
		t.Fatal("payload mismatch")
	}
}

func TestCloseEOF(t *testing.T) {
	client, server, cleanup := pair(t)
	defer cleanup()

	if _, err := client.Write([]byte("bye")); err != nil {
		t.Fatalf("client write: %v", err)
	}
	if err := client.Close(); err != nil {
		t.Fatalf("client close: %v", err)
	}

	if err := server.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil {
		t.Fatalf("SetReadDeadline: %v", err)
	}
	got, err := io.ReadAll(server)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(got) != "bye" {
		t.Fatalf("read %q, want %q", got, "bye")
	}
}

func TestReadDeadline(t *testing.T) {
	client, _, cleanup := pair(t)
	defer cleanup()

	if err := client.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil {
		t.Fatalf("SetReadDeadline: %v", err)
	}
	buf := make([]byte, 4)
	if _, err := client.Read(buf); !os.IsTimeout(err) {
		t.Fatalf("Read returned %v, want a timeout", err)
	}
}

func TestLocalNameEscape(t *testing.T) {
	root := t.TempDir()
	storage, err := newLocalStorage(root)
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}

	if err := storage.Put(context.Background(), "../escaped", []byte("x")); err != nil {
		t.Fatalf("Put: %v", err)
	}
	if _, err := os.Stat(root + "/escaped"); err != nil {
		t.Fatalf("name was not clamped inside the root: %v", err)
	}
}

func TestLocalMissingObject(t *testing.T) {
	storage, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}

	if _, err := storage.Get(context.Background(), "nothing/here"); err != errNotFound {
		t.Fatalf("Get returned %v, want errNotFound", err)
	}
	names, err := storage.List(context.Background(), "nothing")
	if err != nil {
		t.Fatalf("List: %v", err)
	}
	if len(names) != 0 {
		t.Fatalf("List returned %v, want none", names)
	}
}

func TestResumeAfterIdle(t *testing.T) {
	client, server, cleanup := pair(t)
	defer cleanup()

	if _, err := client.Write([]byte("first")); err != nil {
		t.Fatalf("client write: %v", err)
	}
	expectRead(t, server, "first")

	time.Sleep(200 * time.Millisecond)

	if _, err := client.Write([]byte("second")); err != nil {
		t.Fatalf("client write: %v", err)
	}
	expectRead(t, server, "second")
}

func TestParseEntry(t *testing.T) {
	cases := []struct {
		name string
		seq  int64
		ok   bool
	}{
		{"000000000.seg", 0, true},
		{"000000042.seg", 42, true},
		{"000000007.end", 7, true},
		{"000000001.tmp", 0, false},
		{"notanumber.seg", 0, false},
		{"000000001", 0, false},
	}
	for _, c := range cases {
		seq, ok := parseEntry(c.name)
		if ok != c.ok || (ok && seq != c.seq) {
			t.Fatalf("parseEntry(%q) = %d, %v; want %d, %v", c.name, seq, ok, c.seq, c.ok)
		}
	}
}

func TestParamDefaults(t *testing.T) {
	p := paramsFromConfig(&Config{})
	if p.segmentBytes != defaultSegmentBytes || p.flushInterval != defaultFlushInterval {
		t.Fatalf("defaults not applied: %+v", p)
	}

	p = paramsFromConfig(&Config{SegmentBytes: 1 << 30, PollIntervalMs: 400, MaxPollIntervalMs: 100})
	if p.segmentBytes != maxSegmentBytes {
		t.Fatalf("segmentBytes is %d, want %d", p.segmentBytes, maxSegmentBytes)
	}
	if p.maxPollInterval < p.minPollInterval {
		t.Fatalf("maxPollInterval %v below minPollInterval %v", p.maxPollInterval, p.minPollInterval)
	}
}

func waitFor(t *testing.T, what string, done func() bool) {
	t.Helper()

	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if done() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

func newTestListener(t *testing.T, folder string) *Listener {
	t.Helper()

	storage, err := newLocalStorage(folder)
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}
	return &Listener{
		ctx:       context.Background(),
		storage:   storage,
		params:    paramsFromConfig(&Config{SessionTtlSeconds: 1}),
		active:    make(map[string]bool),
		handled:   make(map[string]time.Time),
		idleSince: make(map[string]time.Time),
	}
}

func TestCollectAbandoned(t *testing.T) {
	folder := t.TempDir()
	listener := newTestListener(t, folder)

	if err := listener.storage.Put(context.Background(), uplinkPrefix("dead")+"/000000000.seg", []byte("x")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	if err := listener.collect(); err != nil {
		t.Fatalf("collect: %v", err)
	}
	names, _ := listener.storage.List(context.Background(), streamsDir)
	if len(names) != 1 {
		t.Fatalf("first pass removed the session, got %v", names)
	}

	listener.idleSince["dead"] = time.Now().Add(-2 * time.Second)
	if err := listener.collect(); err != nil {
		t.Fatalf("collect: %v", err)
	}
	names, _ = listener.storage.List(context.Background(), streamsDir)
	if len(names) != 0 {
		t.Fatalf("abandoned session still there, got %v", names)
	}
}

func TestCollectKeepsActive(t *testing.T) {
	folder := t.TempDir()
	listener := newTestListener(t, folder)
	listener.active["live"] = true

	if err := listener.storage.Put(context.Background(), uplinkPrefix("live")+"/000000000.seg", []byte("x")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	listener.idleSince["live"] = time.Now().Add(-2 * time.Second)
	if err := listener.collect(); err != nil {
		t.Fatalf("collect: %v", err)
	}
	names, _ := listener.storage.List(context.Background(), streamsDir)
	if len(names) != 1 {
		t.Fatalf("collected an active session, got %v", names)
	}
}

func TestParseAnnounce(t *testing.T) {
	session, at, ok := parseAnnounce("1757000000123456789-abc123")
	if !ok || session != "abc123" || at.UnixNano() != 1757000000123456789 {
		t.Fatalf("parseAnnounce returned %q, %v, %v", session, at.UnixNano(), ok)
	}
	for _, bad := range []string{"abc123", "-abc123", "1757000000-", "notanumber-abc"} {
		if _, _, ok := parseAnnounce(bad); ok {
			t.Fatalf("parseAnnounce accepted %q", bad)
		}
	}
}

func TestStaleAnnounce(t *testing.T) {
	folder := t.TempDir()
	listener := newTestListener(t, folder)

	ctx := context.Background()
	stale := announceName("ghost", time.Now().Add(-time.Hour))
	if err := listener.storage.Put(ctx, stale, nil); err != nil {
		t.Fatalf("Put: %v", err)
	}
	if err := listener.storage.Put(ctx, uplinkPrefix("ghost")+"/000000000.seg", []byte("x")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	accepted, err := listener.acceptPending(ctx)
	if err != nil {
		t.Fatalf("acceptPending: %v", err)
	}
	if accepted {
		t.Fatal("accepted a stale announcement")
	}

	waitFor(t, "the stale announcement to be removed", func() bool {
		names, _ := listener.storage.List(ctx, sessionsDir)
		return len(names) == 0
	})
	waitFor(t, "the stale session data to be removed", func() bool {
		names, _ := listener.storage.List(ctx, streamsDir)
		return len(names) == 0
	})
}

func TestFreshAnnounce(t *testing.T) {
	folder := t.TempDir()
	listener := newTestListener(t, folder)
	listener.addConn = func(conn stat.Connection) { conn.Close() }

	ctx := context.Background()
	if err := listener.storage.Put(ctx, announceName("fresh", time.Now()), nil); err != nil {
		t.Fatalf("Put: %v", err)
	}

	accepted, err := listener.acceptPending(ctx)
	if err != nil {
		t.Fatalf("acceptPending: %v", err)
	}
	if !accepted {
		t.Fatal("did not accept a fresh announcement")
	}
}

func TestAnnouncePrecision(t *testing.T) {
	at := time.Unix(1757000000, int64(900*time.Millisecond))
	entry := strings.TrimPrefix(announceName("abc123", at), sessionsDir+"/")

	session, parsed, ok := parseAnnounce(entry)
	if !ok || session != "abc123" {
		t.Fatalf("parseAnnounce(%q) returned %q, %v", entry, session, ok)
	}
	if !parsed.Equal(at) {
		t.Fatalf("timestamp came back as %v, want %v", parsed, at)
	}
}

func TestRecentAnnounceTTL(t *testing.T) {
	folder := t.TempDir()
	listener := newTestListener(t, folder)
	listener.addConn = func(conn stat.Connection) { conn.Close() }

	ctx := context.Background()
	recent := time.Now().Add(-900 * time.Millisecond)
	if err := listener.storage.Put(ctx, announceName("recent", recent), nil); err != nil {
		t.Fatalf("Put: %v", err)
	}

	accepted, err := listener.acceptPending(ctx)
	if err != nil {
		t.Fatalf("acceptPending: %v", err)
	}
	if !accepted {
		t.Fatal("dropped an announcement younger than the TTL")
	}
}

func TestMissingSegment(t *testing.T) {
	storage, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20, HoleTimeoutMs: 200})
	if err := storage.Put(ctx, objectName("hole", 1, segSuffix), []byte("second")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	reader := newWALReader(ctx, storage, "hole", p)
	select {
	case _, ok := <-reader.ch:
		if ok {
			t.Fatal("delivered data past a missing segment")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("reader did not give up on a missing segment")
	}

	err = reader.Err()
	if err == nil || err == io.EOF {
		t.Fatalf("Err returned %v, want a failure", err)
	}
}

func TestIdleStreamWaits(t *testing.T) {
	storage, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20, HoleTimeoutMs: 100})
	reader := newWALReader(ctx, storage, "idle", p)

	time.Sleep(400 * time.Millisecond)
	if err := storage.Put(ctx, objectName("idle", 0, segSuffix), []byte("late")); err != nil {
		t.Fatalf("Put: %v", err)
	}

	select {
	case data, ok := <-reader.ch:
		if !ok {
			t.Fatalf("the reader gave up on an idle stream: %v", reader.Err())
		}
		if string(data) != "late" {
			t.Fatalf("read %q, want %q", data, "late")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("reader missed a late segment")
	}
}

func TestFailureMarker(t *testing.T) {
	storage, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20})
	if err := storage.Put(ctx, objectName("broken", 0, segSuffix), []byte("first")); err != nil {
		t.Fatalf("Put: %v", err)
	}
	if err := storage.Put(ctx, objectName("broken", 1, errSuffix), nil); err != nil {
		t.Fatalf("Put: %v", err)
	}

	reader := newWALReader(ctx, storage, "broken", p)

	select {
	case data, ok := <-reader.ch:
		if !ok || string(data) != "first" {
			t.Fatalf("want the segment before the marker, got %q %v", data, ok)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("reader did not deliver the first segment")
	}

	select {
	case _, ok := <-reader.ch:
		if ok {
			t.Fatal("delivered data past the failure marker")
		}
	case <-time.After(5 * time.Second):
		t.Fatal("reader did not stop on the failure marker")
	}

	if err := reader.Err(); err == nil || err == io.EOF {
		t.Fatalf("Err returned %v, want a failure", err)
	}
}

type inlineOnlyStorage struct {
	Storage
	gets int64
}

func (s *inlineOnlyStorage) List(ctx context.Context, prefix string) ([]Entry, error) {
	return []Entry{{Name: "000000000" + segSuffix, Inline: []byte("carried by the listing")}}, nil
}

func (s *inlineOnlyStorage) Get(ctx context.Context, name string) ([]byte, error) {
	atomic.AddInt64(&s.gets, 1)
	return nil, errNotFound
}

func (s *inlineOnlyStorage) Delete(ctx context.Context, name string) error {
	return nil
}

func TestInlinePayload(t *testing.T) {
	base, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}
	storage := &inlineOnlyStorage{Storage: base}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	reader := newWALReader(ctx, storage, "inline", paramsFromConfig(&Config{PollIntervalMs: 5}))
	select {
	case data, ok := <-reader.ch:
		if !ok {
			t.Fatalf("reader stopped: %v", reader.Err())
		}
		if string(data) != "carried by the listing" {
			t.Fatalf("read %q, want the inline payload", data)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("reader did not deliver the inline payload")
	}

	if got := atomic.LoadInt64(&storage.gets); got != 0 {
		t.Fatalf("called Get %d times for an inline payload", got)
	}
}

// runSessions opens n sessions at once, both ends in this process, passes data
// both ways over each and closes them
func runSessions(t *testing.T, streamSettings *internet.MemoryStreamConfig, n int) {
	t.Helper()

	accepted := make(chan stat.Connection, 1)
	listener, err := Serve(context.Background(), net.LocalHostIP, net.Port(0), streamSettings, func(conn stat.Connection) {
		accepted <- conn
	})
	if err != nil {
		t.Fatalf("Serve: %v", err)
	}
	defer listener.Close()

	var ends []stat.Connection
	defer func() {
		for _, end := range ends {
			end.Close()
		}
	}()
	for range n {
		client, err := Dial(context.Background(), net.Destination{}, streamSettings)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		ends = append(ends, client)
		select {
		case peer := <-accepted:
			ends = append(ends, peer)
		case <-time.After(testPatience):
			t.Fatal("listener did not accept the session")
		}
	}
	for i := 0; i < len(ends); i += 2 {
		if _, err := ends[i].Write([]byte("ping")); err != nil {
			t.Fatalf("client write: %v", err)
		}
		expectRead(t, ends[i+1], "ping")
		if _, err := ends[i+1].Write([]byte("pong")); err != nil {
			t.Fatalf("server write: %v", err)
		}
		expectRead(t, ends[i], "pong")
	}
}

// sessionEnd is the session and end of it that a request to the fake store
// comes from, or "" for the listener and the token
func sessionEnd(r *http.Request) string {
	name := strings.TrimPrefix(r.URL.Path, "/folder/")
	if r.Method == "PROPFIND" {
		name = r.URL.Query().Get("prefix")
	}
	write := r.Method == http.MethodPut
	parts := strings.Split(name, flatSeparator)
	switch {
	case len(parts) > 2 && parts[0] == streamsDir:
		// each end writes its own direction and reads the other
		if (parts[2] == uplinkDir) == write {
			return parts[1] + " client"
		}
		return parts[1] + " server"
	case len(parts) == 2 && parts[0] == sessionsDir && write:
		_, session, _ := strings.Cut(parts[1], "-")
		return session + " client"
	}
	return ""
}

// with "maxConcurrency" 1 every end of a session has a connection of its own,
// with "maxConnections" the ends share that many, and the sessions after them
// take the same ones again. The listener and the token use none of them, and
// the sessions get the token once
func TestXmuxConnections(t *testing.T) {
	const n = 3
	for _, c := range []struct {
		name   string
		xmux   *splithttp.XmuxConfig
		conns  int
		shared bool
	}{
		{"maxConcurrency", &splithttp.XmuxConfig{MaxConcurrency: &splithttp.RangeConfig{From: 1, To: 1}}, 2 * n, false},
		{"maxConnections", &splithttp.XmuxConfig{MaxConnections: &splithttp.RangeConfig{From: 2, To: 2}}, 2, true},
	} {
		t.Run(c.name, func(t *testing.T) {
			resetSharedStorage()
			t.Cleanup(resetSharedStorage)

			store := &fakeStore{objects: make(map[string][]byte), needAuth: "Bearer tok-fake"}
			var (
				mu    sync.Mutex
				conns = make(map[string]map[string]bool)
			)
			server := newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, func(w http.ResponseWriter, r *http.Request) {
				end := sessionEnd(r)
				mu.Lock()
				if conns[end] == nil {
					conns[end] = make(map[string]bool)
				}
				conns[end][r.RemoteAddr] = true
				mu.Unlock()
				store.handle(w, r)
			})
			store.server = server.Server
			settings := templateSettings(store, map[string]interface{}{
				"type":     "oauth2",
				"tokenUrl": server.URL + "/token",
				"form":     map[string]interface{}{"grant_type": "refresh_token", "refresh_token": "{secret0}"},
				"header":   map[string]interface{}{"Authorization": "Bearer {token}"},
			}, []string{"refresh"})
			settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
			settings.ProtocolSettings.(*Config).Xmux = c.xmux

			runSessions(t, settings, n)
			mu.Lock()
			first := conns
			conns = make(map[string]map[string]bool)
			mu.Unlock()
			dialed := server.conns.Load()
			runSessions(t, settings, n)
			if got := server.conns.Load(); got != dialed {
				t.Errorf("the next %d sessions dialed %d connections", n, got-dialed)
			}

			owner := make(map[string]string)
			for end, addrs := range first {
				if end == "" {
					continue
				}
				if len(addrs) != 1 {
					t.Errorf("%s used %d connections", end, len(addrs))
				}
				for addr := range addrs {
					if owner[addr] != "" && !c.shared {
						t.Errorf("%s and %s shared a connection", end, owner[addr])
					}
					owner[addr] = end
				}
			}
			if len(owner) != c.conns {
				t.Errorf("%d sessions had %d connections, want %d", 2*n, len(owner), c.conns)
			}
			for addr := range first[""] {
				if owner[addr] != "" {
					t.Errorf("the listener or the token went over the connection of %s", owner[addr])
				}
			}
			store.mu.Lock()
			defer store.mu.Unlock()
			if store.tokens != 1 {
				t.Errorf("the sessions got the token %d times", store.tokens)
			}
		})
	}
}

func TestDriveXmuxConnections(t *testing.T) {
	drive := newFakeDrive(t)
	server := newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, drive.server.Config.Handler.ServeHTTP)
	driveTokenURL = server.URL + "/token"
	driveFilesURL = server.URL + "/files"
	driveUploadURL = server.URL + "/upload"
	settings := driveSettings()
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	settings.ProtocolSettings.(*Config).Xmux = &splithttp.XmuxConfig{MaxConnections: &splithttp.RangeConfig{From: 3, To: 3}}

	runSessions(t, settings, 3)
	if got := server.conns.Load(); got != 4 {
		t.Errorf("6 ends of sessions with 3 connections and the listener took %d connections", got)
	}

	drive.mu.Lock()
	defer drive.mu.Unlock()
	if drive.tokens != 1 {
		t.Errorf("6 ends of sessions got the token %d times", drive.tokens)
	}
}

// a session knows the hosts that answered the storage over h1, the polling of
// the listener for one, so it sends requests in a row over one connection, not
// first over one for h2 that then stays idle, and parallel ones over more
func TestSessionKnowsHTTP1Hosts(t *testing.T) {
	check := func(name string, server *tlsServer, storage Storage) {
		put := func(s Storage, names ...string) {
			for _, n := range names {
				if err := s.Put(context.Background(), "streams/abc/c2s/"+n, []byte("x")); err != nil {
					t.Errorf("%s: Put: %v", name, err)
				}
			}
		}
		put(storage, "a", "b")
		before := server.conns.Load()
		session := sessionStorage(storage)
		defer session.Close()
		put(session, "c", "d")
		if n := server.conns.Load() - before; n != 1 {
			t.Errorf("%s: 2 requests in a row of a session over h1 took %d connections", name, n)
		}

		before = server.conns.Load()
		var wg sync.WaitGroup
		for _, n := range []string{"e", "f", "g", "h"} {
			wg.Go(func() { put(session, n) })
		}
		wg.Wait()
		if server.conns.Load() == before {
			t.Errorf("%s: 4 parallel requests of a session over h1 took one connection", name)
		}
	}
	slow := func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			time.Sleep(20 * time.Millisecond)
			next(w, r)
		}
	}

	store := &fakeStore{objects: make(map[string][]byte)}
	server := newTLSServer(t, &gotls.Config{NextProtos: []string{"http/1.1"}}, slow(store.handle))
	store.server = server.Server
	settings := templateSettings(store, map[string]interface{}{"type": "none"}, nil)
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	template, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
	if err != nil {
		t.Fatalf("newTemplateStorage: %v", err)
	}
	check("template", server, template)

	drive := newFakeDrive(t)
	server = newTLSServer(t, &gotls.Config{NextProtos: []string{"http/1.1"}}, slow(drive.server.Config.Handler.ServeHTTP))
	driveTokenURL = server.URL + "/token"
	driveFilesURL = server.URL + "/files"
	settings = driveSettings()
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	google, err := newDriveStorage(settings, settings.ProtocolSettings.(*Config))
	if err != nil {
		t.Fatalf("newDriveStorage: %v", err)
	}
	check("drive", server, google)
}

// with "hMaxRequestTimes" a session moves on to a new connection after that
// many requests, and the old one closes once the requests on it end, not before
func TestXmuxRequestLimit(t *testing.T) {
	store := &fakeStore{objects: make(map[string][]byte)}
	arrived := make(chan struct{}, 1)
	release := make(chan struct{})
	server := newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			arrived <- struct{}{}
			<-release
		}
		store.handle(w, r)
	})
	store.server = server.Server
	settings := templateSettings(store, map[string]interface{}{"type": "none"}, nil)
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	settings.ProtocolSettings.(*Config).Xmux = &splithttp.XmuxConfig{HMaxRequestTimes: &splithttp.RangeConfig{From: 3, To: 3}}
	storage, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
	if err != nil {
		t.Fatalf("newTemplateStorage: %v", err)
	}
	session := sessionStorage(storage)
	defer session.Close()
	put := func(names ...string) {
		for _, name := range names {
			if err := session.Put(context.Background(), "streams/abc/c2s/"+name, []byte("x")); err != nil {
				t.Fatalf("Put: %v", err)
			}
		}
	}

	put("a")
	got := make(chan error, 1)
	go func() {
		_, err := session.Get(context.Background(), "streams/abc/c2s/a")
		got <- err
	}()
	<-arrived
	put("b", "c")
	if n := server.conns.Load(); n != 2 {
		t.Errorf("4 requests with 3 to a connection took %d connections", n)
	}
	if server.closed.Load() != 0 {
		t.Error("the first connection closed under a request")
	}
	close(release)
	if err := <-got; err != nil {
		t.Errorf("Get on the first connection: %v", err)
	}
	waitFor(t, "the first connection to close", func() bool { return server.closed.Load() == 1 })

	put("d", "e", "f")
	if n := server.conns.Load(); n != 3 {
		t.Errorf("7 requests with 3 to a connection took %d connections", n)
	}
	waitFor(t, "the second connection to close", func() bool { return server.closed.Load() == 2 })
}

// a session that fails to announce itself lets go of its connection, which the
// next one takes with "maxConcurrency" 1
func TestFailedDialReleasesConnection(t *testing.T) {
	resetSharedStorage()
	t.Cleanup(resetSharedStorage)

	store := &fakeStore{objects: make(map[string][]byte)}
	var refuse atomic.Bool
	refuse.Store(true)
	server := newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, func(w http.ResponseWriter, r *http.Request) {
		if refuse.Load() {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		store.handle(w, r)
	})
	store.server = server.Server
	settings := templateSettings(store, map[string]interface{}{"type": "none"}, nil)
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	settings.ProtocolSettings.(*Config).Xmux = &splithttp.XmuxConfig{MaxConcurrency: &splithttp.RangeConfig{From: 1, To: 1}}

	if _, err := Dial(context.Background(), net.Destination{}, settings); err == nil {
		t.Fatal("Dial announced a session that the service refused")
	}
	refuse.Store(false)
	conn, err := Dial(context.Background(), net.Destination{}, settings)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	if n := server.conns.Load(); n != 1 {
		t.Errorf("a failed Dial and the next took %d connections", n)
	}
}

// a session gets the token over a connection of the storage, not over one of
// XMUX
func TestSessionTokenConnection(t *testing.T) {
	check := func(name string, server *tlsServer, storage Storage) {
		session := sessionStorage(storage)
		defer session.Close()
		if err := session.Put(context.Background(), "sessions/abc", nil); err != nil {
			t.Fatalf("%s: Put: %v", name, err)
		}
		if n := server.conns.Load(); n != 2 {
			t.Errorf("%s: the token and a Put of a session took %d connections, want 2", name, n)
		}
	}

	store := &fakeStore{objects: make(map[string][]byte), needAuth: "Bearer tok-fake"}
	server := newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, store.handle)
	store.server = server.Server
	settings := templateSettings(store, map[string]interface{}{
		"type":     "oauth2",
		"tokenUrl": server.URL + "/token",
		"header":   map[string]interface{}{"Authorization": "Bearer {token}"},
	}, nil)
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	template, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
	if err != nil {
		t.Fatalf("newTemplateStorage: %v", err)
	}
	check("template", server, template)

	drive := newFakeDrive(t)
	server = newTLSServer(t, &gotls.Config{NextProtos: chromeALPN}, drive.server.Config.Handler.ServeHTTP)
	driveTokenURL = server.URL + "/token"
	driveFilesURL = server.URL + "/files"
	settings = driveSettings()
	settings.SecuritySettings = server.settings(&tls.Config{}).SecuritySettings
	google, err := newDriveStorage(settings, settings.ProtocolSettings.(*Config))
	if err != nil {
		t.Fatalf("newDriveStorage: %v", err)
	}
	check("drive", server, google)
}

// local storage has no connections to keep apart
func TestLocalSessionStorage(t *testing.T) {
	storage, err := newLocalStorage(t.TempDir())
	if err != nil {
		t.Fatalf("newLocalStorage: %v", err)
	}
	if sessionStorage(storage) != Storage(storage) {
		t.Fatal("a session of local storage got a storage of its own")
	}
}