1

from qmsk.net.socket.socket cimport *

     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)

   254

   255

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

   256

        """

   257

            Transmit a message to the connected remote endpoint.

   258

   259

                buf         - the data to send

   260

                flags       - (optional) MSG_* flags to send with

   261

   262

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

   263

        """

   264

   265

        cdef void *buf_ptr

   266

        cdef size_t buf_len

   267

        cdef libc.ssize_t ret

   268

   269

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

   270

   271

        # send()

   272

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

   273

   274

        if ret < 0 :

   275

            raise_errno('send')

   276

   277

        else :

   278

            return ret

   279

   280

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

   281

        """

   282

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

   283

            specified, and this acts like send()

   284

   285

                buf         - the data to send

   286

                flags       - (optional) MSG_* flags to send with

   287

                addr        - (optional) target address

   288

   289

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

   290

        """

   291

   292

        cdef void *buf_ptr

   293

        cdef size_t buf_len

   294

        cdef libc.ssize_t ret

   295

   296

        cdef platform.sockaddr *sa

   297

        cdef platform.socklen_t sa_len

   298

   299

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

   300

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

   301

   302

        # send()

   303

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

   304

   305

        if ret < 0 :

   306

            raise_errno('sendto')

   307

   308

        else :

   309

            return ret

   310

   311

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

   312

        """

   313

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

   314

            extra parameters.

   315

   316

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

   317

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

   318

                control     - (optional) control message to transmit

   319

                flags       - (optional) MSG_* flags to send with

   320

   321

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

   322

        """

   323

   324

        cdef libc.ssize_t ret

   325

        cdef libc.iovec *iovec

   326

        cdef platform.msghdr msg

   327

   328

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

   329

   330

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

   331

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

   332

   333

        # iov

   334

        if iov :

   335

            iov = tuple(iov)

   336

   337

            # numerb of bufs = number of iovecs

   338

            msg.msg_iovlen = len(iov)

   339

   340

            # alloca the required number of iovec's

   341

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

   342

   343

            # fill in the iovecs

   344

            for i, buf in enumerate(iov) :

   345

                iovec = &msg.msg_iov[i]

   346

   347

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

   348

   349

        # sendmsg()

   350

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

   351

   352

        if ret < 0 :

   353

            raise_errno('sendmsg')

   354

   355

        else :

   356

            return ret

   357

   358

    def write (self, object buf) :

   359

        """

   360

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

   361

   362

                buf         - the data to send

   363

   364

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

   365

        """

   366

   367

        cdef void *buf_ptr

   368

        cdef size_t buf_len

   369

        cdef libc.ssize_t ret

   370

   371

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

   372

   373

        # send()

   374

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

   375

   376

        if ret < 0 :

   377

            raise_errno('write')

   378

   379

        else :

   380

            return ret

   381

   382

    def writev (self, iov) :

   383

        """

   384

            Write data to a socket from multiple read-buffers.

   385

   386

                iov         - sequence of read-buffers to transmit

   387

   388

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

   389

        """

   390

   391

        # iov

   392

        cdef libc.iovec *iov_list = NULL

   393

        cdef size_t iov_count = 0

   394

        cdef libc.iovec *iovec

   395

   396

        iov = tuple(iov)

   397

   398

        # numerb of bufs = number of iovecs

   399

        iov_count = len(iov)

   400

   401

        # alloca the required number of iovec's

   402

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

   403

   404

        # fill in the iovecs

   405

        for i, buf in enumerate(iov) :

   406

            iovec = &iov_list[i]

   407

   408

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

   409

   410

        # sendmsg()

   411

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

   412

   413

        if ret < 0 :

   414

            raise_errno('writev')

   415

   416

        else :

   417

            return ret

   418

   419

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

   420

        """

   421

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

   422

   423

                len         - size of buffer to use for recv

   424

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

   425

   426

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

   427

        """

   428

   429

        cdef sockbuf sb

   430

        cdef libc.ssize_t ret

   431

   432

        # alloc the recv buffer

   433

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

   434

   435

        try :

   436

            # recv()

   437

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

   438

   439

            if ret < 0 :

   440

                raise_errno('recv')

   441

   442

            # truncate to correct length

   443

            return sockbuf_truncate(&sb, ret)

   444

   445

        finally :

   446

            sockbuf_deinit(&sb)

   447

   448

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

   449

        """

   450

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

   451

   452

                len         - size of buffer to use for recv

   453

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

   454

   455

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

   456

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

   457

                src_addr    - the source address the message was recieved from

   458

        """

   459

   460

        cdef sockbuf sb

   461

        cdef libc.ssize_t ret

   462

        cdef object str

   463

        cdef sockaddr src_addr

   464

   465

        # prep the sockaddr that we will return

   466

        cdef platform.sockaddr_storage ss

   467

        cdef platform.socklen_t ss_len = sizeof(ss)

   468

   469

        # alloc recv buf

   470

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

   471

   472

        try :

   473

            # recvfrom()

   474

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

   475

   476

            if ret < 0 :

   477

                raise_errno('recv')

   478

   479

            # truncate

   480

            str = sockbuf_truncate(&sb, ret)

   481

   482

            # prep the new addr

   483

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

   484

   485

            return str, src_addr

   486

   487

        finally :

   488

            sockbuf_deinit(&sb)

   489

   490

    def recvmsg (self, bint recv_addr = True, object iov_lens = None, size_t control_len = 0, int flags = 0) :

   491

        """

   492

            Recieve a message along with some extra data.

   493

   494

                recv_addr   - ask for and return a sockaddr for the source address of the message?

   495

                iov_lens    - (optional) sequence of buffer sizes to use for the message data iov

   496

                control_len - (optional) amount of auxiliary data to recieve

   497

                flags       - (optional) flags to pass to recvmsg()