peer-at-code-framework/src/dev/peerat/framework/TomcatServlet.java

package dev.peerat.framework;

import java.io.IOException;
import java.io.OutputStream;
import java.lang.reflect.Method;

import javax.servlet.ServletConfig;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.jose4j.jwt.consumer.InvalidJwtException;

public class TomcatServlet extends HttpServlet{
	
	private static Router<?> router;
	private static Locker<TomcatWebsocketContainer> sessions;
	
	static void setRouter(Router<?> router){
		TomcatServlet.router = router;
	}
	
	static {
		sessions = new Locker<>();
		sessions.listenAsync((container) -> {
			try{
				Thread.sleep(1000);
				WebSocketWriter writer = new WebSocketWriter(new TomcatHttpWriter(container.getResponse()));
				writer.write("{\"test\":4}");
				writer.flush();
			}catch(Exception e){
				e.printStackTrace();
			}
		});
	}
	
	@Override
	public void init(ServletConfig config) throws ServletException{
		try{
			String value = config.getInitParameter("router");
			Class<?> clazz = Class.forName(value);
			for(Method method : clazz.getDeclaredMethods()){
				if(method.getName().equals("main")) method.invoke(null, new Object[] {null});
			}
		}catch(Exception ex){
			throw new RuntimeException(ex);
		}
	}
	
	public void doGet(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		if(request.getHeader("Sec-WebSocket-Key") != null){
			Context context = null;
			TomcatHttpReader reader = null;
			TomcatHttpWriter writer = null;
			try{
				reader = new TomcatHttpReader(request);
				writer = new TomcatHttpWriter(response);
				context = new Context(RequestType.GET, request.getPathInfo(), isLogin(router, RequestType.GET, reader, writer), writer, router.getDefaultHeaders(RequestType.GET));	
				Response resp = router.execTomcat(context, reader, writer);
				OutputStream output = response.getOutputStream();
				writer.flush(output);
				writer.free(output);
				WebSocketReader wreader = new WebSocketReader(reader);
				sessions.setValue(new TomcatWebsocketContainer(wreader, response));
			}catch(InvalidJwtException e){
			}catch(Exception e){
				response.setHeader("exce", e.getMessage());
				if(context != null && context.getResponseCode() == 0){
					try{
						context.response(500);
					}catch(Exception ex){}
				}
				throw new IOException(e);
			}
			if(context != null) router.getLogger().setValue(context);
			return;
		}
		doIt(RequestType.GET, request, response);
	}
	
	public void doDelete(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.DELETE, request, response);
	}
	
	public void doHead(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.HEAD, request, response);
	}
	
	public void doOptions(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.OPTIONS, request, response);
	}
	
	public void doPost(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.POST, request, response);
	}
	
	public void doPut(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.PUT, request, response);
	}
	
	public void doTrace(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		doIt(RequestType.TRACE, request, response);
	}
	
	private void doIt(RequestType type, HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
		Context context = null;
		TomcatHttpReader reader = null;
		TomcatHttpWriter writer = null;
		try{
			reader = new TomcatHttpReader(request);
			writer = new TomcatHttpWriter(response);
			context = new Context(type, request.getPathInfo(), isLogin(router, type, reader, writer), writer, router.getDefaultHeaders(type));	
			router.exec(context, reader, writer);
			writer.flush(response.getOutputStream());
		}catch(InvalidJwtException e){
		}catch(Exception e){
			if(context != null && context.getResponseCode() == 0){
				try{
					context.response(500);
				}catch(Exception ex){}
			}
			throw new IOException(e);
		}
		if(context != null) router.getLogger().setValue(context);
	}
	
	private User isLogin(Router<?> router, RequestType type, HttpReader reader, HttpWriter writer) throws Exception{
		String auth = reader.getHeader("Authorization");
		if(auth == null) return null;
		auth = auth.substring(7);
		try{
			return router.getUser(auth);
		}catch(InvalidJwtException e){
			try{
				writer.response(401, router.getDefaultHeaders(type));
				writer.flush();
				writer.close();
			}catch(Exception ex){}
			throw e;
		}
	}
	
	private static class TomcatWebsocketContainer{
		
		private WebSocketReader reader;
		private HttpServletResponse response;
		
		public TomcatWebsocketContainer(WebSocketReader reader, HttpServletResponse response){
			this.reader = reader;
			this.response = response;
		}
		
		public WebSocketReader getReader(){
			return reader;
		}
		
		public HttpServletResponse getResponse(){
			return response;
		}
		
	}
}