Merge pull request #136 from basho/ss-cs-copy

slfritchie · slfritchie · commit eb93d6301fb6 · 2014-09-19T14:56:12.000+09:00
Add PUT Copy operation to CS driver
diff --git a/examples/cs.config.sample b/examples/cs.config.sample
@@ -13,12 +13,24 @@
 
 {driver, basho_bench_driver_cs}.
 
-%% Replace this with a user you have created.
-%% Instructions to create a user are here:
-%% https://github.com/basho/riak_cs/wiki/Creating-a-User
+%% Two ways of user configuration
+%% 1. Set cs_access_key and cs_secret_key for existing user.
+%%    Replace this with a user you have created.
+%%    Instructions to create a user are here:
+%%    https://github.com/basho/riak_cs/wiki/Creating-a-User
 {cs_access_key, "ZG7SS3ZPECF-8LZOEBMA"}.
 {cs_secret_key, "21HoIRdeO617nJrIbam9mKH2MBxmcsMEwESvmQ=="}.
-{cs_bucket, "test"}. % create this with s3cmd before running basho_bench!
+
+%% 2. Let the driver create or retrieve user
+%%    In this case, you should
+%%    - setup riak-cs with {admin_auth_enabled, false} and
+%%    - comment out cs_access_key and cs_secret_key above
+%%    cs_display_name is used to identify user.
+%% {cs_display_name, "test-user"}.
+
+%% CS Bucket, it will be created if needed
+{cs_bucket, "test-bucket"}.
+
 {cs_disconnect_frequency, 5}.    % # ops before disconnecting HTTP socket
 {cs_raw_ip, "s3.amazonaws.com"}. % DO NOT CHANGE
 {cs_raw_port, 80}.               % DO NOT CHANGE
@@ -46,3 +58,4 @@
 %% bad idea: {operations, [{insert, 1}, {get, 1}]}.
 {operations, [{insert, 1}]}.
 %{operations, [{get, 1}]}.
+%{operations, [{delete, 1}]}.
diff --git a/src/basho_bench_driver_cs.erl b/src/basho_bench_driver_cs.erl
@@ -120,10 +120,15 @@ new(ID) ->
     lager:log(info, self(), "ID ~p Proxy host ~p TCP port ~p\n",
               [ID, ProxyH, ProxyP]),
 
-    {ok, #state{client_id=ID, hosts=Hosts,
-                bucket = basho_bench_config:get(cs_bucket, "test"),
-                report_fun = ReportFun,
-                http_proxy_host = ProxyH, http_proxy_port = ProxyP}}.
+    State = #state{client_id=ID, hosts=Hosts,
+                   bucket = basho_bench_config:get(cs_bucket, "test"),
+                   report_fun = ReportFun,
+                   http_proxy_host = ProxyH, http_proxy_port = ProxyP},
+    case ID of
+        1 -> ok = setup_user_and_bucket(State);
+        _ -> ok
+    end,
+    {ok, State}.
 
 %% This module does some crazy stuff, but it's there for a reason.
 %% The reason is that basho_bench is expecting the run() function to
