improved sample using jwt

Author: yorek

app.py

@@ -8,23 +8,47 @@
 from threading import Lock
 from tenacity import *
 import logging
+import jwt
 
 # Initialize Flask
 app = Flask(__name__)
 
 # Setup Flask Restful framework
 api = Api(app)
 parser = reqparse.RequestParser()
-parser.add_argument('X-UserName', location='headers')
+parser.add_argument('Authorization', location='headers')
 
 # Set connection string
 application_name = ";APP={0}".format(socket.gethostname())  
 connection_string = os.environ['SQLAZURECONNSTR_RLS'] + application_name
 
 class Queryable(Resource):
-    def get(self):  
-        args = parser.parse_args()
-        result = self.executeQueryJson("get", args["X-UserName"])   
+    def __authorize(self):
+        encoded = ""
+        user_hash_id = 0
+        
+        request_args = parser.parse_args()
+        authorization = request_args["Authorization"]
+        tokens = authorization.split()
+
+        if tokens[0].lower() == "bearer":
+            encoded = tokens[1]
+        else:
+            raise Exception("Wrong authorization schema")
+        
+        try:
+            secure_payload = jwt.decode(encoded, 'mySUPERs3cr3t', algorithms=['HS256'])        
+            user_hash_id = int(secure_payload["user-hash-id"])
+        except jwt.InvalidSignatureError:
+            user_hash_id = 0
+        except:
+            raise
+
+        return user_hash_id
+
+    def get(self):      
+        user_hash_id = self.__authorize()
+        result = self.executeQueryJson("get", user_hash_id)   
         return result, 200
 
     @retry(stop=stop_after_attempt(3), wait=wait_fixed(10), retry=retry_if_exception_type(pyodbc.OperationalError), after=after_log(app.logger, logging.DEBUG))
@@ -39,11 +63,13 @@ def executeQueryJson(self, verb, username, payload=None):
             cursor = conn.cursor()
 
             # set session context info, used by Row-Level Security
-            cursor.execute(f"EXEC sys.sp_set_session_context @key=N'username', @value=?, @read_only=1;", username)
+            cursor.execute(f"EXEC sys.sp_set_session_context @key=N'user-hash-id', @value=?, @read_only=1;", username)                    
 
             if payload:
+                print("EXEC %s %s" % (procedure, json.dumps(payload)))
                 cursor.execute(f"EXEC {procedure} ?", json.dumps(payload))
             else:
+                print("EXEC %s" % procedure)
                 cursor.execute(f"EXEC {procedure}")
 
             result = cursor.fetchone()

requirements.txt

@@ -2,4 +2,5 @@ pyodbc
 flask
 flask-restful
 tenacity
-python-dotenv
+python-dotenv
+PyJWT

sql/SetupRLS.sql renamed to sql/00-SetupRLS.sql

@@ -7,6 +7,13 @@ drop function if exists rls.fn_SecurityPredicate;
 drop table if exists rls.SensitiveDataPermissions;
 drop table if exists dbo.EvenMoreSensitiveData;
 drop table if exists dbo.SensitiveData;
+drop procedure if exists web.get_sensitivedata
+drop procedure if exists web.get_evenmoresensitivedata
+go
+
+drop schema if exists web;
+go
+create schema web;
 go
 
 drop schema if exists rls;
@@ -38,13 +45,16 @@ go
 
 create table rls.SensitiveDataPermissions
 (
-    UserName sysname not null,
+    UserHashId bigint not null,
     SensitiveDataId int not null foreign key references dbo.SensitiveData(Id),
     HasAccess bit not null default(0),
-    constraint pk__rls_SensitiveDataPermissions primary key clustered (UserName, SensitiveDataId)
+    constraint pk__rls_SensitiveDataPermissions primary key nonclustered (UserHashId, SensitiveDataId)
 )
 go
 
+create clustered index ixc on rls.SensitiveDataPermissions (SensitiveDataId)
+go
+
 insert into dbo.SensitiveData values
 (1, 'Davide', 'Mauri', '{"SuperPowers":"Fly"}'),
 (2, 'John', 'Doe', '{"SuperPowers":"Laser Eyes"}')
@@ -59,8 +69,8 @@ insert into dbo.EvenMoreSensitiveData values
 go
 
 insert into rls.SensitiveDataPermissions values
-('damauri', 1, 1),
-('jdoe', 2, 1)
+(-6134311, 1, 1),
+(1225328053, 2, 1)
 go
 
 create function rls.fn_securitypredicate(@SensitiveDataId int)  
@@ -73,32 +83,13 @@ select
 from
     rls.SensitiveDataPermissions
 where
-    database_principal_id() = database_principal_id('MiddleTierUser')
+    (database_principal_id() = database_principal_id('MiddleTierUser') or is_member('db_owner') = 1)
 and     
-    UserName = session_context(N'username')
+    UserHashId = cast(session_context(N'user-hash-id') as bigint)
 and
     SensitiveDataId = @SensitiveDataId
 go
 
-select * from dbo.SensitiveData
-go
-
-create security policy rls.SensitiveDataPolicy
-add filter predicate rls.fn_SecurityPredicate(Id) on dbo.SensitiveData,
-add filter predicate rls.fn_SecurityPredicate(SensitiveDataId) on dbo.EvenMoreSensitiveData
-with (state = on);  
-
-select * from dbo.SensitiveData
-select * from dbo.EvenMoreSensitiveData
-go
-
-exec sys.sp_set_session_context @key = N'username', @value = 'damauri', @read_only = 0;  
-go
-
-exec sys.sp_set_session_context @key = N'username', @value = 'jdoe', @read_only = 0;  
-go
-
-
 create or alter procedure web.get_sensitivedata
 as
 select
@@ -136,3 +127,9 @@ from
 for
 	json path
 go
+
+create security policy rls.SensitiveDataPolicy
+add filter predicate rls.fn_SecurityPredicate(Id) on dbo.SensitiveData,
+add filter predicate rls.fn_SecurityPredicate(SensitiveDataId) on dbo.EvenMoreSensitiveData
+with (state = off);  
+

sql/01-Test.sql

@@ -0,0 +1,32 @@
+select * from dbo.SensitiveData
+select * from dbo.EvenMoreSensitiveData
+go
+
+alter security policy rls.SensitiveDataPolicy
+with (state = on)
+go
+
+select * from dbo.SensitiveData
+select * from dbo.EvenMoreSensitiveData
+go
+
+exec sys.sp_set_session_context @key = N'user-hash-id', @value = -6134311, @read_only = 0;  
+go
+
+select * from dbo.SensitiveData
+select * from dbo.EvenMoreSensitiveData
+go
+
+exec sys.sp_set_session_context @key = N'user-hash-id', @value = 1225328053, @read_only = 0;  
+go
+
+select * from dbo.SensitiveData
+select * from dbo.EvenMoreSensitiveData
+go
+
+exec sys.sp_set_session_context @key = 'user-hash-id', @value = 0, @read_only = 0;  
+go
+
+select * from dbo.SensitiveData
+select * from dbo.EvenMoreSensitiveData
+go