aerospike 2.20.1 → 2.22.0

data/lib/aerospike/command/admin_command.rb

@@ -18,17 +18,22 @@ module Aerospike
 
   private
   # Commands
-  AUTHENTICATE    = 0
-  CREATE_USER     = 1
-  DROP_USER       = 2
-  SET_PASSWORD    = 3
-  CHANGE_PASSWORD = 4
-  GRANT_ROLES     = 5
-  REVOKE_ROLES    = 6
-  #CREATE_ROLE = 8
-  QUERY_USERS = 9
-  #QUERY_ROLES =  10
-  LOGIN = 20
+  AUTHENTICATE      = 0
+  CREATE_USER       = 1
+  DROP_USER         = 2
+  SET_PASSWORD      = 3
+  CHANGE_PASSWORD   = 4
+  GRANT_ROLES       = 5
+  REVOKE_ROLES      = 6
+  QUERY_USERS       = 9
+  CREATE_ROLE       = 10
+  DROP_ROLE         = 11
+  GRANT_PRIVILEGES  = 12
+  REVOKE_PRIVILEGES = 13
+  SET_WHITELIST     = 14
+  SET_QUOTAS        = 15
+  QUERY_ROLES       = 16
+  LOGIN             = 20
 
   # Field IDs
   USER           = 0
@@ -36,8 +41,17 @@ module Aerospike
   OLD_PASSWORD   = 2
   CREDENTIAL     = 3
   CLEAR_PASSWORD = 4
+  SESSION_TOKEN  = 5
+  SESSION_TTL    = 6
   ROLES          = 10
-  #PRIVILEGES =  11
+  ROLE           = 11
+  PRIVILEGES     = 12
+  ALLOWLIST      = 13
+  READ_QUOTA     = 14
+  WRITE_QUOTA    = 15
+  READ_INFO      = 16
+  WRITE_INFO     = 17
+  CONNECTIONS    = 18
 
   # Misc
   MSG_VERSION  = 2
@@ -55,34 +69,6 @@ module Aerospike
       @data_offset =  8
     end
 
-    def authenticate(conn, user, password)
-      begin
-        set_authenticate(user, password)
-        conn.write(@data_buffer, @data_offset)
-        conn.read(@data_buffer, HEADER_SIZE)
-
-        result = @data_buffer.read(RESULT_CODE)
-
-        # read the rest of the buffer
-        size = @data_buffer.read_int64(0)
-        length = (size & 0xFFFFFFFFFFFF) - HEADER_REMAINING
-        conn.read(@data_buffer, length)
-
-        raise Exceptions::Aerospike.new(result, "Authentication failed") if result != 0
-      ensure
-        Buffer.put(@data_buffer)
-      end
-    end
-
-    def set_authenticate(user, password)
-      write_header(LOGIN, 2)
-      write_field_str(USER, user)
-      write_field_bytes(CREDENTIAL, password)
-      write_size
-
-      return @data_offset
-    end
-
     def create_user(cluster, policy, user, password, roles)
       write_header(CREATE_USER, 3)
       write_field_str(USER, user)
@@ -126,6 +112,61 @@ module Aerospike
       execute_command(cluster, policy)
     end
 
+    def create_role(cluster, policy, role_name, privileges = [], allowlist = [], read_quota = 0, write_quota = 0)
+      field_count = 1
+      field_count += 1 if privileges.size > 0
+      field_count += 1 if allowlist.size > 0
+      field_count += 1 if read_quota > 0
+      field_count += 1 if write_quota > 0
+
+      write_header(CREATE_ROLE, field_count)
+      write_field_str(ROLE, role_name)
+
+      write_privileges(privileges) if privileges.size > 0
+      write_allowlist(allowlist) if allowlist.size > 0
+
+      write_field_uint32(READ_QUOTA, read_quota) if read_quota > 0
+      write_field_uint32(WRITE_QUOTA, write_quota) if write_quota > 0
+
+      execute_command(cluster, policy)
+    end
+
+    def drop_role(cluster, policy, role)
+      write_header(DROP_ROLE, 1)
+      write_field_str(ROLE, role)
+      execute_command(cluster, policy)
+    end
+
+    def grant_privileges(cluster, policy, role, privileges)
+      write_header(GRANT_PRIVILEGES, 2)
+      write_field_str(ROLE, role)
+      write_privileges(privileges)
+      execute_command(cluster, policy)
+    end
+
+    def revoke_privileges(cluster, policy, role, privileges)
+      write_header(REVOKE_PRIVILEGES, 2)
+      write_field_str(ROLE, role)
+      write_privileges(privileges)
+      execute_command(cluster, policy)
+    end
+
+    def set_allowlist(cluster, policy, role, allowlist = [])
+      field_count = 1
+      field_count += 1 if allowlist.size > 0
+      write_header(SET_WHITELIST, field_count)
+      write_allowlist(allowlist) if allowlist.size > 0
+      execute_command(cluster, policy)
+    end
+
+    def set_quotas(cluster, policy, role, read_quota, write_quota)
+      write_header(SET_QUOTAS, 3)
+      write_field_str(ROLE, role)
+      write_field_uint32(READ_QUOTA, read_quota)
+      write_field_uint32(WRITE_QUOTA, write_quota)
+      execute_command(cluster, policy)
+    end
+
     def query_user(cluster, policy, user)
       # TODO: Remove the workaround in the future
       sleep(0.010)
