1

from qmsk.net.socket.socket cimport *

     2

from qmsk.net.socket.addr cimport sockaddr, build_sockaddr

     3

     4

cimport qmsk.net.socket.platform as platform

     5

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

     6

     7

from qmsk.net.py cimport raise_errno

     8

     9

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

    10

    if addr is not None :

    11

        addr._get_sockaddr(sa_ptr, sa_len)

    12

    13

    elif optional :

    14

        sa_ptr[0] = NULL

    15

        sa_len[0] = 0

    16

    17

    else :

    18

        raise ValueError(addr)

    19

    20

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

    21

    cdef libc.ssize_t tmp_len

    22

    23

    if buf is not None :

    24

        # XXX: test that except works right

    25

        # XXX: this complains about const...

    26

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

    27

    28

        # XXX: ensure that this is >= 0

    29

        buf_len[0] = tmp_len

    30

    31

    elif optional :

    32

        buf_ptr[0] = NULL

    33

        buf_len[0] = 0

    34

    35

    else :

    36

        raise ValueError(buf)

    37

    38

## helper for socket.recv*

    39

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

    40

cdef struct sockbuf :

    41

    py.PyObject *str

    42

    43

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

    44

    """

    45

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

    46

        contents as a char*.

    47

    48

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

    49

    """

    50

    51

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

    52

    53

    return py.PyString_AS_STRING(buf.str)

    54

    55

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

    56

    """

    57

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

    58

    """

    59

    60

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

    61

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

    62

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

    63

        py._PyString_Resize(&buf.str, len)

    64

    65

    return <object> buf.str

    66

    67

cdef void sockbuf_deinit (sockbuf *buf) :

    68

    """

    69

        Release the PyObject.

    70

    71

        This is safe if the sockbuf was initialized to zero.

    72

    """

    73

    74

    py.Py_XDECREF(buf.str)

    75

    76

# XXX: do some GIL-releasin'

    77

