Api change

frontend/archive/src/main/java/org/pytorch/serve/archive/model/InvalidKeyException.java

@@ -0,0 +1,32 @@
+package org.pytorch.serve.archive.model;
+
+public class InvalidKeyException extends ModelException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs an {@code InvalidKeyException} with the specified detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     */
+    public InvalidKeyException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs an {@code InvalidKeyException} with the specified detail message and cause.
+     *
+     * <p>Note that the detail message associated with {@code cause} is <i>not</i> automatically
+     * incorporated into this exception's detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     * @param cause The cause (which is saved for later retrieval by the {@link #getCause()}
+     *     method). (A null value is permitted, and indicates that the cause is nonexistent or
+     *     unknown.)
+     */
+    public InvalidKeyException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

frontend/archive/src/main/java/org/pytorch/serve/archive/model/KeyTimeOutException.java

@@ -0,0 +1,32 @@
+package org.pytorch.serve.archive.model;
+
+public class KeyTimeOutException extends ModelException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs an {@code KeyTimeOutException} with the specified detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     */
+    public KeyTimeOutException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs an {@code KeyTimeOutException} with the specified detail message and cause.
+     *
+     * <p>Note that the detail message associated with {@code cause} is <i>not</i> automatically
+     * incorporated into this exception's detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     * @param cause The cause (which is saved for later retrieval by the {@link #getCause()}
+     *     method). (A null value is permitted, and indicates that the cause is nonexistent or
+     *     unknown.)
+     */
+    public KeyTimeOutException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

frontend/server/src/main/java/org/pytorch/serve/ModelServer.java

