Add capability to set headers to outgoing response

3 files changed, 120 insertions, 8 deletions

diff --git a/library/cpp/actors/http/http.cpp b/library/cpp/actors/http/http.cpp
index f5d7f5cf5c..a8e2c09922 100644
--- a/library/cpp/actors/http/http.cpp
+++ b/library/cpp/actors/http/http.cpp
@@ -452,7 +452,7 @@ THttpIncomingResponse::THttpIncomingResponse(THttpOutgoingRequestPtr request)
     : Request(request)
 {}
 
-THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status, TStringBuf message, TStringBuf contentType, TStringBuf body, TInstant lastModified) {
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateIncompleteResponse(TStringBuf status, TStringBuf message) {
     TStringBuf version = Version;
     if (version != "1.0" && version != "1.1") {
         version = "1.1";
@@ -462,19 +462,30 @@ THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status,
     if (!Endpoint->WorkerName.empty()) {
         response->Set("X-Worker-Name", Endpoint->WorkerName);
     }
-    if (!contentType.empty() && !body.empty()) {
-        response->Set<&THttpResponse::ContentType>(contentType);
+    return response;
+}
+
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateIncompleteResponse(TStringBuf status, TStringBuf message, const THeaders& headers) {
+    THttpOutgoingResponsePtr response = CreateIncompleteResponse(status, message);
+    response->Set(headers);
+    return response;
+}
+
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateIncompleteResponse(TStringBuf status, TStringBuf message, const THeaders& headers, TStringBuf body) {
+    THttpOutgoingResponsePtr response = CreateIncompleteResponse(status, message, headers);
+    if (!response->ContentType.empty() && !body.empty()) {
         if (!Endpoint->CompressContentTypes.empty()) {
-            contentType = contentType.Before(';');
+            TStringBuf contentType = response->ContentType.Before(';');
             Trim(contentType, ' ');
             if (Count(Endpoint->CompressContentTypes, contentType) != 0) {
                 response->EnableCompression();
             }
         }
     }
-    if (lastModified) {
-        response->Set<&THttpResponse::LastModified>(lastModified.FormatGmTime("%a, %d %b %Y %H:%M:%S GMT"));
-    }
+    return response;
+}
+
+void THttpIncomingRequest::FinishResponse(THttpOutgoingResponsePtr& response, TStringBuf body) {
     if (response->IsNeedBody() || !body.empty()) {
         if (Method == "HEAD") {
             response->Set<&THttpResponse::ContentLength>(ToString(body.size()));
@@ -482,9 +493,37 @@ THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status,
             response->SetBody(body);
         }
     }
+}
+
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status, TStringBuf message) {
+    THttpOutgoingResponsePtr response = CreateIncompleteResponse(status, message);
+    FinishResponse(response);
+    return response;
+}
+
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status, TStringBuf message, const THeaders& headers) {
+    THttpOutgoingResponsePtr response = CreateIncompleteResponse(status, message, headers);
+    FinishResponse(response);
+    return response;
+}
+
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status, TStringBuf message, const THeaders& headers, TStringBuf body) {
+    THttpOutgoingResponsePtr response = CreateIncompleteResponse(status, message, headers, body);
+    FinishResponse(response, body);
     return response;
 }
 
+THttpOutgoingResponsePtr THttpIncomingRequest::CreateResponse(TStringBuf status, TStringBuf message, TStringBuf contentType, TStringBuf body, TInstant lastModified) {
+    NHttp::THeadersBuilder headers;
+    if (!contentType.empty() && !body.empty()) {
+        headers.Set("Content-Type", contentType);
+    }
+    if (lastModified) {
+        headers.Set("Last-Modified", lastModified.FormatGmTime("%a, %d %b %Y %H:%M:%S GMT"));
+    }
+    return CreateResponse(status, message, headers, body);
+}
+
 THttpIncomingRequestPtr THttpIncomingRequest::Duplicate() {
     THttpIncomingRequestPtr request = new THttpIncomingRequest(*this);
     request->Reparse();
diff --git a/library/cpp/actors/http/http.h b/library/cpp/actors/http/http.h
index 295b8f97b5..b7227f6ca7 100644
--- a/library/cpp/actors/http/http.h
+++ b/library/cpp/actors/http/http.h
@@ -741,14 +741,24 @@ public:
     THttpOutgoingResponsePtr CreateResponseNotFound(TStringBuf html = TStringBuf(), TStringBuf contentType = "text/html"); // 404
     THttpOutgoingResponsePtr CreateResponseServiceUnavailable(TStringBuf html = TStringBuf(), TStringBuf contentType = "text/html"); // 503
     THttpOutgoingResponsePtr CreateResponseGatewayTimeout(TStringBuf html = TStringBuf(), TStringBuf contentType = "text/html"); // 504
+    THttpOutgoingResponsePtr CreateResponse(TStringBuf status, TStringBuf message);
+    THttpOutgoingResponsePtr CreateResponse(TStringBuf status, TStringBuf message, const THeaders& headers);
+    THttpOutgoingResponsePtr CreateResponse(TStringBuf status, TStringBuf message, const THeaders& headers, TStringBuf body);
     THttpOutgoingResponsePtr CreateResponse(
             TStringBuf status,
             TStringBuf message,
-            TStringBuf contentType = TStringBuf(),
+            TStringBuf contentType,
             TStringBuf body = TStringBuf(),
             TInstant lastModified = TInstant());
 
+    THttpOutgoingResponsePtr CreateIncompleteResponse(TStringBuf status, TStringBuf message);
+    THttpOutgoingResponsePtr CreateIncompleteResponse(TStringBuf status, TStringBuf message, const THeaders& headers);
+    THttpOutgoingResponsePtr CreateIncompleteResponse(TStringBuf status, TStringBuf message, const THeaders& headers, TStringBuf body);
+
     THttpIncomingRequestPtr Duplicate();
+
+private:
+    void FinishResponse(THttpOutgoingResponsePtr& response, TStringBuf body = TStringBuf());
 };
 
 class THttpIncomingResponse;