cdef class socket :

    78

    79

    def __cinit__ (self) :

    80

        """

    81

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

    82

        """

    83

    84

        self.fd = -1

    85

    86

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

    87

        """

    88

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

    89

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

    90

    91

                family      - one of AF_*

    92

                socktype    - one of SOCK_*

    93

                protocol    - one of IPPROTO_* or zero to select default

    94

        """

    95

    96

        if fd >= 0 :

    97

            # given fd

    98

            self.fd = fd

    99

   100

        else :

   101

            # socket()

   102

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

   103

   104

        # trap

   105

        if self.fd < 0 :

   106

            raise_errno('socket')

   107

   108

    def fileno (self) :

   109

        """

   110

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

   111

        """

   112

   113

        return self.fd

   124

   125

    def bind (self, sockaddr addr) :

   126

        """

   127

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

   128

            compatible address family.

   129

   130

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

   131

                              ephemeral port.

   132

        """

   133

   134

        cdef platform.sockaddr *sa

   135

        cdef platform.socklen_t sa_len

   136

   137

        # XXX: require non-NULL addr?

   138

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

   139

   140

        # bind()

   141

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

   142

            raise_errno('bind')

   143

   144

    def listen (self, int backlog) :

   145

        """

   146

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

   147

            requests using sock.accept().

   148

   149

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

   150

   151

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

   152

        """

   153

   154

        # listen()

   155

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

   156

            raise_errno('listen')

   157

   158

    def connect (self, sockaddr addr) :

   159

        """

   160

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

   161

            should be of the same or a compatible address family.

   162

   163

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

   164

   165

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

   166

            ephemeral port.

   167

   168

                addr        - the remote address to connect to.

   169

        """

   170

   171

        cdef platform.sockaddr *sa

   172

        cdef platform.socklen_t sa_len

   173

   174

        # XXX: require non-NULL addr?

   175

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

   176

   177

        # connect()

   178

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

   179

            raise_errno('connect')

   180

   181

    def accept (self) :

   182

        """

   183

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

   184

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

   185

            (.listen()).

   186

   187

            This returns a (sock, src_addr) tuple:

   188

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

   189

                src_addr    - the remote address of the incoming connection

   190

        """

   191

   192

        # prep the sockaddr that we will return

   193

        cdef platform.sockaddr_storage ss

   194

        cdef platform.socklen_t ss_len = sizeof(ss)

   195

   196

        # accept()

   197

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

   198

   199

        if sock_fd < 0 :

   200

            raise_errno('accept')

   201

   202

        try :

   203

            # prep the new socket

   204

            sock_obj = socket(fd=sock_fd)

   205

   206

        except :

   207

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

   208

            platform.close(sock_fd)

   209

   210

            raise

   211

   212

        # prep the new addr

   213

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

   214

   215

        return sock_obj, src_addr

   216

   217

    def getsockname (self) :

   218

        """

   219

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

   220

   221

            Returns a sockaddr object.

   222

        """

   223

   224

        # prep the sockaddr that we will return

   225

        cdef platform.sockaddr_storage ss

   226

        cdef platform.socklen_t ss_len = sizeof(ss)

   227

   228

        # getsockname()

   229

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

   230

            raise_errno('getsockname')

   231

   232

        # build the new sockaddr

   233

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

   234

   235

    def getpeername (self) :

   236

        """

   237

            Get the remote address this socket is currently connected to.

   238

   239

            Returns a sockaddr object.

   240

        """

   241

   242

        # prep the sockaddr that we will return

   243

        cdef platform.sockaddr_storage ss

   244

        cdef platform.socklen_t ss_len = sizeof(ss)

   245

   246

        # getpeername()

   247

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

   248

            raise_errno('getpeername')

   249

   250

        # build the new sockaddr

   251

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

   252

   253

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

   254

        """

   255

            Transmit a message to the connected remote endpoint.

   256

   257

                buf         - the data to send

   258

                flags       - (optional) MSG_* flags to send with

   259

   260

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

   261

        """

   262

   263

        cdef void *buf_ptr

   264

        cdef size_t buf_len

   265

        cdef libc.ssize_t ret

   266

   267

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

   268

   269

        # send()

   270

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

   271

   272

        if ret < 0 :

   273

            raise_errno('send')

   274

   275

        else :

   276

            return ret

   277

   278

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

   279

        """

   280

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

   281

            specified, and this acts like send()

   282

   283

                buf         - the data to send

   284

                flags       - (optional) MSG_* flags to send with

   285

                addr        - (optional) target address

   286

   287

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

   288

        """

   289

   290

        cdef void *buf_ptr

   291

        cdef size_t buf_len

   292

        cdef libc.ssize_t ret

   293

   294

        cdef platform.sockaddr *sa

   295

        cdef platform.socklen_t sa_len

   296

   297

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

   298

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

   299

   300

        # send()

   301

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

   302

   303

        if ret < 0 :

   304

            raise_errno('sendto')

   305

   306

        else :

   307

            return ret

   308

   309

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

   310

        """

   311

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

   312

            extra parameters.

   313

   314

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

   315

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

   316

                control     - (optional) control message to transmit

   317

                flags       - (optional) MSG_* flags to send with

   318

   319

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

   320

        """

   321

   322

        cdef libc.ssize_t ret

   323

        cdef libc.iovec *iovec

   324

        cdef platform.msghdr msg

   325

   326

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

   327

   328

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

   329

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

   330

   331

        # iov

   332

        if iov :

   333

            iov = tuple(iov)

   334

   335

            # numerb of bufs = number of iovecs

   336

            msg.msg_iovlen = len(iov)

   337

   338

            # alloca the required number of iovec's

   339

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

   340

   341

            # fill in the iovecs

   342

            for i, buf in enumerate(iov) :

   343

                iovec = &msg.msg_iov[i]

   344

   345

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

   346

   347

        # sendmsg()

   348

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

   349

   350

        if ret < 0 :

   351

            raise_errno('sendmsg')

   352

   353

        else :

   354

            return ret

   355

   356

    def write (self, object buf) :

   357

        """

   358

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

   359

   360

                buf         - the data to send

   361

   362

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

   363

        """

   364

   365

        cdef void *buf_ptr

   366

        cdef size_t buf_len

   367

        cdef libc.ssize_t ret

   368

   369

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

   370

   371

        # send()

   372

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

   373

   374

        if ret < 0 :

   375

            raise_errno('write')

   376

   377

        else :

   378

            return ret

   379

   380

    def writev (self, iov) :

   381

        """

   382

            Write data to a socket from multiple read-buffers.

   383

   384

                iov         - sequence of read-buffers to transmit

   385

   386

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

   387

        """

   388

   389

        # iov

   390

        cdef libc.iovec *iov_list = NULL

   391

        cdef size_t iov_count = 0

   392

        cdef libc.iovec *iovec

   393

   394

        iov = tuple(iov)

   395

   396

        # numerb of bufs = number of iovecs

   397

        iov_count = len(iov)

   398

   399

        # alloca the required number of iovec's

   400

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

   401

   402

        # fill in the iovecs

   403

        for i, buf in enumerate(iov) :

   404

            iovec = &iov_list[i]

   405

   406

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

   407

   408

        # sendmsg()

   409

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

   410

   411

        if ret < 0 :

   412

            raise_errno('writev')

   413

   414

        else :

   415

            return ret

   416

   417

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

   418

        """

   419

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

   420

   421

                len         - size of buffer to use for recv

   422

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

   423

   424

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

   425

        """

   426

   427

        cdef sockbuf sb

   428

        cdef libc.ssize_t ret

   429

   430

        # alloc the recv buffer

   431

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

   432

   433

        try :

   434

            # recv()

   435

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

   436

   437

            if ret < 0 :

   438

                raise_errno('recv')

   439

   440

            # truncate to correct length

   441

            return sockbuf_truncate(&sb, ret)

   442

   443

        finally :

   444

            sockbuf_deinit(&sb)

   445

   446

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

   447

        """

   448

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

   449

   450

                len         - size of buffer to use for recv

   451

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

   452

   453

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

   454

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

   455

                src_addr    - the source address the message was recieved from

   456

        """

   457

   458

        cdef sockbuf sb

   459

        cdef libc.ssize_t ret

   460

        cdef object str

   461

        cdef sockaddr src_addr

   462

   463

        # prep the sockaddr that we will return

   464

        cdef platform.sockaddr_storage ss

   465

        cdef platform.socklen_t ss_len = sizeof(ss)

   466

   467

        # alloc recv buf

   468

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

   469

   470

        try :

   471

            # recvfrom()

   472

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

   473

   474

            if ret < 0 :

   475

                raise_errno('recv')

   476

   477

            # truncate

   478

            str = sockbuf_truncate(&sb, ret)

   479

   480

            # prep the new addr

   481

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

   482

   483

            return str, src_addr

   484

   485

        finally :

   486

            sockbuf_deinit(&sb)

   487

   488

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

   489

        """

   490

            Recieve a message along with some extra data.

   491

   492

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

   493

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

   494

                control_len - (optional) amount of auxiliary data to recieve

   495

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

   496

   497

            Returns a (name, iovs, control, flags) tuple :