def test_transform_backends_from_config():
input = yaml.safe_load('''
hosts:
- host: www.mango.com
servers:
- localhost:8081
- localhost:8082
- host: www.apple.com
servers:
- localhost:9081
- localhost:9082
paths:
- path: /mango
servers:
- localhost:8081
- localhost:8082
- path: /apple
servers:
- localhost:9081
- localhost:9082
''')
output = transform_backends_from_config(input)
assert list(output.keys()) == [
'www.mango.com', 'www.apple.com', '/mango', '/apple'
]
assert output['www.mango.com'][0] == Server('localhost:8081')
assert output['www.mango.com'][1] == Server('localhost:8082')
assert output['www.apple.com'][0] == Server('localhost:9081')
assert output['www.apple.com'][1] == Server('localhost:9082')
assert output['/mango'][0] == Server('localhost:8081')
assert output['/mango'][1] == Server('localhost:8082')
assert output['/apple'][0] == Server('localhost:9081')
assert output['/apple'][1] == Server('localhost:9082')
def test_transform_backends_from_config():
input = yaml.safe_load('''
hosts:
- host: www.mango.com
servers:
- localhost:8081
- localhost:8082
- host: www.apple.com
servers:
- localhost:9081
- localhost:9082
paths:
- path: /mango
servers:
- localhost:8081
- localhost:8082
- path: /apple
servers:
- localhost:9081
- localhost:9082
''')
output = transform_backends_from_config(input)
assert list(output.keys()) == [
"www.mango.com", "www.apple.com", "/mango", "/apple"
]
assert output["www.mango.com"][0] == Server("localhost:8081")
assert output["www.mango.com"][1] == Server("localhost:8082")
assert output["www.apple.com"][0] == Server("localhost:9081")
assert output["www.apple.com"][1] == Server("localhost:9082")
assert output["/mango"][0] == Server("localhost:8081")
assert output["/mango"][1] == Server("localhost:8082")
assert output["/apple"][0] == Server("localhost:9081")
assert output["/apple"][1] == Server("localhost:9082")
def test_healthcheck():
config = yaml.safe_load('''
hosts:
- host: www.mango.com
servers:
- localhost:8081
- localhost:8888
- host: www.apple.com
servers:
- localhost:9081
- localhost:4444
''')
register = healthcheck(transform_backends_from_config(config))
assert register['www.apple.com'][0].healthy
assert not register['www.apple.com'][1].healthy
assert register['www.mango.com'][0].healthy
assert not register['www.mango.com'][1].healthy
def test_healthcheck():
config = yaml.safe_load("""
hosts:
- host: www.mango.com
servers:
- localhost:8081
- localhost:8888
- host: www.apple.com
servers:
- localhost:9081
- localhost:4444
""")
register = healthcheck(transform_backends_from_config(config))
assert register["www.apple.com"][0].healthy == True
assert register["www.apple.com"][1].healthy == False
assert register["www.mango.com"][0].healthy == True
assert register["www.mango.com"][1].healthy == False
def test_healthcheck():
config = yaml.safe_load(
"""
hosts:
- host: www.anthrax.com
servers:
- localhost:8081
- localhost:8888
- host: www.metallica.com
servers:
- localhost:9081
- localhost:4444
"""
)
register = healthcheck(transform_backends_from_config(config))
assert register["www.metallica.com"][0].healthy
assert not register["www.metallica.com"][1].healthy
assert register["www.anthrax.com"][0].healthy
assert not register["www.anthrax.com"][1].healthy
def test_transform_backends_from_config():
input = yaml.safe_load(
"""
hosts:
- host: www.anthrax.com
servers:
- localhost:8081
- localhost:8082
- host: www.metallica.com
servers:
- localhost:9081
- localhost:9082
paths:
- path: /anthrax
servers:
- localhost:8081
- localhost:8082
- path: /metallica
servers:
- localhost:9081
- localhost:9082
"""
)
output = transform_backends_from_config(input)
assert list(output.keys()) == [
"www.anthrax.com",
"www.metallica.com",
"/anthrax",
"/metallica",
]
assert output["www.anthrax.com"][0] == Server("localhost:8081")
assert output["www.anthrax.com"][1] == Server("localhost:8082")
assert output["www.metallica.com"][0] == Server("localhost:9081")
assert output["www.metallica.com"][1] == Server("localhost:9082")
assert output["/anthrax"][0] == Server("localhost:8081")
assert output["/anthrax"][1] == Server("localhost:8082")
assert output["/metallica"][0] == Server("localhost:9081")
assert output["/metallica"][1] == Server("localhost:9082")
from flask import Flask, request
import requests, random
from utils import load_configuration, transform_backends_from_config, get_healthy_server, process_rules, process_rewrite_rules, process_firewall_rules_flag
from tasks import healthcheck
import sys
loadbalancer = Flask(__name__)
config = load_configuration('loadbalancer.yaml')
register = transform_backends_from_config(config)
@loadbalancer.route("/")
@loadbalancer.route("/<path>")
def router(path="/"):
updated_register = healthcheck(register)
host_header = request.headers["Host"]
if not process_firewall_rules_flag(
config, host_header, request.environ["REMOTE_ADDR"], "/" + path):
return "Forbidden", 403
for entry in config["hosts"]:
if host_header == entry["host"]:
healthy_server = get_healthy_server(entry["host"],
updated_register)
if not healthy_server:
return "No Backends servers available", 503
headers = process_rules(config, host_header,
{k: v
for k, v in request.headers.items()},
"header")
params = process_rules(config, host_header,
{k: v
