livekit/pkg/service/utils_test.go

// Copyright 2023 LiveKit, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service_test

import (
	"bytes"
	"compress/gzip"
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/stretchr/testify/require"

	"github.com/livekit/livekit-server/pkg/service"
)

func redisClientDocker(t testing.TB) *redis.Client {
	addr := runRedis(t)
	cli := redis.NewClient(&redis.Options{
		Addr: addr,
	})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := cli.Ping(ctx).Err(); err != nil {
		_ = cli.Close()
		t.Fatal(err)
	}
	t.Cleanup(func() {
		_ = cli.Close()
	})
	return cli
}

func redisClient(t testing.TB) *redis.Client {
	cli := redis.NewClient(&redis.Options{
		Addr: "localhost:6379",
	})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := cli.Ping(ctx).Err()
	if err == nil {
		t.Cleanup(func() {
			_ = cli.Close()
		})
		return cli
	}
	_ = cli.Close()
	t.Logf("local redis not available: %v", err)

	t.Logf("starting redis in docker")
	return redisClientDocker(t)
}

func TestIsValidDomain(t *testing.T) {
	list := map[string]bool{
		"turn.myhost.com":  true,
		"turn.google.com":  true,
		"https://host.com": false,
		"turn://host.com":  false,
	}
	for key, result := range list {
		service.IsValidDomain(key)
		require.Equal(t, service.IsValidDomain(key), result)
	}
}

func compress(t *testing.T, payload []byte) []byte {
	t.Helper()
	var buf bytes.Buffer
	gw, err := gzip.NewWriterLevel(&buf, gzip.BestCompression)
	require.NoError(t, err)
	_, err = gw.Write(payload)
	require.NoError(t, err)
	require.NoError(t, gw.Close())
	return buf.Bytes()
}

func TestDecompressGzip(t *testing.T) {
	t.Run("small payload", func(t *testing.T) {
		out, err := service.DecompressGzip(compress(t, []byte("hello world")))
		require.NoError(t, err)
		require.Equal(t, []byte("hello world"), out)
	})

	t.Run("payload exactly at cap", func(t *testing.T) {
		raw := make([]byte, http.DefaultMaxHeaderBytes)
		out, err := service.DecompressGzip(compress(t, raw))
		require.NoError(t, err)
		require.Len(t, out, http.DefaultMaxHeaderBytes)
	})

	t.Run("payload one byte over capd", func(t *testing.T) {
		raw := make([]byte, http.DefaultMaxHeaderBytes+1)
		_, err := service.DecompressGzip(compress(t, raw))
		require.ErrorIs(t, err, service.ErrGzipTooLarge)
	})

	t.Run("gzip decompression bomb", func(t *testing.T) {
		// 100 MB of zeros
		raw := make([]byte, 100<<20)
		compressed := compress(t, raw)
		require.Less(t, len(compressed), 1<<20,
			"sanity: bomb input should compress dramatically")

		_, err := service.DecompressGzip(compressed)
		require.ErrorIs(t, err, service.ErrGzipTooLarge)
	})

	t.Run("malformed gzip compression", func(t *testing.T) {
		_, err := service.DecompressGzip([]byte("not gzip data"))
		require.Error(t, err)
		require.Contains(t, err.Error(), "cannot read decompressed")
	})
}

// URL.Path and URL.RawPath must stay in step: once they disagree,
// EscapedPath() re-encodes from Path and every escape is lost.
func TestRemoveDoubleSlashes(t *testing.T) {
	cases := []struct {
		name   string
		target string
		want   string
	}{
		{"doubled slash is trimmed", "//agents/a/d/x", "/agents/a/d/x"},
		{"single slash untouched", "/agents/a/d/x", "/agents/a/d/x"},
		{"doubled slash keeps the encoded tail", "//agents/a/d/f%2Fg", "/agents/a/d/f%2Fg"},
		{"encoded name survives", "//agents/LODHA%20Vayam/d/x", "/agents/LODHA%20Vayam/d/x"},
		{"leading encoded slash is not doubled", "/%2Fagents/a/d/x", "/%2Fagents/a/d/x"},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodGet, c.target, nil)
			var got string
			service.RemoveDoubleSlashes(httptest.NewRecorder(), r, func(_ http.ResponseWriter, r *http.Request) {
				got = r.URL.EscapedPath()
			})
			require.Equal(t, c.want, got)
		})
	}
}

func TestGetClientIP(t *testing.T) {
	cases := []struct {
		name    string
		headers map[string]string
		want    string
	}{
		{"remote address", nil, "192.0.2.1"},
		{"cloudflare header wins", map[string]string{"CF-Connecting-IP": "203.0.113.7", "X-Forwarded-For": "198.51.100.2"}, "203.0.113.7"},
		{"single forwarded address", map[string]string{"X-Forwarded-For": "198.51.100.2"}, "198.51.100.2"},
		{"forwarded through two proxies", map[string]string{"X-Forwarded-For": "198.51.100.2, 10.0.0.5"}, "198.51.100.2"},
		{"forwarded list without spaces", map[string]string{"X-Forwarded-For": "198.51.100.2,10.0.0.5,10.0.0.6"}, "198.51.100.2"},
		{"forwarded ipv6", map[string]string{"X-Forwarded-For": "2001:db8::1, 10.0.0.5"}, "2001:db8::1"},
		{"real ip header", map[string]string{"X-Real-IP": "198.51.100.9"}, "198.51.100.9"},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodGet, "/rtc", nil)
			r.RemoteAddr = "192.0.2.1:54321"
			for k, v := range c.headers {
				r.Header.Set(k, v)
			}
			require.Equal(t, c.want, service.GetClientIP(r))
		})
	}
}