Databricks L5 (SSE) Technical Phone Screen | Cleared
Anonymous User
10423

Implement a firewall feature considering following points:

  1. Some rules are defined ("ALLOW", "192.168.1.0/24"), ("ALLOW", "1.2.3.4"), ("DENY", "10.5.12.10/28")
  2. Only allow requests from IP addresses that follow the rule. All other requests should be denied.

I implemented this system like following:

class Firewall {
	Map<CIDRBlock, String> map;
	
	public Firewall() {
		this.map = new HashMap<>();
	}
	
	public void addRule(String cidr, String mode) {
		this.map.put(new CIDRBlock(cidr), mode);
	}
	
	public boolean allow(String ip) throws UnknownHostException {
		InetAddress input = InetAddress.getByName(ip);
		
		for (CIDRBlock cidr : this.map.keySet()) {
			String mode = this.map.get(cidr);
			if (cidr.contains(input)) {
				return (mode.equals("ALLOW")) ? true : false;
			}
		}
		return false;
	}
}

class CIDRBlock {
	InetAddress networkAddress;
	int prefixLength;
	int mask;
	
	public CIDRBlock(String cidr) throws UnknownHostException {
		String[] parts = cidr.split("/");
        this.networkAddress = InetAddress.getByName(parts[0]);
        this.prefixLength = Integer.parseInt(parts[1]);
        this.mask = (0xFFFFFFFF << (32 - this.prefixLength));
	}
	
	public boolean contains(InetAddress address) throws UnknownHostException {
        byte[] networkBytes = this.networkAddress.getAddress();
        byte[] addressBytes = address.getAddress();

        int networkInt = byteArrayToInt(networkBytes);
        int addressInt = byteArrayToInt(addressBytes);

        return (networkInt & mask) == (addressInt & mask);
    }

    private int byteArrayToInt(byte[] bytes) {
        int result = 0;
        for (byte b : bytes) {
            result = (result << 8) | (b & 0xFF);
        }
        return result;
    }
}
Comments (2)