Merge pull request #19 from Climbgunks/issue-18

db2 query data > 32k fixes
nakagami · Feb 6, 2024 · c11db06 · c11db06
2 parents e547b53 + ebfd5fc
commit c11db06
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/drda/connection.py b/drda/connection.py
@@ -48,7 +48,19 @@ def _parse_response(self):
         more_data = False
         while True:
             while chained:
-                dss_type, chained, correlation_id, code_point, obj = ddm.read_dss(self.sock, self.db_type)
+                dss_type, chained, correlation_id, code_point, obj, more_data = ddm.read_dss(self.sock, self.db_type)
+                while more_data:
+                    # server is waiting for us to request more query data
+                    # may want to check code_point here
+                    ddm.write_request_dss(
+                        self.sock,
+                        ddm.packCNTQRY(
+                            self.pkgid, self.pkgcnstkn, self.pkgsn, self.database
+                        ),
+                        1, False, True
+                    )
+                    _X_dss_type, _X_chained, _X_correlation_id, _X_xcode_point, extra_obj, more_data = ddm.read_dss(self.sock,self.db_type)
+                    obj += extra_obj
                 if code_point == cp.SQLERRRM:
                     err_msg = ddm.parse_reply(obj).get(cp.SRVDGN)
                 elif code_point == cp.SQLCARD:
@@ -106,7 +118,7 @@ def _parse_accsecrd(self):
         secmec = sectkn = None
         chained = True
         while chained:
-            dss_type, chained, correlation_id, code_point, obj = ddm.read_dss(self.sock, self.db_type)
+            dss_type, chained, correlation_id, code_point, obj, more_data = ddm.read_dss(self.sock, self.db_type)
             if code_point == cp.ACCSECRD:
                 while len(obj):
                     ln = int.from_bytes(obj[:2], byteorder='big')

diff --git a/drda/ddm.py b/drda/ddm.py
@@ -33,7 +33,7 @@
 def _recv_from_sock(sock, nbytes, max_attempts=16):
     n = nbytes
     attempts = 0
-    received = b''
+    received = bytearray() # mutable
     while n > 0 and attempts < max_attempts:
         bs = sock.recv(n)
         if len(bs) > 0:
@@ -241,15 +241,20 @@ def read_dss(sock, db_type):
     correlation_id = int.from_bytes(b[4:6],  byteorder='big')
     obj_ln = int.from_bytes(_recv_from_sock(sock, 2), byteorder='big')
     code_point = int.from_bytes(_recv_from_sock(sock, 2), byteorder='big')
+    more_data = False
 
     if dss_ln == 0xFFFF:
         assert code_point == 0x241B     # QRYDTA
         if db_type == 'db2':
-            assert obj_ln == 32772      # ???
-            obj = _recv_from_sock(sock, 32757)   # ???
+            assert obj_ln == 32772      # 0x8004 protocol magic
+            obj = _recv_from_sock(sock, 32757)   # 0x7fff - 6 - 4
+            # !! assumes there is only 1 additional "page".. not sure what controls this
+            # !! worried it depends on QRYBLKSZ (which is 65535 below)
             next_ln = int.from_bytes(_recv_from_sock(sock, 2), byteorder='big')
             extra = _recv_from_sock(sock, next_ln-2)
             obj += extra
+            if next_ln == 0x7ffe:
+                more_data = True
         elif db_type == 'derby':
             assert obj_ln == 32776      # ???
             ln = int.from_bytes(_recv_from_sock(sock, 4), byteorder='big')
@@ -264,7 +269,7 @@ def read_dss(sock, db_type):
             raise ConnectionError("invalid DSS packet from socket")
         assert len(obj) == (obj_ln - 4)
 
-    return dss_type, chained, correlation_id, code_point, obj
+    return dss_type, chained, correlation_id, code_point, obj, more_data
 
 
 def write_request_dss(sock, o, cur_id, next_dss_has_same_id, last_packet):