Repository: googleapis/google-auth-library-python

Issue:
google.oauth2.service_account.Credential objects are not pickleable if cryptography is installed
#### Environment details

  - OS: MacOS 13.5.2
  - Python version: 3.9
  - pip version: 23.1.2
  - `google-auth` version: 2.22.0

#### Steps to reproduce

Install `cryptography`

Run:
```
import pickle
from google.oauth2 import service_account

credentials = service_account.Credentials.from_service_account_file("credentials.json")

pickle.dumps(credentials)
```

The object in question seems to be `credentials._signer` which is either a pickleable pure Python object if `cryptography` is not installed, or else an unpickleable `google.auth.crypt._cryptography_rsa.RSASigner` if it is. Specifically, the signer._key object is of type `cryptography.hazmat.backends.openssl.rsa._RSAPrivateKey`.

This conversation on SO seems related: https://stackoverflow.com/questions/39321606/cant-pickle-an-rsa-key-to-send-over-a-socket

This is impacting the Storage SDK's multiprocessing capability: https://github.com/googleapis/python-storage/issues/1116 despite efforts to work around it.


Discussion:
I'm not sure how this should be implemented.. Any object that is pickled would need to reconstruct the signer object to work, which requires the key's PEM file. Put another way, it seems to me that the the object needs to be recreated each time.

@arithmetic1728 @sai-sunder-s what do you think?

The base OAuth2.0 credential object does support pickling https://github.com/googleapis/google-auth-library-python/blob/main/google/oauth2/credentials.py#L156.


Maybe we could store a buffer of the key's bytes, and then recreate the signer that way?
@clundin25 yes I think we could store the key bytes and recreate the signer

Resolving patch:
diff --git a/google/auth/crypt/_cryptography_rsa.py b/google/auth/crypt/_cryptography_rsa.py
--- a/google/auth/crypt/_cryptography_rsa.py
+++ b/google/auth/crypt/_cryptography_rsa.py
@@ -134,3 +134,18 @@ def from_string(cls, key, key_id=None):
             key, password=None, backend=_BACKEND
         )
         return cls(private_key, key_id=key_id)
+
+    def __getstate__(self):
+        """Pickle helper that serializes the _key attribute."""
+        state = self.__dict__.copy()
+        state["_key"] = self._key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        return state
+
+    def __setstate__(self, state):
+        """Pickle helper that deserializes the _key attribute."""
+        state["_key"] = serialization.load_pem_private_key(state["_key"], None)
+        self.__dict__.update(state)
diff --git a/google/auth/crypt/es256.py b/google/auth/crypt/es256.py
--- a/google/auth/crypt/es256.py
+++ b/google/auth/crypt/es256.py
@@ -158,3 +158,18 @@ def from_string(cls, key, key_id=None):
             key, password=None, backend=_BACKEND
         )
         return cls(private_key, key_id=key_id)
+
+    def __getstate__(self):
+        """Pickle helper that serializes the _key attribute."""
+        state = self.__dict__.copy()
+        state["_key"] = self._key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        return state
+
+    def __setstate__(self, state):
+        """Pickle helper that deserializes the _key attribute."""
+        state["_key"] = serialization.load_pem_private_key(state["_key"], None)
+        self.__dict__.update(state)


Test patch:
diff --git a/system_tests/secrets.tar.enc b/system_tests/secrets.tar.enc
Binary files a/system_tests/secrets.tar.enc and b/system_tests/secrets.tar.enc differ
diff --git a/tests/crypt/test__cryptography_rsa.py b/tests/crypt/test__cryptography_rsa.py
--- a/tests/crypt/test__cryptography_rsa.py
+++ b/tests/crypt/test__cryptography_rsa.py
@@ -14,6 +14,7 @@
 
 import json
 import os
+import pickle
 
 from cryptography.hazmat.primitives.asymmetric import rsa
 import pytest  # type: ignore
@@ -159,3 +160,17 @@ def test_from_service_account_file(self):
 
         assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
         assert isinstance(signer._key, rsa.RSAPrivateKey)
+
+    def test_pickle(self):
+        signer = _cryptography_rsa.RSASigner.from_service_account_file(
+            SERVICE_ACCOUNT_JSON_FILE
+        )
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, rsa.RSAPrivateKey)
+
+        pickled_signer = pickle.dumps(signer)
+        signer = pickle.loads(pickled_signer)
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, rsa.RSAPrivateKey)
diff --git a/tests/crypt/test_es256.py b/tests/crypt/test_es256.py
--- a/tests/crypt/test_es256.py
+++ b/tests/crypt/test_es256.py
@@ -15,6 +15,7 @@
 import base64
 import json
 import os
+import pickle
 
 from cryptography.hazmat.primitives.asymmetric import ec
 import pytest  # type: ignore
@@ -141,3 +142,15 @@ def test_from_service_account_file(self):
 
         assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
         assert isinstance(signer._key, ec.EllipticCurvePrivateKey)
+
+    def test_pickle(self):
+        signer = es256.ES256Signer.from_service_account_file(SERVICE_ACCOUNT_JSON_FILE)
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, ec.EllipticCurvePrivateKey)
+
+        pickled_signer = pickle.dumps(signer)
+        signer = pickle.loads(pickled_signer)
+
+        assert signer.key_id == SERVICE_ACCOUNT_INFO[base._JSON_FILE_PRIVATE_KEY_ID]
+        assert isinstance(signer._key, ec.EllipticCurvePrivateKey)