s3_glue.cpp

sleuthit-2.09 一个磁盘的工具集

#define CURLINFO_RESPONSE_CODE CURLINFO_HTTP_CODE
#endif

class response_buffer *request(string method,string path,string query,time_t expires,
			       const char *sendbuf,size_t sendbuflen,
			       const s3headers *extraheaders)
{
    /* Note: this function is not threadsafe */
    static bool curl_initted = false;
    if(!curl_initted){
	curl_global_init(CURL_GLOBAL_ALL);	
	curl_initted=true;
    }

    int retry_count=0;
    class response_buffer *b = 0;
    class buffer *h = 0;
    do {

	if(s3_debug>1) printf("==================================================\n");
	if(s3_debug && retry_count>0) printf("=== S3 RETRY %d ===\n",retry_count);
	
	CURL *c = curl_easy_init();
	struct curl_slist *headers=NULL;

	if(expires==0){
	    /* Add the Date: field to the header */
	    struct tm tm;
	    time_t t = time(0);
	    char date[64];
	    strftime(date,sizeof(date),"Date: %a, %d %b %Y %X GMT",gmtime_r(&t,&tm));
	    headers = curl_slist_append(headers, date);
	}

	/* Add the extra headers */
	while(extraheaders  && extraheaders[0].name){
	    int len = strlen(extraheaders[0].name)+strlen(extraheaders[0].value)+4;
	    char *buf = (char *)alloca(len);
	    snprintf(buf,len,"%s: %s",extraheaders[0].name,extraheaders[0].value);
	    headers = curl_slist_append(headers, buf);
	    extraheaders++;
	}

	string url = aws_base_url + path;
	string canonical_str     = canonical_string(method,path,headers,expires);
	string encoded_canonical = encode(aws_secret_access_key,canonical_str);

	if(expires==0){
	    /* Create an Authorization header */

	    char authorization[96];
	
	    snprintf(authorization,sizeof(authorization),"Authorization: AWS %s:%s",
		     aws_access_key_id,encoded_canonical.c_str());
	    headers = curl_slist_append(headers, authorization);
	    curl_easy_setopt(c, CURLOPT_HTTPHEADER, headers);
	}

	if(expires){
	    /* Add authorization to the URL*/
	    if(query.size()>0) query += "&";
	    query += "Signature=" + quote_plus(encoded_canonical);
	    query += "&Expires=" + itos(expires);
	    query += "&AWSAccessKeyId=" + string(aws_access_key_id);
	}

	if(query.size()>0){
	    url += "?" + query;
	}

	if(b) delete b;
	b = new response_buffer();
	memset(b->ETag,0,sizeof(b->ETag));
	if(s3_debug>1) curl_easy_setopt(c,CURLOPT_VERBOSE,1);
	if(method != "GET"){
	    curl_easy_setopt(c,CURLOPT_CUSTOMREQUEST,method.c_str());
	}

	if(method == "HEAD"){
	    curl_easy_setopt(c,CURLOPT_NOBODY,1);
	}

	/* Queries that take longer than an hour should timeout */
	curl_easy_setopt(c,CURLOPT_TIMEOUT,60*60);

	/* Disable DNS cache */
	curl_easy_setopt(c,CURLOPT_DNS_CACHE_TIMEOUT,0); // per amazon specification
	curl_easy_setopt(c,CURLOPT_WRITEFUNCTION,buffer_write);
	curl_easy_setopt(c,CURLOPT_WRITEDATA,b); // fourth argument
	curl_easy_setopt(c,CURLOPT_URL,url.c_str());

	/* Are we sending data */
	class buffer *sendbuffer = 0;
	if(sendbuf){
	    sendbuffer = new buffer(sendbuf,sendbuflen);
	    curl_easy_setopt(c,CURLOPT_READFUNCTION,buffer_read);
	    curl_easy_setopt(c,CURLOPT_READDATA,sendbuffer);
	    curl_easy_setopt(c,CURLOPT_UPLOAD,1);
	    curl_easy_setopt(c,CURLOPT_INFILESIZE,sendbuflen);
	    //fprintf(stderr,"***** sendbuflen= %d %qd\n",sizeof(sendbuflen),sendbuflen);
	}

	/* Make provisions to get the response headers */
	if(h) delete h;
	h = new buffer();
	curl_easy_setopt(c,CURLOPT_HEADERFUNCTION,buffer_write);
	curl_easy_setopt(c,CURLOPT_WRITEHEADER,h); // fourth argument

	/* Make provisions for getting the headers */

	int success = curl_easy_perform(c);

	if(sendbuffer){
	    delete sendbuffer;
	    sendbuffer = 0;
	    if(success==0) s3_bytes_written += sendbuflen;
	}

	s3_bytes_read += h->len;
	s3_bytes_read += b->len;

	// CURL API says do not assume NULL terminate, so terminate it
	h->write("\000",1);			
	curl_easy_getinfo(c,CURLINFO_RESPONSE_CODE,&b->result);

	/* Now clean up */
	s3_request_retry_count = retry_count;
	if(headers) curl_slist_free_all(headers);
	curl_easy_cleanup(c);

	/* Process the results */
	if(success!=0){
	    delete h;
	    delete b;
	    s3_request_retry_count = retry_count;
	    return 0;			// internal CURL error
	}
	if(s3_debug>2){
	    printf("Header results:\n");
	    h->print();
	    printf("Data results:\n");
	    b->print();
	    printf("\n");
	}
    } while(b->result==500 && ++retry_count<s3_retry_max);

    if(b->result==404) errno=ENOENT;

    /* Pull out the headers */
    char *line,*brkt;
    for(line = strtok_r(h->base,"\r\n",&brkt);
	line;
	line = strtok_r(NULL,"\r\n",&brkt)){
	char *cc = strchr(line,':');
	if(cc){
	    *cc++ = '\000';
	    while(*cc && isspace(*cc)) cc++;
	    b->rheaders[line] = cc;
	}
    }

    /* Find the ETag in the header and put in the buffer */
    const char *e = b->rheaders["ETag"].c_str();
    if(strlen(e)==34){
	for(int i=0;i<16;i++){
	    b->ETag[i] = (hexval(e[i*2+1])<<4) + hexval(e[i*2+2]);
	}
    }

    delete h;				// we don't care about it
    if(s3_debug>1) printf(".\n\n");
    return b;
}

