UrlsWorker/src/main/java/eu/openaire/urls_worker/controllers/GeneralController.java

package eu.openaire.urls_worker.controllers;

import eu.openaire.urls_worker.components.AssignmentsHandler;
import eu.openaire.urls_worker.components.plugins.PublicationsRetrieverPlugin;
import eu.openaire.urls_worker.util.UriBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import javax.servlet.http.HttpServletRequest;
import java.io.File;
import java.net.UnknownHostException;
import java.util.regex.Pattern;


@RestController
@RequestMapping("")
public class GeneralController {

    private static final Logger logger = LoggerFactory.getLogger(GeneralController.class);

    private final String controllerIp;
    private final String workerReportsDirPath;
    private final String workerId;


    private static final Pattern DOMAIN_DETECTOR = Pattern.compile("^.*[a-zA-Z].*$");

    public GeneralController(@Value("${info.controllerIp}") String controllerIp, @Value("${workerReportsDirPath}") String workerReportsDirPath, @Value("${info.workerId}") String workerId)
    {
        if ( DOMAIN_DETECTOR.matcher(controllerIp).matches() ) {
            try {
                this.controllerIp = java.net.InetAddress.getByName(controllerIp).getHostAddress();
            } catch (UnknownHostException uhe) {
                String errorMsg = "The domain given for the Controller (" + controllerIp + ") is unknown to the world! So its IP cannot be retrieved!";
                logger.error(errorMsg);
                throw new RuntimeException(errorMsg);
            }
        } else
            this.controllerIp = controllerIp;

        this.workerReportsDirPath = workerReportsDirPath;
        this.workerId = workerId;
    }


    @GetMapping("isAlive")
    public ResponseEntity<?> isWorkerAlive() {
        logger.info("Received an \"isAlive\" request.");
        return ResponseEntity.ok().build();
    }


    public static boolean shouldShutdownWorker = false;

    @PostMapping("shutdownWorker")
    public ResponseEntity<?> shutdownWorkerGracefully(HttpServletRequest request)
    {
        String initMsg = "Received a \"shutdownWorker\" request. ";
        ResponseEntity<?> responseEntity = passSecurityChecks(request, initMsg);
        if ( responseEntity != null )
            return responseEntity;

        String finalMsg = "";
        if ( shouldShutdownWorker )
            finalMsg = "The worker has already received a \"shutdownWorker\" (which was not canceled afterwards).";
        else
            shouldShutdownWorker = true;

        finalMsg += "The worker will shutdown, after finishing current work.";
        logger.info(initMsg + finalMsg);
        return ResponseEntity.ok().body(finalMsg + "\n");
    }


    @PostMapping("cancelShutdownWorker")
    public ResponseEntity<?> cancelShutdownWorkerGracefully(HttpServletRequest request)
    {
        String initMsg = "Received a \"cancelShutdownWorker\" request. ";
        ResponseEntity<?> responseEntity = passSecurityChecks(request, initMsg);
        if ( responseEntity != null )
            return responseEntity;

        shouldShutdownWorker = false;
        String finalMsg = "Any previous \"shutdownWorker\"-request is canceled. The \"maxAssignmentsBatchesToHandleBeforeShutdown\" will still be honored (if it's set).";
        logger.info(initMsg + finalMsg);
        return ResponseEntity.ok().body(finalMsg + "\n");
    }


    @GetMapping("getHandledAssignmentsCounts")
    public ResponseEntity<?> getHandledAssignmentsCounts()
    {
        return ResponseEntity.ok(AssignmentsHandler.assignmentsNumsHandled);
    }


    @PostMapping("addReportResultToWorker/{assignmentsCounter}")
    public ResponseEntity<?> addReportResultToWorker(@PathVariable long assignmentsCounter, @RequestBody(required=false) String errorMsg)
    {
        String directoryPath = PublicationsRetrieverPlugin.assignmentsBasePath + "assignments_" + assignmentsCounter + "_fullTexts";
        File dir = new File(directoryPath);
        if ( ! dir.isDirectory() ) {
            logger.error("The \"addReportResultToWorker\"-endpoint was called for an unknown \"assignmentsCounter\": " + assignmentsCounter);
            return ResponseEntity.notFound().build();
        }

        if ( errorMsg == null ) {
            logger.info("The Controller successfully handled the WorkerReport, for assignments_" + assignmentsCounter + ". The worker-report and all full-text files associated with it, will be deleted.");
            FullTextsController.deleteAssignmentsDirectory(assignmentsCounter, dir);
            FullTextsController.deleteFile(this.workerReportsDirPath + this.workerId + "_assignments_" + assignmentsCounter + "_report.json");
        } else
            logger.error("The Controller failed to handle the WorkerReport, for assignments_" + assignmentsCounter + ". The error is:\n" + errorMsg);

        return ResponseEntity.ok().build();
    }


    public ResponseEntity<?> passSecurityChecks(HttpServletRequest request, String initMsg)
    {
        if ( request == null ) {
            logger.error(initMsg + "The \"HttpServletRequest\" is null!");
            return ResponseEntity.internalServerError().build();
        }
        String remoteAddr = request.getHeader("X-FORWARDED-FOR");   // This retrieves the original IP address, if the request passes through a proxy server.
        if ( remoteAddr == null )
            remoteAddr = request.getRemoteAddr();

        if ( ! (remoteAddr.equals("127.0.0.1") || remoteAddr.equals(UriBuilder.ip) || remoteAddr.equals(controllerIp)) ) {
            logger.error(initMsg + "The request came from another IP: " + remoteAddr + " | while this worker has the IP: " + UriBuilder.ip);
            return ResponseEntity.status(HttpStatus.FORBIDDEN).build();
        }
        return null;    // The checks are passing.
    }

}