1

from qmsk.net.socket.socket cimport *

     2

from qmsk.net.socket._address cimport sockaddr

     3

from qmsk.net.socket.address cimport build_sockaddr

     4

     5

cimport qmsk.net.socket.platform as platform

     6

cimport qmsk.net.libc as libc, qmsk.net.py as py

     7

     8

from qmsk.net.py cimport raise_errno

     9

    10

cdef parse_sockaddr (platform.sockaddr **sa_ptr, platform.socklen_t *sa_len, sockaddr addr, int optional = 0) :

    11

    if addr is not None :

    12

        addr._get_sockaddr(sa_ptr, sa_len)

    13

    14

    elif optional :

    15

        sa_ptr[0] = NULL

    16

        sa_len[0] = 0

    17

    18

    else :

    19

        raise ValueError(addr)

    20

    21

cdef parse_buf (void **buf_ptr, size_t *buf_len, object buf, int optional = 0) :

    22

    cdef libc.ssize_t tmp_len

    23

    24

    if buf is not None :

    25

        # XXX: test that except works right

    26

        # XXX: this complains about const...

    27

        py.PyObject_AsCharBuffer(buf, <char **> buf_ptr, &tmp_len)

    28

    29

        # XXX: ensure that this is >= 0

    30

        buf_len[0] = tmp_len

    31

    32

    elif optional :

    33

        buf_ptr[0] = NULL

    34

        buf_len[0] = 0

    35

    36

    else :

    37

        raise ValueError(buf)

    38

    39

## helper for socket.recv*

    40

# XXX: make sure these don't leak the PyString in case of errors...

    41

cdef struct sockbuf :

    42

    py.PyObject *str

    43

    44

cdef char* sockbuf_init (sockbuf *buf, size_t len) :

    45

    """

    46

        Initialize the sockbuf to contain a PyString that can hold `len` bytes, and return a pointer into its

    47

        contents as a char*.

    48

    49

        Note that this requires use of a try-finally with sockbuf_deinit...

    50

    """

    51

    52

    buf.str = py.PyString_FromStringAndSize(NULL, len)

    53

    54

    return py.PyString_AS_STRING(buf.str)

    55

    56

cdef object sockbuf_truncate (sockbuf *buf, size_t len) :

    57

    """

    58

        Truncate the given sockbuf's PyString to the given length, and return the PyObject*

    59

    """

    60

    61

    # optimize for the no-need-to-resize case

    62

    # this also fixes behaviour for zero-length strings (heh), since they're interned and can't be resized

    63

    if len != py.PyString_GET_SIZE(buf.str) :

    64

        py._PyString_Resize(&buf.str, len)

    65

    66

    return <object> buf.str

    67

    68

cdef void sockbuf_deinit (sockbuf *buf) :

    69

    """

    70

        Release the PyObject.

    71

    72

        This is safe if the sockbuf was initialized to zero.

    73

    """

    74

    75

    py.Py_XDECREF(buf.str)

    76

    77

# XXX: do some GIL-releasin'

    78

cdef class socket :

    79

    80

    def __cinit__ (self) :

    81

        """

    82

            Initialize the socket to set fd to -1, so that we dont't try and close stdin too often :)

    83

        """

    84

    85

        self.fd = -1

    86

    87

    def __init__ (self, int family = platform.AF_INET, int socktype = platform.SOCK_STREAM, int protocol = 0, int fd = -1) :

    88

        """

    89

            Create a new socket endpoint with the given family/domain, socktype and optionally, specific protocol,

    90

            unless the fd argument is given as >= 0, in which case it used directly.

    91

    92

                family      - one of AF_*

    93

                socktype    - one of SOCK_*

    94

                protocol    - one of IPPROTO_* or zero to select default

    95

        """

    96

    97

        if fd >= 0 :

    98

            # given fd

    99

            self.fd = fd

   100

   101

        else :

   102

            # socket()

   103

            self.fd = platform.socket(family, socktype, protocol)

   104

   105

        # trap

   106

        if self.fd < 0 :

   107

            raise_errno('socket')

   108

   109

    def fileno (self) :

   110

        """

   111

            Returns the OS-level file desriptor for this socket as an integer

   112

        """

   113

   114

        return self.fd

   115

   116

#    # XXX: does this *really* belong here?

   117

#    def setblocking (self, bint blocking = True) :

   118

