github.com/QuangHoangHao/kafka-go@v0.4.36/sasl/scram/scram.go (about)

     1  package scram
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha256"
     6  	"crypto/sha512"
     7  	"hash"
     8  
     9  	"github.com/QuangHoangHao/kafka-go/sasl"
    10  	"github.com/xdg/scram"
    11  )
    12  
    13  // Algorithm determines the hash function used by SCRAM to protect the user's
    14  // credentials.
    15  type Algorithm interface {
    16  	// Name returns the algorithm's name, e.g. "SCRAM-SHA-256"
    17  	Name() string
    18  
    19  	// Hash returns a new hash.Hash.
    20  	Hash() hash.Hash
    21  }
    22  
    23  type sha256Algo struct{}
    24  
    25  func (sha256Algo) Name() string {
    26  	return "SCRAM-SHA-256"
    27  }
    28  
    29  func (sha256Algo) Hash() hash.Hash {
    30  	return sha256.New()
    31  }
    32  
    33  type sha512Algo struct{}
    34  
    35  func (sha512Algo) Name() string {
    36  	return "SCRAM-SHA-512"
    37  }
    38  
    39  func (sha512Algo) Hash() hash.Hash {
    40  	return sha512.New()
    41  }
    42  
    43  var (
    44  	SHA256 Algorithm = sha256Algo{}
    45  	SHA512 Algorithm = sha512Algo{}
    46  )
    47  
    48  type mechanism struct {
    49  	algo   Algorithm
    50  	client *scram.Client
    51  }
    52  
    53  type session struct {
    54  	convo *scram.ClientConversation
    55  }
    56  
    57  // Mechanism returns a new sasl.Mechanism that will use SCRAM with the provided
    58  // Algorithm to securely transmit the provided credentials to Kafka.
    59  //
    60  // SCRAM-SHA-256 and SCRAM-SHA-512 were added to Kafka in 0.10.2.0.  These
    61  // mechanisms will not work with older versions.
    62  func Mechanism(algo Algorithm, username, password string) (sasl.Mechanism, error) {
    63  	hashGen := scram.HashGeneratorFcn(algo.Hash)
    64  	client, err := hashGen.NewClient(username, password, "")
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	return &mechanism{
    70  		algo:   algo,
    71  		client: client,
    72  	}, nil
    73  }
    74  
    75  func (m *mechanism) Name() string {
    76  	return m.algo.Name()
    77  }
    78  
    79  func (m *mechanism) Start(ctx context.Context) (sasl.StateMachine, []byte, error) {
    80  	convo := m.client.NewConversation()
    81  	str, err := convo.Step("")
    82  	if err != nil {
    83  		return nil, nil, err
    84  	}
    85  	return &session{convo: convo}, []byte(str), nil
    86  }
    87  
    88  func (s *session) Next(ctx context.Context, challenge []byte) (bool, []byte, error) {
    89  	str, err := s.convo.Step(string(challenge))
    90  	return s.convo.Done(), []byte(str), err
    91  }