response_buffer *get_url(const char *url)
{
    int retry_count = 0;
    response_buffer *b = new response_buffer();
    do {
	CURL *c = curl_easy_init();
	curl_easy_setopt(c,CURLOPT_WRITEFUNCTION,buffer_write);
	curl_easy_setopt(c,CURLOPT_WRITEDATA,b);
	curl_easy_setopt(c,CURLOPT_URL,url);
	int success = curl_easy_perform(c);
	curl_easy_getinfo(c,CURLINFO_RESPONSE_CODE,&b->result);
	curl_easy_cleanup(c);
    } while(b->result!=200 && ++retry_count<s3_retry_max);
    s3_request_retry_count = retry_count;
    return b;
}




class s3_result *list_buckets()
{
    time_t expires = time(0)+60;
    expires = 0;
    class response_buffer *b = request("GET","","",expires,0,0,0);

    class s3_result *r = xml_extract_response(b);
    delete b;
    return r;
}

class s3_result *list_bucket(string bucket,string prefix,string marker,int max_keys)
{
    string query;


    if(prefix.size()>0) query += "prefix=" + prefix;
    if(marker.size()>0){
	if(query.size()>0) query += "&";
	query += "marker=" + marker;
    }
    if(max_keys>0){
	if(query.size()>0) query += "&";;
	query += "max-keys=" + itos(max_keys);
    }
    class response_buffer *b = request("GET",bucket,query,0,0,0,0);
    if(!b) return 0;
    class s3_result *r = xml_extract_response(b);
    delete b;
    return r;
}

/* 
 * af_hexbuf:
 * Turn a binay string into a hex string, optionally with spaces.
 */

