tailscale/cmd/tsconnect/driveprobe/driveprobe_test.go

// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

package driveprobe

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func multistatus(hrefs ...string) string {
	var b strings.Builder
	b.WriteString(`<?xml version="1.0" encoding="UTF-8"?><D:multistatus xmlns:D="DAV:">`)
	for _, h := range hrefs {
		fmt.Fprintf(&b, `<D:response><D:href>%s</D:href><D:propstat><D:prop><D:resourcetype><D:collection/></D:resourcetype></D:prop><D:status>HTTP/1.1 200 OK</D:status></D:propstat></D:response>`, h)
	}
	b.WriteString(`</D:multistatus>`)
	return b.String()
}

func TestHasChild(t *testing.T) {
	tests := []struct {
		name string
		body string
		want bool
	}{
		{"root only, prefix stripped", multistatus("/"), false},
		{"root only, prefix kept", multistatus("/v0/drive/"), false},
		{"root only, no trailing slash", multistatus("/v0/drive"), false},
		{"one share, prefix stripped", multistatus("/", "/docs"), true},
		{"one share, prefix kept", multistatus("/v0/drive/", "/v0/drive/docs"), true},
		{"absolute urls", multistatus("http://100.1.2.3:1234/v0/drive/", "http://100.1.2.3:1234/v0/drive/docs"), true},
		{"percent-encoded share name", multistatus("/v0/drive/", "/v0/drive/my%20share"), true},
		{"unicode share name", multistatus("/v0/drive/", "/v0/drive/%E6%97%A5%E6%9C%AC"), true},
		{"empty multistatus", multistatus(), false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := hasChild(strings.NewReader(tt.body), drivePath)
			if err != nil {
				t.Fatalf("hasChild: %v", err)
			}
			if got != tt.want {
				t.Errorf("hasChild = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestHasChildMalformed(t *testing.T) {
	body := strings.TrimSuffix(multistatus("/v0/drive/", "/v0/drive/docs"), "</D:multistatus>")
	// Truncation after a child href still answers the question.
	got, err := hasChild(strings.NewReader(body), drivePath)
	if err != nil {
		t.Fatalf("hasChild: %v", err)
	}
	if !got {
		t.Error("hasChild = false on a truncated body that already listed a share")
	}

	if _, err := hasChild(strings.NewReader("<D:multistatus"), drivePath); err == nil {
		t.Error("hasChild on malformed XML: want error, got nil")
	}
}

// serveDrive returns a server answering PROPFIND at the taildrive root with the
// given body and status, and records the requests it saw.
func serveDrive(t *testing.T, status int, body string) (*httptest.Server, *[]*http.Request) {
	t.Helper()
	var mu sync.Mutex
	var reqs []*http.Request
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		reqs = append(reqs, r)
		mu.Unlock()
		w.WriteHeader(status)
		io := []byte(body)
		w.Write(io)
	}))
	t.Cleanup(srv.Close)
	return srv, &reqs
}

func TestHasShares(t *testing.T) {
	srv, reqs := serveDrive(t, http.StatusMultiStatus, multistatus("/v0/drive/", "/v0/drive/docs"))
	got, err := HasShares(context.Background(), srv.Client(), srv.URL)
	if err != nil {
		t.Fatalf("HasShares: %v", err)
	}
	if !got {
		t.Error("HasShares = false, want true")
	}

	if len(*reqs) != 1 {
		t.Fatalf("got %d requests, want 1", len(*reqs))
	}
	r := (*reqs)[0]
	if r.Method != "PROPFIND" {
		t.Errorf("method = %q, want PROPFIND", r.Method)
	}
	if r.URL.Path != drivePath {
		t.Errorf("path = %q, want %q", r.URL.Path, drivePath)
	}
	if d := r.Header.Get("Depth"); d != "1" {
		t.Errorf("Depth = %q, want 1", d)
	}
}

func TestHasSharesTrailingSlashInPeerURL(t *testing.T) {
	srv, reqs := serveDrive(t, http.StatusMultiStatus, multistatus("/"))
	if _, err := HasShares(context.Background(), srv.Client(), srv.URL+"/"); err != nil {
		t.Fatalf("HasShares: %v", err)
	}
	if p := (*reqs)[0].URL.Path; p != drivePath {
		t.Errorf("path = %q, want %q", p, drivePath)
	}
}

func TestHasSharesNonMultiStatus(t *testing.T) {
	// Taildrive disabled, or we hold no cap on that peer.
	for _, status := range []int{http.StatusNotFound, http.StatusForbidden, http.StatusOK} {
		srv, _ := serveDrive(t, status, "")
		got, err := HasShares(context.Background(), srv.Client(), srv.URL)
		if err == nil {
			t.Errorf("status %d: want error, got nil", status)
		}
		if got {
			t.Errorf("status %d: HasShares = true", status)
		}
	}
}

func TestHasSharesCancelled(t *testing.T) {
	// The handler must also unblock on release: a client-side cancel does not
	// reliably reach the server's request context, and Close waits for it.
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-r.Context().Done():
		case <-release:
		}
	}))
	defer srv.Close()
	defer close(release)

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	if _, err := HasShares(ctx, srv.Client(), srv.URL); err == nil {
		t.Error("want error on cancelled probe, got nil")
	}
}

func TestHasSharesMultiOrderAndFailures(t *testing.T) {
	withShares, _ := serveDrive(t, http.StatusMultiStatus, multistatus("/v0/drive/", "/v0/drive/docs"))
	noShares, _ := serveDrive(t, http.StatusMultiStatus, multistatus("/v0/drive/"))
	refused, _ := serveDrive(t, http.StatusNotFound, "")

	urls := []string{noShares.URL, withShares.URL, refused.URL, "http://127.0.0.1:1/dead", withShares.URL}
	want := []bool{false, true, false, false, true}

	var logs atomic.Int32
	got := HasSharesMulti(context.Background(), withShares.Client(), urls, func(string, ...any) {
		logs.Add(1)
	})

	for i := range want {
		if got[i] != want[i] {
			t.Errorf("result[%d] = %v, want %v (%s)", i, got[i], want[i], urls[i])
		}
	}
	// The 404 peer and the dead address are both reported, not swallowed.
	if n := logs.Load(); n != 2 {
		t.Errorf("logged %d probe failures, want 2", n)
	}
}

func TestHasSharesMultiRunsInParallel(t *testing.T) {
	const delay = 100 * time.Millisecond
	var inFlight, peak atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := inFlight.Add(1)
		for {
			old := peak.Load()
			if n <= old || peak.CompareAndSwap(old, n) {
				break
			}
		}
		time.Sleep(delay)
		inFlight.Add(-1)
		w.WriteHeader(http.StatusMultiStatus)
		w.Write([]byte(multistatus("/v0/drive/", "/v0/drive/docs")))
	}))
	defer srv.Close()

	urls := make([]string, maxProbes)
	for i := range urls {
		urls[i] = srv.URL
	}

	start := time.Now()
	got := HasSharesMulti(context.Background(), srv.Client(), urls, func(string, ...any) {})
	elapsed := time.Since(start)

	for i, ok := range got {
		if !ok {
			t.Errorf("result[%d] = false, want true", i)
		}
	}
	if elapsed >= delay*time.Duration(len(urls)) {
		t.Errorf("probes serialized: %v for %d probes of %v each", elapsed, len(urls), delay)
	}
	if peak.Load() < 2 {
		t.Errorf("peak concurrency = %d, want >= 2", peak.Load())
	}
}

func TestHasSharesMultiLimitsConcurrency(t *testing.T) {
	var inFlight, peak atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		n := inFlight.Add(1)
		for {
			old := peak.Load()
			if n <= old || peak.CompareAndSwap(old, n) {
				break
			}
		}
		time.Sleep(20 * time.Millisecond)
		inFlight.Add(-1)
		w.WriteHeader(http.StatusMultiStatus)
		w.Write([]byte(multistatus("/v0/drive/")))
	}))
	defer srv.Close()

	urls := make([]string, maxProbes*3)
	for i := range urls {
		urls[i] = srv.URL
	}
	client := &http.Client{Transport: &http.Transport{MaxConnsPerHost: 0}}
	HasSharesMulti(context.Background(), client, urls, func(string, ...any) {})

	if peak.Load() > maxProbes {
		t.Errorf("peak concurrency = %d, want <= %d", peak.Load(), maxProbes)
	}
}