content: /*
   * Copyright 2021-2024 Aklivity Inc.
   *
   * Aklivity licenses this file 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.
   */
  package io.aklivity.zilla.runtime.binding.kafka.internal.stream;
  
  import static io.aklivity.zilla.runtime.binding.kafka.internal.KafkaConfiguration.KAFKA_CLIENT_ID_DEFAULT;
  
  import java.nio.charset.StandardCharsets;
  import java.security.MessageDigest;
  import java.security.SecureRandom;
  import java.util.Arrays;
  import java.util.Base64;
  import java.util.List;
  import java.util.Map;
  import java.util.function.LongUnaryOperator;
  import java.util.function.Supplier;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  
  import javax.crypto.Mac;
  import javax.crypto.spec.SecretKeySpec;
  
  import org.agrona.DirectBuffer;
  import org.agrona.LangUtil;
  import org.agrona.MutableDirectBuffer;
  import org.agrona.collections.LongLongConsumer;
  import org.agrona.collections.Object2ObjectHashMap;
  import org.agrona.concurrent.UnsafeBuffer;
  
  import io.aklivity.zilla.runtime.binding.kafka.config.KafkaSaslConfig;
  import io.aklivity.zilla.runtime.binding.kafka.config.KafkaServerConfig;
  import io.aklivity.zilla.runtime.binding.kafka.identity.KafkaClientIdSupplier;
  import io.aklivity.zilla.runtime.binding.kafka.internal.KafkaConfiguration;
  import io.aklivity.zilla.runtime.binding.kafka.internal.config.KafkaScramMechanism;
  import io.aklivity.zilla.runtime.binding.kafka.internal.events.KafkaEventContext;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.String16FW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.RequestHeaderFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.ResponseHeaderFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.sasl.SaslAuthenticateRequestFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.sasl.SaslAuthenticateResponseFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.sasl.SaslHandshakeMechanismResponseFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.sasl.SaslHandshakeRequestFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.codec.sasl.SaslHandshakeResponseFW;
  import io.aklivity.zilla.runtime.binding.kafka.internal.types.stream.DataFW;
  import io.aklivity.zilla.runtime.engine.EngineContext;
  
  public abstract class KafkaClientSaslHandshaker
  {
      private static final short SASL_HANDSHAKE_API_KEY = 17;
      private static final short SASL_HANDSHAKE_API_VERSION = 1;
  
      private static final short SASL_AUTHENTICATE_API_KEY = 36;
      private static final short SASL_AUTHENTICATE_API_VERSION = 1;
      private static final int ERROR_SASL_AUTHENTICATION_FAILED = 58;
      private static final int ERROR_NONE = 0;
      private static final int ERROR_CLUSTER_AUTHORIZATION_FAILED = 31;
      private static final int ERROR_UNSUPPORTED_VERSION = 35;
  
      private static final String CLIENT_KEY = "Client Key";
      private static final String SERVER_KEY = "Server Key";
      private static final String PRINTABLE = "[\\x21-\\x7E&&[^,]]+";
      private static final String BASE64_CHAR = "[a-zA-Z0-9/+]";
      private static final String BASE64 = String.format("(?:%s{4})*(?:%s{3}=|%s{2}==)?", BASE64_CHAR, BASE64_CHAR, BASE64_CHAR);
      private static final Pattern SASL_SCRAM_SERVER_FIRST_MESSAGE = Pattern.compile(String.format(
              "r=(?<nonce>%s),s=(?<salt>%s),i=(?<iterations>[0-9]+)", PRINTABLE, BASE64));
      private static final Pattern SASL_SCRAM_SERVER_FINAL_MESSAGE = Pattern.compile(String.format(
              "(?:v=(?<verifier>%s))", BASE64));
      private static final String AUTH_MESSAGE_FORMAT = "n=%s,r=%s,r=%s,s=%s,i=%s,c=%s,r=%s";
      private static final byte[] SASL_SCRAM_CHANNEL_BINDING = "n,,".getBytes(StandardCharsets.US_ASCII);
      private static final byte[] SASL_SCRAM_USERNAME = "n=".getBytes(StandardCharsets.US_ASCII);
      private static final byte[] SASL_SCRAM_RANDOM = ",r=".getBytes(StandardCharsets.US_ASCII);
      private static final byte[] SASL_SCRAM_CHANNEL = "c=".getBytes(StandardCharsets.US_ASCII);
      private static final byte[] SASL_SCRAM_SALT_PASSWORD = ",p=".getBytes(StandardCharsets.US_ASCII);
      private static final String SASL_SCRAM_CHANNEL_RANDOM = Base64.getEncoder().encodeToString(SASL_SCRAM_CHANNEL_BINDING);
  
      private static final String16FW KAFKA_CLIENT_ID_DEFAULT_VALUE = new String16FW(KAFKA_CLIENT_ID_DEFAULT);
  
      private final RequestHeaderFW.Builder requestHeaderRW = new RequestHeaderFW.Builder();
      private final SaslHandshakeRequestFW.Builder saslHandshakeRequestRW = new SaslHandshakeRequestFW.Builder();
      private final SaslAuthenticateRequestFW.Builder saslAuthenticateRequestRW = new SaslAuthenticateRequestFW.Builder();
  
      private final ResponseHeaderFW responseHeaderRO = new ResponseHeaderFW();
      private final SaslHandshakeResponseFW saslHandshakeResponseRO = new SaslHandshakeResponseFW();
      private final SaslHandshakeMechanismResponseFW saslHandshakeMechanismResponseRO = new SaslHandshakeMechanismResponseFW();
      private final SaslAuthenticateResponseFW saslAuthenticateResponseRO = new SaslAuthenticateResponseFW();
      private final KafkaEventContext event;
  
      private KafkaSaslClientDecoder decodeSaslPlainAuthenticate = this::decodeSaslPlainAuthenticate;
      private KafkaSaslClientDecoder decodeSaslScramAuthenticateFirst = this::decodeSaslScramAuthenticateFirst;
      private KafkaSaslClientDecoder decodeSaslScramAuthenticateFinal = this::decodeSaslScramAuthenticateFinal;
  
      private final SecureRandom random = new SecureRandom();
  
      private final MutableDirectBuffer scramBuffer = new UnsafeBuffer(new byte[1024]);
  
      private MessageDigest messageDigest;
      private Mac mac;
      private Supplier<String> nonceSupplier;
      private KafkaScramMechanism mechanism;
      private Matcher serverResponseMatcher;
      private byte[] result, ui, prev;
  
      private final Map<KafkaServerConfig, String16FW> clientIdsByServer;
  
      protected final KafkaClientIdSupplier clientIdSupplier;
      protected final LongUnaryOperator supplyInitialId;
      protected final LongUnaryOperator supplyReplyId;
      protected final MutableDirectBuffer writeBuffer;
  
      public KafkaClientSaslHandshaker(
          KafkaConfiguration config,
          EngineContext context)
      {
          this.clientIdSupplier = KafkaClientIdSupplier.instantiate(config);
          this.supplyInitialId = context::supplyInitialId;
          this.supplyReplyId = context::supplyReplyId;
          this.writeBuffer = new UnsafeBuffer(new byte[context.writeBuffer().capacity()]);
          this.nonceSupplier = config.nonceSupplier();
          this.clientIdsByServer = new Object2ObjectHashMap<>();
          this.event = new KafkaEventContext(context);
      }
  
      public abstract class KafkaSaslClient
      {
          protected final KafkaSaslConfig sasl;
          protected final long originId;
          protected final long routedId;
          protected final KafkaServerConfig server;
          protected final String16FW clientId;
  
          protected long initialId;
          protected long replyId;
  
          protected int nextRequestId;
  
          private int decodableResponseBytes;
          private int decodableMechanisms;
          private String clientNonce;
          private String serverNonce;
          private String salt;
          private String iterationCount;
          byte[] saltedPassword;
          byte[] authMessage;
          private LongLongConsumer encodeSaslAuthenticate;
          private KafkaSaslClientDecoder decodeSaslAuthenticate;
  
          protected KafkaSaslClient(
              List<KafkaServerConfig> servers,
              KafkaSaslConfig sasl,
              long originId,
              long routedId)
          {
              this(servers != null && !servers.isEmpty()
                      ? servers.get(random.nextInt(servers.size()))
                      : null,
                  sasl, originId, routedId);
          }
  
          protected KafkaSaslClient(
              KafkaServerConfig server,
              KafkaSaslConfig sasl,
              long originId,
              long routedId)
          {
              this.server = server;
              this.sasl = sasl;
              this.originId = originId;
              this.routedId = routedId;
              this.clientId = supplyClientId(server);
              this.initialId = supplyInitialId.applyAsLong(routedId);
              this.replyId = supplyReplyId.applyAsLong(initialId);
          }
  
          protected final void doEncodeSaslHandshakeRequest(
              long traceId,
              long budgetId)
          {
              final MutableDirectBuffer encodeBuffer = writeBuffer;
              final int encodeOffset = DataFW.FIELD_OFFSET_PAYLOAD;
              final int encodeLimit = encodeBuffer.capacity();
  
              int encodeProgress = encodeOffset;
  
              final RequestHeaderFW requestHeader = requestHeaderRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                      .length(0)
                      .apiKey(SASL_HANDSHAKE_API_KEY)
                      .apiVersion(SASL_HANDSHAKE_API_VERSION)
                      .correlationId(0)
                      .clientId(clientId)
                      .build();
  
              encodeProgress = requestHeader.limit();
  
              final String mechanism = sasl.mechanism;
              final SaslHandshakeRequestFW saslRequest = saslHandshakeRequestRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                      .mechanism(mechanism.toUpperCase())
                      .build();
  
              encodeProgress = saslRequest.limit();
  
              final int requestId = nextRequestId++;
              final int requestSize = encodeProgress - encodeOffset - RequestHeaderFW.FIELD_OFFSET_API_KEY;
  
              requestHeaderRW.wrap(encodeBuffer, requestHeader.offset(), requestHeader.limit())
                      .length(requestSize)
                      .apiKey(requestHeader.apiKey())
                      .apiVersion(requestHeader.apiVersion())
                      .correlationId(requestId)
                      .clientId(requestHeader.clientId())
                      .build();
  
              if (KafkaConfiguration.DEBUG)
              {
                  System.out.format("[0x%016x] SASL HANDSHAKE %s\n", replyId, mechanism);
              }
  
              doNetworkData(traceId, budgetId, encodeBuffer, encodeOffset, encodeProgress);
  
              doDecodeSaslHandshakeResponse(traceId);
          }
  
          protected final void doEncodeSaslAuthenticateRequest(
                  long traceId,
                  long budgetId)
          {
              encodeSaslAuthenticate.accept(traceId, budgetId);
          }
  
          private void doEncodeSaslPlainAuthenticateRequest(
              long traceId,
              long budgetId)
          {
              final MutableDirectBuffer encodeBuffer = writeBuffer;
              final int encodeOffset = DataFW.FIELD_OFFSET_PAYLOAD;
              final int encodeLimit = encodeBuffer.capacity();
  
              int encodeProgress = encodeOffset;
  
              final RequestHeaderFW requestHeader = requestHeaderRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                      .length(0)
                      .apiKey(SASL_AUTHENTICATE_API_KEY)
                      .apiVersion(SASL_AUTHENTICATE_API_VERSION)
                      .correlationId(0)
                      .clientId(clientId)
                      .build();
  
              encodeProgress = requestHeader.limit();
  
              final String username = sasl.username;
              final String password = sasl.password;
  
              final SaslAuthenticateRequestFW authenticateRequest =
                      saslAuthenticateRequestRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                              .authBytes(a -> a.put((b, o, l) ->
                              {
                                  int p = o;
                                  b.putByte(p++, (byte) 0);
                                  p += b.putStringWithoutLengthUtf8(p, username);
                                  b.putByte(p++, (byte) 0);
                                  p += b.putStringWithoutLengthUtf8(p, password);
                                  return p - o;
                              }))
                              .build();
              encodeProgress = authenticateRequest.limit();
  
              final int requestId = nextRequestId++;
              final int requestSize = encodeProgress - encodeOffset - RequestHeaderFW.FIELD_OFFSET_API_KEY;
  
              requestHeaderRW.wrap(encodeBuffer, requestHeader.offset(), requestHeader.limit())
                      .length(requestSize)
                      .apiKey(requestHeader.apiKey())
                      .apiVersion(requestHeader.apiVersion())
                      .correlationId(requestId)
                      .clientId(requestHeader.clientId())
                      .build();
  
              if (KafkaConfiguration.DEBUG)
              {
                  System.out.format("[0x%016x] SASL AUTHENTICATE %s\n", replyId, username);
              }
  
              doNetworkData(traceId, budgetId, encodeBuffer, encodeOffset, encodeProgress);
  
              doDecodeSaslAuthenticateResponse(traceId);
          }
  
          private void doEncodeSaslScramFirstAuthenticateRequest(
                  long traceId,
                  long budgetId)
          {
              final MutableDirectBuffer encodeBuffer = writeBuffer;
              final int encodeOffset = DataFW.FIELD_OFFSET_PAYLOAD;
              final int encodeLimit = encodeBuffer.capacity();
  
              int encodeProgress = encodeOffset;
  
              final RequestHeaderFW requestHeader = requestHeaderRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                      .length(0)
                      .apiKey(SASL_AUTHENTICATE_API_KEY)
                      .apiVersion(SASL_AUTHENTICATE_API_VERSION)
                      .correlationId(0)
                      .clientId(clientId)
                      .build();
  
              encodeProgress = requestHeader.limit();
  
              final String username = sasl.username;
              clientNonce = nonceSupplier.get();
  
              int scramBytes = 0;
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_CHANNEL_BINDING);
              scramBytes += SASL_SCRAM_CHANNEL_BINDING.length;
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_USERNAME);
              scramBytes += SASL_SCRAM_USERNAME.length;
              scramBytes += scramBuffer.putStringWithoutLengthUtf8(scramBytes, username);
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_RANDOM);
              scramBytes += SASL_SCRAM_RANDOM.length;
              scramBuffer.putBytes(scramBytes, clientNonce.getBytes());
              scramBytes += clientNonce.getBytes().length;
  
              final SaslAuthenticateRequestFW authenticateRequest =
                      saslAuthenticateRequestRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                              .authBytes(scramBuffer, 0, scramBytes)
                              .build();
  
              encodeProgress = authenticateRequest.limit();
  
              final int requestId = nextRequestId++;
              final int requestSize = encodeProgress - encodeOffset - RequestHeaderFW.FIELD_OFFSET_API_KEY;
  
              requestHeaderRW.wrap(encodeBuffer, requestHeader.offset(), requestHeader.limit())
                      .length(requestSize)
                      .apiKey(requestHeader.apiKey())
                      .apiVersion(requestHeader.apiVersion())
                      .correlationId(requestId)
                      .clientId(requestHeader.clientId())
                      .build();
  
              if (KafkaConfiguration.DEBUG)
              {
                  System.out.format("[0x%016x] SASL AUTHENTICATE %s\n", replyId, username);
              }
  
              doNetworkData(traceId, budgetId, encodeBuffer, encodeOffset, encodeProgress);
  
              doDecodeSaslAuthenticateResponse(traceId);
          }
  
          private void doEncodeSaslScramFinalAuthenticateRequest(
                  long traceId,
                  long budgetId)
          {
              mechanism = KafkaScramMechanism.forMechanismName(sasl.mechanism.toUpperCase());
              try
              {
                  messageDigest = MessageDigest.getInstance(mechanism.hashAlgorithm());
                  mac = Mac.getInstance(mechanism.macAlgorithm());
              }
              catch (Exception e)
              {
                  LangUtil.rethrowUnchecked(e);
              }
              saltedPassword = hi(sasl.password.getBytes(StandardCharsets.US_ASCII),
                      Base64.getDecoder().decode(toBytes(salt)),
                      Integer.parseInt(iterationCount));
              authMessage = toBytes(String.format(AUTH_MESSAGE_FORMAT,
                      sasl.username, clientNonce, serverNonce, salt, iterationCount, SASL_SCRAM_CHANNEL_RANDOM, serverNonce));
  
              int scramBytes = 0;
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_CHANNEL);
              scramBytes += SASL_SCRAM_CHANNEL.length;
              scramBytes += scramBuffer.putStringWithoutLengthUtf8(scramBytes, SASL_SCRAM_CHANNEL_RANDOM);
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_RANDOM);
              scramBytes += SASL_SCRAM_RANDOM.length;
              scramBytes += scramBuffer.putStringWithoutLengthUtf8(scramBytes, serverNonce);
              scramBuffer.putBytes(scramBytes, SASL_SCRAM_SALT_PASSWORD);
              scramBytes += SASL_SCRAM_SALT_PASSWORD.length;
              scramBytes += scramBuffer.putStringWithoutLengthUtf8(scramBytes, clientProof(saltedPassword, authMessage));
  
              final MutableDirectBuffer encodeBuffer = writeBuffer;
              final int encodeOffset = DataFW.FIELD_OFFSET_PAYLOAD;
              final int encodeLimit = encodeBuffer.capacity();
  
              int encodeProgress = encodeOffset;
  
              final RequestHeaderFW requestHeader = requestHeaderRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                      .length(0)
                      .apiKey(SASL_AUTHENTICATE_API_KEY)
                      .apiVersion(SASL_AUTHENTICATE_API_VERSION)
                      .correlationId(0)
                      .clientId(clientId)
                      .build();
  
              encodeProgress = requestHeader.limit();
  
              final SaslAuthenticateRequestFW authenticateRequest =
                      saslAuthenticateRequestRW.wrap(encodeBuffer, encodeProgress, encodeLimit)
                              .authBytes(scramBuffer, 0, scramBytes)
                              .build();
  
              encodeProgress = authenticateRequest.limit();
  
              final int requestId = nextRequestId++;
              final int requestSize = encodeProgress - encodeOffset - RequestHeaderFW.FIELD_OFFSET_API_KEY;
  
              requestHeaderRW.wrap(encodeBuffer, requestHeader.offset(), requestHeader.limit())
                      .length(requestSize)
                      .apiKey(requestHeader.apiKey())
                      .apiVersion(requestHeader.apiVersion())
                      .correlationId(requestId)
                      .clientId(requestHeader.clientId())
                      .build();
  
              doNetworkData(traceId, budgetId, encodeBuffer, encodeOffset, encodeProgress);
  
              doDecodeSaslAuthenticateResponse(traceId);
          }
  
          protected final void onDecodeResponseErrorCode(
              long traceId,
              long bindingId,
              int apiKey,
              int apiVersion,
              int errorCode)
          {
              switch (errorCode)
              {
              case ERROR_CLUSTER_AUTHORIZATION_FAILED -> event.clusterAuthorizationFailed(traceId, bindingId, apiKey, apiVersion);
              case ERROR_UNSUPPORTED_VERSION -> event.apiVersionRejected(traceId, bindingId, apiKey, apiVersion);
              }
          }
  
          protected abstract void doNetworkData(
              long traceId,
              long budgetId,
              DirectBuffer buffer,
              int offset,
              int limit);
  
          protected abstract void onDecodeSaslHandshakeResponse(
              long traceId,
              long authorization,
              int errorCode);
  
          protected abstract void onDecodeSaslAuthenticateResponse(
              long traceId,
              long authorization,
              int errorCode);
  
          protected abstract void onDecodeSaslResponse(
              long traceId);
  
          protected abstract void doDecodeSaslHandshakeResponse(
                  long traceId);
  
          protected abstract void doDecodeSaslAuthenticateResponse(
                  long traceId);
  
          protected abstract void doDecodeSaslAuthenticate(
              long traceId);
  
          protected abstract void doDecodeSaslHandshake(
              long traceId);
  
          protected abstract void doDecodeSaslHandshakeMechanisms(
              long traceId);
  
          protected abstract void doDecodeSaslHandshakeMechansim(
              long traceId);
      }
  
      @FunctionalInterface
      private interface KafkaSaslClientEncoder
      {
          void encode(
              KafkaSaslClient client,
              long replyId,
              long traceId,
              long budgetId);
      }
  
      @FunctionalInterface
      private interface KafkaSaslClientDecoder
      {
          int decode(
              KafkaSaslClient client,
              long traceId,
              long authorization,
              long budgetId,
              int reserved,
              MutableDirectBuffer buffer,
              int offset,
              int progress,
              int limit);
      }
  
      protected final int decodeSaslHandshakeResponse(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final ResponseHeaderFW responseHeader = responseHeaderRO.tryWrap(buffer, progress, limit);
              if (responseHeader != null)
              {
                  progress = responseHeader.limit();
                  client.decodableResponseBytes = responseHeader.length();
                  client.doDecodeSaslHandshake(traceId);
              }
          }
  
          return progress;
      }
  
      protected final int decodeSaslHandshake(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final SaslHandshakeResponseFW saslResponse = saslHandshakeResponseRO.tryWrap(buffer, progress, limit);
  
              if (saslResponse != null)
              {
                  final int errorCode = saslResponse.errorCode();
  
                  progress = saslResponse.limit();
  
                  client.decodableResponseBytes -= saslResponse.sizeof();
                  assert client.decodableResponseBytes >= 0;
                  if (errorCode == ERROR_NONE)
                  {
                      switch (client.sasl.mechanism)
                      {
                      case "plain" :
                          client.encodeSaslAuthenticate = client::doEncodeSaslPlainAuthenticateRequest;
                          client.decodeSaslAuthenticate = decodeSaslPlainAuthenticate;
                          break;
                      case "scram-sha-1" :
                      case "scram-sha-256" :
                      case "scram-sha-512" :
                          client.encodeSaslAuthenticate = client::doEncodeSaslScramFirstAuthenticateRequest;
                          client.decodeSaslAuthenticate = decodeSaslScramAuthenticateFirst;
                          break;
                      }
                  }
                  client.onDecodeSaslHandshakeResponse(traceId, authorization, errorCode);
  
                  client.decodableMechanisms = saslResponse.mechanismCount();
                  client.doDecodeSaslHandshakeMechanisms(traceId);
              }
          }
  
          return progress;
      }
  
      protected final int decodeSaslHandshakeMechanisms(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          if (client.decodableMechanisms == 0)
          {
              client.onDecodeSaslResponse(traceId);
              client.doDecodeSaslAuthenticateResponse(traceId);
          }
          else
          {
              client.doDecodeSaslHandshakeMechansim(traceId);
          }
  
          return progress;
      }
  
      protected final int decodeSaslHandshakeMechanism(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          decode:
          if (length != 0)
          {
              final SaslHandshakeMechanismResponseFW mechanism = saslHandshakeMechanismResponseRO.tryWrap(buffer, progress, limit);
              if (mechanism == null)
              {
                  break decode;
              }
  
              progress = mechanism.limit();
  
              client.decodableResponseBytes -= mechanism.sizeof();
              assert client.decodableResponseBytes >= 0;
  
              client.decodableMechanisms--;
              assert client.decodableMechanisms >= 0;
  
              client.doDecodeSaslHandshakeMechanisms(traceId);
          }
  
          return progress;
      }
  
      protected final int decodeSaslAuthenticateResponse(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final ResponseHeaderFW responseHeader = responseHeaderRO.tryWrap(buffer, progress, limit);
              if (responseHeader != null)
              {
                  progress = responseHeader.limit();
                  client.decodableResponseBytes = responseHeader.length();
                  client.doDecodeSaslAuthenticate(traceId);
              }
          }
  
          return progress;
      }
  
      protected final int decodeSaslAuthenticate(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          MutableDirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          return client.decodeSaslAuthenticate.decode(
                  client, traceId, authorization,
                  budgetId, reserved, buffer,
                  offset, progress, limit);
      }
  
      private int decodeSaslPlainAuthenticate(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final SaslAuthenticateResponseFW authenticateResponse = saslAuthenticateResponseRO.tryWrap(buffer, progress, limit);
  
              if (authenticateResponse != null)
              {
                  final int errorCode = authenticateResponse.errorCode();
                  if (errorCode != ERROR_NONE)
                  {
                      event.authorizationFailed(traceId, client.originId, client.sasl.username);
                  }
  
                  progress = authenticateResponse.limit();
  
                  client.decodableResponseBytes -= authenticateResponse.sizeof();
                  assert client.decodableResponseBytes >= 0;
  
                  client.onDecodeSaslResponse(traceId);
                  client.onDecodeSaslAuthenticateResponse(traceId, authorization, errorCode);
              }
          }
  
          return progress;
      }
  
      private int decodeSaslScramAuthenticateFirst(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final SaslAuthenticateResponseFW authenticateResponse = saslAuthenticateResponseRO.tryWrap(buffer, progress, limit);
  
              if (authenticateResponse != null)
              {
                  final int errorCode = authenticateResponse.errorCode();
                  if (errorCode != ERROR_NONE)
                  {
                      event.authorizationFailed(traceId, client.originId, client.sasl.username);
                  }
  
                  progress = authenticateResponse.limit();
  
                  client.decodableResponseBytes -= authenticateResponse.sizeof();
                  assert client.decodableResponseBytes >= 0;
  
                  DirectBuffer serverFirstResponse = authenticateResponse.authBytes().value();
                  String serverFirstMessage = serverFirstResponse.getStringWithoutLengthUtf8(0,
                          serverFirstResponse.capacity());
  
                  serverResponseMatcher = SASL_SCRAM_SERVER_FIRST_MESSAGE.matcher(serverFirstMessage);
                  if (serverResponseMatcher.matches())
                  {
                      client.serverNonce = serverResponseMatcher.group("nonce");
                      client.salt = serverResponseMatcher.group("salt");
                      client.iterationCount = serverResponseMatcher.group("iterations");
                  }
                  if (errorCode == ERROR_NONE && client.serverNonce != null &&
                          client.serverNonce.startsWith(client.clientNonce))
                  {
                      client.encodeSaslAuthenticate = client::doEncodeSaslScramFinalAuthenticateRequest;
                      client.decodeSaslAuthenticate = decodeSaslScramAuthenticateFinal;
                      client.onDecodeSaslResponse(traceId);
                      client.onDecodeSaslHandshakeResponse(traceId, authorization, ERROR_NONE);
                  }
                  else
                  {
                      client.onDecodeSaslResponse(traceId);
                      client.onDecodeSaslAuthenticateResponse(traceId, authorization, ERROR_SASL_AUTHENTICATION_FAILED);
                  }
              }
          }
  
          return progress;
      }
  
      private int decodeSaslScramAuthenticateFinal(
          KafkaSaslClient client,
          long traceId,
          long authorization,
          long budgetId,
          int reserved,
          DirectBuffer buffer,
          int offset,
          int progress,
          int limit)
      {
          final int length = limit - progress;
  
          if (length != 0)
          {
              final SaslAuthenticateResponseFW authenticateResponse = saslAuthenticateResponseRO.tryWrap(buffer, progress, limit);
  
              if (authenticateResponse != null)
              {
                  final int errorCode = authenticateResponse.errorCode();
  
                  progress = authenticateResponse.limit();
  
                  client.decodableResponseBytes -= authenticateResponse.sizeof();
                  assert client.decodableResponseBytes >= 0;
  
                  byte[] serverKey = serverKey(client.saltedPassword);
                  byte[] serverSignature = hmac(serverKey, client.authMessage);
                  DirectBuffer serverFinalResponse = authenticateResponse.authBytes().value();
                  String serverFinalMessage = serverFinalResponse.getStringWithoutLengthUtf8(0,
                          serverFinalResponse.capacity());
                  serverResponseMatcher.reset();
                  serverResponseMatcher = SASL_SCRAM_SERVER_FINAL_MESSAGE.matcher(serverFinalMessage);
                  if (serverResponseMatcher.matches())
                  {
                      serverFinalMessage = serverResponseMatcher.group("verifier");
                  }
                  if (!Arrays.equals(Base64.getDecoder().decode(serverFinalMessage),
                          serverSignature))
                  {
                      client.onDecodeSaslResponse(traceId);
                      client.onDecodeSaslAuthenticateResponse(traceId, authorization, ERROR_SASL_AUTHENTICATION_FAILED);
                  }
                  else
                  {
                      client.onDecodeSaslResponse(traceId);
                      client.onDecodeSaslAuthenticateResponse(traceId, authorization, errorCode);
                  }
              }
          }
  
          return progress;
      }
  
      private String16FW supplyClientId(
          KafkaServerConfig server)
      {
          return server != null
              ? clientIdsByServer.computeIfAbsent(server, this::createClientId)
              : KAFKA_CLIENT_ID_DEFAULT_VALUE;
      }
  
      private String16FW createClientId(
          KafkaServerConfig server)
      {
          String clientId = clientIdSupplier.get(server.host);
          return clientId != null ? new String16FW(clientId) : KAFKA_CLIENT_ID_DEFAULT_VALUE;
      }
  
      public byte[] hmac(byte[] key, byte[] bytes)
      {
          try
          {
              mac.init(new SecretKeySpec(key, mac.getAlgorithm()));
          }
          catch (Exception e)
          {
              LangUtil.rethrowUnchecked(e);
          }
          return mac.doFinal(bytes);
      }
  
      public byte[] xor(byte[] first, byte[] second)
      {
          result = new byte[first.length];
          if (first.length == second.length)
          {
              for (int i = 0; i < result.length; i++)
              {
                  result[i] = (byte) (first[i] ^ second[i]);
              }
          }
          return result;
      }
  
      public byte[] hi(byte[] str, byte[] salt, int iterations)
      {
          try
          {
              mac.init(new SecretKeySpec(str, mac.getAlgorithm()));
          }
          catch (Exception e)
          {
              LangUtil.rethrowUnchecked(e);
          }
          mac.update(salt);
          result = prev = mac.doFinal(new byte[]{0, 0, 0, 1});
          for (int i = 2; i <= iterations; i++)
          {
              ui = hmac(str, prev);
              result = xor(result, ui);
              prev = ui;
          }
          return result;
      }
  
      public byte[] hash(byte[] str)
      {
          return messageDigest.digest(str);
      }
  
      public byte[] clientKey(byte[] saltedPassword)
      {
          return hmac(saltedPassword, toBytes(CLIENT_KEY));
      }
  
      public byte[] serverKey(byte[] saltedPassword)
      {
          return hmac(saltedPassword, toBytes(SERVER_KEY));
      }
  
      public String clientProof(byte[] saltedPassword, byte[] authMessage)
      {
          byte[] clientKey = clientKey(saltedPassword);
          return Base64.getEncoder().encodeToString(xor(clientKey,
                  hmac(hash(clientKey), authMessage)));
      }
  
      public byte[] toBytes(String str)
      {
          return str.getBytes(StandardCharsets.UTF_8);
      }
  }
  