admin_command.go

// Copyright 2013-2015 Aerospike, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use acmd 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.

package aerospike

import (
	"time"

	"github.com/THE108/aerospike-client-go/pkg/bcrypt"
	. "github.com/THE108/aerospike-client-go/types"
	Buffer "github.com/THE108/aerospike-client-go/utils/buffer"
)

const (
	// Commands
	_AUTHENTICATE    byte = 0
	_CREATE_USER     byte = 1
	_DROP_USER       byte = 2
	_SET_PASSWORD    byte = 3
	_CHANGE_PASSWORD byte = 4
	_GRANT_ROLES     byte = 5
	_REVOKE_ROLES    byte = 6
	_REPLACE_ROLES   byte = 7
	//_CREATE_ROLE byte = 8;
	_QUERY_USERS byte = 9
	//_QUERY_ROLES byte =  10;

	// Field IDs
	_USER         byte = 0
	_PASSWORD     byte = 1
	_OLD_PASSWORD byte = 2
	_CREDENTIAL   byte = 3
	_ROLES        byte = 10
	//_PRIVILEGES byte =  11;

	// Misc
	_MSG_VERSION int64 = 0
	_MSG_TYPE    int64 = 2

	_HEADER_SIZE      int = 24
	_HEADER_REMAINING int = 16
	_RESULT_CODE      int = 9
	_QUERY_END        int = 50
)

type AdminCommand struct {
	dataBuffer []byte
	dataOffset int
}

func newAdminCommand() *AdminCommand {
	return &AdminCommand{
		dataBuffer: bufPool.Get(),
		dataOffset: 8,
	}
}

func (acmd *AdminCommand) authenticate(conn *Connection, user string, password []byte) error {

	acmd.setAuthenticate(user, password)
	if _, err := conn.Write(acmd.dataBuffer[:acmd.dataOffset]); err != nil {
		return err
	}

	if _, err := conn.Read(acmd.dataBuffer, _HEADER_SIZE); err != nil {
		return err
	}

	result := acmd.dataBuffer[_RESULT_CODE]
	if result != 0 {
		return NewAerospikeError(ResultCode(result), "Authentication failed")
	}

	bufPool.Put(acmd.dataBuffer)

	return nil
}

func (acmd *AdminCommand) setAuthenticate(user string, password []byte) int {
	acmd.writeHeader(_AUTHENTICATE, 2)
	acmd.writeFieldStr(_USER, user)
	acmd.writeFieldBytes(_CREDENTIAL, password)
	acmd.writeSize()

	return acmd.dataOffset
}

func (acmd *AdminCommand) createUser(cluster *Cluster, policy *AdminPolicy, user string, password []byte, roles []string) error {
	acmd.writeHeader(_CREATE_USER, 3)
	acmd.writeFieldStr(_USER, user)
	acmd.writeFieldBytes(_PASSWORD, password)
	acmd.writeRoles(roles)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) dropUser(cluster *Cluster, policy *AdminPolicy, user string) error {
	acmd.writeHeader(_DROP_USER, 1)
	acmd.writeFieldStr(_USER, user)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) setPassword(cluster *Cluster, policy *AdminPolicy, user string, password []byte) error {
	acmd.writeHeader(_SET_PASSWORD, 2)
	acmd.writeFieldStr(_USER, user)
	acmd.writeFieldBytes(_PASSWORD, password)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) changePassword(cluster *Cluster, policy *AdminPolicy, user string, password []byte) error {
	acmd.writeHeader(_CHANGE_PASSWORD, 3)
	acmd.writeFieldStr(_USER, user)
	acmd.writeFieldBytes(_OLD_PASSWORD, cluster.Password())
	acmd.writeFieldBytes(_PASSWORD, password)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) grantRoles(cluster *Cluster, policy *AdminPolicy, user string, roles []string) error {
	acmd.writeHeader(_GRANT_ROLES, 2)
	acmd.writeFieldStr(_USER, user)
	acmd.writeRoles(roles)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) revokeRoles(cluster *Cluster, policy *AdminPolicy, user string, roles []string) error {
	acmd.writeHeader(_REVOKE_ROLES, 2)
	acmd.writeFieldStr(_USER, user)
	acmd.writeRoles(roles)
	return acmd.executeCommand(cluster, policy)
}

