s3ext: support to check if a key exists on S3

Signed-off-by: NAdam Lee <ali@pivotal.io>
Signed-off-by: NHaozhou Wang <hawang@pivotal.io>

gpAux/extensions/gps3ext/include/restful_service.h

@@ -19,6 +19,16 @@ enum ResponseStatus {
     RESPONSE_ERROR,  // server error (server return code is not 200)
 };
 
+typedef long ResponseCode;
+
+struct UploadData {
+    UploadData(const vector<uint8_t>& buff) : buffer(buff), currentPosition(0) {
+    }
+
+    const vector<uint8_t>& buffer;
+    uint64_t currentPosition;
+};
+
 class Response {
    public:
     Response() : status(RESPONSE_FAIL) {
@@ -82,6 +92,9 @@ class RESTfulService {
 
     virtual Response put(const string& url, HTTPHeaders& headers, const map<string, string>& params,
                          const vector<uint8_t>& data) = 0;
+
+    virtual ResponseCode head(const string& url, HTTPHeaders& headers,
+                              const map<string, string>& params) = 0;
 };
 
 #endif /* INCLUDE_RESTFUL_SERVICE_H_ */

gpAux/extensions/gps3ext/include/s3interface.h

@@ -76,6 +76,11 @@ class S3Interface {
                                                    const S3Credential& cred) {
         throw std::runtime_error("Default implementation must not be called.");
     }
+
+    virtual bool checkKeyExistence(const string& keyUrl, const string& region,
+                                   const S3Credential& cred) {
+        throw std::runtime_error("Default implementation must not be called.");
+    }
 };
 
 class S3Service : public S3Interface {
@@ -91,6 +96,8 @@ class S3Service : public S3Interface {
     S3CompressionType checkCompressionType(const string& keyUrl, const string& region,
                                            const S3Credential& cred);
 
+    bool checkKeyExistence(const string& keyUrl, const string& region, const S3Credential& cred);
+
     // following two functions are exposed publicly for UT tests
     void setRESTfulService(RESTfulService* restfullService) {
         this->restfulService = restfullService;

gpAux/extensions/gps3ext/include/s3restful_service.h

@@ -8,8 +8,11 @@ class S3RESTfulService : public RESTfulService {
     S3RESTfulService();
     virtual ~S3RESTfulService();
     Response get(const string& url, HTTPHeaders& headers, const map<string, string>& params);
+
     Response put(const string& url, HTTPHeaders& headers, const map<string, string>& params,
                  const vector<uint8_t>& data);
+
+    ResponseCode head(const string& url, HTTPHeaders& headers, const map<string, string>& params);
 };
 
 #endif /* INCLUDE_S3RESTFUL_SERVICE_H_ */

gpAux/extensions/gps3ext/src/s3interface.cpp

@@ -48,7 +48,7 @@ Response S3Service::getResponseWithRetries(const string &url, HTTPHeaders &heade
                                            const map<string, string> &params, uint64_t retries) {
     while (retries--) {
         // declare response here to leverage RVO (Return Value Optimization)
-        Response response = restfulService->get(url, headers, params);
+        Response response = this->restfulService->get(url, headers, params);
         if (response.isSuccess() || (retries == 0)) {
             return response;
         };
@@ -387,6 +387,26 @@ S3CompressionType S3Service::checkCompressionType(const string &keyUrl, const st
     return S3_COMPRESSION_PLAIN;
 }
 
+bool S3Service::checkKeyExistence(const string &keyUrl, const string &region,
+                                  const S3Credential &cred) {
+    HTTPHeaders headers;
+    map<string, string> params;
+    UrlParser parser(keyUrl);
+
+    headers.Add(HOST, parser.getHost());
+    headers.Add(X_AMZ_CONTENT_SHA256, "UNSIGNED-PAYLOAD");
+
+    SignRequestV4("HEAD", &headers, region, parser.getPath(), "", cred);
+
+    ResponseCode code = this->restfulService->head(keyUrl, headers, params);
+
+    if (code == 200 || code == 206) {
+        return true;
+    }
+
+    return false;
+}
+
 ListBucketResult::~ListBucketResult() {
     vector<BucketContent *>::iterator i;
     for (i = this->contents.begin(); i != this->contents.end(); i++) {

gpAux/extensions/gps3ext/src/s3restful_service.cpp

@@ -5,6 +5,7 @@
 #include <inttypes.h>
 #include <map>
 #include <string>
+#include <string.h>
 
 #include "gpcommon.h"
 #include "s3http_headers.h"
@@ -38,15 +39,7 @@ size_t RESTfulServiceWriteFuncCallback(char *ptr, size_t size, size_t nmemb, voi
     return realsize;
 }
 
-struct UploadData {
-    UploadData(const vector<uint8_t> &buff) : buffer(buff), currentPosition(0) {
-    }
-
-    const vector<uint8_t> &buffer;
-    uint64_t currentPosition;
-};
-
-// curl's read function callback.
+// curl's reading function callback.
 size_t RESTfulServiceReadFuncCallback(char *ptr, size_t size, size_t nmemb, void *userp) {
     if (QueryCancelPending) {
         return -1;
@@ -56,16 +49,16 @@ size_t RESTfulServiceReadFuncCallback(char *ptr, size_t size, size_t nmemb, void
     uint64_t dataLeft = data->buffer.size() - data->currentPosition;
 
     size_t requestedSize = size * nmemb;
-    size_t copyItemNum = requestedSize < dataLeft ? nmemb : (dataLeft / size);
-    size_t copySize = copyItemNum * size;
+    size_t copiedItemNum = requestedSize < dataLeft ? nmemb : (dataLeft / size);
+    size_t copiedDataSize = copiedItemNum * size;
 
-    if (copySize == 0) return 0;
+    if (copiedDataSize == 0) return 0;
 
-    memcpy(ptr, data->buffer.data() + data->currentPosition, copySize);
+    memcpy(ptr, data->buffer.data() + data->currentPosition, copiedDataSize);
 
-    data->currentPosition += copySize;
+    data->currentPosition += copiedDataSize;
 
-    return copyItemNum;
+    return copiedItemNum;
 }
 
 // get() will execute HTTP GET RESTful API with given url/headers/params,
@@ -146,7 +139,7 @@ Response S3RESTfulService::put(const string &url, HTTPHeaders &headers,
 
     headers.CreateList();
 
-    /* configs for download */
+    /* options for downloading */
     curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
     curl_easy_setopt(curl, CURLOPT_FORBID_REUSE, 1L);
     curl_easy_setopt(curl, CURLOPT_NOSIGNAL, 1L);
@@ -154,7 +147,7 @@ Response S3RESTfulService::put(const string &url, HTTPHeaders &headers,
     curl_easy_setopt(curl, CURLOPT_WRITEDATA, (void *)&response);
     curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, RESTfulServiceWriteFuncCallback);
 
-    /* configs for upload */
+    /* options for uploading */
     UploadData uploadData(data);
     curl_easy_setopt(curl, CURLOPT_READDATA, (void *)&uploadData);
     curl_easy_setopt(curl, CURLOPT_READFUNCTION, RESTfulServiceReadFuncCallback);
@@ -208,3 +201,49 @@ Response S3RESTfulService::put(const string &url, HTTPHeaders &headers,
 
     return response;
 }
+
+// head() will execute HTTP HEAD RESTful API with given url/headers/params, and return the HTTP
+// response code.
+//
+// Currently, this method only return the HTTP code, will be extended if needed in the future
+// implementation.
+ResponseCode S3RESTfulService::head(const string &url, HTTPHeaders &headers,
+                                    const map<string, string> &params) {
+    ResponseCode responseCode = -1;
+
+    CURL *curl = curl_easy_init();
+    CHECK_OR_DIE_MSG(curl != NULL, "%s", "Failed to create curl handler");
+
+    headers.CreateList();
+
+    curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
+    curl_easy_setopt(curl, CURLOPT_FORBID_REUSE, 1L);
+    curl_easy_setopt(curl, CURLOPT_NOSIGNAL, 1L);
+    curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers.GetList());
+
+    curl_easy_setopt(curl, CURLOPT_CUSTOMREQUEST, "HEAD");
+    curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);
+
+    map<string, string>::const_iterator iter = params.find("debug");
+    if (iter != params.end() && iter->second == "true") {
+        curl_easy_setopt(curl, CURLOPT_VERBOSE, 1L);
+    }
+
+    if (s3ext_debug_curl) {
+        curl_easy_setopt(curl, CURLOPT_VERBOSE, 1L);
+    }
+
+    CURLcode res = curl_easy_perform(curl);
+
+    if (res != CURLE_OK) {
+        S3ERROR("curl_easy_perform() failed: %s", curl_easy_strerror(res));
+    } else {
+        // Get the HTTP response status code from HTTP header
+        curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &responseCode);
+    }
+
+    curl_easy_cleanup(curl);
+    headers.FreeList();
+
+    return responseCode;
+}

gpAux/extensions/gps3ext/test/mock_classes.h

@@ -31,6 +31,9 @@ class MockS3RESTfulService : public S3RESTfulService {
    public:
     MOCK_METHOD3(get, Response(const string &url, HTTPHeaders &headers,
                                const map<string, string> &params));
+
+    MOCK_METHOD3(head, ResponseCode(const string &url, HTTPHeaders &headers,
+                               const map<string, string> &params));
 };
 
 class XMLGenerator {

gpAux/extensions/gps3ext/test/s3interface_test.cpp

@@ -393,3 +393,35 @@ TEST_F(S3ServiceTest, fetchDataWithResponseError) {
                      "https://s3-us-west-2.amazonaws.com/s3test.pivotal.io/whatever", region, cred),
                  std::runtime_error);
 }
+
+TEST_F(S3ServiceTest, HeadResponse200) {
+    string url = "https://s3-us-west-2.amazonaws.com/s3test.pivotal.io/whatever";
+
+    EXPECT_CALL(mockRestfulService, head(_, _, _)).WillOnce(Return(200));
+
+    EXPECT_TRUE(s3service->checkKeyExistence(url, this->region, this->cred));
+}
+
+TEST_F(S3ServiceTest, HeadResponse206) {
+    string url = "https://s3-us-west-2.amazonaws.com/s3test.pivotal.io/whatever";
+
+    EXPECT_CALL(mockRestfulService, head(_, _, _)).WillOnce(Return(206));
+
+    EXPECT_TRUE(s3service->checkKeyExistence(url, this->region, this->cred));
+}
+
+TEST_F(S3ServiceTest, HeadResponse404) {
+    string url = "https://s3-us-west-2.amazonaws.com/s3test.pivotal.io/whatever";
+
+    EXPECT_CALL(mockRestfulService, head(_, _, _)).WillOnce(Return(404));
+
+    EXPECT_FALSE(s3service->checkKeyExistence(url, this->region, this->cred));
+}
+
+TEST_F(S3ServiceTest, HeadResponse403) {
+    string url = "https://s3-us-west-2.amazonaws.com/s3test.pivotal.io/whatever";
+
+    EXPECT_CALL(mockRestfulService, head(_, _, _)).WillOnce(Return(403));
+
+    EXPECT_FALSE(s3service->checkKeyExistence(url, this->region, this->cred));
+}

gpAux/extensions/gps3ext/test/s3restful_service_test.cpp

@@ -167,3 +167,40 @@ TEST(S3RESTfulService, DISABLED_PutToDummyServer) {
     EXPECT_EQ(RESPONSE_OK, resp.getStatus());
     EXPECT_TRUE(compareVector(data, resp.getRawData()));
 }
+
+TEST(S3RESTfulService, HeadWithoutURL) {
+    HTTPHeaders headers;
+    map<string, string> params;
+    string url;
+    S3RESTfulService service;
+
+    ResponseCode code = service.head(url, headers, params);
+
+    EXPECT_EQ(-1, code);
+}
+
+TEST(S3RESTfulService, HeadWithWrongURL) {
+    HTTPHeaders headers;
+    map<string, string> params;
+    string url;
+    S3RESTfulService service;
+
+    url = "https://www.bing.com/pivotal.html";
+
+    ResponseCode code = service.head(url, headers, params);
+
+    EXPECT_EQ(404, code);
+}
+
+TEST(S3RESTfulService, HeadWithCorrectURL) {
+    HTTPHeaders headers;
+    map<string, string> params;
+    string url;
+    S3RESTfulService service;
+
+    url = "https://www.bing.com/";
+
+    ResponseCode code = service.head(url, headers, params);
+
+    EXPECT_EQ(200, code);
+}