@@ -152,6 +193,32 @@ module Aerospike
       end
     end
 
+    def query_role(cluster, policy, role)
+      # TODO: Remove the workaround in the future
+      sleep(0.010)
+
+      list = []
+      begin
+        write_header(QUERY_ROLES, 1)
+        write_field_str(ROLE, role)
+        list = read_roles(cluster, policy)
+        return (list.is_a?(Array) && list.length > 0 ? list.first : nil)
+      ensure
+        Buffer.put(@data_buffer)
+      end
+    end
+
+    def query_roles(cluster, policy)
+      # TODO: Remove the workaround in the future
+      sleep(0.010)
+      begin
+        write_header(QUERY_ROLES, 0)
+        return read_roles(cluster, policy)
+      ensure
+        Buffer.put(@data_buffer)
+      end
+    end
+
     def write_roles(roles)
       offset = @data_offset + FIELD_HEADER_SIZE
       @data_buffer.write_byte(roles.length.ord, offset)
@@ -174,6 +241,54 @@ module Aerospike
       @data_buffer.write_int64(size, 0)
     end
 
+    def write_privileges(privileges)
+      offset = @data_offset
+      @data_offset += FIELD_HEADER_SIZE
+      write_byte(privileges.size)
+
+      for privilege in privileges
+        write_byte(privilege.to_code)
+        if privilege.can_scope?
+          if privilege.set_name.to_s.size > 0 && privilege.namespace.to_s.size == 0
+            raise Aerospike::Exceptions::Aerospike.new(Aerospike::ResultCode::INVALID_PRIVILEGE, "Admin privilege #{privilege.namespace} has a set scope with an empty namespace")
+          end
+
+          write_str(privilege.namespace.to_s)
+          write_str(privilege.set_name.to_s)
+        else
+          if privilege.set_name.to_s.bytesize > 0 || privilege.namespace.to_s.bytesize > 0
+            raise Aerospike::Exceptions::Aerospike.new(Aerospike::ResultCode::INVALID_PRIVILEGE, "Admin global privilege #{privilege} can't have a namespace or set")
+          end
+        end
+      end
+
+      size = @data_offset - offset - FIELD_HEADER_SIZE
+      @data_offset = offset
+      write_field_header(PRIVILEGES, size)
+      @data_offset += size
+    end
+
+    def write_allowlist(allowlist)
+      offset = @data_offset
+      @data_offset += FIELD_HEADER_SIZE
+
+      comma = false
+      for addr in allowlist
+        if comma
+          write_byte(",")
+        else
+          comma = true
+        end
+
+        @data_offset += @data_buffer.write_binary(addr, @data_offset)
+      end
+
+      size = @data_offset - offset - FIELD_HEADER_SIZE
+      @data_offset = offset
+      write_field_header(ALLOWLIST, size)
+      @data_offset += size
+    end
+
     def write_header(command, field_count)
       # Authenticate header is almost all zeros
       i = @data_offset
@@ -186,12 +301,27 @@ module Aerospike
       @data_offset += 16
     end
 
+    def write_byte(b)
+      @data_offset += @data_buffer.write_byte(b, @data_offset)
+    end
+
+    def write_str(str)
+      @data_offset += @data_buffer.write_byte(str.bytesize, @data_offset)
+      @data_offset += @data_buffer.write_binary(str, @data_offset)
+    end
+
     def write_field_str(id, str)
       len = @data_buffer.write_binary(str, @data_offset+FIELD_HEADER_SIZE)
       write_field_header(id, len)
       @data_offset += len
     end
 