func (acmd *AdminCommand) queryUser(cluster *Cluster, policy *AdminPolicy, user string) (*UserRoles, error) {
	// TODO: Remove the workaround in the future
	time.Sleep(time.Millisecond * 10)
	defer bufPool.Put(acmd.dataBuffer)

	acmd.writeHeader(_QUERY_USERS, 1)
	acmd.writeFieldStr(_USER, user)
	list, err := acmd.readUsers(cluster, policy)
	if err != nil {
		return nil, err
	}

	if len(list) > 0 {
		return list[0], nil
	}

	return nil, nil
}

func (acmd *AdminCommand) queryUsers(cluster *Cluster, policy *AdminPolicy) ([]*UserRoles, error) {
	// TODO: Remove the workaround in the future
	time.Sleep(time.Millisecond * 10)
	defer bufPool.Put(acmd.dataBuffer)

	acmd.writeHeader(_QUERY_USERS, 0)
	list, err := acmd.readUsers(cluster, policy)
	if err != nil {
		return nil, err
	}
	return list, nil
}

func (acmd *AdminCommand) writeRoles(roles []string) {
	offset := acmd.dataOffset + int(_FIELD_HEADER_SIZE)
	acmd.dataBuffer[offset] = byte(len(roles))
	offset++

	for _, role := range roles {
		len := copy(acmd.dataBuffer[offset+1:], role)
		acmd.dataBuffer[offset] = byte(len)
		offset += len + 1
	}

	size := offset - acmd.dataOffset - int(_FIELD_HEADER_SIZE)
	acmd.writeFieldHeader(_ROLES, size)
	acmd.dataOffset = offset
}

func (acmd *AdminCommand) writeSize() {
	// Write total size of message which is the current offset.
	var size = int64(acmd.dataOffset-8) | (_MSG_VERSION << 56) | (_MSG_TYPE << 48)
	Buffer.Int64ToBytes(size, acmd.dataBuffer, 0)
}

func (acmd *AdminCommand) writeHeader(command byte, fieldCount int) {
	// Authenticate header is almost all zeros
	for i := acmd.dataOffset; i < acmd.dataOffset+16; i++ {
		acmd.dataBuffer[i] = 0
	}
	acmd.dataBuffer[acmd.dataOffset+2] = command
	acmd.dataBuffer[acmd.dataOffset+3] = byte(fieldCount)
	acmd.dataOffset += 16
}

func (acmd *AdminCommand) writeFieldStr(id byte, str string) {
	len := copy(acmd.dataBuffer[acmd.dataOffset+int(_FIELD_HEADER_SIZE):], str)
	acmd.writeFieldHeader(id, len)
	acmd.dataOffset += len
}

func (acmd *AdminCommand) writeFieldBytes(id byte, bytes []byte) {
	copy(acmd.dataBuffer[acmd.dataOffset+int(_FIELD_HEADER_SIZE):], bytes)
	acmd.writeFieldHeader(id, len(bytes))
	acmd.dataOffset += len(bytes)
}

func (acmd *AdminCommand) writeFieldHeader(id byte, size int) {
	Buffer.Int32ToBytes(int32(size+1), acmd.dataBuffer, acmd.dataOffset)
	acmd.dataOffset += 4
	acmd.dataBuffer[acmd.dataOffset] = id
	acmd.dataOffset++
}

func (acmd *AdminCommand) executeCommand(cluster *Cluster, policy *AdminPolicy) error {
	// TODO: Remove the workaround in the future
	defer time.Sleep(time.Millisecond * 10)

	defer bufPool.Put(acmd.dataBuffer)

	acmd.writeSize()
	node, err := cluster.GetRandomNode()
	if err != nil {
		return nil
	}
	timeout := 1 * time.Second
	if policy != nil && policy.Timeout > 0 {
		timeout = policy.Timeout
	}

	conn, err := node.GetConnection(timeout)
	if err != nil {
		return err
	}

	if _, err := conn.Write(acmd.dataBuffer[:acmd.dataOffset]); err != nil {
		node.InvalidateConnection(conn)
		return err
	}

	if _, err := conn.Read(acmd.dataBuffer, _HEADER_SIZE); err != nil {
		node.InvalidateConnection(conn)
		return err
	}

	node.PutConnection(conn)

	result := acmd.dataBuffer[_RESULT_CODE]
	if result != 0 {
		return NewAerospikeError(ResultCode(result))
	}

	return nil
}