diff --git a/library/cpp/actors/http/http_ut.cpp b/library/cpp/actors/http/http_ut.cpp
index 91921c9120..e06de07867 100644
--- a/library/cpp/actors/http/http_ut.cpp
+++ b/library/cpp/actors/http/http_ut.cpp
@@ -230,6 +230,69 @@ Y_UNIT_TEST_SUITE(HttpProxy) {
         UNIT_ASSERT_VALUES_EQUAL(requestData, "GET /data/url HTTP/1.1\r\nHost: www.yandex.ru\r\nAccept: */*\r\nCookie: cookie1=123456; cookie2=45678;\r\n");
     }
 
+    Y_UNIT_TEST(BasicRenderOutgoingResponse) {
+        NHttp::THttpIncomingRequestPtr request = new NHttp::THttpIncomingRequest();
+        EatWholeString(request, "GET /test HTTP/1.1\r\nHost: test\r\nSome-Header: 32344\r\n\r\n");
+
+        NHttp::THttpOutgoingResponsePtr httpResponseOk = request->CreateResponseOK("response ok");
+        UNIT_ASSERT_EQUAL(httpResponseOk->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseOk->Status, "200");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseOk->Message, "OK");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseOk->ContentType, "text/html");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseOk->Body, "response ok");
+
+        NHttp::THttpOutgoingResponsePtr httpResponseBadRequest = request->CreateResponseBadRequest();
+        UNIT_ASSERT_EQUAL(httpResponseBadRequest->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseBadRequest->Status, "400");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseBadRequest->Message, "Bad Request");
+        UNIT_ASSERT(httpResponseBadRequest->ContentType.empty());
+        UNIT_ASSERT(httpResponseBadRequest->Body.empty());
+
+        NHttp::THttpOutgoingResponsePtr httpResponseNotFound = request->CreateResponseNotFound();
+        UNIT_ASSERT_EQUAL(httpResponseNotFound->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseNotFound->Status, "404");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseNotFound->Message, "Not Found");
+        UNIT_ASSERT(httpResponseNotFound->ContentType.empty());
+        UNIT_ASSERT(httpResponseNotFound->Body.empty());
+
+        NHttp::THttpOutgoingResponsePtr httpResponseServiceUnavailable = request->CreateResponseServiceUnavailable();
+        UNIT_ASSERT_EQUAL(httpResponseServiceUnavailable->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseServiceUnavailable->Status, "503");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseServiceUnavailable->Message, "Service Unavailable");
+        UNIT_ASSERT(httpResponseServiceUnavailable->ContentType.empty());
+        UNIT_ASSERT(httpResponseServiceUnavailable->Body.empty());
+
+        NHttp::THttpOutgoingResponsePtr httpResponseGatewayTimeout = request->CreateResponseGatewayTimeout("gateway timeout body");
+        UNIT_ASSERT_EQUAL(httpResponseGatewayTimeout->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseGatewayTimeout->Status, "504");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseGatewayTimeout->Message, "Gateway Timeout");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseGatewayTimeout->ContentType, "text/html");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseGatewayTimeout->Body, "gateway timeout body");
+
+        NHttp::THttpOutgoingResponsePtr httpIncompleteResponse = request->CreateIncompleteResponse("401", "Unauthorized");
+        UNIT_ASSERT_EQUAL(httpIncompleteResponse->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Header);
+        UNIT_ASSERT_STRINGS_EQUAL(httpIncompleteResponse->Status, "401");
+        UNIT_ASSERT_STRINGS_EQUAL(httpIncompleteResponse->Message, "Unauthorized");
+
+        NHttp::THttpOutgoingResponsePtr httpResponse = request->CreateResponse("401", "Unauthorized");
+        UNIT_ASSERT_EQUAL(httpResponse->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponse->Status, "401");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponse->Message, "Unauthorized");
+
+        NHttp::THeadersBuilder headers;
+        NHttp::TCookiesBuilder cookies;
+        cookies.Set("cookie1", "123456");
+        headers.Set("Set-Cookie", cookies.Render());
+        headers.Set("Location", "http://www.yandex.ru/data/url");
+
+        NHttp::THttpOutgoingResponsePtr httpResponseRedirect = request->CreateResponse("302", "Found", headers);
+        UNIT_ASSERT_EQUAL(httpResponseRedirect->Stage, NHttp::THttpOutgoingResponse::ERenderStage::Done);
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseRedirect->Status, "302");
+        UNIT_ASSERT_STRINGS_EQUAL(httpResponseRedirect->Message, "Found");
+        UNIT_ASSERT_STRING_CONTAINS(httpResponseRedirect->Headers, "Set-Cookie: cookie1=123456;");
+        UNIT_ASSERT_STRING_CONTAINS(httpResponseRedirect->Headers, "Location: http://www.yandex.ru/data/url");
+    }
+
     Y_UNIT_TEST(BasicRunning4) {
         NActors::TTestActorRuntimeBase actorSystem;
         TPortManager portManager;