Add support for limiting the number of redirects

bas-vk · mrtazz · commit 731a2129dd25 · 2018-06-11T11:27:12.000+02:00
diff --git a/README.md b/README.md
@@ -65,6 +65,8 @@ conn->SetUserAgent("foo/cool");
 
 // enable following of redirects (default is off)
 conn->FollowRedirects(true);
+// and limit the number of redirects (default is -1, unlimited)
+conn->FollowRedirects(true, 3);
 
 // set headers
 RestClient::HeaderFields headers;
diff --git a/include/restclient-cpp/connection.h b/include/restclient-cpp/connection.h
@@ -79,6 +79,8 @@ class Connection {
       *  Member 'timeout' contains the configured timeout
       *  @var Info::followRedirects
       *  Member 'followRedirects' contains whether or not to follow redirects
+      *  @var Info::maxRedirects
+      *  Member 'maxRedirects' contains the maximum number of redirect to follow (-1 unlimited)
       *  @var Info::basicAuth
       *  Member 'basicAuth' contains information about basic auth
       *  @var basicAuth::username
@@ -105,6 +107,7 @@ class Connection {
       RestClient::HeaderFields headers;
       int timeout;
       bool followRedirects;
+      int maxRedirects;
       bool noSignal;
       struct {
         std::string username;
@@ -138,6 +141,9 @@ class Connection {
     // set whether to follow redirects
     void FollowRedirects(bool follow);
 
+    // set whether to follow redirects (-1 for unlimited)
+    void FollowRedirects(bool follow, int maxRedirects);
+
     // set custom user agent
     // (this will result in the UA "foo/cool restclient-cpp/VERSION")
     void SetUserAgent(const std::string& userAgent);
@@ -191,6 +197,7 @@ class Connection {
     RestClient::HeaderFields headerFields;
     int timeout;
     bool followRedirects;
+    int maxRedirects;
     bool noSignal;
     struct {
       std::string username;
diff --git a/source/connection.cc b/source/connection.cc
@@ -34,6 +34,7 @@ RestClient::Connection::Connection(const std::string& baseUrl)
   this->baseUrl = baseUrl;
   this->timeout = 0;
   this->followRedirects = false;
+  this->maxRedirects = -1l;
   this->noSignal = false;
 }
 
@@ -57,6 +58,7 @@ RestClient::Connection::GetInfo() {
   ret.headers = this->GetHeaders();
   ret.timeout = this->timeout;
   ret.followRedirects = this->followRedirects;
+  ret.maxRedirects = this->maxRedirects;
   ret.noSignal = this->noSignal;
   ret.basicAuth.username = this->basicAuth.username;
   ret.basicAuth.password = this->basicAuth.password;
@@ -120,6 +122,19 @@ RestClient::Connection::GetHeaders() {
 void
 RestClient::Connection::FollowRedirects(bool follow) {
   this->followRedirects = follow;
+  this->maxRedirects = -1l;
+}
+
+/**
+ * @brief configure whether to follow redirects on this connection
+ *
+ * @param follow - boolean whether to follow redirects
+ * @param maxRedirects - int indicating the maximum number of redirect to follow (-1 unlimited)
+ */
+void
+RestClient::Connection::FollowRedirects(bool follow, int maxRedirects) {
+  this->followRedirects = follow;
+  this->maxRedirects = maxRedirects;
 }
 
 /**
@@ -327,6 +342,8 @@ RestClient::Connection::performCurlRequest(const std::string& uri) {
   // set follow redirect
   if (this->followRedirects == true) {
     curl_easy_setopt(this->curlHandle, CURLOPT_FOLLOWLOCATION, 1L);
+    curl_easy_setopt(this->curlHandle, CURLOPT_MAXREDIRS,
+                     static_cast<int64_t>(this->maxRedirects));
   }
 
   if (this->noSignal) {
diff --git a/test/test_connection.cc b/test/test_connection.cc
@@ -185,6 +185,21 @@ TEST_F(ConnectionTest, TestFollowRedirect)
   EXPECT_EQ(200, res.code);
 }
 
+TEST_F(ConnectionTest, TestFollowRedirectLimited)
+{
+  RestClient::Response res = conn->get("/redirect/2");
+  EXPECT_EQ(302, res.code);
+  conn->FollowRedirects(true, 1);
+  res = conn->get("/redirect/2");
+  EXPECT_EQ(-1, res.code);
+  conn->FollowRedirects(true, 2);
+  res = conn->get("/redirect/2");
+  EXPECT_EQ(200, res.code);
+  conn->FollowRedirects(true, -1);
+  res = conn->get("/redirect/2");
+  EXPECT_EQ(200, res.code);
+}
+
 TEST_F(ConnectionTest, TestGetInfoFromRedirect)
 {
   conn->FollowRedirects(true);
