forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
"""Low level socket IO for docker API."""
|
|
import struct
|
|
import logging
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SOCK_BUF_SIZE = 1024
|
|
|
|
class DockerSocket:
|
|
"""Wrapper around docker API's socket object. Can be used as a context manager."""
|
|
|
|
_timeout: int = 5
|
|
|
|
|
|
def __init__(self, socket, timeout: int = _timeout):
|
|
self.socket = socket
|
|
self.socket._sock.settimeout(timeout)
|
|
# self.socket._sock.setblocking(False)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self.close()
|
|
|
|
def close(self):
|
|
logger.debug("closing socket...")
|
|
self.socket._sock.shutdown(2) # 2 = SHUT_RDWR
|
|
self.socket._sock.close()
|
|
self.socket.close()
|
|
|
|
def sendall(self, data: bytes) -> None:
|
|
self.socket._sock.sendall(data)
|
|
|
|
def setblocking(self, flag: bool) -> None:
|
|
self.socket._sock.setblocking(flag)
|
|
|
|
def recv(self) -> Any:
|
|
"""Wrapper for socket.recv that does buffured read."""
|
|
|
|
# NOTE: this is optional as a bonus
|
|
# TODO: Recv with TTY enabled
|
|
#
|
|
# When the TTY setting is enabled in POST /containers/create, the stream
|
|
# is not multiplexed. The data exchanged over the hijacked connection is
|
|
# simply the raw data from the process PTY and client's stdin.
|
|
|
|
# header := [8]byte{STREAM_TYPE, 0, 0, 0, SIZE1, SIZE2, SIZE3, SIZE4}
|
|
# STREAM_TYPE can be:
|
|
#
|
|
# 0: stdin (is written on stdout)
|
|
# 1: stdout
|
|
# 2: stderr
|
|
# SIZE1, SIZE2, SIZE3, SIZE4 are the four bytes of the uint32 size encoded as
|
|
# big endian.
|
|
#
|
|
# Following the header is the payload, which is the specified number of bytes of
|
|
# STREAM_TYPE.
|
|
#
|
|
# The simplest way to implement this protocol is the following:
|
|
#
|
|
# - Read 8 bytes.
|
|
# - Choose stdout or stderr depending on the first byte.
|
|
# - Extract the frame size from the last four bytes.
|
|
# - Read the extracted size and output it on the correct output.
|
|
# - Goto 1.
|
|
|
|
chunks = []
|
|
# try:
|
|
# self.socket._sock.recv(8)
|
|
# except BlockingIOError as e:
|
|
# raise ValueError("incomplete read from container output")
|
|
|
|
while True:
|
|
header = b''
|
|
try:
|
|
# strip the header
|
|
# the first recv is blocking to wait for the container to start
|
|
header = self.socket._sock.recv(8)
|
|
except BlockingIOError:
|
|
# logger.debug("[header] blocking IO")
|
|
break
|
|
|
|
self.socket._sock.setblocking(False)
|
|
|
|
if header == b'':
|
|
break
|
|
stream_type, size = struct.unpack("!BxxxI", header)
|
|
|
|
payload = b''
|
|
while size:
|
|
chunk = b''
|
|
try:
|
|
chunk = self.socket._sock.recv(min(size, SOCK_BUF_SIZE))
|
|
except BlockingIOError:
|
|
# logger.debug("[body] blocking IO")
|
|
break
|
|
if chunk == b'':
|
|
raise ValueError("incomplete read from container output")
|
|
payload += chunk
|
|
size -= len(chunk)
|
|
chunks.append((stream_type, payload))
|
|
# try:
|
|
# msg = self.socket._sock.recv(SOCK_BUF_SIZE)
|
|
# chunk += msg
|
|
# except BlockingIOError as e:
|
|
# break
|
|
|
|
return chunks
|