From fa499a6a53e328b6693ac5f354f1ad2660ff8dc2 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Mon, 5 Sep 2022 12:35:06 +0100
Subject: [PATCH] Improve verification of proto parsing across Cdm, RemoteCdm
 and Device

This ensures that a partially parsing input (because of optional flags in the proto) does not get past any verification checks.

This prevents issues like an invalid License Challenging from getting an exception later down the line, as well as possibility of it also passing that check by pure luck, resulting in hard to debug issues.
---
 pywidevine/cdm.py       |  2 ++
 pywidevine/device.py    | 33 +++++++++++++++++++++++++++------
 pywidevine/remotecdm.py |  7 ++++++-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/pywidevine/cdm.py b/pywidevine/cdm.py
index d3aa9f8..7c6696e 100644
--- a/pywidevine/cdm.py
+++ b/pywidevine/cdm.py
@@ -360,6 +360,8 @@ class Cdm:
             signed_message = SignedMessage()
             try:
                 signed_message.ParseFromString(license_message)
+                if signed_message.SerializeToString() != license_message:
+                    raise DecodeError(license_message)
             except DecodeError as e:
                 raise InvalidLicenseMessage(f"Could not parse license_message as a SignedMessage, {e}")
             license_message = signed_message
diff --git a/pywidevine/device.py b/pywidevine/device.py
index 4e29351..6a02cbd 100644
--- a/pywidevine/device.py
+++ b/pywidevine/device.py
@@ -110,20 +110,37 @@ class Device:
         self.client_id = ClientIdentification()
         try:
             self.client_id.ParseFromString(client_id)
-        except DecodeError:
-            raise ValueError("Failed to parse client_id as a ClientIdentification")
+            if self.client_id.SerializeToString() != client_id:
+                raise DecodeError("partial parse")
+        except DecodeError as e:
+            raise ValueError(f"Failed to parse client_id as a ClientIdentification, {e}")
 
         self.vmp = FileHashes()
         if self.client_id.vmp_data:
             try:
                 self.vmp.ParseFromString(self.client_id.vmp_data)
-            except DecodeError:
-                raise ValueError("Failed to parse Client ID's VMP data as a FileHashes")
+                if self.vmp.SerializeToString() != self.client_id.vmp_data:
+                    raise DecodeError("partial parse")
+            except DecodeError as e:
+                raise ValueError(f"Failed to parse Client ID's VMP data as a FileHashes, {e}")
 
         signed_drm_certificate = SignedDrmCertificate()
-        signed_drm_certificate.ParseFromString(self.client_id.token)
         drm_certificate = DrmCertificate()
-        drm_certificate.ParseFromString(signed_drm_certificate.drm_certificate)
+
+        try:
+            signed_drm_certificate.ParseFromString(self.client_id.token)
+            if signed_drm_certificate.SerializeToString() != self.client_id.token:
+                raise DecodeError("partial parse")
+        except DecodeError as e:
+            raise DecodeError(f"Failed to parse the Signed DRM Certificate of the Client ID, {e}")
+
+        try:
+            drm_certificate.ParseFromString(signed_drm_certificate.drm_certificate)
+            if drm_certificate.SerializeToString() != signed_drm_certificate.drm_certificate:
+                raise DecodeError("partial parse")
+        except DecodeError as e:
+            raise DecodeError(f"Failed to parse the DRM Certificate of the Client ID, {e}")
+
         self.system_id = drm_certificate.system_id
 
     def __repr__(self) -> str:
@@ -190,6 +207,8 @@ class Device:
             if data.vmp:
                 try:
                     vmp.ParseFromString(data.vmp)
+                    if vmp.SerializeToString() != data.vmp:
+                        raise DecodeError("partial parse")
                 except DecodeError as e:
                     raise DecodeError(f"Failed to parse VMP data as FileHashes, {e}")
                 data.vmp = vmp
@@ -197,6 +216,8 @@ class Device:
                 client_id = ClientIdentification()
                 try:
                     client_id.ParseFromString(data.client_id)
+                    if client_id.SerializeToString() != data.client_id:
+                        raise DecodeError("partial parse")
                 except DecodeError as e:
                     raise DecodeError(f"Failed to parse VMP data as FileHashes, {e}")
 
diff --git a/pywidevine/remotecdm.py b/pywidevine/remotecdm.py
index 4ff602f..6d9e743 100644
--- a/pywidevine/remotecdm.py
+++ b/pywidevine/remotecdm.py
@@ -175,8 +175,11 @@ class RemoteCdm(Cdm):
         r = r["data"]
 
         try:
+            challenge = base64.b64decode(r["challenge_b64"])
             license_message = SignedMessage()
-            license_message.ParseFromString(base64.b64decode(r["challenge_b64"]))
+            license_message.ParseFromString(challenge)
+            if license_message.SerializeToString() != challenge:
+                raise DecodeError("partial parse")
         except DecodeError as e:
             raise InvalidLicenseMessage(f"Failed to parse license request, {e}")
 
@@ -196,6 +199,8 @@ class RemoteCdm(Cdm):
             signed_message = SignedMessage()
             try:
                 signed_message.ParseFromString(license_message)
+                if signed_message.SerializeToString() != license_message:
+                    raise DecodeError("partial parse")
             except DecodeError as e:
                 raise InvalidLicenseMessage(f"Could not parse license_message as a SignedMessage, {e}")
             license_message = signed_message