Merge pull request #55 from codeocean/master

Metadata service compatible with imdsv2

install.sh

@@ -83,6 +83,9 @@ Options
     -s, --initial-size  SIZE_GB
                         Initial size of the volume in GB. (Default: 200)
                         Only used if --initial-device is NOT specified.
+
+    -i, --imdsv2        
+                        Enable imdsv2 for instance metadata API requests.
     
 EOF
 )
@@ -106,8 +109,6 @@ BASEDIR=$(dirname $0)
 
 . ${BASEDIR}/shared/utils.sh
 
-initialize
-
 # parse options
 PARAMS=""
 while (( "$#" )); do
@@ -160,6 +161,10 @@ while (( "$#" )); do
             MOUNTPOINT=$2
             shift 2
             ;;
+        -i|--imdsv2)
+            IMDSV2="true"
+            shift 1
+            ;;
         -h|--help)
             echo "$USAGE"
             exit
@@ -180,6 +185,8 @@ done
 
 eval set -- "$PARAMS"
 
+initialize
+
 # for backwards compatibility evaluate positional parameters like previous 2.0.x and 2.1.x releases
 # this will be removed in the future
 if [ ! -z "$PARAMS" ]; then

shared/utils.sh

@@ -28,10 +28,22 @@
 #  IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 #  POSSIBILITY OF SUCH DAMAGE.
 
+function get_metadata() {
+    local key=$1
+    local metadata_ip='169.254.169.254'
+
+    if [ ! -z "$IMDSV2" ]; then
+        local token=$(curl -s -X PUT "http://$metadata_ip/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 60")
+        local token_wrapper='-H "X-aws-ec2-metadata-token: $token"'
+    fi
+
+    echo `curl -s $token_wrapper http://$metadata_ip/latest/meta-data/$key`
+}
+
 function initialize() {
-    export AWS_AZ=$(curl -s  http://169.254.169.254/latest/meta-data/placement/availability-zone/)
+    export AWS_AZ=$(get_metadata placement/availability-zone)
     export AWS_REGION=$(echo ${AWS_AZ} | sed -e 's/[a-z]$//')
-    export INSTANCE_ID=$(curl -s  http://169.254.169.254/latest/meta-data/instance-id)
+    export INSTANCE_ID=$(get_metadata instance-id)
     export EBS_AUTOSCALE_CONFIG_FILE=/etc/ebs-autoscale.json
 }
 
@@ -51,11 +63,6 @@ function get_config_value() {
     jq -r $filter $EBS_AUTOSCALE_CONFIG_FILE
 }
 
-function get_metadata() {
-    local key=$1
-    echo `curl -s http://169.254.169.254/latest/meta-data/$key`
-}
-
 function logthis() {
     echo "[`date`] $1" >> $(get_config_value .logging.log_file)
 }