+    def write_field_uint32(id, val)
+      len = @data_buffer.write_uint32(val, @data_offset+FIELD_HEADER_SIZE)
+      write_field_header(id, len)
+      @data_offset += len
+    end
+
     def write_field_bytes(id, bytes)
       @data_buffer.write_binary(bytes, @data_offset+FIELD_HEADER_SIZE)
       write_field_header(id, bytes.bytesize)
@@ -251,7 +381,7 @@ module Aerospike
         raise e
       end
 
-      raise Exceptions::Aerospike.new(result) if status > 0
+      raise Exceptions::Aerospike.new(status) if status > 0
 
       return list
     end
@@ -292,7 +422,7 @@ module Aerospike
           return (result_code == QUERY_END ? -1 : result_code)
         end
 
-        userRoles = UserRoles.new
+        user_roles = UserRoles.new
         field_count = @data_buffer.read(@data_offset+3)
         @data_offset += HEADER_REMAINING
 
@@ -306,10 +436,17 @@ module Aerospike
 
           case id
           when USER
-            userRoles.user = @data_buffer.read(@data_offset, len)
+            user_roles.user = @data_buffer.read(@data_offset, len)
             @data_offset += len
           when ROLES
-            parse_roles(userRoles)
+            parse_roles(user_roles)
+          when READ_INFO
+            user_roles.read_info = parse_info
+          when WRITE_INFO
+            user_roles.write_info = parse_info
+          when CONNECTIONS
+            user_roles.conns_in_use = @data_buffer.read_int32(@data_offset)
+            @data_offset += len
           else
             @data_offset += len
           end
@@ -317,19 +454,19 @@ module Aerospike
           i = i.succ
         end
 
-        next if userRoles.user == "" && userRoles.roles == nil
+        next if user_roles.user == "" && user_roles.roles == nil
 
-        userRoles.roles = [] if userRoles.roles == nil
-        list << userRoles
+        user_roles.roles = [] if user_roles.roles == nil
+        list << user_roles
       end
 
       return 0, list
     end
 
-    def parse_roles(userRoles)
+    def parse_roles(user_roles)
       size = @data_buffer.read(@data_offset)
       @data_offset += 1
-      userRoles.roles = []
+      user_roles.roles = []
 
       i = 0
       while i < size
@@ -337,17 +474,188 @@ module Aerospike
         @data_offset += 1
         role = @data_buffer.read(@data_offset, len)
         @data_offset += len
-        userRoles.roles << role
+        user_roles.roles << role
 
         i = i.succ
       end
     end
 
-    SALT = '$2a$10$7EqJtq98hPqEX7fNZaFWoO'
-    def self.hash_password(password)
-      # Hashing the password with the cost of 10, with a static salt
-      return BCrypt::Engine.hash_secret(password, SALT, :cost => 10)
+    def parse_info
+      size = @data_buffer.read(@data_offset)
+      @data_offset += 1
+      list = []
+
+      i = 0
+      while i < size
+        val = @data_buffer.read_int32(@data_offset)
+        @data_offset += 4
+        list << val
+
+        i = i.succ
+      end
+
+      list
     end
