Ensure Riak TS tables and queries are unicode

Throw unicode_error when invalid data is used for a table name or query

src/riakc_ts.erl

@@ -81,11 +81,7 @@
         query_common(Pid, Query, Interpolations, Cover, Options).
 
 query_common(Pid, Query, Interpolations, Cover, Options)
-  when is_pid(Pid), is_list(Query) ->
-    QueryBin = unicode:characters_to_binary(Query),
-    query_common(Pid, QueryBin, Interpolations, Cover, Options);
-query_common(Pid, Query, Interpolations, Cover, Options)
-  when is_pid(Pid), is_binary(Query) ->
+  when is_pid(Pid) ->
     Msg0 = riakc_ts_query_operator:serialize(Query, Interpolations),
     Msg1 = Msg0#tsqueryreq{cover_context = Cover},
     Msg = {Msg1, {msgopts, Options}},
@@ -98,10 +94,12 @@ query_common(Pid, Query, Interpolations, Cover, Options)
 -spec get_coverage(pid(), table_name(), QueryText::iolist()) ->
                           {ok, Entries::[term()]} | {error, term()}.
 get_coverage(Pid, Table, Query) ->
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
+    {ok, Q} = riakc_utils:characters_to_unicode_binary(Query),
     Message =
-        #tscoveragereq{'query' = #tsinterpolation{base = iolist_to_binary(Query)},
+        #tscoveragereq{'query' = #tsinterpolation{base = Q},
                        replace_cover = undefined,
-                       table = iolist_to_binary(Table)},
+                       table = T},
     case server_call(Pid, Message) of
         {ok, Entries} ->
             {ok, riak_pb_ts_codec:decode_cover_list(Entries)};
@@ -119,11 +117,13 @@ replace_coverage(Pid, Table, Query, Cover) ->
                        OtherCover::list(binary())) ->
                               {ok, Entries::[term()]} | {error, term()}.
 replace_coverage(Pid, Table, Query, Cover, Other) ->
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
+    {ok, Q} = riakc_utils:characters_to_unicode_binary(Query),
     Message =
-        #tscoveragereq{'query' = #tsinterpolation{base = iolist_to_binary(Query)},
+        #tscoveragereq{'query' = #tsinterpolation{base = Q},
                        replace_cover = Cover,
                        unavailable_cover = Other,
-                       table = iolist_to_binary(Table)},
+                       table = T},
     case server_call(Pid, Message) of
         {ok, Entries} ->
             {ok, riak_pb_ts_codec:decode_cover_list(Entries)};
@@ -174,7 +174,8 @@ put(Pid, Table, Measurements, Options)
 delete(Pid, Table, Key, Options)
   when is_pid(Pid), (is_binary(Table) orelse is_list(Table)),
        is_list(Key), is_list(Options) ->
-    Message = #tsdelreq{table   = iolist_to_binary(Table),
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
+    Message = #tsdelreq{table   = T,
                         key     = riak_pb_ts_codec:encode_cells_non_strict(Key),
                         vclock  = proplists:get_value(vclock, Options),
                         timeout = proplists:get_value(timeout, Options)},
@@ -223,9 +224,9 @@ stream_list_keys(Pid, Table, Timeout) when is_integer(Timeout) ->
     stream_list_keys(Pid, Table, [{timeout, Timeout}]);
 stream_list_keys(Pid, Table, Options)
   when is_pid(Pid), (is_binary(Table) orelse is_list(Table)), is_list(Options) ->
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
     ReqTimeout = proplists:get_value(timeout, Options),
-    Req = #tslistkeysreq{table   = iolist_to_binary(Table),
-                         timeout = ReqTimeout},
+    Req = #tslistkeysreq{table = T, timeout = ReqTimeout},
     ReqId = riakc_pb_socket:mk_reqid(),
     gen_server:call(Pid, {req, Req, ?DEFAULT_PB_TIMEOUT, {ReqId, self()}}, infinity).
 

src/riakc_ts_get_operator.erl

@@ -33,12 +33,12 @@
 
 
 serialize(Table, Key, true) ->
-    #tsgetreq{table = iolist_to_binary(Table),
-              key   = Key};
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
+    #tsgetreq{table = T, key = Key};
 serialize(Table, Key, false) ->
+    {ok, T} = riakc_utils:characters_to_unicode_binary(Table),
     SerializedKey = riak_pb_ts_codec:encode_cells_non_strict(Key),
