Add authentication for websockets

api/example.env

@@ -5,6 +5,8 @@ MYSQL_URL=user:password@tcp(localhost)/dbname
 API_ENDPOINT=http://localhost:4001/
 API_LISTEN_PORT=4001
 
+PUBLIC_ENDPOINT=http://localhost:1234/
+
 # Set to 1 to set server in development mode, this enables graphql playground
 # Remove this if running in production
 DEVELOPMENT=1

api/go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/go-sql-driver/mysql v1.5.0
 	github.com/golang-migrate/migrate v3.5.4+incompatible
 	github.com/gorilla/mux v1.7.4
+	github.com/gorilla/websocket v1.4.1
 	github.com/h2non/filetype v1.0.10
 	github.com/joho/godotenv v1.3.0
 	github.com/lib/pq v1.3.0

api/go.sum

@@ -22,6 +22,8 @@ github.com/gorilla/mux v1.7.4 h1:VuZ8uybHlWmqV03+zRzdwKL4tUnIp1MAQtp1mIFE1bc=
 github.com/gorilla/mux v1.7.4/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
 github.com/gorilla/websocket v1.2.0 h1:VJtLvh6VQym50czpZzx07z/kw9EgAxI3x1ZB8taTMQQ=
 github.com/gorilla/websocket v1.2.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
+github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
+github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/h2non/filetype v1.0.10 h1:z+SJfnL6thYJ9kAST+6nPRXp1lMxnOVbMZHNYHMar0s=
 github.com/h2non/filetype v1.0.10/go.mod h1:isekKqOuhMj+s/7r3rIeTErIRy4Rub5uBWHfvMusLMU=
 github.com/hashicorp/golang-lru v0.5.0 h1:CL2msUPvZTLb5O648aiLNJw3hnBxN2+1Jq8rCOH9wdo=

api/graphql/auth/auth.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"regexp"
 
+	"github.com/99designs/gqlgen/handler"
 	"github.com/viktorstrate/photoview/api/graphql/models"
 )
 
@@ -32,16 +33,14 @@ func Middleware(db *sql.DB) func(http.Handler) http.Handler {
 				return
 			}
 
