WsController.java

package org.flasby.ws;

import java.security.Principal;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;

import org.flasby.entity.Present;
import org.flasby.entity.repository.PresentsRepository;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.event.EventListener;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.messaging.simp.annotation.SendToUser;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Controller;
import org.springframework.web.socket.messaging.AbstractSubProtocolEvent;
import org.springframework.web.socket.messaging.SessionConnectEvent;
import org.springframework.web.socket.messaging.SessionDisconnectEvent;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.log4j.Log4j2;

@Log4j2
@Controller
public class WsController {

    @Autowired
    private SimpMessagingTemplate messagingTemplate;

    @Autowired
    private PresentsRepository pr;

    // @Getter
    public static class UserInfo {

        public UserInfo(String name, Collection<? extends GrantedAuthority> auths) {
            this.name = name;
            this.auths = Collections.unmodifiableCollection(auths);
        }

        private final String name;
        private final Collection<? extends GrantedAuthority> auths;

        public String getName() {
            return name;
        }

        public Collection<? extends GrantedAuthority> getAuths() {
            return Collections.unmodifiableCollection(auths);
        }

    }

    @Data
    @AllArgsConstructor
    public static class Ping {
        @SuppressWarnings("unused") // Needed by Jackson
        private Ping() {
        }

        private int ping;
    }

    @Data
    @AllArgsConstructor
    public static class RowId {
        @SuppressWarnings("unused") // Needed by Jackson
        private RowId() {
        }

        private long id;
    }

    @Data
    @AllArgsConstructor
    public static class PresentData {
        @SuppressWarnings("unused") // Needed by Jackson
        private PresentData() {
        }

        private String forUser;
        private List<Present> myRequests;
        private List<Present> allOthersRequest;
        private List<Present> myClaims;
    }

    @Data
    @AllArgsConstructor
    public static class Pong {
        private final int count;
        private final long time;
    }

    private Map<String, UserInfo> connectedUsers = new HashMap<>();

    public WsController() {
    }

    static String getUserName(AbstractSubProtocolEvent event) {
        if ( event != null ) {
            if ( event.getUser() != null ){
                return event.getUser().getName();
            } else {
                log.warn("Missing user on event");
                return "";
            }
        } else {
            log.error("Null Protocol Event");
            return "";
        }
    }


    @EventListener
    public void onDisconnectEvent(SessionDisconnectEvent event) {
        log.info("Client with username {} disconnected", getUserName(event));
        connectedUsers.remove(getUserName(event));

        // Logout the current user
        SecurityContextHolder.getContext().setAuthentication(null);

        // // Expire the user from the current websocket session
        // SessionState ss = sessionManager.expire(event.getSessionId());
    }

    @EventListener
    public void onConnectEvent(SessionConnectEvent event) {
        log.info("Client with username {} connected", getUserName(event));
        Authentication auth = (Authentication) event.getUser();

        connectedUsers.put(getUserName(event), new UserInfo(auth.getName(), auth.getAuthorities()));
        // Logout the current user
        SecurityContextHolder.getContext().setAuthentication(null);

        log.info("Connected Principals " + Arrays.toString(connectedUsers.keySet().toArray()));
        // // Expire the user from the current websocket session
        // SessionState ss = sessionManager.expire(event.getSessionId());
    }

    @MessageMapping("/whoami")
    @SendToUser("/queue/youare")
    public UserInfo whoAmI(Principal user) {
        Authentication auth = (Authentication) user;
        messagingTemplate.convertAndSendToUser(auth.getName(), "/queue/version",
                "" + this.getClass().getPackage().getImplementationVersion());
        return new UserInfo(auth.getName(), auth.getAuthorities());
    }

    @MessageMapping("/ping")
    @SendToUser("/queue/pong")
    public Pong ping(Ping ping) {
        return new Pong(ping.getPing(), System.currentTimeMillis());
    }

