feat(security): handle list mechanism response

src/main/java/com/xiaomi/infra/pegasus/rpc/async/Negotiation.java

@@ -18,33 +18,35 @@
  */
 package com.xiaomi.infra.pegasus.rpc.async;
 
+import com.xiaomi.infra.pegasus.apps.negotiation_request;
 import com.xiaomi.infra.pegasus.apps.negotiation_response;
 import com.xiaomi.infra.pegasus.apps.negotiation_status;
 import com.xiaomi.infra.pegasus.base.blob;
 import com.xiaomi.infra.pegasus.base.error_code;
 import com.xiaomi.infra.pegasus.operator.negotiation_operator;
 import com.xiaomi.infra.pegasus.rpc.ReplicationException;
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import javax.security.auth.Subject;
-import javax.security.sasl.Sasl;
 import org.slf4j.Logger;
 
 public class Negotiation {
   private static final Logger logger = org.slf4j.LoggerFactory.getLogger(Negotiation.class);
+  // Because negotiation message is always the first rpc sent to pegasus server,
+  // which will cost much more time. so we set negotiation timeout to 10s here
+  private static final int negotiationTimeoutMS = 10000;
+  private static final List<String> expectedMechanisms =
+      new ArrayList<>(Collections.singletonList("GSSAPI"));
+
   private negotiation_status status;
   private ReplicaSession session;
-  private String serviceName; // used for SASL authentication
-  private String serviceFqdn; // name used for SASL authentication
-  private final HashMap<String, Object> props = new HashMap<String, Object>();
-  private final Subject subject;
+  public SaslWrapper saslWrapper;
 
   public Negotiation(
-      ReplicaSession session, Subject subject, String serviceName, String serviceFqdn) {
+      ReplicaSession session, Subject subject, String serviceName, String serviceFQDN) {
+    this.saslWrapper = new SaslWrapper(subject, serviceName, serviceFQDN);
     this.session = session;
-    this.subject = subject;
-    this.serviceName = serviceName;
-    this.serviceFqdn = serviceFqdn;
-    this.props.put(Sasl.QOP, "auth");
   }
 
   public void start() {
@@ -53,11 +55,14 @@ public void start() {
   }
 
   public void send(negotiation_status status, blob msg) {
-    // TODO: send negotiation message, using RecvHandler to handle the corresponding response.
+    negotiation_request request = new negotiation_request(status, msg);
+    negotiation_operator operator = new negotiation_operator(request);
+    session.asyncSend(
+        operator, new RecvHandler(operator), negotiationTimeoutMS, /* isBackupRequest */ false);
   }
 
   private class RecvHandler implements Runnable {
-    negotiation_operator op;
+    private negotiation_operator op;
 
     RecvHandler(negotiation_operator op) {
       this.op = op;
@@ -72,6 +77,7 @@ public void run() {
         handleResponse();
       } catch (Exception e) {
         logger.error("Negotiation failed", e);
+        negotiationFailed();
       }
     }
 
@@ -81,19 +87,58 @@ private void handleResponse() throws Exception {
         throw new Exception("RecvHandler received a null response, abandon it");
       }
 
-      switch (resp.status) {
-        case SASL_LIST_MECHANISMS_RESP:
-        case SASL_SELECT_MECHANISMS_RESP:
-        case SASL_CHALLENGE:
-        case SASL_SUCC:
+      switch (status) {
+        case SASL_LIST_MECHANISMS:
+          onRecvMechanisms(resp);
+          break;
+        case SASL_SELECT_MECHANISMS:
+        case SASL_INITIATE:
+        case SASL_CHALLENGE_RESP:
+          // TBD(zlw):
           break;
         default:
-          throw new Exception("Received an unexpected response, status " + resp.status);
+          throw new Exception("unexpected negotiation status: " + resp.status);
+      }
+    }
+  }
+
+  public void onRecvMechanisms(negotiation_response response) throws Exception {
+    checkStatus(response.status, negotiation_status.SASL_LIST_MECHANISMS_RESP);
+
+    String[] matchMechanism = new String[1];
+    matchMechanism[0] = getMatchMechanism(new String(response.msg.data));
+    blob msg = new blob(saslWrapper.init(matchMechanism));
+
+    status = negotiation_status.SASL_SELECT_MECHANISMS;
+    send(status, msg);
+  }
+
+  public String getMatchMechanism(String respString) {
+    String matchMechanism = new String();
+    String[] serverSupportMechanisms = respString.split(",");
+    for (String serverSupportMechanism : serverSupportMechanisms) {
+      if (expectedMechanisms.contains(serverSupportMechanism)) {
+        matchMechanism = serverSupportMechanism;
+        break;
       }
     }
+
+    return matchMechanism;
+  }
+
+  public void checkStatus(negotiation_status status, negotiation_status expected_status)
+      throws Exception {
+    if (status != expected_status) {
+      throw new Exception("status is " + status + " while expect " + expected_status);
+    }
+  }
+
+  private void negotiationFailed() {
+    status = negotiation_status.SASL_AUTH_FAIL;
+    session.closeSession();
   }
 
-  public negotiation_status get_status() {
+  public negotiation_status getStatus() {
     return status;
   }
 }

src/main/java/com/xiaomi/infra/pegasus/rpc/async/ReplicaSession.java

@@ -341,6 +341,10 @@ void tryNotifyFailureWithSeqID(int seqID, error_types errno, boolean isTimeoutTa
   }
 
   private void write(final RequestEntry entry, VolatileFields cache) {
+    if (!interceptorManager.onSendMessage(this, entry)) {
+      return;
+    }
+
     cache
         .nettyChannel
         .writeAndFlush(entry)
@@ -381,6 +385,26 @@ public void run() {
         TimeUnit.MILLISECONDS);
   }
 
+  // return value:
+  //   true  - pend succeed
+  //   false - pend failed
+  public boolean tryPendRequest(RequestEntry entry) {
+    // double check. the first one doesn't lock the lock.
+    // Because negotiationSucceed only transfered from false to true.
+    // So if it is true now, it will not change in the later.
+    // But if it is false now, maybe it will change soon. So we should use lock to protect it.
+    if (!this.negotiationSucceed) {
+      synchronized (negotiationPendingSend) {
+        if (!this.negotiationSucceed) {
+          negotiationPendingSend.offer(entry);
+          return true;
+        }
+      }
+    }
+
+    return false;
+  }
+
   final class DefaultHandler extends SimpleChannelInboundHandler<RequestEntry> {
     @Override
     public void channelInactive(ChannelHandlerContext ctx) throws Exception {
@@ -445,6 +469,8 @@ static final class VolatileFields {
   private Bootstrap boot;
   private EventLoopGroup rpcGroup;
   private ReplicaSessionInterceptorManager interceptorManager;
+  private boolean negotiationSucceed;
+  final Queue<RequestEntry> negotiationPendingSend = new LinkedList<>();
 
   // Session will be actively closed if all the rpcs across `sessionResetTimeWindowMs`
   // are timed out, in that case we suspect that the server is unavailable.

src/main/java/com/xiaomi/infra/pegasus/rpc/async/SaslWrapper.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF 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 com.xiaomi.infra.pegasus.rpc.async;
+
+import java.security.PrivilegedActionException;
+import java.security.PrivilegedExceptionAction;
+import java.util.HashMap;
+import javax.security.auth.Subject;
+import javax.security.sasl.Sasl;
+import javax.security.sasl.SaslClient;
+
+public class SaslWrapper {
+  private SaslClient saslClient;
+  private Subject subject;
+  private String serviceName;
+  private String serviceFQDN;
+  private HashMap<String, Object> properties = new HashMap<>();
+
+  public SaslWrapper(Subject subject, String serviceName, String serviceFQDN) {
+    this.subject = subject;
+    this.serviceName = serviceName;
+    this.serviceFQDN = serviceFQDN;
+    this.properties.put(Sasl.QOP, "auth");
+  }
+
+  public byte[] init(String[] mechanims) throws PrivilegedActionException {
+    return Subject.doAs(
+        subject,
+        (PrivilegedExceptionAction<byte[]>)
+            () -> {
+              saslClient =
+                  Sasl.createSaslClient(
+                      mechanims, null, serviceName, serviceFQDN, properties, null);
+              return saslClient.getMechanismName().getBytes();
+            });
+  }
+}

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/ReplicaSessionInterceptor.java

@@ -23,4 +23,7 @@
 public interface ReplicaSessionInterceptor {
   // The behavior when a rpc session is connected.
   void onConnected(ReplicaSession session);
+
+  // The behavior when rpc session is sending a message.
+  boolean onSendMessage(ReplicaSession session, final ReplicaSession.RequestEntry entry);
 }

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/ReplicaSessionInterceptorManager.java

@@ -39,4 +39,12 @@ public void onConnected(ReplicaSession session) {
       interceptor.onConnected(session);
     }
   }
+
+  public boolean onSendMessage(ReplicaSession session, final ReplicaSession.RequestEntry entry) {
+    boolean ret = true;
+    for (ReplicaSessionInterceptor interceptor : interceptors) {
+      ret &= interceptor.onSendMessage(session, entry);
+    }
+    return ret;
+  }
 }

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/SecurityReplicaSessionInterceptor.java

@@ -19,6 +19,7 @@
 package com.xiaomi.infra.pegasus.rpc.interceptor;
 
 import com.sun.security.auth.callback.TextCallbackHandler;
+import com.xiaomi.infra.pegasus.operator.negotiation_operator;
 import com.xiaomi.infra.pegasus.rpc.async.Negotiation;
 import com.xiaomi.infra.pegasus.rpc.async.ReplicaSession;
 import javax.security.auth.Subject;
@@ -41,7 +42,15 @@ public SecurityReplicaSessionInterceptor(String serviceName, String serviceFqdn)
     this.serviceFqdn = serviceFqdn;
 
     try {
+      // The LoginContext class provides the basic methods used to authenticate subjects, and
+      // provides a way to develop an application independent of the underlying authentication
+      // technology. For more details:
+      // https://docs.oracle.com/javase/7/docs/technotes/guides/security/jaas/JAASRefGuide.html#LoginContext
+
+      // let the LoginContext instantiate a new Subject
       loginContext = new LoginContext("client", new TextCallbackHandler());
+
+      // authenticate the Subject
       loginContext.login();
 
       subject = loginContext.getSubject();
@@ -59,4 +68,13 @@ public void onConnected(ReplicaSession session) {
     Negotiation negotiation = new Negotiation(session, subject, serviceName, serviceFqdn);
     negotiation.start();
   }
+
+  public boolean onSendMessage(ReplicaSession session, final ReplicaSession.RequestEntry entry) {
+    // tryPendRequest returns false means that the negotiation is succeed now
+    return isNegotiationRequest(entry) || !session.tryPendRequest(entry);
+  }
+
+  private boolean isNegotiationRequest(final ReplicaSession.RequestEntry entry) {
+    return entry.op instanceof negotiation_operator;
+  }
 }

src/test/java/com/xiaomi/infra/pegasus/rpc/async/NegotiationTest.java

@@ -18,21 +18,82 @@
  */
 package com.xiaomi.infra.pegasus.rpc.async;
 
-import static com.xiaomi.infra.pegasus.apps.negotiation_status.SASL_LIST_MECHANISMS;
 import static org.mockito.ArgumentMatchers.any;
 
+import com.xiaomi.infra.pegasus.apps.negotiation_response;
+import com.xiaomi.infra.pegasus.apps.negotiation_status;
+import com.xiaomi.infra.pegasus.base.blob;
+import javax.security.auth.Subject;
 import org.junit.Assert;
 import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
 import org.mockito.Mockito;
 
 public class NegotiationTest {
+  private Negotiation negotiation = new Negotiation(null, new Subject(), "", "");
+
   @Test
   public void testStart() {
-    Negotiation negotiation = new Negotiation(null, null, "", "");
     Negotiation mockNegotiation = Mockito.spy(negotiation);
 
     Mockito.doNothing().when(mockNegotiation).send(any(), any());
     mockNegotiation.start();
-    Assert.assertEquals(mockNegotiation.get_status(), SASL_LIST_MECHANISMS);
+    Assert.assertEquals(mockNegotiation.getStatus(), negotiation_status.SASL_LIST_MECHANISMS);
+  }
+
+  @Test
+  public void tetGetMatchMechanism() {
+    String matchMechanism = negotiation.getMatchMechanism("GSSAPI,ABC");
+    Assert.assertEquals(matchMechanism, "GSSAPI");
+
+    matchMechanism = negotiation.getMatchMechanism("TEST,ABC");
+    Assert.assertEquals(matchMechanism, "");
+  }
+
+  @Test
+  public void testCheckStatus() {
+    negotiation_status expectedStatus = negotiation_status.SASL_LIST_MECHANISMS;
+
+    Assertions.assertDoesNotThrow(
+        () -> negotiation.checkStatus(negotiation_status.SASL_LIST_MECHANISMS, expectedStatus));
+
+    Assertions.assertThrows(
+        Exception.class,
+        () ->
+            negotiation.checkStatus(negotiation_status.SASL_LIST_MECHANISMS_RESP, expectedStatus));
+  }
+
+  @Test
+  public void testRecvMechanisms() {
+    Negotiation mockNegotiation = Mockito.spy(negotiation);
+    SaslWrapper mockSaslWrapper = Mockito.mock(SaslWrapper.class);
+    mockNegotiation.saslWrapper = mockSaslWrapper;
+
+    Mockito.doNothing().when(mockNegotiation).send(any(), any());
+    Assertions.assertDoesNotThrow(
+        () -> {
+          Mockito.when(mockNegotiation.saslWrapper.init(any())).thenReturn(new byte[0]);
+        });
+
+    // normal case
+    Assertions.assertDoesNotThrow(
+        () -> {
+          negotiation_response response =
+              new negotiation_response(
+                  negotiation_status.SASL_LIST_MECHANISMS_RESP, new blob(new byte[0]));
+          mockNegotiation.onRecvMechanisms(response);
+          Assert.assertEquals(
+              mockNegotiation.getStatus(), negotiation_status.SASL_SELECT_MECHANISMS);
+        });
+
+    // deal with wrong response.status
+    Assertions.assertThrows(
+        Exception.class,
+        () -> {
+          negotiation_response response =
+              new negotiation_response(
+                  negotiation_status.SASL_LIST_MECHANISMS, new blob(new byte[0]));
+          mockNegotiation.onRecvMechanisms(response);
+        });
   }
 }