-    #tsgetreq{table = iolist_to_binary(Table),
-              key   = SerializedKey}.
+    #tsgetreq{table = T, key = SerializedKey}.
 
 deserialize({error, {Code, Message}}) when is_integer(Code), is_list(Message) ->
     {error, {Code, iolist_to_binary(Message)}};

src/riakc_ts_put_operator.erl

@@ -35,12 +35,14 @@
 %% As of 2015-11-05, columns parameter is ignored, Riak TS
 %% expects the full set of fields in each element of Data.
 serialize(TableName, Measurements, true) ->
-    #tsputreq{table   = iolist_to_binary(TableName),
+    {ok, T} = riakc_utils:characters_to_unicode_binary(TableName),
+    #tsputreq{table   = T,
               columns = [],
               rows    = Measurements};
 serialize(TableName, Measurements, false) ->
+    {ok, T} = riakc_utils:characters_to_unicode_binary(TableName),
     SerializedRows = riak_pb_ts_codec:encode_rows_non_strict(Measurements),
-    #tsputreq{table   = TableName,
+    #tsputreq{table   = T,
               columns = [],
               rows    = SerializedRows}.
 

src/riakc_ts_query_operator.erl

@@ -31,9 +31,12 @@
 -export([serialize/2,
          deserialize/1, deserialize/2]).
 
-serialize(QueryText, Interpolations) ->
+
+serialize(QueryText, Interpolations)
+  when is_binary(QueryText) orelse is_list(QueryText) ->
+    {ok, Q} = riakc_utils:characters_to_unicode_binary(QueryText),
     Content = #tsinterpolation{
-                 base           = iolist_to_binary(QueryText),
+                 base           = Q,
                  interpolations = serialize_interpolations(Interpolations)},
     #tsqueryreq{'query' = Content}.
 

src/riakc_utils.erl

@@ -22,7 +22,7 @@
 
 -module(riakc_utils).
 
--export([wait_for_list/1]).
+-export([wait_for_list/1, characters_to_unicode_binary/1]).
 
 %% @doc Wait for the results of a listing operation
 wait_for_list(ReqId) ->
@@ -33,3 +33,17 @@ wait_for_list(ReqId, Acc) ->
         {ReqId, {error, Reason}} -> {error, Reason};
         {ReqId, {_, Res}} -> wait_for_list(ReqId, [Res|Acc])
     end.
+
+%% @doc Convert to unicode binary with informative errors
+%% @throws {unicode_error, ErrMsg}
+characters_to_unicode_binary(String) ->
+    case unicode:characters_to_binary(String) of
+        {incomplete, Encoded, Rest} ->
+            ErrMsg = io_lib:format("Incomplete unicode data provided. Encoded: ~p Rest: ~p", [Encoded, Rest]),
+            throw({unicode_error, ErrMsg});
+        {error, Encoded, Rest} ->
+            ErrMsg = io_lib:format("Unicode encoding error. Encoded: ~p Rest: ~p", [Encoded, Rest]),
+            throw({unicode_error, ErrMsg});
+        Binary ->
+            {ok, Binary}
+    end.

test/riakc_utils_tests.erl

@@ -0,0 +1,34 @@
+%% -------------------------------------------------------------------
+%%
+%% riakc_utils_tests
+%%
+%% Copyright (c) 2016 Basho Technologies, Inc.  All Rights Reserved.
+%%
+%% This file is provided to you under the Apache License,
+%% Version 2.0 (the "License"); you may not use this file
+%% except in compliance with the License.  You may obtain
+%% a copy of the License at
+%%
+%%   http://www.apache.org/licenses/LICENSE-2.0
+%%
+%% Unless required by applicable law or agreed to in writing,
+%% software distributed under the License is distributed on an
+%% "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+%% KIND, either express or implied.  See the License for the
+%% specific language governing permissions and limitations
+%% under the License.
+%%
+%% -------------------------------------------------------------------
+-ifdef(TEST).
+
+-module(riakc_utils_tests).
+
+-compile(export_all).
+
+-include_lib("eunit/include/eunit.hrl").
+
+bad_unicode_binary_test() ->
+    S = <<"\xa0\xa1">>,
+    ?assertThrow({unicode_error, _Msg}, riakc_utils:characters_to_unicode_binary(S)).
+
+-endif.