Chatto/src/main/java/org/ros/chatto/controller/RegistrationController.java

package org.ros.chatto.controller;

import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;

import javax.imageio.ImageIO;
import javax.validation.Valid;

import org.ros.chatto.dto.UserRegistrationDTO;
import org.ros.chatto.service.CaptchaService;
import org.ros.chatto.service.UserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.validation.BindingResult;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.ModelAttribute;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

@Controller
@RequiredArgsConstructor
@Slf4j
public class RegistrationController {

	@Autowired
	private final UserService userService;

	@Autowired
	private final CaptchaService captchaService;

	// FIXME must change this to a timeout base cache otherwise memory leak!
	private final Map<Long, String> captchaMap = new ConcurrentHashMap<>();

	@GetMapping("/registration")
	public String registrationForm(Model model) {
		UserRegistrationDTO userRegistrationDTO = new UserRegistrationDTO();
		String captchaText = captchaService.getRandomText();
		userRegistrationDTO.setCaptchaText(captchaText);
		log.debug("captcha text = {}", captchaText);
		Long captchaID = ThreadLocalRandom.current().nextLong();
		userRegistrationDTO.setCaptchaID(captchaID);
		captchaMap.put(captchaID, captchaText);
		model.addAttribute("userRegistrationDTO", userRegistrationDTO);
		return "registration";
	}

	@PostMapping("/perform_registration")
	public String performRegistration(
			@ModelAttribute("userRegistrationDTO") @Valid UserRegistrationDTO userRegistrationDTO,
			BindingResult bindingResult) {
		if (bindingResult.hasErrors()) {
			log.warn("Registration input has errors!");
			return "redirect:registration?error";
		}
		if (userService.getUser(userRegistrationDTO.getUserName()).isPresent()) {
			return "redirect:registration?error&duplicate=true";
		}

		log.debug("Captcha text from user input = {}", userRegistrationDTO.getCaptchaInput());
		log.debug("Captcha text from captcha map = {}", captchaMap.get(userRegistrationDTO.getCaptchaID()));
		if (userRegistrationDTO.getCaptchaInput().equals(captchaMap.get(userRegistrationDTO.getCaptchaID()))) {
			log.info("Registration captcha equal success");
			userService.createUser(userRegistrationDTO);
			return "redirect:registration?success";
		} else {
			log.warn("Registration captcha equal fail");
			return "redirect:registration?error&captchaError=true";
		}
	}

	@GetMapping(value = "/img/captcha/{image_id}", produces = MediaType.IMAGE_PNG_VALUE)
	public ResponseEntity<byte[]> getImage(@PathVariable("image_id") Long imageId) throws IOException {

		final String captchaText = captchaMap.get(imageId);
		final HttpHeaders headers = new HttpHeaders();
		headers.setContentType(MediaType.IMAGE_PNG);
		BufferedImage captchaBufferedImage = captchaService.createCaptchaImage(captchaText);

		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		ImageIO.write(captchaBufferedImage, "png", baos);
		byte[] imageBytes = baos.toByteArray();

		return new ResponseEntity<byte[]>(imageBytes, headers, HttpStatus.OK);
	}
}