Scope OSX Eggroll session reuse checks

java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/PutBatchSinkUtil.java

@@ -22,7 +22,7 @@
 import java.util.concurrent.TimeUnit;
 
 public class PutBatchSinkUtil {
-    public static LoadingCache<String, ErSession> sessionCache =
+    public static LoadingCache<SessionCacheKey, ErSession> sessionCache =
 
 
             CacheBuilder.newBuilder()
@@ -31,10 +31,10 @@ public class PutBatchSinkUtil {
                     .concurrencyLevel(100)
                     .recordStats()
                     .softValues()
-                    .build(new CacheLoader<String, ErSession>() {
+                    .build(new CacheLoader<SessionCacheKey, ErSession>() {
                                @Override
-                               public ErSession load(String sessionId) throws Exception {
-                                   return new ErSession(sessionId, false);
+                               public ErSession load(SessionCacheKey sessionCacheKey) throws Exception {
+                                   return new ErSession(sessionCacheKey.getSessionId(), false);
                                }
 
 

java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/SessionCacheKey.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2019 The FATE Authors. All Rights Reserved.
+ *
+ * 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.
+ */
+package org.fedai.osx.broker.eggroll;
+
+import com.webank.ai.eggroll.api.networking.proxy.Proxy;
+
+import java.util.Objects;
+
+public class SessionCacheKey {
+
+    private final String sessionId;
+    private final String taskId;
+    private final String rollSiteSessionId;
+    private final String name;
+    private final String tag;
+    private final String srcRole;
+    private final String srcPartyId;
+    private final String dstRole;
+    private final String dstPartyId;
+
+    public SessionCacheKey(String sessionId, Proxy.Metadata metadata, ErRollSiteHeader rollSiteHeader) {
+        this.sessionId = sessionId;
+        this.taskId = metadata == null ? "" : metadata.getTask().getTaskId();
+        this.rollSiteSessionId = rollSiteHeader.getRollSiteSessionId();
+        this.name = rollSiteHeader.getName();
+        this.tag = rollSiteHeader.getTag();
+        this.srcRole = rollSiteHeader.getSrcRole();
+        this.srcPartyId = rollSiteHeader.getSrcPartyId();
+        this.dstRole = rollSiteHeader.getDstRole();
+        this.dstPartyId = rollSiteHeader.getDstPartyId();
+    }
+
+    public SessionCacheKey(String sessionId, ErRollSiteHeader rollSiteHeader) {
+        this(sessionId, Proxy.Metadata.getDefaultInstance(), rollSiteHeader);
+    }
+
+    public String getSessionId() {
+        return sessionId;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (!(o instanceof SessionCacheKey)) {
+            return false;
+        }
+        SessionCacheKey that = (SessionCacheKey) o;
+        return Objects.equals(sessionId, that.sessionId)
+                && Objects.equals(taskId, that.taskId)
+                && Objects.equals(rollSiteSessionId, that.rollSiteSessionId)
+                && Objects.equals(name, that.name)
+                && Objects.equals(tag, that.tag)
+                && Objects.equals(srcRole, that.srcRole)
+                && Objects.equals(srcPartyId, that.srcPartyId)
+                && Objects.equals(dstRole, that.dstRole)
+                && Objects.equals(dstPartyId, that.dstPartyId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(sessionId, taskId, rollSiteSessionId, name, tag, srcRole, srcPartyId, dstRole, dstPartyId);
+    }
+
+    @Override
+    public String toString() {
+        return "SessionCacheKey{"
+                + "sessionId='" + sessionId + '\''
+                + ", taskId='" + taskId + '\''
+                + ", rollSiteSessionId='" + rollSiteSessionId + '\''
+                + ", name='" + name + '\''
+                + ", tag='" + tag + '\''
+                + ", srcRole='" + srcRole + '\''
+                + ", srcPartyId='" + srcPartyId + '\''
+                + ", dstRole='" + dstRole + '\''
+                + ", dstPartyId='" + dstPartyId + '\''
+                + '}';
+    }
+}

java/osx/osx-broker/src/main/java/org/fedai/osx/broker/grpc/QueuePushReqStreamObserver.java

@@ -45,6 +45,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.net.InetAddress;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.*;
@@ -90,12 +91,17 @@ public void setForwardPushReqSO(StreamObserver<Proxy.Packet> forwardPushReqSO) {
     }
 
     public void init(Proxy.Packet packet) throws Exception {
+        String inboundSourcePartyId = context.getSrcNodeId();
+        String inboundDestinationPartyId = context.getDesNodeId();
         TransferUtil.assableContextFromProxyPacket(context, packet);
         Proxy.Metadata metadata = packet.getHeader();
         String desPartyId = context.getDesNodeId();
         String srcPartyId = context.getSrcNodeId();
         ByteString encodedRollSiteHeader = metadata.getExt();
         rsHeader = ErRollSiteHeader.parseFromPb(Transfer.RollSiteHeader.parseFrom(encodedRollSiteHeader));
+        checkInboundParty("source party", inboundSourcePartyId, srcPartyId);
+        checkInboundParty("destination party", inboundDestinationPartyId, desPartyId);
+        checkRollSiteHeader(metadata, rsHeader);
         Integer partitionId = rsHeader.getPartitionId();
         brokerTag = "putBatch-" + rsHeader.getRsKey("#", "__rsk") + "-" + partitionId;
         context.setSessionId(rsHeader.getRollSiteSessionId());
@@ -184,12 +190,14 @@ private void initEggroll(OsxContext context, Proxy.Packet firstRequest) {
         metadata = firstRequest.getHeader();
         String oneLineStringMetadata = ToStringUtils.toOneLineString(metadata);
         context.setActionType(ActionType.PUSH_EGGPAIR.name());
+        checkInboundSource(context);
         String rsKey = rsHeader.getRsKey("#", "__rsk");
         String sessionId = String.join("_", rsHeader.getRollSiteSessionId(), rsHeader.getDstRole(), rsHeader.getDstPartyId());
         context.setSessionId(sessionId);
+        SessionCacheKey sessionCacheKey = new SessionCacheKey(sessionId, metadata, rsHeader);
         ErSession session = null;
         try {
-            session = PutBatchSinkUtil.sessionCache.get(sessionId);
+            session = PutBatchSinkUtil.sessionCache.get(sessionCacheKey);
         } catch (ExecutionException e) {
             logger.error("get session error ", e);
         }
@@ -259,6 +267,57 @@ private void initEggroll(OsxContext context, Proxy.Packet firstRequest) {
         putBatchSinkPushReqSO = stub.send(new PutBatchSinkPushRespSO(metadata, commandFuture, backRespSO, finishLatch, routerInfo));
     }
 
+    private void checkInboundSource(OsxContext context) {
+        RouterInfo sourceRouterInfo = routerService.route(context.getDesNodeId(), context.getDesComponent(), context.getSrcNodeId(), context.getSrcComponent());
+        if (sourceRouterInfo == null) {
+            throw new NoRouterInfoException("no router is found for source party " + context.getSrcNodeId());
+        }
+        String sourceIp = context.getSourceIp();
+        if (StringUtils.isNotEmpty(sourceIp)
+                && StringUtils.isNotEmpty(sourceRouterInfo.getHost())
+                && !isSameHost(sourceIp, sourceRouterInfo.getHost())) {
+            throw new ParameterException("request source does not match route table for source party " + context.getSrcNodeId());
+        }
+    }
+
+    private void checkInboundParty(String fieldName, String inboundValue, String packetValue) {
+        if (StringUtils.isNotEmpty(inboundValue)
+                && StringUtils.isNotEmpty(packetValue)
+                && !inboundValue.equals(packetValue)) {
+            throw new ParameterException("inbound " + fieldName + " does not match packet header");
+        }
+    }
+
+    private boolean isSameHost(String sourceIp, String routeHost) {
+        if (sourceIp.equals(routeHost)) {
+            return true;
+        }
+        try {
+            return InetAddress.getByName(sourceIp).equals(InetAddress.getByName(routeHost));
+        } catch (Exception e) {
+            logger.warn("unable to resolve source ip {} or route host {}", sourceIp, routeHost);
+            return false;
+        }
+    }
+
+    private void checkRollSiteHeader(Proxy.Metadata metadata, ErRollSiteHeader rollSiteHeader) {
+        if (rollSiteHeader == null) {
+            throw new ParameterException("rollSiteHeader is null");
+        }
+        checkHeaderField("srcRole", metadata.getSrc().getRole(), rollSiteHeader.getSrcRole());
+        checkHeaderField("srcPartyId", metadata.getSrc().getPartyId(), rollSiteHeader.getSrcPartyId());
+        checkHeaderField("dstRole", metadata.getDst().getRole(), rollSiteHeader.getDstRole());
+        checkHeaderField("dstPartyId", metadata.getDst().getPartyId(), rollSiteHeader.getDstPartyId());
+    }
+
+    private void checkHeaderField(String fieldName, String metadataValue, String rollSiteHeaderValue) {
+        if (StringUtils.isNotEmpty(metadataValue)
+                && StringUtils.isNotEmpty(rollSiteHeaderValue)
+                && !metadataValue.equals(rollSiteHeaderValue)) {
+            throw new ParameterException("metadata " + fieldName + " does not match rollSiteHeader");
+        }
+    }
+
 
     @Override
     public void onNext(Proxy.Packet value) {

java/osx/osx-broker/src/main/java/org/fedai/osx/broker/util/ContextUtil.java

@@ -80,7 +80,14 @@ public static void assableContextFromInbound(OsxContext context) {
     public static OsxContext buildFateContext(Protocol protocol) {
         OsxContext context = new OsxContext();
         context.setProtocol(protocol);
-        //  context.setSourceIp(ContextPrepareInterceptor.sourceIp.get() != null ? ContextPrepareInterceptor.sourceIp.get().toString() : "");
+        context.setSrcNodeId(CONTEXTKEY_FROM_NODE_ID.get(io.grpc.Context.current()));
+        context.setSrcInstId(CONTEXTKEY_FROM_INST_ID.get(io.grpc.Context.current()));
+        context.setDesNodeId(CONTEXTKEY_TARGET_NODE_ID.get(io.grpc.Context.current()));
+        context.setDesInstId(CONTEXTKEY_TARGET_INST_ID.get(io.grpc.Context.current()));
+        Object sourceIp = CONTEXTKEY_SOURCEIP.get(io.grpc.Context.current());
+        if (sourceIp != null) {
+            context.setSourceIp(sourceIp.toString());
+        }
 
         return context;
     }

java/osx/osx-broker/src/test/java/org/fedai/osx/broker/eggroll/SessionCacheKeyTest.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2019 The FATE Authors. All Rights Reserved.
+ *
+ * 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.
+ */
+package org.fedai.osx.broker.eggroll;
+
+import com.webank.ai.eggroll.api.networking.proxy.Proxy;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class SessionCacheKeyTest {
+
+    @Test
+    public void distinguishesSourcePartyForSameBackendSession() {
+        SessionCacheKey guest9999 = new SessionCacheKey("sess-42_host_10000", header("guest", "9999", "table_x", "batch_0"));
+        SessionCacheKey guest1234 = new SessionCacheKey("sess-42_host_10000", header("guest", "1234", "table_x", "batch_0"));
+
+        Assert.assertNotEquals(guest9999, guest1234);
+    }
+
+    @Test
+    public void distinguishesRollSiteNameAndTagForSameBackendSession() {
+        SessionCacheKey tableX = new SessionCacheKey("sess-42_host_10000", header("guest", "9999", "table_x", "batch_0"));
+        SessionCacheKey tableY = new SessionCacheKey("sess-42_host_10000", header("guest", "9999", "table_y", "batch_1"));
+
+        Assert.assertNotEquals(tableX, tableY);
+    }
+
+    @Test
+    public void distinguishesTaskIdForSameBackendSession() {
+        ErRollSiteHeader header = header("guest", "9999", "table_x", "batch_0");
+        SessionCacheKey taskA = new SessionCacheKey("sess-42_host_10000", metadata("job-train-1"), header);
+        SessionCacheKey taskB = new SessionCacheKey("sess-42_host_10000", metadata("job-train-9"), header);
+
+        Assert.assertNotEquals(taskA, taskB);
+    }
+
+    @Test
+    public void matchesSameRollSiteContextForSameBackendSession() {
+        SessionCacheKey first = new SessionCacheKey("sess-42_host_10000", header("guest", "9999", "table_x", "batch_0"));
+        SessionCacheKey second = new SessionCacheKey("sess-42_host_10000", header("guest", "9999", "table_x", "batch_0"));
+
+        Assert.assertEquals(first, second);
+        Assert.assertEquals(first.hashCode(), second.hashCode());
+    }
+
+    private ErRollSiteHeader header(String srcRole, String srcPartyId, String name, String tag) {
+        ErRollSiteHeader header = new ErRollSiteHeader();
+        header.setRollSiteSessionId("sess-42");
+        header.setName(name);
+        header.setTag(tag);
+        header.setSrcRole(srcRole);
+        header.setSrcPartyId(srcPartyId);
+        header.setDstRole("host");
+        header.setDstPartyId("10000");
+        return header;
+    }
+
+    private Proxy.Metadata metadata(String taskId) {
+        return Proxy.Metadata.newBuilder()
+                .setTask(Proxy.Task.newBuilder().setTaskId(taskId))
+                .build();
+    }
+}