From df9449207dbc3c43b49b9c315c86e8329ad519da Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 18 May 2017 16:51:28 +0530
Subject: [PATCH] Start work on testing restrictions

---
 src/calibre/srv/ajax.py       |  5 ++++-
 src/calibre/srv/handler.py    |  5 +++++
 src/calibre/srv/tests/ajax.py | 38 +++++++++++++++++++++++++++++++++--
 src/calibre/srv/tests/base.py |  1 +
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/calibre/srv/ajax.py b/src/calibre/srv/ajax.py
index 83ad6766e0..34ab439934 100644
--- a/src/calibre/srv/ajax.py
+++ b/src/calibre/srv/ajax.py
@@ -554,7 +554,10 @@ def search_result(ctx, rd, db, query, num, offset, sort, sort_order, vl=''):
 @endpoint('/ajax/search/{library_id=None}', postprocess=json)
 def search(ctx, rd, library_id):
     '''
-    Return the books (as list of ids) matching the specified search query.
+    Return the books matching the specified search query.
+    The returned object is a dict with the field book_ids which
+    is a list of matched book ids. For all the other fields in the object, see
+    :func:`search_result`.
 
     Optional: ?num=100&offset=0&sort=title&sort_order=asc&query=&vl=
     '''
diff --git a/src/calibre/srv/handler.py b/src/calibre/srv/handler.py
index 60b33569c7..e29a40d60d 100644
--- a/src/calibre/srv/handler.py
+++ b/src/calibre/srv/handler.py
@@ -137,6 +137,7 @@ class Context(object):
 
     def search(self, request_data, db, query, vl=''):
         restrict_to_ids = self.get_effective_book_ids(db, request_data, vl)
+        query = query or ''
         key = query, restrict_to_ids
         with self.lock:
             cache = self.library_broker.search_caches[db.server_library_id]
@@ -181,3 +182,7 @@ class Handler(object):
 
     def close(self):
         self.router.ctx.library_broker.close()
+
+    @property
+    def ctx(self):
+        return self.router.ctx
diff --git a/src/calibre/srv/tests/ajax.py b/src/calibre/srv/tests/ajax.py
index 732ac3c5d2..f81662a957 100644
--- a/src/calibre/srv/tests/ajax.py
+++ b/src/calibre/srv/tests/ajax.py
@@ -6,14 +6,17 @@ from __future__ import (unicode_literals, division, absolute_import,
 __license__ = 'GPL v3'
 __copyright__ = '2015, Kovid Goyal <kovid at kovidgoyal.net>'
 
-import httplib, zlib, json
+import httplib, zlib, json, base64, os
 from functools import partial
 from urllib import urlencode
+from httplib import OK, NOT_FOUND
 
 from calibre.srv.tests.base import LibraryBaseTest
 
 
-def make_request(conn, url, headers={}, prefix='/ajax'):
+def make_request(conn, url, headers={}, prefix='/ajax', username=None, password=None):
+    if username and password:
+        headers[b'Authorization'] = b'Basic ' + base64.standard_b64encode((username + ':' + password).encode('utf-8'))
     conn.request('GET', prefix + url, headers=headers)
     r = conn.getresponse()
     data = r.read()
@@ -78,3 +81,34 @@ class ContentTest(LibraryBaseTest):
             r, data = request('/search?' + urlencode({'query': 'tags:"=Tag One"', 'vl':'1'}))
             self.ae(set(data['book_ids']), {2})
     # }}}
+
+    def test_srv_restrictions(self):
+        ' Test that virtual lib. + search restriction works on all end points'
+        with self.create_server(auth=True, auth_mode='basic') as server:
+            db = server.handler.router.ctx.library_broker.get(None)
+            db.set_pref('virtual_libraries', {'1':'id:1', '12':'id:1 or id:2'})
+            db.set_field('tags', {1: ['present'], 3: ['missing']})
+            server.handler.ctx.user_manager.add_user('12', 'test', restriction={
+                'library_restrictions':{os.path.basename(db.backend.library_path): 'id:1 or id:2'}})
+            conn = server.connect()
+            url_for = server.handler.router.url_for
+
+            def r(path, status=OK):
+                r, data = make_request(conn, path, username='12', password='test', prefix='')
+                self.assertEqual(status, r.status)
+                return data
+            ok = r
+            nf = partial(r, status=NOT_FOUND)
+
+            ok(url_for('/ajax/book', book_id=1))
+            nf(url_for('/ajax/book', book_id=3))
+            data = ok(url_for('/ajax/books'))
+            self.assertIsNone(data['3'])
+            for i in '12':
+                self.assertIsNotNone(data[i])
+            self.assertEqual(set(r('/ajax/search')['book_ids']), {1,2})
+            self.assertEqual(set(r('/ajax/search?query=id:2')['book_ids']), {2})
+            self.assertEqual(set(r('/ajax/search?vl=1')['book_ids']), {1})
+
+            nf(url_for('/book-manifest', book_id=3, fmt='x'))
+            nf(url_for('/book-file', book_id=3, fmt='x', size=1, mtime=1, name='x'))
diff --git a/src/calibre/srv/tests/base.py b/src/calibre/srv/tests/base.py
index 85795aa2cd..ddc0f19aa6 100644
--- a/src/calibre/srv/tests/base.py
+++ b/src/calibre/srv/tests/base.py
@@ -96,6 +96,7 @@ class TestServer(Thread):
         kwargs['shutdown_timeout'] = kwargs.get('shutdown_timeout', 0.1)
         kwargs['listen_on'] = kwargs.get('listen_on', 'localhost')
         kwargs['port'] = kwargs.get('port', 0)
+        kwargs['userdb'] = kwargs.get('userdb', ':memory:')
 
     def run(self):
         try: