Warp/server/internal/federation/media_test.go

package federation

import (
	"context"
	"errors"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/http/httptest"
	"net/textproto"
	"strings"
	"testing"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
)

// mediaServer is another server's media endpoints, as a test wants them.
func mediaServer(t *testing.T, mux *http.ServeMux) *httptest.Server {
	t.Helper()
	srv := httptest.NewTLSServer(mux)
	t.Cleanup(srv.Close)
	return srv
}

func mediaClient(t *testing.T, srv *httptest.Server, allowPrivate bool) *Client {
	t.Helper()
	key, err := signing.Generate("")
	if err != nil {
		t.Fatal(err)
	}
	return NewClient(ClientOptions{
		ServerName: "warp.example",
		Key:        key,
		TLSConfig:  srv.Client().Transport.(*http.Transport).TLSClientConfig.Clone(),
		Enabled:    func() bool { return true },
		Resolve: func(context.Context, string) (Resolved, error) {
			_, port, _ := net.SplitHostPort(srv.Listener.Addr().String())
			return Resolved{Addresses: []string{srv.Listener.Addr().String()}, Host: "example.com:" + port}, nil
		},
		AllowPrivateMediaRedirects: allowPrivate,
	})
}

// writeMultipart answers the way the federation media endpoints do.
func writeMultipart(w http.ResponseWriter, contentHeader textproto.MIMEHeader, content string) {
	mw := multipart.NewWriter(w)
	w.Header().Set("Content-Type", "multipart/mixed; boundary="+mw.Boundary())
	meta, _ := mw.CreatePart(textproto.MIMEHeader{"Content-Type": {"application/json"}})
	_, _ = meta.Write([]byte("{}"))
	part, _ := mw.CreatePart(contentHeader)
	_, _ = part.Write([]byte(content))
	_ = mw.Close()
}

func readMedia(t *testing.T, m MediaResponse) string {
	t.Helper()
	defer func() { _ = m.Body.Close() }()
	data, err := io.ReadAll(m.Body)
	if err != nil {
		t.Fatal(err)
	}
	return string(data)
}

func TestDownloadMediaMultipart(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /_matrix/federation/v1/media/download/{mediaID}", func(w http.ResponseWriter, r *http.Request) {
		if !strings.HasPrefix(r.Header.Get("Authorization"), "X-Matrix ") {
			http.Error(w, "unsigned", http.StatusUnauthorized)
			return
		}
		if r.URL.Query().Get("timeout_ms") != "5000" {
			http.Error(w, "no timeout", http.StatusBadRequest)
			return
		}
		writeMultipart(w, textproto.MIMEHeader{
			"Content-Type":        {"text/plain"},
			"Content-Disposition": {`inline; filename="a.txt"`},
		}, "content of "+r.PathValue("mediaID"))
	})
	srv := mediaServer(t, mux)
	c := mediaClient(t, srv, false)

	m, err := c.DownloadMedia(t.Context(), "remote.example", "abc", 5*time.Second)
	if err != nil {
		t.Fatalf("DownloadMedia: %v", err)
	}
	if m.ContentType != "text/plain" || m.ContentDisposition != `inline; filename="a.txt"` {
		t.Errorf("headers = %q, %q", m.ContentType, m.ContentDisposition)
	}
	if got := readMedia(t, m); got != "content of abc" {
		t.Errorf("content = %q", got)
	}
}

func TestDownloadMediaFallsBackToTheLegacyEndpoint(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("/_matrix/federation/", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusNotFound)
		_, _ = w.Write([]byte(`{"errcode":"M_UNRECOGNIZED","error":"Unrecognized request"}`))
	})
	mux.HandleFunc("GET /_matrix/media/v3/download/{origin}/{mediaID}", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Query().Get("allow_remote") != "false" {
			http.Error(w, "would fetch from elsewhere", http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write([]byte("legacy " + r.PathValue("origin") + "/" + r.PathValue("mediaID")))
	})
	srv := mediaServer(t, mux)
	c := mediaClient(t, srv, false)

	m, err := c.DownloadMedia(t.Context(), "old.example", "xyz", 0)
	if err != nil {
		t.Fatalf("DownloadMedia: %v", err)
	}
	if got := readMedia(t, m); got != "legacy old.example/xyz" || m.ContentType != "image/png" {
		t.Errorf("legacy download = %q (%s)", got, m.ContentType)
	}
}

func TestDownloadMediaNotFoundIsNotRetried(t *testing.T) {
	legacyAsked := false
	mux := http.NewServeMux()
	mux.HandleFunc("/_matrix/federation/", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusNotFound)
		_, _ = w.Write([]byte(`{"errcode":"M_NOT_FOUND","error":"no such media"}`))
	})
	mux.HandleFunc("/_matrix/media/", func(w http.ResponseWriter, _ *http.Request) { legacyAsked = true })
	c := mediaClient(t, mediaServer(t, mux), false)

	_, err := c.DownloadMedia(t.Context(), "remote.example", "gone", 0)
	if !IsNotFound(err) {
		t.Fatalf("DownloadMedia = %v, want a 404", err)
	}
	if legacyAsked {
		t.Error("a server that knows the endpoint and lacks the media was asked again")
	}
}

func TestMediaThumbnailRedirect(t *testing.T) {
	var target string
	mux := http.NewServeMux()
	mux.HandleFunc("GET /_matrix/federation/v1/media/thumbnail/{mediaID}", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		if q.Get("width") != "96" || q.Get("height") != "96" || q.Get("method") != "crop" {
			http.Error(w, "bad size", http.StatusBadRequest)
			return
		}
		writeMultipart(w, textproto.MIMEHeader{"Location": {target}}, "")
	})
	mux.HandleFunc("GET /cdn/thumb", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "image/jpeg")
		_, _ = w.Write([]byte("from the cdn"))
	})
	srv := mediaServer(t, mux)
	target = srv.URL + "/cdn/thumb"
	req := ThumbnailRequest{Width: 96, Height: 96, Method: "crop"}

	m, err := mediaClient(t, srv, true).MediaThumbnail(t.Context(), "remote.example", "abc", req, 0)
	if err != nil {
		t.Fatalf("MediaThumbnail: %v", err)
	}
	if got := readMedia(t, m); got != "from the cdn" || m.ContentType != "image/jpeg" {
		t.Errorf("redirected thumbnail = %q (%s)", got, m.ContentType)
	}

	// The same redirect into a private address is refused when private
	// redirects are not allowed — which is everywhere but tests.
	_, err = mediaClient(t, srv, false).MediaThumbnail(t.Context(), "remote.example", "abc", req, 0)
	if !errors.Is(err, errPrivateAddress) {
		t.Errorf("redirect to loopback = %v, want errPrivateAddress", err)
	}
}

func TestMediaNeedsFederation(t *testing.T) {
	c := NewClient(ClientOptions{ServerName: "warp.example"})
	if _, err := c.DownloadMedia(t.Context(), "remote.example", "abc", 0); !errors.Is(err, ErrDisabled) {
		t.Errorf("DownloadMedia with federation off = %v, want ErrDisabled", err)
	}
}

func TestIsPublic(t *testing.T) {
	for addr, want := range map[string]bool{
		"8.8.8.8":         true,
		"2001:4860::8888": true,
		"127.0.0.1":       false,
		"10.1.2.3":        false,
		"192.168.0.1":     false,
		"169.254.169.254": false,
		"::1":             false,
		"fe80::1":         false,
		"fd00::1":         false,
		"0.0.0.0":         false,
	} {
		if got := isPublic(net.ParseIP(addr)); got != want {
			t.Errorf("isPublic(%s) = %v, want %v", addr, got, want)
		}
	}
}