OIDC needs pass tracing headers (#19653)oidc-1.1.2

2 files changed, 65 insertions, 44 deletions

diff --git a/ydb/mvp/oidc_proxy/oidc_protected_page.cpp b/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
index a397199753b..ee41ca32ffc 100644
--- a/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
+++ b/ydb/mvp/oidc_proxy/oidc_protected_page.cpp
@@ -39,14 +39,7 @@ void THandlerSessionServiceCheck::HandleProxy(NHttp::TEvHttpProxy::TEvHttpIncomi
         if (NeedSendSecureHttpRequest(response)) {
             return SendSecureHttpRequest(response);
         }
-        NHttp::THeadersBuilder headers = GetResponseHeaders(response);
-        TStringBuf contentType = headers.Get("Content-Type").NextTok(';');
-        if (contentType == "text/html") {
-            TString newBody = FixReferenceInHtml(response->Body, response->GetRequest()->Host);
-            return ReplyAndPassAway(Request->CreateResponse(response->Status, response->Message, headers, newBody));
-        } else {
-            return ReplyAndPassAway(Request->CreateResponse(response->Status, response->Message, headers, response->Body));
-        }
+        return ReplyAndPassAway(CreateProxiedResponse(response));
     } else {
         static constexpr size_t MAX_LOGGED_SIZE = 1024;
         BLOG_D("Can not process request to protected resource:\n" << event->Get()->Request->GetObfuscatedData().substr(0, MAX_LOGGED_SIZE));
@@ -135,17 +128,7 @@ bool THandlerSessionServiceCheck::IsAuthorizedRequest(TStringBuf authHeader) {
 
 void THandlerSessionServiceCheck::ForwardUserRequest(TStringBuf authHeader, bool secure) {
     BLOG_D("Forward user request bypass OIDC");
-    NHttp::THttpOutgoingRequestPtr httpRequest = NHttp::THttpOutgoingRequest::CreateRequest(Request->Method, ProtectedPageUrl);
-    ForwardRequestHeaders(httpRequest);
-    if (!authHeader.empty()) {
-        httpRequest->Set(AUTH_HEADER_NAME, authHeader);
-    }
-    if (Request->HaveBody()) {
-        httpRequest->SetBody(Request->Body);
-    }
-    if (RequestedPageScheme.empty()) {
-        httpRequest->Secure = secure;
-    }
+    auto httpRequest = CreateProxiedRequest(authHeader, secure);
 
     auto requestEvent = std::make_unique<NHttp::TEvHttpProxy::TEvHttpOutgoingRequest>(httpRequest);
     requestEvent->Timeout = TDuration::Seconds(120);
@@ -190,7 +173,9 @@ TString THandlerSessionServiceCheck::FixReferenceInHtml(TStringBuf html, TString
     return FixReferenceInHtml(result, host, findString);
 }
 
-void THandlerSessionServiceCheck::ForwardRequestHeaders(NHttp::THttpOutgoingRequestPtr& request) const {
+NHttp::THttpOutgoingRequestPtr THandlerSessionServiceCheck::CreateProxiedRequest(TStringBuf authHeader, bool secure) const {
+    auto request = NHttp::THttpOutgoingRequest::CreateRequest(Request->Method, ProtectedPageUrl);
+
     static const TVector<TStringBuf> HEADERS_WHITE_LIST = {
         "Connection",
         "Accept-Language",
@@ -201,7 +186,10 @@ void THandlerSessionServiceCheck::ForwardRequestHeaders(NHttp::THttpOutgoingRequ
         "Sec-Fetch-User",
         "Upgrade-Insecure-Requests",
         "Content-Type",
-        "Origin"
+        "Origin",
+        "X-Trace-Verbosity",
+        "X-Want-Trace",
+        "traceparent"
     };
     NHttp::THeadersBuilder headers(Request->Headers);
     for (const auto& header : HEADERS_WHITE_LIST) {
@@ -210,9 +198,23 @@ void THandlerSessionServiceCheck::ForwardRequestHeaders(NHttp::THttpOutgoingRequ
         }
     }
     request->Set("Accept-Encoding", "deflate");
+
+    if (!authHeader.empty()) {
+        request->Set(AUTH_HEADER_NAME, authHeader);
+    }
+    if (Request->HaveBody()) {
+        request->SetBody(Request->Body);
+    }
+    if (RequestedPageScheme.empty()) {
+        request->Secure = secure;
+    }
+
+    return request;
 }
 
-NHttp::THeadersBuilder THandlerSessionServiceCheck::GetResponseHeaders(const NHttp::THttpIncomingResponsePtr& response) {
+NHttp::THeadersBuilder THandlerSessionServiceCheck::ProxyResponseHeaders(const NHttp::THttpIncomingResponsePtr& response) {
+    NHttp::THeadersBuilder headers(response->Headers);
+    NHttp::THeadersBuilder outHeaders;
     static const TVector<TStringBuf> HEADERS_WHITE_LIST = {
         "Content-Type",
         "Connection",
@@ -221,20 +223,37 @@ NHttp::THeadersBuilder THandlerSessionServiceCheck::GetResponseHeaders(const NHt
         "Access-Control-Allow-Origin",
         "Access-Control-Allow-Credentials",
         "Access-Control-Allow-Headers",
-        "Access-Control-Allow-Methods"
+        "Access-Control-Allow-Methods",
+        "traceresponse"
     };
-    NHttp::THeadersBuilder headers(response->Headers);
-    NHttp::THeadersBuilder resultHeaders;
     for (const auto& header : HEADERS_WHITE_LIST) {
         if (headers.Has(header)) {
-            resultHeaders.Set(header, headers.Get(header));
+            outHeaders.Set(header, headers.Get(header));
         }
     }
+
     static const TString LOCATION_HEADER_NAME = "Location";
     if (headers.Has(LOCATION_HEADER_NAME)) {
-        resultHeaders.Set(LOCATION_HEADER_NAME, GetFixedLocationHeader(headers.Get(LOCATION_HEADER_NAME)));
+        outHeaders.Set(LOCATION_HEADER_NAME, GetFixedLocationHeader(headers.Get(LOCATION_HEADER_NAME)));
     }
-    return resultHeaders;
+
+    return outHeaders;
+}
+
+TString THandlerSessionServiceCheck::ProxyResponseBody(const NHttp::THttpIncomingResponsePtr& response) {
+    NHttp::THeadersBuilder headers(response->Headers);
+    TStringBuf contentType = headers.Get("Content-Type").NextTok(';');
+    if (contentType == "text/html") {
+        return FixReferenceInHtml(response->Body, response->GetRequest()->Host);
+    } else {
+        return TString(response->Body);
+    }
+}
+
+NHttp::THttpOutgoingResponsePtr THandlerSessionServiceCheck::CreateProxiedResponse(const NHttp::THttpIncomingResponsePtr& response) {
+    auto outHeaders = ProxyResponseHeaders(response);
+    auto outBody = ProxyResponseBody(response);
+    return Request->CreateResponse(response->Status, response->Message, outHeaders, outBody);
 }
 
 void THandlerSessionServiceCheck::SendSecureHttpRequest(const NHttp::THttpIncomingResponsePtr& response) {
@@ -246,21 +265,21 @@ void THandlerSessionServiceCheck::SendSecureHttpRequest(const NHttp::THttpIncomi
 
 TString THandlerSessionServiceCheck::GetFixedLocationHeader(TStringBuf location) {
     TStringBuf scheme, host, uri;
-        NHttp::CrackURL(ProtectedPageUrl, scheme, host, uri);
-        if (location.StartsWith("//")) {
-            return TStringBuilder() << '/' << (scheme.empty() ? "" : TString(scheme) + "://") << location.SubStr(2);
-        } else if (location.StartsWith('/')) {
-            return TStringBuilder() << '/'
-                                    << (scheme.empty() ? "" : TString(scheme) + "://")
-                                    << host << location;
-        } else {
-            TStringBuf locScheme, locHost, locUri;
-            NHttp::CrackURL(location, locScheme, locHost, locUri);
-            if (!locScheme.empty()) {
-                return TStringBuilder() << '/' << location;
-            }
+    NHttp::CrackURL(ProtectedPageUrl, scheme, host, uri);
+    if (location.StartsWith("//")) {
+        return TStringBuilder() << '/' << (scheme.empty() ? "" : TString(scheme) + "://") << location.SubStr(2);
+    } else if (location.StartsWith('/')) {
+        return TStringBuilder() << '/'
+                                << (scheme.empty() ? "" : TString(scheme) + "://")
+                                << host << location;
+    } else {
+        TStringBuf locScheme, locHost, locUri;
+        NHttp::CrackURL(location, locScheme, locHost, locUri);
+        if (!locScheme.empty()) {
+            return TStringBuilder() << '/' << location;
         }
-        return TString(location);
+    }
+    return TString(location);
 }
 
 NHttp::THttpOutgoingResponsePtr THandlerSessionServiceCheck::CreateResponseForbiddenHost() {
diff --git a/ydb/mvp/oidc_proxy/oidc_protected_page.h b/ydb/mvp/oidc_proxy/oidc_protected_page.h
index 41022702213..f8f1bb31d22 100644
--- a/ydb/mvp/oidc_proxy/oidc_protected_page.h
+++ b/ydb/mvp/oidc_proxy/oidc_protected_page.h
@@ -56,7 +56,7 @@ protected:
     virtual bool NeedSendSecureHttpRequest(const NHttp::THttpIncomingResponsePtr& response) const = 0;
 
     bool CheckRequestedHost();
-    void ForwardRequestHeaders(NHttp::THttpOutgoingRequestPtr& request) const;
+    NHttp::THttpOutgoingRequestPtr CreateProxiedRequest(TStringBuf authHeader, bool secure) const;
     void ReplyAndPassAway(NHttp::THttpOutgoingResponsePtr httpResponse);
 
     static bool IsAuthorizedRequest(TStringBuf authHeader);
@@ -64,7 +64,9 @@ protected:
     static TString FixReferenceInHtml(TStringBuf html, TStringBuf host);
 
 private:
-    NHttp::THeadersBuilder GetResponseHeaders(const NHttp::THttpIncomingResponsePtr& response);
+    NHttp::THeadersBuilder ProxyResponseHeaders(const NHttp::THttpIncomingResponsePtr& response);
+    TString ProxyResponseBody(const NHttp::THttpIncomingResponsePtr& response);
+    NHttp::THttpOutgoingResponsePtr CreateProxiedResponse(const NHttp::THttpIncomingResponsePtr& response);
     void SendSecureHttpRequest(const NHttp::THttpIncomingResponsePtr& response);
     TString GetFixedLocationHeader(TStringBuf location);
     NHttp::THttpOutgoingResponsePtr CreateResponseForbiddenHost();