    @MessageMapping("/subscribe")
    @SendToUser("/queue/initialList")
    public PresentData subscribe( Principal me) throws JsonProcessingException {

        System.out.println();
        log.info("Connected Principals " + Arrays.toString(connectedUsers.keySet().toArray()));
        System.out.println();

        System.err.println("Handling request for /subscribe");
        // pushDataEverywhere("changed", getPresentsAsJSON());
        System.err.println("Now sending Present List in response to request for /subscribe");

        return new PresentData(
                me.getName(),
                pr.getMyRequests(me.getName()),
                pr.getRequestsExceptMine(me.getName()),
                pr.getMyClaims(me.getName()));
    }

    List<Present> getPresents() {
        List<Present> presents = pr.findAll();// new ArrayList<>();
        Collections.sort(presents, new Comparator<Present>() {
            @Override
            public int compare(Present arg0, Present arg1) {
                if (arg0.getName().equals(arg1.getName())) {
                    return arg0.getDescription().compareTo(arg1.getDescription());
                }
                return arg0.getName().compareTo(arg1.getName());
            }
        });
        return presents;
    }

    String getPresentsAsJSON() throws JsonProcessingException {
        List<Present> presents = getPresents();
        ObjectWriter ow = new ObjectMapper().writer();// .withDefaultPrettyPrinter();
        return ow.writeValueAsString(presents);
    }

    @MessageMapping("/unclaim")
    public void unclaim(RowId row, Principal user) throws JsonProcessingException {
        System.err.println("Unclaim " + row.getId());
        Optional<Present> p = pr.findById(row.getId());
        p.ifPresent(new Consumer<Present>() {
            public void accept(Present present) {
                present.setClaimed(false);
                present.setClaimedBy("");
                pr.save(present);
            };
        });
        sendAllLists();
    }

    @MessageMapping("/claim")
    public void claim(RowId row, Principal user) throws JsonProcessingException {
        System.err.println("Claim " + row.getId());
        Optional<Present> p = pr.findById(row.getId());
        p.ifPresent(new Consumer<Present>() {
            public void accept(Present present) {
                present.setClaimed(true);
                present.setClaimedBy(user.getName());
                pr.save(present);
            };
        });
        sendAllLists();
    }

    private void sendAllLists() throws JsonProcessingException {

        for (UserInfo connectedUser : connectedUsers.values()) {
            System.out.println();
            System.out.println(Arrays.toString(pr.getMyRequests(connectedUser.getName()).toArray()));
            System.out.println();
            System.out.println(Arrays.toString(pr.getMyClaims(connectedUser.getName()).toArray()));
            System.out.println();
            System.out.println(Arrays.toString(pr.getRequestsExceptMine(connectedUser.getName()).toArray()));
            System.out.println();

            PresentData pb = new PresentData(
                    connectedUser.getName(),
                    pr.getMyRequests(connectedUser.getName()),
                    pr.getRequestsExceptMine(connectedUser.getName()),
                    pr.getMyClaims(connectedUser.getName()));

            messagingTemplate.convertAndSendToUser(connectedUser.getName(), "/queue/changed", pb);
            ObjectWriter ow = new ObjectMapper().writer().withDefaultPrettyPrinter();
            log.info("Sending for user:" + connectedUser.getName() + " " + ow.writeValueAsString(pb));
        }
    }

    @MessageMapping("/askfor")
    public void askFor(String mywish, Principal me) throws JsonProcessingException {
        log.info("New Present: " + pr.save(new Present(me.getName(), mywish)));
        sendAllLists();
    }

    @MessageMapping("/log")
    public void log(String msg) {
        log.info("LOG from GUI: " + msg);
    }

    public void sendErrorTo(String errorMessage, Principal user) {
        messagingTemplate.convertAndSendToUser(user.getName(), "/queue/error", errorMessage);
    }

    public void pushDataEverywhere(String topicName, String json) {
        log.info("Pushing to topic " + topicName + " - " + json);
        messagingTemplate.convertAndSend("/topic/" + topicName, json);
    }
}