From 0e129cf8239b07151566af33a6244b8b3844efb4 Mon Sep 17 00:00:00 2001
From: Sumner Evans <sumner@beeper.com>
Date: Wed, 28 Feb 2024 14:45:24 -0700
Subject: [PATCH] auth: add ability to authenticate downloads with JWT tokens

Signed-off-by: Sumner Evans <sumner@beeper.com>
---
 go.mod  |  7 ++++--
 go.sum  |  2 ++
 main.go | 74 +++++++++++++++++++++++++++++++++++++++++++--------------
 3 files changed, 63 insertions(+), 20 deletions(-)

diff --git a/go.mod b/go.mod
index 171d0ed..bb99e27 100644
--- a/go.mod
+++ b/go.mod
@@ -1,5 +1,8 @@
 module github.com/matrix-org/rageshake
 
-go 1.19
+go 1.22
 
-require gopkg.in/yaml.v2 v2.2.2
+require (
+	github.com/golang-jwt/jwt/v5 v5.2.0
+	gopkg.in/yaml.v2 v2.2.2
+)
diff --git a/go.sum b/go.sum
index e936db1..019b3fa 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/golang-jwt/jwt/v5 v5.2.0 h1:d/ix8ftRUorsN+5eMIlF4T6J8CAt9rch3My2winC1Jw=
+github.com/golang-jwt/jwt/v5 v5.2.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
diff --git a/main.go b/main.go
index f9925b2..fa05120 100644
--- a/main.go
+++ b/main.go
@@ -26,6 +26,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/golang-jwt/jwt/v5"
 	"gopkg.in/yaml.v2"
 )
 
@@ -37,6 +38,8 @@ type config struct {
 	BugsUser string `yaml:"listings_auth_user"`
 	BugsPass string `yaml:"listings_auth_pass"`
 
+	BugsJWTSecret string `yaml:"listings_jwt_secret"`
+
 	// External URI to /api
 	APIPrefix string `yaml:"api_prefix"`
 
@@ -47,15 +50,59 @@ type config struct {
 	WebhookURL string `yaml:"webhook_url"`
 }
 
-func basicAuth(handler http.Handler, username, password, realm string) http.Handler {
+const (
+	rageshakeIssuer = "com.beeper.rageshake"
+	apiServerIssuer = "com.beeper.api-server"
+)
+
+func basicAuthOrJWTAuthenticated(handler http.Handler, username, password, realm string, jwtSecret []byte) http.Handler {
+	if (username == "" || password == "") && len(jwtSecret) == 0 {
+		panic("Either username or password for basic auth must be set, or JWT secret must be set, or both")
+	}
+
+	unauthorized := func(w http.ResponseWriter) {
+		w.WriteHeader(401)
+		w.Write([]byte("Unauthorised.\n"))
+	}
+
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		user, pass, ok := r.BasicAuth() // pull creds from the request
-
-		// check user and pass securely
-		if !ok || subtle.ConstantTimeCompare([]byte(user), []byte(username)) != 1 || subtle.ConstantTimeCompare([]byte(pass), []byte(password)) != 1 {
+		if !ok && len(jwtSecret) == 0 { // if no basic auth and no JWT auth, return unauthorized
+			unauthorized(w)
+			return
+		} else if !ok { // if no basic auth, try to do JWT auth
+			token, err := jwt.ParseWithClaims(r.URL.Query().Get("tok"), &jwt.RegisteredClaims{}, func(token *jwt.Token) (any, error) {
+				// Don't forget to validate the alg is what you expect:
+				if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
+					return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
+				}
+				return jwtSecret, nil
+			})
+			if err != nil {
+				log.Printf("Error parsing JWT: %v", err)
+				unauthorized(w)
+				return
+			}
+
+			claims, ok := token.Claims.(*jwt.RegisteredClaims)
+			if !token.Valid || !ok {
+				log.Printf("Token invalid or claims not RegisteredClaims: %v", err)
+				unauthorized(w)
+				return
+			} else if claims.Issuer != rageshakeIssuer && claims.Issuer != apiServerIssuer {
+				log.Printf("Token issuer not rageshake or API server: %s", claims.Issuer)
+				unauthorized(w)
+				return
+			} else if claims.Subject != r.URL.Path {
+				log.Printf("Token subject (%s) not the request path (%s)", claims.Subject, r.URL.Path)
+				unauthorized(w)
+				return
+			}
+
+			log.Printf("Valid token from %s for accessing %s", claims.Issuer, claims.Subject)
+		} else if subtle.ConstantTimeCompare([]byte(user), []byte(username)) != 1 || subtle.ConstantTimeCompare([]byte(pass), []byte(password)) != 1 { // check user and pass securely
 			w.Header().Set("WWW-Authenticate", `Basic realm="`+realm+`"`)
-			w.WriteHeader(401)
-			w.Write([]byte("Unauthorised.\n"))
+			unauthorized(w)
 			return
 		}
 
@@ -99,19 +146,10 @@ func main() {
 
 	// serve files under "bugs"
 	ls := &logServer{"bugs"}
-	fs := http.StripPrefix("/api/listing/", ls)
-
-	// set auth if env vars exist
-	usr := cfg.BugsUser
-	pass := cfg.BugsPass
-	if usr == "" || pass == "" {
-		fmt.Println("No listings_auth_user/pass configured. No authentication is running for /api/listing")
-	} else {
-		fs = basicAuth(fs, usr, pass, "Riot bug reports")
-	}
-	http.Handle("/api/listing/", fs)
+	fs := basicAuthOrJWTAuthenticated(ls, cfg.BugsUser, cfg.BugsPass, "Riot bug reports", []byte(cfg.BugsJWTSecret))
+	http.Handle("/api/listing/", http.StripPrefix("/api/listing/", fs))
 
-	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
 		fmt.Fprint(w, "ok")
 	})