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

package dev.peerat.framework;

import java.lang.reflect.Method;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.Key;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.regex.Pattern;

import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.SSLServerSocketFactory;

import org.jose4j.jwa.AlgorithmConstraints.ConstraintType;
import org.jose4j.jwk.RsaJsonWebKey;
import org.jose4j.jwk.RsaJwkGenerator;
import org.jose4j.jwk.JsonWebKey.OutputControlLevel;
import org.jose4j.jws.AlgorithmIdentifiers;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwt.consumer.JwtConsumer;
import org.jose4j.jwt.consumer.JwtConsumerBuilder;
import org.jose4j.lang.JoseException;

public class Router<U extends User>{
	
	public static void main(String[] args) {
		
	}
	
	private Locker<Context> logger;
	private Locker<Exception> exceptions;
	private RouteMapper<U>[] mappers;
	private List<RouteInterceptor> interceptors;
	private Response noFileFound;
	private RsaJsonWebKey rsaJsonWebKey;
	private JwtConsumer jwtConsumer;
	private Consumer<JwtClaims> claims;
	private Function<JwtClaims, U> userCreator;
	private String[][] headers;
	private ServerSocket serverSocket;

	public Router(){
		this.logger = new Locker<>();
		this.exceptions = new Locker<>();
		int types = RequestType.values().length;
		this.mappers = new RouteMapper[types];
		this.interceptors = new ArrayList<>();
		for(RequestType type : RequestType.values()) this.mappers[type.ordinal()] = new RouteMapper<>(this);
		this.headers = new String[types][0];
	}
	
	public Router<U> configureJwt(Consumer<JwtConsumerBuilder> consumer, Consumer<JwtClaims> claims, Function<JwtClaims, U> userCreator) throws Exception{
		this.rsaJsonWebKey = RsaJwkGenerator.generateJwk(2048);
		configureJwtWithKey(consumer, claims, userCreator, this.rsaJsonWebKey.getKey());
		return this;
	}
	
	public Router<U> configureJwt(Consumer<JwtConsumerBuilder> consumer, Consumer<JwtClaims> claims, Function<JwtClaims, U> userCreator, Map<String, Object> keyParams) throws Exception{
		this.rsaJsonWebKey = new RsaJsonWebKey(keyParams);
		configureJwtWithKey(consumer, claims, userCreator, this.rsaJsonWebKey.getKey());
		return this;
	}
	
	private void configureJwtWithKey(Consumer<JwtConsumerBuilder> consumer, Consumer<JwtClaims> claims, Function<JwtClaims, U> userCreator, Key key) throws Exception{
		JwtConsumerBuilder builder = new JwtConsumerBuilder()
			.setRequireExpirationTime()
			.setAllowedClockSkewInSeconds(30)
			.setVerificationKey(key)
			.setJwsAlgorithmConstraints(ConstraintType.PERMIT, AlgorithmIdentifiers.RSA_USING_SHA256);
				
			consumer.accept(builder);
				
			this.jwtConsumer = builder.build();
			this.claims = claims;
			this.userCreator = userCreator;
	}
	
	public Map<String, Object> exportJwtKey(){
		return this.rsaJsonWebKey.toParams(OutputControlLevel.INCLUDE_PRIVATE);
	}
	
	public Router<U> addDefaultHeaders(RequestType type, String... headers){
		String[] origin = this.headers[type.ordinal()]; 
		String[] copy = new String[origin.length+headers.length];
		System.arraycopy(origin, 0, copy, 0, origin.length);
		System.arraycopy(headers, 0, copy, origin.length, headers.length);
		this.headers[type.ordinal()] = copy;
		return this;
	}
	
	public Router<U> register(Response response){
		try{
			Method method = response.getClass().getDeclaredMethod("exec",
					Response.class.getDeclaredMethods()[0].getParameterTypes());
			Route route = method.getAnnotation(Route.class);

			this.mappers[route.type().ordinal()].register(response, method, route, Pattern.compile(route.path()));
		}catch(Exception e){
			throw new IllegalArgumentException(e);
		}
		return this;
	}

	public Router<U> setDefault(Response response){
		this.noFileFound = response;
		return this;
	}
	
	public Router<U> addInterceptor(RouteInterceptor interceptor){
		this.interceptors.add(interceptor);
		return this;
	}
	
	public Router<U> activeReOrdering(){
		for(RouteMapper<?> mapper : this.mappers) mapper.activeReOrdering();
		return this;
	}
	
	String[] getDefaultHeaders(RequestType type){
		return  this.headers[type.ordinal()];
	}

	void exec(Context context, HttpReader reader, HttpWriter writer) throws Exception{
		if(this.mappers[context.getType().ordinal()].exec(context, reader, writer, this.interceptors)) return;
		if(noFileFound != null) noFileFound.exec(null, context, reader, writer);
	}

	public Router<U> configureSSL(String keyStore, String keyStorePassword){
		System.setProperty("javax.net.ssl.keyStore", keyStore);
		System.setProperty("javax.net.ssl.keyStorePassword", keyStorePassword);
		return this;
	}
	
	public U getUser(String token) throws InvalidJwtException{
		return this.userCreator.apply(this.jwtConsumer.processToClaims(token));
	}

	public String createAuthUser(U user) throws JoseException{
		JwtClaims claims = new JwtClaims();
		claims.setGeneratedJwtId(); // a unique identifier for the token
		claims.setIssuedAtToNow(); // when the token was issued/created (now)
		claims.setNotBeforeMinutesInThePast(2); // time before which the token is not yet valid (2 minutes ago)
		this.claims.accept(claims);

		user.write(claims);

		JsonWebSignature jws = new JsonWebSignature();
		jws.setPayload(claims.toJson());
		jws.setKey(rsaJsonWebKey.getPrivateKey());
		jws.setKeyIdHeaderValue(rsaJsonWebKey.getKeyId());
		jws.setAlgorithmHeaderValue(AlgorithmIdentifiers.RSA_USING_SHA256);
		return jws.getCompactSerialization();
	}
	
	public Locker<Context> getLogger(){
		return this.logger;
	}
	
	public Locker<Exception> getExceptionLogger(){
		return this.exceptions;
	}
	
	public void listen(int port, boolean ssl) throws Exception{
		if (ssl) { // Not needed with the use of a proxy
			try {
				SSLServerSocketFactory ssf = (SSLServerSocketFactory) SSLServerSocketFactory.getDefault();
				serverSocket = (SSLServerSocket) ssf.createServerSocket(port);

				while (!serverSocket.isClosed()) {
					Socket socket = serverSocket.accept();
					Client<U> client = new Client<>(socket, this);
					client.start();
				}
			} catch (Exception e) {
				e.printStackTrace();
			} finally {
				stop();
			}
		} else {
			try{
				serverSocket = new ServerSocket(port);
				while (!serverSocket.isClosed()) {
					Socket socket = serverSocket.accept();
					Client<U> client = new Client<>(socket, this);
					client.start();
				}
			} catch (Exception e) {
				e.printStackTrace();
			} finally {
				stop();
			}
		}
	}
	
	public void stop(){
		if(serverSocket == null) return;
		try {
			serverSocket.close();
			serverSocket = null;
		}catch(Exception e){
			e.printStackTrace();
		}
	}
	
}