#        """

   119

#            Control the OS-level nonblocking-IO mode flag for this socket.

   120

#

   121

#                blocking    - True for normal blocking operation, False to use non-blocking operation

   122

#        """

   123

#        

   124

#        # fcntl magic

   125

#        libc.fcntl_set_flag(self.fd, libc.O_NONBLOCK, not blocking)

   126

   127

    def bind (self, sockaddr addr) :

   128

        """

   129

            Bind this socket to the given local socket address. The given sockaddr should be of the same or a

   130

            compatible address family.

   131

   132

                addr        - the local address to bind to. The port may be zero to let the system choose an unused

   133

                              ephemeral port.

   134

        """

   135

   136

        cdef platform.sockaddr *sa

   137

        cdef platform.socklen_t sa_len

   138

   139

        # XXX: require non-NULL addr?

   140

        parse_sockaddr(&sa, &sa_len, addr, 1)

   141

   142

        # bind()

   143

        if platform.bind(self.fd, sa, sa_len) :

   144

            raise_errno('bind')

   145

   146

    def listen (self, int backlog) :

   147

        """

   148

            Listen for connections, marking this socket as a passive socket, which can accept incoming connection

   149

            requests using sock.accept().

   150

   151

            It is customary to call .bind() before .listen().

   152

   153

                backlog     - maximum number of pending connections (those not yet .accept()'d).

   154

        """

   155

   156

        # listen()

   157

        if platform.listen(self.fd, backlog) :

   158

            raise_errno('listen')

   159

   160

    def connect (self, sockaddr addr) :

   161

        """

   162

            Initiate a connection, connecting this socket to the remote endpoint specified by `addr`. The given sockaddr

   163

            should be of the same or a compatible address family.

   164

   165

            If the socket is in non-blocking mode, this will presumeably return errno.EINPROGRESS.

   166

   167

            If the socket has not yet been bound (using .bind()), the system will pick an appropriate local address and

   168

            ephemeral port.

   169

   170

                addr        - the remote address to connect to.

   171

        """

   172

   173

        cdef platform.sockaddr *sa

   174

        cdef platform.socklen_t sa_len

   175

   176

        # XXX: require non-NULL addr?

   177

        parse_sockaddr(&sa, &sa_len, addr, 1)

   178

   179

        # connect()

   180

        if platform.connect(self.fd, sa, sa_len) :

   181

            raise_errno('connect')

   182

   183

    def accept (self) :

   184

        """

   185

            Accept a connection, dequeueing the first pending connection and returning a new sock object for it. This

   186

            socket must be a connection-based socket (SOCK_STREAM/SOCK_SEQPACKET) and in the passive listening mode

   187

            (.listen()).

   188

   189

            This returns a (sock, src_addr) tuple:

   190

                sock        - the newly created sock, corresponding to the incoming connection

   191

                src_addr    - the remote address of the incoming connection

   192

        """

   193

   194

        # prep the sockaddr that we will return

   195

        cdef platform.sockaddr_storage ss

   196

        cdef platform.socklen_t ss_len = sizeof(ss)

   197

   198

        # accept()

   199

        cdef socket_t sock_fd = platform.accept(self.fd, <platform.sockaddr *> &ss, &ss_len)

   200

   201

        if sock_fd < 0 :

   202

            raise_errno('accept')

   203

   204

        try :

   205

            # prep the new socket

   206

            sock_obj = socket(fd=sock_fd)

   207

   208

        except :

   209

            # XXX: don't leak the socket fd? How does socket.__init__ handle this?

   210

            platform.close(sock_fd)

   211

   212

            raise

   213

   214

        # prep the new addr

   215

        cdef sockaddr src_addr = build_sockaddr(<platform.sockaddr *> &ss, ss_len)

   216

   217

        return sock_obj, src_addr

   218

   219

    def getsockname (self) :

   220

        """

   221

            Get the local address this socket is currently bound to. This can be set using bind(), or automatically.

   222

   223

            Returns a sockaddr object.

   224

        """

   225

   226

        # prep the sockaddr that we will return

   227

        cdef platform.sockaddr_storage ss

   228

        cdef platform.socklen_t ss_len = sizeof(ss)

   229

   230

        # getsockname()

   231

        if platform.getsockname(self.fd, <platform.sockaddr *> &ss, &ss_len) :

   232

            raise_errno('getsockname')

   233

   234

        # build the new sockaddr

   235

        return build_sockaddr(<platform.sockaddr *> &ss, ss_len)

   236

   237

    def getpeername (self) :

   238

        """

   239

            Get the remote address this socket is currently connected to.

   240

   241

            Returns a sockaddr object.

   242

        """

   243

   244

        # prep the sockaddr that we will return

   245

        cdef platform.sockaddr_storage ss

   246

        cdef platform.socklen_t ss_len = sizeof(ss)

   247

   248

        # getpeername()

   249

        if platform.getpeername(self.fd, <platform.sockaddr *> &ss, &ss_len) :

   250

            raise_errno('getpeername')

   251

   252

        # build the new sockaddr

   253

        return build_sockaddr(<platform.sockaddr *> &ss, ss_len)

   272

   273

    def send (self, object buf, int flags = 0) :

   274

        """

   275

            Transmit a message to the connected remote endpoint.

   276

   277

                buf         - the data to send

   278

                flags       - (optional) MSG_* flags to send with

   279

   280

            Returns the number of bytes sent, which may be less than the length of buf.

   281

        """

   282

   283

        cdef void *buf_ptr

   284

        cdef size_t buf_len

   285

        cdef libc.ssize_t ret

   286

   287

        parse_buf(&buf_ptr, &buf_len, buf, 0)

   288

   289

        # send()

   290

        ret = platform.send(self.fd, buf_ptr, buf_len, flags)

   291

   292

        if ret < 0 :

   293

            raise_errno('send')

   294

   295

        else :

   296

            return ret

   297

   298

    def sendto (self, object buf, int flags = 0, sockaddr addr = None) :

   299

        """

   300

            Transmit a message to the given remote endpoint. If this socket is connected, the addr must not be

   301

            specified, and this acts like send()

   302

   303

                buf         - the data to send

   304

                flags       - (optional) MSG_* flags to send with

   305

                addr        - (optional) target address

   306

   307

            Returns the number of bytes sent, which may be less than the length of buf.

   308

        """

   309

   310

        cdef void *buf_ptr

   311

        cdef size_t buf_len

   312

        cdef libc.ssize_t ret

   313

   314

        cdef platform.sockaddr *sa

   315

        cdef platform.socklen_t sa_len

   316

   317

        parse_sockaddr(&sa, &sa_len, addr, 1)

   318

        parse_buf(&buf_ptr, &buf_len, buf, 0)

   319

   320

        # send()

   321

        ret = platform.sendto(self.fd, buf_ptr, buf_len, flags, sa, sa_len)

   322

   323

        if ret < 0 :

   324

            raise_errno('sendto')

   325

   326

        else :

   327

            return ret

   328

   329

    def sendmsg (self, sockaddr addr = None, iov = None, control = None, int flags = 0) :

   330

        """

   331

            Transmit an extended message to the given remote endpoint (or default for connected sockets) with the given

   332

            extra parameters.

   333

   334

                addr        - (optional) destination address (struct msghdr::msg_name)

   335

                iov         - (optional) sequence of read-buffers to transmit

   336

                control     - (optional) control message to transmit

   337

                flags       - (optional) MSG_* flags to send with

   338

   339

            Returns the number of bytes sent, which may be less than the total length of iov.

   340

        """

   341

   342

        cdef libc.ssize_t ret

   343

        cdef libc.iovec *iovec

   344

        cdef platform.msghdr msg

   345

   346

        libc.memset(&msg, 0, sizeof(msg))

   347

   348

        parse_sockaddr(<platform.sockaddr **> &msg.msg_name, &msg.msg_namelen, addr, 1)

   349

        parse_buf(&msg.msg_control, &msg.msg_controllen, control, 1)

   350

   351

        # iov

   352

        if iov :

   353

            iov = tuple(iov)

   354

   355

            # numerb of bufs = number of iovecs

   356

            msg.msg_iovlen = len(iov)

   357

   358

            # alloca the required number of iovec's

   359

            msg.msg_iov = <libc.iovec *> libc.alloca(msg.msg_iovlen * sizeof(libc.iovec))

   360

   361

            # fill in the iovecs

   362

            for i, buf in enumerate(iov) :

   363

                iovec = &msg.msg_iov[i]

   364

   365

                parse_buf(&iovec.iov_base, &iovec.iov_len, buf, 1)

   366

   367

        # sendmsg()

   368

        ret = platform.sendmsg(self.fd, &msg, flags)

   369

   370

        if ret < 0 :

   371

            raise_errno('sendmsg')

   372

   373

        else :

   374

            return ret

   375

   376

    def write (self, object buf) :

   377

        """

   378

            Write data to socket, mostly equivalent to send() with flags=0.

   379

   380

                buf         - the data to send

   381

   382

            Returns the number of bytes sent, which may be less than the length of buf.

   383

        """

   384

   385

        cdef void *buf_ptr

   386

        cdef size_t buf_len

   387

        cdef libc.ssize_t ret

   388

   389

        parse_buf(&buf_ptr, &buf_len, buf, 0)

   390

   391

        # send()

   392

        ret = libc.write(self.fd, buf_ptr, buf_len)

   393

   394

        if ret < 0 :

   395

            raise_errno('write')

   396

   397

        else :

   398

            return ret

   399

   400

    def writev (self, iov) :

   401

        """

   402

            Write data to a socket from multiple read-buffers.

   403

   404

                iov         - sequence of read-buffers to transmit

   405

   406

            Returns the number of bytes sent, which may be less than the total length of iov.

   407

        """

   408

   409

        # iov

   410

        cdef libc.iovec *iov_list = NULL

   411

        cdef size_t iov_count = 0

   412

        cdef libc.iovec *iovec

   413

   414

        iov = tuple(iov)

   415

   416

        # numerb of bufs = number of iovecs

   417

        iov_count = len(iov)

   418

   419

        # alloca the required number of iovec's

   420

        iov_list = <libc.iovec *> libc.alloca(iov_count * sizeof(libc.iovec))

   421

   422

        # fill in the iovecs

   423

        for i, buf in enumerate(iov) :

   424

            iovec = &iov_list[i]

   425

   426

            parse_buf(&iovec.iov_base, &iovec.iov_len, buf, 1)

   427

   428

        # sendmsg()

   429

        ret = libc.writev(self.fd, iov_list, iov_count)

   430

   431

        if ret < 0 :

   432

            raise_errno('writev')

   433

   434

        else :

   435

            return ret

   436

   437

    def recv (self, size_t len, int flags = 0) :

   438

        """

   439

            Recieve a message, reading and returning at most `len` bytes.

   440

   441

                len         - size of buffer to use for recv

   442

                flags       - (optional) MSG_* flags to use for recv()

   443

   444

            Returns the recieved data as a newly allocated string of the correct length.

   445

        """

   446

   447

        cdef sockbuf sb

   448

        cdef libc.ssize_t ret

   449

   450

        # alloc the recv buffer

   451

        cdef char *buf = sockbuf_init(&sb, len)

   452

   453

        try :

   454

            # recv()

   455

            ret = platform.recv(self.fd, buf, len, flags)

   456

   457

            if ret < 0 :

   458

                raise_errno('recv')

   459

   460

            # truncate to correct length

   461

            return sockbuf_truncate(&sb, ret)

   462

   463

        finally :

   464

            sockbuf_deinit(&sb)

   465

   466

    def recvfrom (self, size_t len, int flags = 0) :

   467

        """

   468

            Recieve a message, reading at most `len` bytes, also returning the source address.

   469

   470

                len         - size of buffer to use for recv

   471

                flags       - (optional) MSG_* flags to use for recvfrom()

   472

   473

            Returns the recieved data and the source address as a (buf, src_addr) tuple :

   474

                buf         - a newly allocated string containing the recieved data, of the correct length

   475

                src_addr    - the source address the message was recieved from

   476

        """

   477

   478

        cdef sockbuf sb

   479

        cdef libc.ssize_t ret

   480

        cdef object str

   481

        cdef sockaddr src_addr

   482

   483

        # prep the sockaddr that we will return

   484

        cdef platform.sockaddr_storage ss

   485

        cdef platform.socklen_t ss_len = sizeof(ss)

   486

   487

        # alloc recv buf

   488

        cdef char *buf = sockbuf_init(&sb, len)

   489

   490

        try :

   491

            # recvfrom()

   492

            ret = platform.recvfrom(self.fd, buf, len, flags, <platform.sockaddr *> &ss, &ss_len)

   493

   494

            if ret < 0 :

   495

                raise_errno('recv')

   496

   497

            # truncate

   498