diff --git a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/PutBatchSinkUtil.java b/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/PutBatchSinkUtil.java index f1a32bbd9d..f35f8e774d 100644 --- a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/PutBatchSinkUtil.java +++ b/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 sessionCache = + public static LoadingCache sessionCache = CacheBuilder.newBuilder() @@ -31,10 +31,10 @@ public class PutBatchSinkUtil { .concurrencyLevel(100) .recordStats() .softValues() - .build(new CacheLoader() { + .build(new CacheLoader() { @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); } diff --git a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/SessionCacheKey.java b/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/eggroll/SessionCacheKey.java new file mode 100644 index 0000000000..ceb118bb36 --- /dev/null +++ b/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 + '\'' + + '}'; + } +} diff --git a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/grpc/QueuePushReqStreamObserver.java b/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/grpc/QueuePushReqStreamObserver.java index 055c9275fd..4a35e6d304 100644 --- a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/grpc/QueuePushReqStreamObserver.java +++ b/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 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) { diff --git a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/util/ContextUtil.java b/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/util/ContextUtil.java index c51a6dfb98..208df28673 100644 --- a/java/osx/osx-broker/src/main/java/org/fedai/osx/broker/util/ContextUtil.java +++ b/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; } diff --git a/java/osx/osx-broker/src/test/java/org/fedai/osx/broker/eggroll/SessionCacheKeyTest.java b/java/osx/osx-broker/src/test/java/org/fedai/osx/broker/eggroll/SessionCacheKeyTest.java new file mode 100644 index 0000000000..c41126524b --- /dev/null +++ b/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(); + } +}