Add logic to secure ims

kaiko-ai · Oct 9, 2024 · d5f5fc5 · d5f5fc5
1 parent fcd7c77
commit d5f5fc5
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 6 deletions.
diff --git a/src/main/java/be/cytomine/api/controller/image/RestSliceInstanceController.java b/src/main/java/be/cytomine/api/controller/image/RestSliceInstanceController.java
@@ -17,9 +17,11 @@
 */
 
 import be.cytomine.api.controller.RestCytomineController;
+import be.cytomine.config.properties.ApplicationProperties;
 import be.cytomine.domain.image.ImageInstance;
 import be.cytomine.domain.image.SliceInstance;
 import be.cytomine.exceptions.ObjectNotFoundException;
+import be.cytomine.security.jwt.JwtTokenGenerator;
 import be.cytomine.service.dto.CropParameter;
 import be.cytomine.service.dto.ImageParameter;
 import be.cytomine.service.dto.WindowParameter;
@@ -35,6 +37,11 @@
 
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 @RestController
 @RequestMapping("/api")
@@ -56,6 +63,8 @@ public class RestSliceInstanceController extends RestCytomineController {
 
     private final ImagePropertiesService imagePropertiesService;
 
+    private final ApplicationProperties applicationProperties;
+
 
     @GetMapping("/imageinstance/{id}/sliceinstance.json")
     public ResponseEntity<String> listByImageInstance(
@@ -64,7 +73,30 @@ public ResponseEntity<String> listByImageInstance(
         log.debug("REST request to list slice instance for image {}", id);
         ImageInstance imageInstance = imageInstanceService.find(id)
                 .orElseThrow(() -> new ObjectNotFoundException("ImageInstance", id));
-        return responseSuccess(sliceInstanceService.list(imageInstance));
+        List<SliceInstance> sliceInstances = sliceInstanceService.list(imageInstance);
+        List<JsonObject> jsonArray = new ArrayList<JsonObject>(); // Use JsonArray to hold the list of JSON objects
+        for (SliceInstance sliceInstance : sliceInstances) {
+            jsonArray.add(addToken(sliceInstance)); 
+        }
+        return responseSuccess(jsonArray);
+    }
+
+    private JsonObject addToken(SliceInstance sliceInstance) {
+        // Injects a temporary token for the slice instance to be able to access the image
+        // directly from the image server
+        JsonObject returnArray = sliceInstance.toJsonObject();
+        Map<String, Object> payload = new HashMap<>();
+        payload.put("filepath", sliceInstance.getPath());
+        // Generate the token and expiry time
+        Map<String, Object> tokenData = JwtTokenGenerator.generateJwtToken(applicationProperties, payload);
+        String token = (String) tokenData.get("token");
+        Instant expiryTime = (Instant) tokenData.get("expiryTime");
+        returnArray.put("temporaryToken", token);
+        return returnArray;
+    }
+
+    private ResponseEntity<String> responseWithToken(SliceInstance sliceInstance) {
+        return responseSuccess(addToken(sliceInstance).toJsonString());
     }
 
 
@@ -74,7 +106,7 @@ public ResponseEntity<String> show(
     ) {
         log.debug("REST request to get slice instance {}", id);
         return sliceInstanceService.find(id)
-                .map(this::responseSuccess)
+                .map(this::responseWithToken)
                 .orElseThrow(() -> new ObjectNotFoundException("SliceInstance", id));
     }
 
@@ -92,7 +124,7 @@ public ResponseEntity<String> getByImageInstanceAndCoordinates(
 
         SliceInstance sliceInstance = sliceInstanceService.find(imageInstance, channel,zStack, time)
                 .orElseThrow(() -> new ObjectNotFoundException("SliceInstance", id + "[" + channel + "-" + zStack + "-" + time + "]"));
-        return responseSuccess(sliceInstance);
+        return responseWithToken(sliceInstance);
     }
 
     @PostMapping("/sliceinstance.json")

diff --git a/src/main/java/be/cytomine/domain/image/SliceInstance.java b/src/main/java/be/cytomine/domain/image/SliceInstance.java
@@ -18,6 +18,7 @@
 
 import be.cytomine.domain.CytomineDomain;
 import be.cytomine.domain.project.Project;
+import be.cytomine.security.jwt.JwtTokenGenerator;
 import be.cytomine.utils.JsonObject;
 import lombok.Getter;
 import lombok.Setter;
@@ -27,6 +28,9 @@
 import javax.persistence.FetchType;
 import javax.persistence.ManyToOne;
 import java.io.Serializable;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
 
 @Entity
@@ -58,9 +62,8 @@ public CytomineDomain buildDomainFromJson(JsonObject json, EntityManager entityM
 
     public static JsonObject getDataFromDomain(CytomineDomain domain) {
         JsonObject returnArray = CytomineDomain.getDataFromDomain(domain);
-        SliceInstance sliceInstance = (SliceInstance) domain;
+        SliceInstance sliceInstance = (SliceInstance) domain;;
         returnArray.put("uploadedFile", Optional.ofNullable(sliceInstance.baseSlice).map(AbstractSlice::getUploadedFileId).orElse(null));
-
         returnArray.put("imageServerUrl", sliceInstance.getImageServerUrl());
         returnArray.put("project", Optional.ofNullable(sliceInstance.project).map(CytomineDomain::getId).orElse(null));
         returnArray.put("baseSlice", Optional.ofNullable(sliceInstance.baseSlice).map(CytomineDomain::getId).orElse(null));

diff --git a/src/main/java/be/cytomine/security/jwt/JwtTokenGenerator.java b/src/main/java/be/cytomine/security/jwt/JwtTokenGenerator.java
@@ -0,0 +1,53 @@
+/*
+ * This file is used to generate a JWT token with a payload, secret key, algorithm, and expiry time.
+ * 
+ * The main purpose is generate temporary access keys for the IMS through a shared secret key.
+ */
+
+package be.cytomine.security.jwt;
+
+import be.cytomine.config.properties.ApplicationProperties;
+
+import java.util.Map;  
+import io.jsonwebtoken.Jwts;
+import io.jsonwebtoken.SignatureAlgorithm;
+import io.jsonwebtoken.security.Keys;
+
+import java.nio.charset.StandardCharsets;
+import java.security.Key;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.Date;
+import java.util.HashMap;  
+
+
+public class JwtTokenGenerator {
+
+    public static Map<String, Object> generateJwtToken(ApplicationProperties applicationProperties,
+                                                      Map<String, Object> payload, 
+                                                      SignatureAlgorithm algorithm, long expiryMinutes) {
+
+        // Get the current time in UTC
+        Instant nowUtc = Instant.now();
+        String secret = applicationProperties.getAuthentication().getJwt().getSecret();
+        // Calculate the expiry time in UTC
+        Instant expiryUtc = nowUtc.plus(expiryMinutes, ChronoUnit.MINUTES);
+        byte[] keyBytes = secret.getBytes(StandardCharsets.UTF_8);
+        Key key = Keys.hmacShaKeyFor(keyBytes);
+        // Generate the JWT token
+        String token = Jwts.builder()
+                .setClaims(payload)
+                .setExpiration(Date.from(expiryUtc))
+                .signWith(key, algorithm)
+                .compact();
+
+        Map<String, Object> result = new HashMap<>();
+        result.put("token", token);
+        result.put("expiryTime", expiryUtc);
+        return result;
+    }
+
+    public static Map<String, Object> generateJwtToken(ApplicationProperties applicationProperties, Map<String, Object> payload) {
+        return generateJwtToken(applicationProperties, payload, SignatureAlgorithm.HS256, 30);
+    }
+}
diff --git a/src/main/java/be/cytomine/service/image/ImageInstanceService.java b/src/main/java/be/cytomine/service/image/ImageInstanceService.java
@@ -341,7 +341,6 @@ public Page<Map<String, Object>> list(SecUser user, List<SearchParameterEntry> s
         String select, from, where, search, sort;
         String request;
 
-        // TODO: Check query security
         select = "SELECT distinct " + imageInstanceAlias + ".* ";
         from = "FROM user_image "+ imageInstanceAlias + " ";
         where = "WHERE user_image_id = :user_id ";

diff --git a/src/main/java/be/cytomine/service/middleware/ImageServerService.java b/src/main/java/be/cytomine/service/middleware/ImageServerService.java
@@ -25,6 +25,7 @@
 import be.cytomine.dto.PimsResponse;
 import be.cytomine.exceptions.*;
 import be.cytomine.repository.middleware.ImageServerRepository;
+import be.cytomine.security.jwt.JwtTokenGenerator;
 import be.cytomine.service.CurrentUserService;
 import be.cytomine.service.ModelService;
 import be.cytomine.service.UrlApi;
@@ -856,6 +857,10 @@ private PimsResponse makeRequest(String httpMethod, String imageServerInternalUr
             for (Map.Entry<String, Object> entry : headers.entrySet()) {
                 requestBuilder.setHeader(entry.getKey(), (String) entry.getValue());
             }
+            Map<String, Object> payload = new HashMap<>();
+            payload.put("application", "core");
+            Map<String, Object> tokenData = JwtTokenGenerator.generateJwtToken(applicationProperties, payload);
+            requestBuilder.setHeader("Authorization", "Bearer "+tokenData.get("token"));
             HttpRequest request = requestBuilder.build();
             HttpResponse<byte[]> response = httpClient.send(request, HttpResponse.BodyHandlers.ofByteArray());
             return processResponse(fullUrl, responseContentType, response);