+
+    def read_roles(cluster, policy)
+      write_size
+      node = cluster.random_node
+
+      timeout = 1
+      timeout = policy.timeout if policy != nil && policy.timeout > 0
+
+      status = -1
+      list = []
+      begin
+        conn = node.get_connection(timeout)
+        conn.write(@data_buffer, @data_offset)
+        status, list = read_role_blocks(conn)
+        node.put_connection(conn)
+      rescue => e
+        conn.close if conn
+        raise e
+      end
+
+      raise Exceptions::Aerospike.new(status) if status > 0
+
+      return list
+    end
+
+    def read_role_blocks(conn)
+      rlist = []
+      status = 0
+      begin
+        while status == 0
+          conn.read(@data_buffer, 8)
+          size = @data_buffer.read_int64(0)
+          receive_size = (size & 0xFFFFFFFFFFFF)
+
+          if receive_size > 0
+            @data_buffer.resize(receive_size) if receive_size > @data_buffer.size
+
+            conn.read(@data_buffer, receive_size)
+            status, list = parse_roles_full(receive_size)
+            rlist.concat(list.to_a)
+          else
+            break
+          end
+        end
+        return status, rlist
+      rescue => e
+        return -1, []
+      end
+    end
+
+    def parse_roles_full(receive_size)
+      @data_offset = 0
+      list = []
+
+      while @data_offset < receive_size
+        result_code = @data_buffer.read(@data_offset+1)
+
+        if result_code != 0
+          return (result_code == QUERY_END ? -1 : result_code)
+        end
+
+        role = Role.new
+        field_count = @data_buffer.read(@data_offset+3)
+        @data_offset += HEADER_REMAINING
+
+        i = 0
+        while i  < field_count
+          len = @data_buffer.read_int32(@data_offset)
+          @data_offset += 4
+          id = @data_buffer.read(@data_offset)
+          @data_offset += 1
+          len -= 1
+
+          case id
+          when ROLE
+            role.name = @data_buffer.read(@data_offset, len).to_s
+            @data_offset += len
+          when PRIVILEGES
+            parse_privileges(role)
+          when ALLOWLIST
+            role.allowlist = parse_allowlist(len)
+          when READ_QUOTA
+            role.read_quota = @data_buffer.read_uint32(@data_offset)
+            @data_offset += len
+          when WRITE_QUOTA
+            role.write_quota = @data_buffer.read_uint32(@data_offset)
+            @data_offset += len
+          else
+            @data_offset += len
+          end
+
+          i = i.succ
+        end
+
+        next if role.name == "" && role.privileges == nil
+
+        role.privileges ||= []
+        list << role
+      end
+
+      return 0, list
+    end
+
+    def parse_privileges(role)
+      size = @data_buffer.read(@data_offset)
+      @data_offset += 1
+      role.privileges = []
+
+      i = 0
+      while i < size
+        priv = Privilege.new
+        priv.code = Privilege.from(@data_buffer.read(@data_offset))
+        @data_offset += 1
+
+        if priv.can_scope?
+          len = @data_buffer.read(@data_offset)
+          @data_offset += 1
+          priv.namespace = @data_buffer.read(@data_offset, len)
+          @data_offset += len
+
+          len = @data_buffer.read(@data_offset)
+          @data_offset += 1
+          priv.set_name = @data_buffer.read(@data_offset, len)
+          @data_offset += len
+        end
+
+        role.privileges << priv
+
+        i = i.succ
+      end
+    end
+
+    def parse_allowlist(len)
+      list = []
+      begn = @data_offset
+      max = begn + len
+
+      while @data_offset < max
+        if @data_buffer.read(@data_offset) == ','
+          l = @data_offset - begn
+          if l > 0
+            s = @data_buffer.read(begn, l)
+            list << s
+          end
+          @data_offset += 1
+          begn = @data_offset
+        else
+          @data_offset += 1
+        end
+      end
+
+      l = @data_offset - begn
+      if l > 0
+        s = @data_buffer.read(begn, l)
+        list << s
+      end
+
+      list      
+    end
+
   end
 end
 

data/lib/aerospike/command/command.rb

@@ -840,12 +840,6 @@ module Aerospike
     end
 
     def size_buffer_sz(size)
-      # Corrupted data streams can result in a hug.length.
-      # Do a sanity check here.
-      if size > Buffer::MAX_BUFFER_SIZE
-        raise Aerospike::Exceptions::Parse.new("Invalid size for buffer: #{size}")
-      end
-
       @data_buffer.resize(size)
     end
 

data/lib/aerospike/command/login_command.rb

