From ada5172dd669545c63efc8761c3fe02e51121808 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <mathieu.velten@beta.gouv.fr>
Date: Tue, 13 Aug 2024 18:06:15 +0200
Subject: [PATCH] Add allowlist

---
 README.md |  2 ++
 main.go   | 38 +++++++++++++++++++++++++++-----------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index a87e95b..4637f79 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,8 @@ $ LIVEKIT_URL="ws://somewhere" LIVEKIT_KEY=devkey LIVEKIT_SECRET=secret go run *
 
 The listening port is configurable via the `LK_JWT_PORT` environment variable and defaults to 8080.
 
+Usage can be limited to a set of specific homeservers by setting `HS_ALLOWLIST` environment variable to a comma-separated list of server names.
+
 ## Disable TLS verification
 
 For testing and debugging (e.g., in the absence of trusted certificates while testing in a lab) you can disable TLS verification for outgoing matrix client connection by setting the environment variable `LIVEKIT_INSECURE_SKIP_VERIFY_TLS` to `YES_I_KNOW_WHAT_I_AM_DOING`.
diff --git a/main.go b/main.go
index 65dae2f..55a6f9f 100644
--- a/main.go
+++ b/main.go
@@ -23,6 +23,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"strings"
 
 	"time"
 
@@ -35,7 +36,8 @@ import (
 
 type Handler struct {
 	key, secret, lk_url string
-	skipVerifyTLS bool
+	skipVerifyTLS       bool
+	homeserverAllowList []string
 }
 
 type OpenIDTokenType struct {
@@ -56,10 +58,18 @@ type SFUResponse struct {
 }
 
 func exchangeOIDCToken(
-	ctx context.Context, token OpenIDTokenType, skipVerifyTLS bool,
-) (*fclient.UserInfo, error) {
+	ctx context.Context, token OpenIDTokenType, skipVerifyTLS bool, homeserverAllowList []string) (*fclient.UserInfo, error) {
 	if token.AccessToken == "" || token.MatrixServerName == "" {
-		return nil, errors.New("Missing parameters in OIDC token")
+		return nil, errors.New("missing parameters in OIDC token")
+	}
+
+	if len(homeserverAllowList) > 0 {
+		for _, allowed_hs := range homeserverAllowList {
+			if allowed_hs == token.MatrixServerName {
+				break
+			}
+		}
+		return nil, fmt.Errorf("homeserver %s not allowed", token.MatrixServerName)
 	}
 
 	if skipVerifyTLS {
@@ -73,7 +83,7 @@ func exchangeOIDCToken(
 	)
 	if err != nil {
 		log.Printf("Failed to look up user info: %v", err)
-		return nil, errors.New("Failed to look up user info")
+		return nil, errors.New("failed to look up user info")
 	}
 	return &userinfo, nil
 }
@@ -130,7 +140,7 @@ func (h *Handler) handle(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		userInfo, err := exchangeOIDCToken(r.Context(), body.OpenIDToken, h.skipVerifyTLS)
+		userInfo, err := exchangeOIDCToken(r.Context(), body.OpenIDToken, h.skipVerifyTLS, h.homeserverAllowList)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
 			err = json.NewEncoder(w).Encode(gomatrix.RespError{
@@ -184,6 +194,11 @@ func main() {
 	secret := os.Getenv("LIVEKIT_SECRET")
 	lk_url := os.Getenv("LIVEKIT_URL")
 
+	homeserverAllowList := []string{}
+	if os.Getenv("HS_ALLOWLIST") != "" {
+		homeserverAllowList = strings.Split(os.Getenv("HS_ALLOWLIST"), ",")
+	}
+
 	// Check if the key, secret or url are empty.
 	if key == "" || secret == "" || lk_url == "" {
 		log.Fatal("LIVEKIT_KEY, LIVEKIT_SECRET and LIVEKIT_URL environment variables must be set")
@@ -194,13 +209,14 @@ func main() {
 		lk_jwt_port = "8080"
 	}
 
-	log.Printf("LIVEKIT_KEY: %s, LIVEKIT_SECRET: %s, LIVEKIT_URL: %s, LK_JWT_PORT: %s", key, secret, lk_url, lk_jwt_port)
+	log.Printf("LIVEKIT_KEY: %s, LIVEKIT_SECRET: __redacted__, LIVEKIT_URL: %s, LK_JWT_PORT: %s, HS_ALLOWLIST: %s", key, lk_url, lk_jwt_port, homeserverAllowList)
 
 	handler := &Handler{
-		key:    key,
-		secret: secret,
-		lk_url: lk_url,
-		skipVerifyTLS: skipVerifyTLS,
+		key:                 key,
+		secret:              secret,
+		lk_url:              lk_url,
+		skipVerifyTLS:       skipVerifyTLS,
+		homeserverAllowList: homeserverAllowList,
 	}
 
 	http.HandleFunc("/sfu/get", handler.handle)