@@ -83,6 +83,11 @@ public static void main(String[] args) {
             ConfigManager.init(arguments);
             ConfigManager configManager = ConfigManager.getInstance();
             PluginsManager.getInstance().initialize();
+            Map<String, ModelServerEndpoint> plugins =
+                    PluginsManager.getInstance().getInferenceEndpoints();
+            if (plugins.containsKey("token")) {
+                configManager.generateKeyFile();
+            }
             MetricCache.init();
             InternalLoggerFactory.setDefaultFactory(Slf4JLoggerFactory.INSTANCE);
             ModelServer modelServer = new ModelServer(configManager);

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/ApiDescriptionRequestHandler.java

@@ -30,7 +30,6 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
-
         if (isApiDescription(segments)) {
             String path = decoder.path();
             if (("/".equals(path) && HttpMethod.OPTIONS.equals(req.method()))

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/InferenceRequestHandler.java

@@ -59,10 +59,13 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
+        ConfigManager configManager = ConfigManager.getInstance();
         if (isInferenceReq(segments)) {
             if (endpointMap.getOrDefault(segments[1], null) != null) {
+                configManager.checkTokenAuthorization(req, false);
                 handleCustomEndpoint(ctx, req, segments, decoder);
             } else {
+                configManager.checkTokenAuthorization(req, true);
                 switch (segments[1]) {
                     case "ping":
                         Runnable r =

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/ManagementRequestHandler.java

@@ -32,6 +32,7 @@
 import org.pytorch.serve.openapi.OpenApiUtils;
 import org.pytorch.serve.servingsdk.ModelServerEndpoint;
 import org.pytorch.serve.util.ApiUtils;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.JsonUtils;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.util.messages.RequestInput;
@@ -61,6 +62,8 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
+        ConfigManager configManager = ConfigManager.getInstance();
+        configManager.checkTokenAuthorization(req, false);
         if (isManagementReq(segments)) {
             if (endpointMap.getOrDefault(segments[1], null) != null) {
                 handleCustomEndpoint(ctx, req, segments, decoder);

frontend/server/src/main/java/org/pytorch/serve/http/api/rest/PrometheusMetricsRequestHandler.java

@@ -24,6 +24,7 @@
 import org.pytorch.serve.archive.model.ModelException;
 import org.pytorch.serve.archive.workflow.WorkflowException;
 import org.pytorch.serve.http.HttpRequestHandlerChain;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.wlm.WorkerInitializationException;
 import org.slf4j.Logger;
@@ -47,6 +48,8 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
+        ConfigManager configManager = ConfigManager.getInstance();
+        configManager.checkTokenAuthorization(req, true);
         if (segments.length >= 2 && "metrics".equals(segments[1])) {
             ByteBuf resBuf = Unpooled.directBuffer();
             List<String> params =

frontend/server/src/main/java/org/pytorch/serve/util/ConfigManager.java

@@ -2,6 +2,8 @@
 
 import com.google.gson.JsonObject;
 import com.google.gson.reflect.TypeToken;
+import io.netty.handler.codec.http.FullHttpRequest;
+import io.netty.handler.codec.http.HttpMethod;
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslContextBuilder;
 import io.netty.handler.ssl.util.SelfSignedCertificate;
@@ -20,11 +22,14 @@
 import java.security.KeyFactory;
 import java.security.KeyStore;
 import java.security.PrivateKey;
+import java.security.SecureRandom;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.PKCS8EncodedKeySpec;
+import java.time.DateTimeException;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Base64;
@@ -36,13 +41,17 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.Option;
 import org.apache.commons.cli.Options;
 import org.apache.commons.io.IOUtils;
+import org.pytorch.serve.archive.model.InvalidKeyException;
+import org.pytorch.serve.archive.model.KeyTimeOutException;
+import org.pytorch.serve.archive.model.ModelException;
 import org.pytorch.serve.metrics.MetricBuilder;
 import org.pytorch.serve.servingsdk.snapshot.SnapshotSerializer;
 import org.pytorch.serve.snapshot.SnapshotSerializerFactory;
@@ -146,6 +155,10 @@ public final class ConfigManager {
     private boolean telemetryEnabled;
     private Logger logger = LoggerFactory.getLogger(ConfigManager.class);
 
+    private SecureRandom secureRandom = new SecureRandom();
+    private Base64.Encoder baseEncoder = Base64.getUrlEncoder();
+    public String keyFileLocation;
+
     private ConfigManager(Arguments args) throws IOException {
         prop = new Properties();
 
@@ -837,6 +850,104 @@ public boolean isSnapshotDisabled() {
         return snapshotDisabled;
     }
 
+    public boolean isTokenExpired(Instant expirationTime) {
+        return !(Instant.now().isBefore(expirationTime));
+    }
+
+    public String generateToken() {
+        byte[] randomBytes = new byte[6];
+        secureRandom.nextBytes(randomBytes);
+        return baseEncoder.encodeToString(randomBytes);
+    }
+
+    public boolean generateKeyFile() throws IOException {
+        String fileData = " ";
+        String absoluteFilePath = getCanonicalPath(".") + "/key_file.txt";
+        keyFileLocation = absoluteFilePath;
+        File file = new File(absoluteFilePath);
+        if (!file.createNewFile() && !file.exists()) {
+            return false;
+        }
+        Integer timeToExpiration = 30; // in minutes
+        fileData =
+                "Management Key: "
+                        + generateToken()
+                        + "\n"
+                        + "Inference Key: "
+                        + generateToken()
+                        + " --- Expiration time: "
+                        + Instant.now().plusSeconds(TimeUnit.MINUTES.toSeconds(timeToExpiration))
+                        + "\n";
+        Files.write(Paths.get("key_file.txt"), fileData.getBytes());
+        return true;
+    }
+
+    public List<String> parseFile(File tsTokenFile) {
+        List<String> parsedTokens = new ArrayList<>();
+        try {
+            InputStream stream = Files.newInputStream(tsTokenFile.toPath());
+            byte[] array = new byte[100];
+            stream.read(array);
+            String data = new String(array);
+            String[] arrOfData = data.split("\n", 2);
+            String[] managementArr = arrOfData[0].split(" ", 3);
+            String[] inferenceArr = arrOfData[1].split(" ", 7);
+            parsedTokens.add(managementArr[2]);
+            parsedTokens.add(inferenceArr[2]);
+            String[] expirationArr = inferenceArr[6].split("\n", 2);
+            parsedTokens.add(expirationArr[0]);
+        } catch (IOException | ArrayIndexOutOfBoundsException e) {
+            System.out.println("Unable to read key file or key file has been modified");
+            return null;
+        }
+        return parsedTokens;
+    }
+
+    public void checkTokenAuthorization(FullHttpRequest req, boolean inferenceRequest)
+            throws ModelException {
+        HttpMethod method = req.method();
+        String filePath = keyFileLocation;
+        if (filePath != null) {
+            File tsTokenFile = new File(filePath);
+            if (tsTokenFile.exists()) {
+                List<String> parsedTokens = parseFile(tsTokenFile);
+                String managementToken = parsedTokens.get(0);
+                String inferenceToken = parsedTokens.get(1);
+                Instant expirationTime = Instant.now();
+                try {
+                    expirationTime = Instant.parse(parsedTokens.get(2));
+                } catch (DateTimeException e) {
+                    e.printStackTrace();
+                    System.out.println("{\n\t\"Error\": Key File has been modified \n}\n");
+                }
+                String tokenBearer = req.headers().get("Authorization");
+                if (tokenBearer == null) {
+                    throw new InvalidKeyException("NO TOKEN PROVIDED");
+                }
+                String[] arrOfStr = tokenBearer.split(" ", 2);
+                if (arrOfStr.length == 1) {
+                    throw new InvalidKeyException("NO TOKEN PROVIDED");
+                }
+                String token = arrOfStr[1];
+                String key = managementToken;
+                if (inferenceRequest) {
+                    key = inferenceToken;
+                }
+
+                if (token.equals(key)) {
+                    if (isTokenExpired(expirationTime) && inferenceRequest) {
+                        throw new KeyTimeOutException("THE CURRENT TOKEN IS EXPIRED");
+                    }
+                    System.out.println("TOKEN AUTHORIZATION WORKED");
+                } else {
+                    throw new InvalidKeyException("TOKEN IS INCORRECT ");
+                }
+            } else {
+                System.out.println("TOKEN AUTHORIZATION IS NOT ENABLED");
+            }
+        }
+    }
+
     public boolean isSSLEnabled(ConnectorType connectorType) {
         String address = prop.getProperty(TS_INFERENCE_ADDRESS, "http://127.0.0.1:8080");
         switch (connectorType) {

frontend/server/src/main/java/org/pytorch/serve/workflow/api/http/WorkflowInferenceRequestHandler.java

@@ -80,6 +80,9 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
+
+        ConfigManager configManager = ConfigManager.getInstance();
+        configManager.checkTokenAuthorization(req, true);
         if ("wfpredict".equalsIgnoreCase(segments[1])) {
             if (segments.length < 3) {
                 throw new ResourceNotFoundException();

frontend/server/src/main/java/org/pytorch/serve/workflow/api/http/WorkflowMgmtRequestHandler.java

@@ -24,6 +24,7 @@
 import org.pytorch.serve.http.MethodNotAllowedException;
 import org.pytorch.serve.http.ResourceNotFoundException;
 import org.pytorch.serve.http.StatusResponse;
+import org.pytorch.serve.util.ConfigManager;
 import org.pytorch.serve.util.JsonUtils;
 import org.pytorch.serve.util.NettyUtils;
 import org.pytorch.serve.wlm.WorkerInitializationException;
@@ -63,6 +64,9 @@ public void handleRequest(
             String[] segments)
             throws ModelException, DownloadArchiveException, WorkflowException,
                     WorkerInitializationException {
+
+        ConfigManager configManager = ConfigManager.getInstance();
+        configManager.checkTokenAuthorization(req, false);
         if (isManagementReq(segments)) {
             if (!"workflows".equals(segments[1])) {
                 throw new ResourceNotFoundException();