@@ -186,6 +191,19 @@ run2(delete, KeyGen, _ValueGen, #state{bucket = Bucket} = State) ->
         {error, Reason} ->
             {error, Reason, S2}
     end;
+run2(copy, KeyGen, _ValueGen, #state{bucket = Bucket} = State) ->
+    {NextHost, S2} = next_host(State),
+    {Host, Port} = NextHost,
+    SourceKey = KeyGen(),
+    DestKey = SourceKey ++ basho_bench_config:get(cs_copy_suffix, "-copy"),
+    Url = url(Host, Port, Bucket, DestKey),
+    Headers = [{"x-amz-copy-source", "/" ++ Bucket ++ "/" ++ SourceKey}],
+    case do_copy({Host, Port}, Url, Headers, State) of
+        ok ->
+            {ok, S2};
+        {error, Reason} ->
+            {error, Reason, S2}
+    end;
 run2(Op, _KeyGen, _ValueGen, State) ->
     {error, {unknown_op, Op}, State}.
 
@@ -278,10 +296,12 @@ insert(KeyGen, ValueGen, {Host, Port}, Bucket, State) ->
     do_put({Host, Port}, Url, [{"content-length", integer_to_list(CL)}], ValueT,
            State).
 
+url(Host, Port, Bucket, undefined) ->
+    UnparsedUrl = lists:concat(["http://", Host, ":", Port, "/", Bucket]),
+    ibrowse_lib:parse_url(UnparsedUrl);
 url(Host, Port, Bucket, Key) ->
     UnparsedUrl = lists:concat(["http://", Host, ":", Port, "/", Bucket, "/", Key]),
-    Url = ibrowse_lib:parse_url(UnparsedUrl),
-    Url.
+    ibrowse_lib:parse_url(UnparsedUrl).
 
 -spec next_host(term()) -> {term(), term()}.
 %% TODO:
@@ -415,6 +435,16 @@ do_delete(Host, Url, Headers, State) ->
             {error, Reason}
     end.
 
+do_copy(Host, Url, Headers, State) ->
+    case send_request(Host, Url, Headers, put, <<>>, proxy_opts(State)) of
+        {ok, "200", _Header, _Body} ->
+            ok;
+        {ok, Code, _Header, _Body} ->
+            {error, {http_error, Code}};
+        {error, Reason} ->
+            {error, Reason}
+    end.
+
 do_get_first_unit(Host, Url, Headers, State) ->
     BufSize = 128*1024,
     Opts = [{max_pipeline_size, 9999999},
@@ -494,7 +524,9 @@ uppercase_verb(put) ->
 uppercase_verb(get) ->
     'GET';
 uppercase_verb(delete) ->
-    'DELETE'.
+    'DELETE';
+uppercase_verb(post) ->
+    'POST'.
 
 to_list(A) when is_atom(A) ->
     atom_to_list(A);
@@ -507,14 +539,126 @@ initiate_request(Host, Url, Headers0, Method, Body, Options) ->
                                'Content-Type', Headers0,
                                'application/octet-stream')),
     Date = httpd_util:rfc1123_date(),
+    Uri = element(7, Url),
     Headers = [{'Content-Type', ContentTypeStr},
                {'Date', Date}|lists:keydelete('Content-Type', 1, Headers0)],
-    Uri = element(7, Url),
-    Sig = stanchion_auth:request_signature(
-            uppercase_verb(Method), Headers, Uri,
-            basho_bench_config:get(cs_secret_key)),
-    AuthStr = ["AWS ", basho_bench_config:get(cs_access_key), ":", Sig],
-    HeadersWithAuth = [{'Authorization', AuthStr}|Headers],
+    HeadersWithAuth =
+        case basho_bench_config:get(cs_access_key, undefined) of
+            undefined ->
+                Headers;
+            AccessKey ->
+                AuthSig = auth_sig(AccessKey, basho_bench_config:get(cs_secret_key),
+                                   uppercase_verb(Method), ContentTypeStr, Date,
+                                   Headers, Uri),
+                [{'Authorization', AuthSig}|Headers]
+        end,
     Timeout = basho_bench_config:get(cs_request_timeout, 5000),
     ibrowse_http_client:send_req(Pid, Url, HeadersWithAuth, Method,
                                  Body, Options, Timeout).
+
+%% S3 utilities
+
+auth_sig(AccessKey, SecretKey, Method, ContentType, Date, Headers, Resource) ->
+    AmzHeaders = lists:filter(fun ({"x-amz-" ++ _, V}) when V =/= undefined -> true; (_) -> false end, Headers),
+    CanonizedAmzHeaders =
+        [[Name, $:, Value, $\n] || {Name, Value} <- lists:sort(AmzHeaders)],
+    StringToSign = [string:to_upper(atom_to_list(Method)), $\n,
+                    "", $\n,  % Content-MD5
+                    ContentType, $\n,
+                    Date, $\n,
+                    CanonizedAmzHeaders,
+                    Resource
+                   ],
+    Signature = base64:encode(stanchion_utils:sha_mac(SecretKey, StringToSign)),
+    ["AWS ", AccessKey, $:, Signature].
+
+%% CS utilities
+
+setup_user_and_bucket(State) ->
+    case basho_bench_config:get(cs_access_key, undefined) of
+        undefined ->
+            DisplayName = basho_bench_config:get(cs_display_name, "test-user"),
+            ok = maybe_create_user(DisplayName, State),
+            {ok, {_DisplayName, KeyId, KeySecret}} = fetch_user_info(DisplayName, State),
+            lager:info("Target User: ~p", [{DisplayName, KeyId, KeySecret}]),
+            ok = basho_bench_config:set(cs_access_key, KeyId),
+            ok = basho_bench_config:set(cs_secret_key, KeySecret);
+        _ ->
+            ok
+    end,
+    ok = maybe_create_bucket(State#state.bucket, State).
+
+maybe_create_user(DisplayName, #state{hosts=Hosts} = State) ->
+    {Host, Port} = hd(Hosts),
+    Json = io_lib:format("{\"email\": \"~s@example.com\", \"name\": \"~s\"}",
+                         [DisplayName, DisplayName]),
+    Url = url(Host, Port, "riak-cs", "user"),
+    Headers = [{'Content-Type', 'application/json'}],
+    case send_request({Host, Port}, Url, Headers, post, Json, proxy_opts(State)) of
+        {ok, "201", _Header, Body} ->
+            lager:debug("User created: ~p~n", [Body]),
+            ok;
+        {ok, "409", _Header, Body} ->
+            lager:debug("User already exists: ~p~n", [Body]),
+            ok;
+        {ok, Code, Header, Body} ->
+            {error, {user_creation, Code, Header, Body}};
+        {error, Reason} ->
+            {error, {user_creation, Reason}}
+    end.
+
+fetch_user_info(DisplayName, State) ->
+    case list_users(State) of
+        {ok, UserList} ->
+            {ok, lists:keyfind(DisplayName, 1, UserList)};
+        {error, Reason} ->
+            {error, Reason}
+    end.
+
+list_users(#state{hosts=Hosts} = State) ->
+    {Host, Port} = hd(Hosts),
+    Url = url(Host, Port, "riak-cs", "users"),
+    case send_request({Host, Port}, Url, [{'Accept', 'application/json'}], get,
+                       [], proxy_opts(State)) of
+        {ok, "200", _Headers, Body} ->
+            {ok, parse_user_info(Body)};
+        {error, Reason} ->
+            {error, {list_users, Reason}}
+    end.
+
+parse_user_info(Output) ->
+    [Boundary | Tokens] = string:tokens(binary_to_list(Output), "\r\n"),
+    parse_user_info(Tokens, Boundary, []).
+
+parse_user_info([_LastToken], _, Users) ->
+    ordsets:from_list(Users);
+parse_user_info(["Content-Type: application/json", RawJson | RestTokens],
+                 Boundary, Users) ->
+    UpdUsers = parse_user_records(RawJson, json) ++ Users,
+    parse_user_info(RestTokens, Boundary, UpdUsers);
+parse_user_info([_ | RestTokens], Boundary, Users) ->
+    parse_user_info(RestTokens, Boundary, Users).
+
+parse_user_records(Output, json) ->
+    JsonData = mochijson2:decode(Output),
+    [begin
+         KeyId = binary_to_list(proplists:get_value(<<"key_id">>, UserJson)),
+         KeySecret = binary_to_list(proplists:get_value(<<"key_secret">>, UserJson)),
+         Name = binary_to_list(proplists:get_value(<<"name">>, UserJson)),
+         {Name, KeyId, KeySecret}
+     end || {struct, UserJson} <- JsonData].
+
+maybe_create_bucket(Bucket, #state{hosts=Hosts} = State) ->
+    {Host, Port} = hd(Hosts),
+    Url = url(Host, Port, Bucket, undefined),
+    case send_request({Host, Port}, Url, [], put, [], proxy_opts(State)) of
+        {ok, "200", _Headers, _Body} ->
+            lager:debug("Bucket created (maybe): ~p~n", [Bucket]),
+            ok;
+        {ok, Code, Header, Body} ->
+            lager:error("Create bucket: ~p~n", [{Code, Header, Body}]),
+            {error, {bucket_creation, Code, Header, Body}};
+        {error, Reason} ->
+            lager:error("Create bucket: ~p~n", [Reason]),
+            {error, {bucket_creation, Reason}}
+    end.
