diff --git a/patroni/api.py b/patroni/api.py index d9a0f527..b99f2cb8 100644 --- a/patroni/api.py +++ b/patroni/api.py @@ -1,3 +1,4 @@ +import base64 import fcntl import json import logging @@ -10,9 +11,36 @@ from threading import Thread logger = logging.getLogger(__name__) +def check_auth(func): + """Decorator function to check authorization header. + + Usage example: + @check_auth + def do_PUT_foo(): + pass + """ + def wrapper(handler): + if handler.check_auth_header(): + return func(handler) + return wrapper + + class RestApiHandler(BaseHTTPRequestHandler): + def send_auth_request(self, body): + self.send_response(401) + self.send_header('WWW-Authenticate', 'Basic realm=\"Patroni\"') + self.send_header('Content-type', 'text/html') + self.end_headers() + + def check_auth_header(self): + auth_header = self.headers.get('Authorization') + status = self.server.check_auth_header(auth_header) + return not status or self.send_auth_request(status) + def do_GET(self): + """Default method for processing all GET requests which can not be routed to other methods""" + response = self.get_postgresql_status() path = '/master' if self.path == '/' else self.path @@ -23,6 +51,31 @@ class RestApiHandler(BaseHTTPRequestHandler): self.end_headers() self.wfile.write(json.dumps(response).encode('utf-8')) + @check_auth + def do_GET_sampleauth(self): + self.send_response(200) + self.send_header('Content-Type', 'text/html') + self.end_headers() + self.wfile.write(b'Hello!') + + def parse_request(self): + """Override parse_request method to enrich basic functionality of `BaseHTTPRequestHandler` class + + Original class can only invoke do_GET, do_POST, do_PUT, etc method implementations if they are defined. + But we would like to have at least some simple routing mechanism, i.e.: + GET /uri1/part2 request should invoke `do_GET_uri1()` + POST /other should invoke `do_POST_other()` + + If the `do__` method does not exists we'll fallback to original behavior.""" + + ret = BaseHTTPRequestHandler.parse_request(self) + if ret: + mname = self.path.lstrip('/').split('/')[0] + mname = self.command + ('_' + mname if mname else '') + if hasattr(self, 'do_' + mname): + self.command = mname + return ret + def get_postgresql_status(self): try: row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'), @@ -52,6 +105,7 @@ class RestApiHandler(BaseHTTPRequestHandler): class RestApiServer(ThreadingMixIn, HTTPServer, Thread): def __init__(self, patroni, config): + self._auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None self.connection_string = 'http://{}/patroni'.format(config.get('connect_address', None) or config['listen']) host, port = config['listen'].split(':') HTTPServer.__init__(self, (host, int(port)), RestApiHandler) @@ -71,3 +125,13 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread): def _set_fd_cloexec(fd): flags = fcntl.fcntl(fd, fcntl.F_GETFD) fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC) + + def check_basic_auth_key(self, key): + return self._auth_key == key + + def check_auth_header(self, auth_header): + if self._auth_key: + if auth_header is None: + return 'no auth header received' + if not auth_header.startswith('Basic ') or not self.check_basic_auth_key(auth_header[6:]): + return 'not authenticated' diff --git a/postgres0.yml b/postgres0.yml index 659a4db2..eb9d10ec 100644 --- a/postgres0.yml +++ b/postgres0.yml @@ -4,6 +4,7 @@ scope: &scope batman restapi: listen: 127.0.0.1:8008 connect_address: 127.0.0.1:8008 + auth: 'username:password' etcd: scope: *scope ttl: *ttl diff --git a/postgres1.yml b/postgres1.yml index bc8b6fd1..52a0bc26 100644 --- a/postgres1.yml +++ b/postgres1.yml @@ -4,6 +4,7 @@ scope: &scope batman restapi: listen: 127.0.0.1:8009 connect_address: 127.0.0.1:8009 + auth: 'username:password' etcd: scope: *scope ttl: *ttl diff --git a/tests/test_api.py b/tests/test_api.py index aecf3df6..46f9e5ae 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -3,13 +3,27 @@ import unittest from patroni.api import RestApiHandler, RestApiServer from six import BytesIO as IO +from six.moves import BaseHTTPServer from test_postgresql import psycopg2_connect +def nop(*args, **kwargs): + pass + + def throws(*args, **kwargs): raise psycopg2.OperationalError() +class Mock_BaseServer__is_shut_down: + + def set(self): + pass + + def clear(self): + pass + + class MockPostgresql: def connection(self): @@ -37,7 +51,7 @@ class MockRequest: class MockRestApiServer(RestApiServer): def __init__(self, Handler, path, *args): - self.patroni = MockPatroni() + super(MockRestApiServer, self).__init__(MockPatroni(), {'listen': '127.0.0.1:8008', 'auth': 'test:test'}) if len(args) > 0: self.query = args[0] Handler(MockRequest(path), ('0.0.0.0', 8080), self) @@ -46,8 +60,20 @@ class MockRestApiServer(RestApiServer): class TestRestApiHandler(unittest.TestCase): def __init__(self, method_name='runTest'): + self.setUp = self.set_up super(TestRestApiHandler, self).__init__(method_name) + def set_up(self): + BaseHTTPServer.HTTPServer.__init__ = nop + RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down() + RestApiServer._BaseServer__shutdown_request = True + RestApiServer.socket = 0 + def test_do_GET(self): MockRestApiServer(RestApiHandler, b'GET /') MockRestApiServer(RestApiHandler, b'GET /', throws) + + def test_do_GET_sampleauth(self): + MockRestApiServer(RestApiHandler, b'GET /sampleauth') + MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization:') + MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization: Basic dGVzdDp0ZXN0') diff --git a/tests/test_patroni.py b/tests/test_patroni.py index a45ebf62..285d0e95 100644 --- a/tests/test_patroni.py +++ b/tests/test_patroni.py @@ -14,6 +14,7 @@ from patroni.etcd import Etcd from patroni import Patroni, main from patroni.zookeeper import ZooKeeper from six.moves import BaseHTTPServer +from test_api import Mock_BaseServer__is_shut_down from test_etcd import Client, etcd_read, etcd_write from test_ha import true, false from test_postgresql import Postgresql, subprocess_call, psycopg2_connect @@ -32,15 +33,6 @@ def time_sleep(*args): raise SleepException() -class Mock_BaseServer__is_shut_down: - - def set(self): - pass - - def clear(self): - pass - - class TestPatroni(unittest.TestCase): def __init__(self, method_name='runTest'):