#define HEXBUF_NO_SPACES 0
#define HEXBUF_SPACE2    0x0001	// space every 2 characters
#define HEXBUF_SPACE4    0x0002	// space every 4 characters
#define HEXBUF_UPPERCASE 0x1000	// uppercase
static const char *hexbuf(char *dst,int dst_len,const unsigned char *bin,int bytes,int flag)
{
    int charcount = 0;
    const char *start = dst;		// remember where the start of the string is
    const char *fmt = (flag & HEXBUF_UPPERCASE) ? "%02X" : "%02x";

    *dst = 0;				// begin with null termination
    while(bytes>0 && dst_len > 3){
	sprintf(dst,fmt,*bin); // convert the next byte
	dst += 2;
	bin += 1;
	dst_len -= 2;
	bytes--;
	charcount++;			// how many characters
	
	bool add_spaces = false;
	if(flag & HEXBUF_SPACE2) add_spaces = true;
	if((flag & HEXBUF_SPACE4) && charcount%2==0){
	    *dst++ = ' ';
	    *dst   = '\000';
	    dst_len -= 1;
	}
    }
    return start;			// return the start
}


/* object_put:
 * Put an object. Make sure that the MD5 of the response matches..
 * Makes a few retry attempts
 * Return 0 if success, -1 if failure. 
 */
int object_put(string bucket,string path,
		  const char *buf,size_t buflen,
		  const struct s3headers *extraheaders)
{
    unsigned char md5[16];
    memset(md5,0,sizeof(md5));
    MD5((const unsigned char *)buf,buflen,md5);
    for(int i=0;i<s3_retry_max;i++){
	s3_object_put_retry_count = i;
	if(i>0){
	    fprintf(stderr,"S3: Attempt to write object '%s' failed. Retrying...\n",
		    path.c_str());
	}

	response_buffer *res = request("PUT",bucket + "/" + path,"",0,buf,buflen,extraheaders);
	if(!res) {
	    fprintf(stderr,"S3 request: No response.\n");
	    continue;
	}
	if(memcmp(res->ETag,md5,16)==0){	/* Check the MD5 of the response */
	    delete res;
	    return 0;
	}
	char buf0[64],buf1[64];
	fprintf(stderr,"S3: Expected ETag '%s' got '%s'\n",
		hexbuf(buf0,sizeof(buf0),md5,16,HEXBUF_SPACE4),
		hexbuf(buf1,sizeof(buf1),res->ETag,16,HEXBUF_SPACE4));
	delete res;
    }
    /* Write failed. Delete the written object and return */
    response_buffer *res = request("DELETE",bucket + "/" + path,"",0,0,0,0);
    if(res) delete res;			
    errno = EIO;
    return -1;
}

int bucket_mkdir(string bucket)
{
    class response_buffer *b =  request("PUT",bucket,"",0,0,0,0);
    int result = b->result;
    delete b;
    switch(result){
    case 409:errno=EEXIST; return -1;
    case 200:errno=0;return 0;
    }
    return -1;			// some unknown error
}

int bucket_rmdir(string bucket)
{
    class response_buffer *b = request("DELETE",bucket,"",0,0,0,0);
    int result = b->result;
    delete b;
    switch(result){
    case 403:errno=EACCES; return -1;
    case 404:errno=ENOENT; return -1;
    case 409:errno=ENOTEMPTY; return -1;
    case 204:errno=0;return 0;		// no content is actually what it gives
    case 200:errno=0;return 0;		// doesn't seem to give this one
    }
    return -1;			// some unknown error
}

class response_buffer *object_get(string bucket,string path,const s3headers *extra_headers)
{
    return request("GET",bucket + "/" + path,"",0,0,0,extra_headers);
}

class response_buffer *object_head(string bucket,string path,const s3headers *extra_headers)
{
    return request("HEAD",bucket + "/" + path,"",0,0,0,extra_headers);
}

int object_rm(string bucket,string path)
{
    class response_buffer *b = request("DELETE",bucket + "/" + path,"",0,0,0,0);
    if(b){
	delete b;
	return 0;
    }
    return -1;
}

}

void s3_audit(int i)
{
    if(i>0 || s3_bytes_written>0 || s3_bytes_read>0){
	fprintf(stderr,"\n");
	fprintf(stderr,"S3 bytes written: %qu\n",s3_bytes_written);
	fprintf(stderr,"S3 bytes read: %qu\n",s3_bytes_read);
    }
}
#else
int main(int argc,char **argv)
{
    fprintf(stderr,"S3 support is not enabled.\n");
    exit(1);
}

#endif /* USE_S3 */