@@ -0,0 +1,164 @@
+# encoding: utf-8
+# Copyright 2014-2020 Aerospike, Inc.
+#
+# Portions may be licensed to Aerospike, Inc. under one or more contributor
+# license agreements.
+#
+# Licensed 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.
+
+require 'aerospike/command/admin_command'
+
+module Aerospike
+
+  private
+
+  attr_reader :session_token, :session_expiration
+
+  class LoginCommand < AdminCommand #:nodoc:
+
+    def login(conn, policy)
+        hashed_pass = LoginCommand.hash_password(policy.password)
+        authenticate(conn, policy, hashed_pass)
+    end
+
+    def authenticate(conn, user, hashed_pass)
+      write_header(LOGIN, 2)
+      write_field_str(USER, policy.user)
+      write_field_bytes(CREDENTIAL, hashed_pass)
+
+      parse_tokens(conn)
+    end
+
+    def authenticate_new(conn, cluster)
+      @data_offset = 8
+      policy = cluster.client_policy
+      case policy.auth_mode
+      when Aerospike::AuthMode::EXTERNAL
+        write_header(LOGIN, 3)
+        write_field_str(USER, policy.user)
+        write_field_bytes(CREDENTIAL, cluster.password)
+        write_field_str(CLEAR_PASSWORD, policy.password)
+      when Aerospike::AuthMode::INTERNAL
+        write_header(LOGIN, 2)
+        write_field_str(USER, policy.user)
+        write_field_bytes(CREDENTIAL, cluster.password)
+      when Aerospike::AuthMode::PKI
+        write_header(LOGIN, 0)
+      else
+        raise Exceptions::Aerospike.new(Aerospike::ResultCode::COMMAND_REJECTED, "Invalid client_policy#auth_mode.")
+      end
+
+      parse_tokens(conn)
+      cluster.session_token = @session_token
+      cluster.session_expiration = @session_expiration
+    end
+
+    def parse_tokens(conn)
+      begin
+        write_size
+        conn.write(@data_buffer, @data_offset)
+        conn.read(@data_buffer, HEADER_SIZE)
+
+        result = @data_buffer.read(RESULT_CODE)
+
+        if result != 0 
+          return if result == Aerospike::ResultCode::SECURITY_NOT_ENABLED
+          raise Exceptions::Aerospike.new(result, "Authentication failed")
+        end
+
+        # read the rest of the buffer
+        size = @data_buffer.read_int64(0)
+        receive_size = (size & 0xFFFFFFFFFFFF) - HEADER_REMAINING
+        field_count = @data_buffer.read(11) & 0xFF
+
+        if receive_size <= 0 || receive_size > @data_buffer.size || field_count <= 0
+          raise Exceptions::Aerospike.new(result, "Node failed to retrieve session token")
+        end
+
+        if @data_buffer.size < receive_size
+          @data_buffer.resize(receive_size)
+        end
+
+        conn.read(@data_buffer, receive_size)
+
+        @data_offset = 0
+        for i in 0...field_count
+          mlen = @data_buffer.read_int32(@data_offset)
+          @data_offset += 4
+          id = @data_buffer.read(@data_offset)
+          @data_offset += 1
+          mlen -= 1
+
+          case id
+          when SESSION_TOKEN
+            # copy the contents of the buffer into a new byte slice
+            @session_token = @data_buffer.read(@data_offset, mlen)
+
+          when SESSION_TTL
+            # Subtract 60 seconds from TTL so client session expires before server session.
+            seconds = @data_buffer.read_int32(@data_offset) - 60
+
+            if seconds > 0
+              @session_expiration = Time.now + (seconds/86400)
+            else
+              Aerospike.logger.warn("Invalid session TTL: #{seconds}")
+              raise Exceptions::Aerospike.new(result, "Node failed to retrieve session token")
+            end
+          end
+
+          @data_offset += mlen
+        end
+
+        if !@session_token
+          raise Exceptions::Aerospike.new(result, "Node failed to retrieve session token")
+        end
+      ensure
+        Buffer.put(@data_buffer)
+      end
+    end
+
+    def authenticate_via_token(conn, cluster)
+      @data_offset = 8
+      policy = cluster.client_policy
+      if policy.auth_mode != Aerospike::AuthMode::PKI
+        write_header(AUTHENTICATE, 2)
+        write_field_str(USER, policy.user)
+      else
+        write_header(AUTHENTICATE, 1)
+      end
+
+      write_field_bytes(SESSION_TOKEN, cluster.session_token) if cluster.session_token
+      write_size
+
+      conn.write(@data_buffer, @data_offset)
+      conn.read(@data_buffer, HEADER_SIZE)
+
+      result = @data_buffer.read(RESULT_CODE)
+      size = @data_buffer.read_int64(0)
+      receive_size = (size & 0xFFFFFFFFFFFF) - HEADER_REMAINING
+      conn.read(@data_buffer, receive_size)
+
+      if result != 0
+        return if result == Aerospike::ResultCode::SECURITY_NOT_ENABLED
+        raise Exceptions::Aerospike.new(result, "Authentication failed")
+      end
+
+      nil
+    end
+
+    SALT = '$2a$10$7EqJtq98hPqEX7fNZaFWoO'
+    def self.hash_password(password)
+      # Hashing the password with the cost of 10, with a static salt
+      return BCrypt::Engine.hash_secret(password, SALT, :cost => 10)
+    end
+  end
+end
+