Add authentication for websockets

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
+	}
+}