From dbf6c48af9aa2f3b41b1146643db6eb0140b38c5 Mon Sep 17 00:00:00 2001
From: Yufei Gu <yufei@apache.org>
Date: Fri, 13 Oct 2017 12:38:58 -0700
Subject: [PATCH] YARN-7270. Fix unsafe casting from long to int for class
 Resource and its sub-classes. (Yufei)

(cherry picked from commit 7a27c2c367518e1bf6ee393391a2f9b412113819)
---
 .../hadoop/yarn/api/records/Resource.java     | 18 +++++++-
 .../hadoop/yarn/api/records/TestResource.java | 43 +++++++++++++++++++
 .../api/records/impl/pb/ResourcePBImpl.java   |  2 +-
 3 files changed, 60 insertions(+), 3 deletions(-)
 create mode 100644 hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
index 0fd41a2f20b..89ca5d62f7e 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/Resource.java
@@ -61,7 +61,7 @@ private static class SimpleResource extends Resource {
     }
     @Override
     public int getMemory() {
-      return (int)memory;
+      return castToIntSafely(memory);
     }
     @Override
     public void setMemory(int memory) {
@@ -77,7 +77,7 @@ public void setMemorySize(long memory) {
     }
     @Override
     public int getVirtualCores() {
-      return (int)vcores;
+      return castToIntSafely(vcores);
     }
     @Override
     public void setVirtualCores(int vcores) {
@@ -206,4 +206,18 @@ public int compareTo(Resource other) {
   public String toString() {
     return "<memory:" + getMemorySize() + ", vCores:" + getVirtualCores() + ">";
   }
+
+  /**
+   * Convert long to int for a resource value safely. This method assumes
+   * resource value is positive.
+   *
+   * @param value long resource value
+   * @return int resource value
+   */
+  protected static int castToIntSafely(long value) {
+    if (value > Integer.MAX_VALUE) {
+      return Integer.MAX_VALUE;
+    }
+    return Long.valueOf(value).intValue();
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java
new file mode 100644
index 00000000000..e0ec3703552
--- /dev/null
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/api/records/TestResource.java
@@ -0,0 +1,43 @@
+/**
+ * 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 org.apache.hadoop.yarn.api.records;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * The class to test {@link Resource}.
+ */
+public class TestResource {
+
+  @Test
+  public void testCastToIntSafely() {
+    assertEquals(0, Resource.castToIntSafely(0));
+    assertEquals(1, Resource.castToIntSafely(1));
+    assertEquals(Integer.MAX_VALUE,
+        Resource.castToIntSafely(Integer.MAX_VALUE));
+
+    assertEquals("Cast to Integer.MAX_VALUE if the long is greater than "
+            + "Integer.MAX_VALUE", Integer.MAX_VALUE,
+        Resource.castToIntSafely(Integer.MAX_VALUE + 1L));
+    assertEquals("Cast to Integer.MAX_VALUE if the long is greater than "
+            + "Integer.MAX_VALUE", Integer.MAX_VALUE,
+        Resource.castToIntSafely(Long.MAX_VALUE));
+  }
+}
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ResourcePBImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ResourcePBImpl.java
index 34109bec0f6..2404a63ea69 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ResourcePBImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ResourcePBImpl.java
@@ -70,7 +70,7 @@ private void maybeInitBuilder() {
   @Override
   @SuppressWarnings("deprecation")
   public int getMemory() {
-    return (int) getMemorySize();
+    return castToIntSafely(this.getMemorySize());
   }
 
   @Override