Merge pull request #14 from orsinium-labs/ws-auth

Support Authorization over WebSocket

Author: orsinium

josh.go

@@ -99,10 +99,6 @@ func Read[T any](t string, r io.Reader) (Data[T], error) {
 
 // Resp is a response type.
 //
-// The generic type T is the type of the data response.
-//
-// If the handler never returns data, use [Void] instead.
-//
 // https://jsonapi.org/format/#document-top-level
 type Resp struct {
 	// The response status code.

middlewares/auth.go

@@ -13,7 +13,16 @@ type AuthValidator[U any] func(string) (U, error)
 //
 // The validators input is the "Bearer" token provided in the "Authorization" request header.
 //
-// If the validator returns an error, that erro is immediately returned
+// If the "Authorization" header is not provided, we also check the "Sec-WebSocket-Protocol"
+// as a necessary workaround for authenticating WebSocket requests from browsers
+// because the browser WebSocket API doesn't allow setting custom headers:
+//
+// https://github.com/whatwg/websockets/issues/16
+//
+// The very first protocol after the "Authorization" protocol is treated as
+// the authorization token.
+//
+// If the validator returns an error, that error is immediately returned
 // as an "Unathorized" response. Otherwise, the returned value (typically, the user
 // or their ID) is added into the request context using [josh.WithSingleton].
 func Auth[U any](v AuthValidator[U], h josh.Handler) josh.Handler {
@@ -31,17 +40,43 @@ func Auth[U any](v AuthValidator[U], h josh.Handler) josh.Handler {
 
 func validateRequest[U any](validator AuthValidator[U], r josh.Req) (U, error) {
 	header := r.Header.Get("Authorization")
-	var def U
 	if header == "" {
-		return def, errors.New("Authorization header not found")
+		// If Authorization header is not provided,
+		// try authenticating it as a WebSocket request.
+		return validateWSRequest(validator, r)
 	}
 	token, hasPrefix := strings.CutPrefix(header, "Bearer ")
 	if !hasPrefix {
+		var def U
 		return def, errors.New("Unsupported Authorization type")
 	}
 	token = strings.TrimSpace(token)
 	if token == "" {
+		var def U
 		return def, errors.New("Authorization token is empty")
 	}
 	return validator(token)
 }
+
+func validateWSRequest[U any](validator AuthValidator[U], r josh.Req) (U, error) {
+	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Sec-WebSocket-Protocol
+	headers := r.Header.Values("Sec-WebSocket-Protocol")
+	var def U
+	if len(headers) == 0 {
+		return def, errors.New("Authorization header not found")
+	}
+
+	foundAuth := false
+	for _, tokens := range headers {
+		for _, token := range strings.Split(tokens, ",") {
+			token = strings.TrimSpace(token)
+			if token == "Authorization" {
+				foundAuth = true
+			} else if foundAuth {
+				return validator(token)
+			}
+		}
+	}
+
+	return def, errors.New("Authorization header not found")
+}

middlewares/auth_test.go

@@ -0,0 +1,99 @@
+package middlewares_test
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/orsinium-labs/josh"
+	"github.com/orsinium-labs/josh/middlewares"
+)
+
+func TestAuth(t *testing.T) {
+	var req josh.Req
+	url := "http://example.com/foo"
+
+	req = httptest.NewRequest("GET", url, nil)
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Authorization", "secret")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Authorization", "Bearer ohno")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Authorization", "Bearer secret")
+	checkAuth(t, req, 200)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "secret")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "secret, Authorization")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "Authorization, ohno")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "Authorization, secret")
+	checkAuth(t, req, 200)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "Authorization")
+	req.Header.Add("Sec-WebSocket-Protocol", "secret")
+	checkAuth(t, req, 200)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "secret")
+	req.Header.Add("Sec-WebSocket-Protocol", "Authorization")
+	checkAuth(t, req, 401)
+
+	req = httptest.NewRequest("GET", url, nil)
+	req.Header.Add("Sec-WebSocket-Protocol", "Authorization")
+	req.Header.Add("Sec-WebSocket-Protocol", "ohno")
+	req.Header.Add("Sec-WebSocket-Protocol", "secret")
+	checkAuth(t, req, 401)
+}
+
+func checkAuth(t *testing.T, req *http.Request, code int) {
+	t.Helper()
+	type User string
+	h := func(r josh.Req) josh.Resp {
+		u, err := josh.GetSingleton[User](r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if u != "Aragorn" {
+			t.Fatalf("got %s, expected Aragorn", u)
+		}
+		return josh.Ok("all is good")
+	}
+	v := func(token string) (User, error) {
+		if token == "secret" {
+			return "Aragorn", nil
+		}
+		return "", fmt.Errorf("bad token: %s", token)
+	}
+	h = middlewares.Auth(v, h)
+	hh := josh.Wrap(h)
+
+	w := httptest.NewRecorder()
+	hh(w, req)
+	resp := w.Result()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("read body: %v", err)
+	}
+	t.Log(string(body))
+	if resp.StatusCode != code {
+		t.Fatalf("got %d, expected %d", resp.StatusCode, code)
+	}
+}