-			regex, _ := regexp.Compile("^Bearer ([a-zA-Z0-9]{24})$")
-			matches := regex.FindStringSubmatch(bearer)
-			if len(matches) != 2 {
+			token, err := TokenFromBearer(&bearer)
+			if err != nil {
+				log.Printf("Invalid bearer format: %s\n", bearer)
 				http.Error(w, "Invalid authorization header format", http.StatusBadRequest)
 				return
 			}
 
-			token := matches[1]
-
-			user, err := models.VerifyTokenAndGetUser(db, token)
+			user, err := models.VerifyTokenAndGetUser(db, *token)
 			if err != nil {
 				log.Printf("Invalid token: %s\n", err)
 				http.Error(w, "Invalid authorization token", http.StatusForbidden)
@@ -58,8 +57,47 @@ func Middleware(db *sql.DB) func(http.Handler) http.Handler {
 	}
 }
 
+func TokenFromBearer(bearer *string) (*string, error) {
+	regex, _ := regexp.Compile("^Bearer ([a-zA-Z0-9]{24})$")
+	matches := regex.FindStringSubmatch(*bearer)
+	if len(matches) != 2 {
+		return nil, errors.New("invalid bearer format")
+	}
+
+	token := matches[1]
+	return &token, nil
+}
+
 // UserFromContext finds the user from the context. REQUIRES Middleware to have run.
 func UserFromContext(ctx context.Context) *models.User {
 	raw, _ := ctx.Value(userCtxKey).(*models.User)
 	return raw
 }
+
+func AuthWebsocketInit(db *sql.DB) func(context.Context, handler.InitPayload) (context.Context, error) {
+	return func(ctx context.Context, initPayload handler.InitPayload) (context.Context, error) {
+
+		bearer, exists := initPayload["Authorization"].(string)
+		if !exists {
+			return ctx, nil
+		}
+
+		token, err := TokenFromBearer(&bearer)
+		if err != nil {
+			log.Printf("Invalid bearer format (websocket): %s\n", bearer)
+			return nil, err
+		}
+
+		user, err := models.VerifyTokenAndGetUser(db, *token)
+		if err != nil {
+			log.Printf("Invalid token in websocket: %s\n", err)
+			return nil, errors.New("invalid authorization token")
+		}
+
+		// put it in context
+		userCtx := context.WithValue(ctx, userCtxKey, user)
+
+		// and return it so the resolvers can see it
+		return userCtx, nil
+	}
+}

api/graphql/models/utils.go

@@ -2,7 +2,6 @@ package models
 
 import (
 	"fmt"
-	"log"
 	"regexp"
 	"strings"
 )
@@ -15,7 +14,6 @@ func (filter *Filter) FormatSQL() (string, error) {
 	result := ""
 
 	if filter.Limit != nil {
-		log.Println("Limit")
 		offset := 0
 		if filter.Offset != nil && *filter.Offset >= 0 {
 			offset = *filter.Offset
@@ -25,7 +23,6 @@ func (filter *Filter) FormatSQL() (string, error) {
 	}
 
 	if filter.OrderBy != nil {
-		log.Println("Order by")
 		order_by := filter.OrderBy
 		match, err := regexp.MatchString("^(\\w+(,\\s*))*\\w+$", strings.TrimSpace(*filter.OrderBy))
 		if err != nil {
@@ -42,6 +39,6 @@ func (filter *Filter) FormatSQL() (string, error) {
 		}
 	}
 
-	log.Printf("Filter: " + result)
+	// log.Printf("Filter: " + result)
 	return result, nil
 }

api/server.go

@@ -9,8 +9,6 @@ import (
 
 	"github.com/gorilla/mux"
 
-	// "github.com/go-chi/chi/middleware"
-	// "github.com/go-chi/cors"
 	"github.com/joho/godotenv"
 
 	"github.com/viktorstrate/photoview/api/database"
@@ -53,14 +51,6 @@ func main() {
 
 	rootRouter.Use(server.CORSMiddleware(devMode))
 
-	// router.Use(cors.New(cors.Options{
-	// 	AllowedOrigins:   []string{"http://localhost:4001", "http://localhost:1234", "*"},
-	// 	AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
-	// 	AllowedHeaders:   []string{"Accept", "Authorization", "Content-Type"},
-	// 	AllowCredentials: true,
-	// 	Debug:            false,
-	// }).Handler)
-
 	graphqlResolver := resolvers.Resolver{Database: db}
 	graphqlDirective := photoview_graphql.DirectiveRoot{}
 	graphqlDirective.IsAdmin = photoview_graphql.IsAdmin(db)
@@ -86,7 +76,13 @@ func main() {
 		})
 	}
 
-	endpointRouter.Handle("/graphql", handler.GraphQL(photoview_graphql.NewExecutableSchema(graphqlConfig), handler.IntrospectionEnabled(devMode)))
+	endpointRouter.Handle("/graphql",
+		handler.GraphQL(photoview_graphql.NewExecutableSchema(graphqlConfig),
+			handler.IntrospectionEnabled(devMode),
+			handler.WebsocketUpgrader(server.WebsocketUpgrader(devMode)),
+			handler.WebsocketInitFunc(auth.AuthWebsocketInit(db)),
+		),
+	)
 
 	photoRouter := endpointRouter.PathPrefix("/photo").Subrouter()
 	routes.RegisterPhotoRoutes(db, photoRouter)

api/server/websocket.go

@@ -0,0 +1,43 @@
+package server
+
+import (
+	"log"
+	"net/http"
+	"net/url"
+	"os"
+
+	"github.com/gorilla/websocket"
+)
+
+func WebsocketUpgrader(devMode bool) websocket.Upgrader {
+	return websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool {
+			if devMode {
+				return true
+			} else {
+				pubEndpoint, err := url.Parse(os.Getenv("PUBLIC_ENDPOINT"))
+				if err != nil {
+					log.Printf("Could not parse API_ENDPOINT environment variable as url: %s", err)
+					return false
+				}
+
+				if r.Header.Get("origin") == "" {
+					return true
+				}
+
+				originURL, err := url.Parse(r.Header.Get("origin"))
+				if err != nil {
+					log.Printf("Could not parse origin header of websocket request: %s", err)
+					return false
+				}
+
+				if pubEndpoint.Host == originURL.Host {
+					return true
+				} else {
+					log.Printf("Not allowing websocket request from %s because it doesn't match PUBLIC_ENDPOINT %s", originURL.Host, pubEndpoint.Host)
+					return false
+				}
+			}
+		},
+	}
+}

ui/src/components/messages/SubscriptionsHook.js

@@ -3,19 +3,22 @@ import PropTypes from 'prop-types'
 import { useSubscription } from 'react-apollo'
 import gql from 'graphql-tag'
 
-const syncSubscription = gql`
-  subscription syncSubscription {
-    scannerStatusUpdate {
-      finished
-      success
-      message
+const notificationSubscription = gql`
+  subscription notificationSubscription {
+    notification {
+      key
+      type
+      header
+      content
       progress
+      positive
+      negative
     }
   }
 `
 
 const SubscriptionsHook = ({ messages, setMessages }) => {
-  const { data, error } = useSubscription(syncSubscription)
+  const { data, error } = useSubscription(notificationSubscription)
 
   useEffect(() => {
     if (error) {