func (acmd *AdminCommand) readUsers(cluster *Cluster, policy *AdminPolicy) ([]*UserRoles, error) {
	acmd.writeSize()
	node, err := cluster.GetRandomNode()
	if err != nil {
		return nil, err
	}
	timeout := 1 * time.Second
	if policy != nil && policy.Timeout > 0 {
		timeout = policy.Timeout
	}

	conn, err := node.GetConnection(timeout)
	if err != nil {
		return nil, err
	}

	if _, err := conn.Write(acmd.dataBuffer[:acmd.dataOffset]); err != nil {
		node.InvalidateConnection(conn)
		return nil, err
	}

	status, list, err := acmd.readUserBlocks(conn)
	if err != nil {
		node.InvalidateConnection(conn)
		return nil, err
	}
	node.PutConnection(conn)

	if status > 0 {
		return nil, NewAerospikeError(ResultCode(status))
	}
	return list, nil
}

func (acmd *AdminCommand) readUserBlocks(conn *Connection) (status int, rlist []*UserRoles, err error) {

	var list []*UserRoles

	for status == 0 {
		if _, err = conn.Read(acmd.dataBuffer, 8); err != nil {
			return -1, nil, err
		}

		size := Buffer.BytesToInt64(acmd.dataBuffer, 0)
		receiveSize := (size & 0xFFFFFFFFFFFF)

		if receiveSize > 0 {
			if receiveSize > int64(len(acmd.dataBuffer)) {
				acmd.dataBuffer = make([]byte, receiveSize)
			}
			if _, err = conn.Read(acmd.dataBuffer, int(receiveSize)); err != nil {
				return -1, nil, err
			}
			status, list, err = acmd.parseUsers(int(receiveSize))
			if err != nil {
				return -1, nil, err
			}
			rlist = append(rlist, list...)
		} else {
			break
		}
	}
	return status, rlist, nil
}

func (acmd *AdminCommand) parseUsers(receiveSize int) (int, []*UserRoles, error) {
	acmd.dataOffset = 0
	list := make([]*UserRoles, 0, 100)

	for acmd.dataOffset < receiveSize {
		resultCode := int(acmd.dataBuffer[acmd.dataOffset+1])

		if resultCode != 0 {
			if resultCode == _QUERY_END {
				return -1, nil, nil
			}
			return resultCode, nil, nil
		}

		userRoles := &UserRoles{}
		fieldCount := int(acmd.dataBuffer[acmd.dataOffset+3])
		acmd.dataOffset += _HEADER_REMAINING

		for i := 0; i < fieldCount; i++ {
			len := int(Buffer.BytesToInt32(acmd.dataBuffer, acmd.dataOffset))
			acmd.dataOffset += 4
			id := acmd.dataBuffer[acmd.dataOffset]
			acmd.dataOffset++
			len--

			if id == _USER {
				userRoles.User = string(acmd.dataBuffer[acmd.dataOffset : acmd.dataOffset+len])
				acmd.dataOffset += len
			} else if id == _ROLES {
				acmd.parseRoles(userRoles)
			} else {
				acmd.dataOffset += len
			}
		}

		if userRoles.User == "" && userRoles.Roles == nil {
			continue
		}

		if userRoles.Roles == nil {
			userRoles.Roles = make([]string, 0)
		}
		list = append(list, userRoles)
	}

	return 0, list, nil
}

func (acmd *AdminCommand) parseRoles(userRoles *UserRoles) {
	size := int(acmd.dataBuffer[acmd.dataOffset])
	acmd.dataOffset++
	userRoles.Roles = make([]string, 0, size)

	for i := 0; i < size; i++ {
		len := int(acmd.dataBuffer[acmd.dataOffset])
		acmd.dataOffset++
		role := string(acmd.dataBuffer[acmd.dataOffset : acmd.dataOffset+len])
		acmd.dataOffset += len
		userRoles.Roles = append(userRoles.Roles, role)
	}
}

func hashPassword(password string) ([]byte, error) {
	// Hashing the password with the cost of 10, with a static salt
	const salt = "$2a$10$7EqJtq98hPqEX7fNZaFWoO"
	hashedPassword, err := bcrypt.Hash(password, salt)
	if err != nil {
		return nil, err
	}
	return []byte(hashedPassword), nil
}