Web优化 2

internal/transport/http/middleware/auth.go

@@ -100,16 +100,26 @@ func UserIDFromContext(ctx context.Context) (uuid.UUID, bool) {
 // because RFC 6750 specifies the literal "Bearer" scheme name and being
 // strict here keeps the parsing surface small.
 //
-// As a fallback, ?token= query parameter is accepted for endpoints that
-// cannot send custom headers (notably browser-native WebSocket connections).
-// Callers must keep token lifetimes short so the URL-borne token is not a
-// long-lived secret should it be logged by intermediaries.
+// For browser-native WebSocket connections, custom HTTP headers cannot be
+// sent during the handshake. Two fallbacks are supported, in order of
+// preference:
+//  1. Sec-WebSocket-Protocol: the client sends a subprotocol entry
+//     "acw-token.<tok>". The token is extracted from the header, which is
+//     not logged by proxies/servers as readily as a URL query parameter.
+//  2. ?token= query parameter (legacy): kept for backward compatibility.
 func extractBearer(r *http.Request) string {
 	h := r.Header.Get("Authorization")
 	const prefix = "Bearer "
 	if strings.HasPrefix(h, prefix) {
 		return strings.TrimSpace(h[len(prefix):])
 	}
+	const wsTokenPrefix = "acw-token."
+	for _, p := range strings.Split(r.Header.Get("Sec-WebSocket-Protocol"), ",") {
+		p = strings.TrimSpace(p)
+		if strings.HasPrefix(p, wsTokenPrefix) {
+			return p[len(wsTokenPrefix):]
+		}
+	}
 	if q := r.URL.Query().Get("token"); q != "" {
 		return q
 	}

internal/transport/http/middleware/auth_test.go

@@ -77,3 +77,18 @@ func TestAuth_BadTokenReturns401(t *testing.T) {
 	h.ServeHTTP(rec, req)
 	require.Equal(t, http.StatusUnauthorized, rec.Code)
 }
+
+func TestAuth_WebSocketTokenFromSubprotocol(t *testing.T) {
+	uid := uuid.New()
+	mw := Auth(fakeAuth{user: uid}, AuthOptions{})
+	h := mw(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+		got, ok := UserIDFromContext(r.Context())
+		require.True(t, ok)
+		require.Equal(t, uid, got)
+	}))
+	req := httptest.NewRequest("GET", "/", nil)
+	req.Header.Set("Sec-WebSocket-Protocol", "acp-v1, acw-token.ws-tok")
+	rec := httptest.NewRecorder()
+	h.ServeHTTP(rec, req)
+	require.Equal(t, http.StatusOK, rec.Code)
+}