;;; udp.ss v1.0 - 23 December 2003 - Category 5
;;;         1.1 - 24 March 2004
;;;               wrapped low-level I/O calls with restart-nonblocking
;;;
;;; An interface to the User Datagram Protocol socket system calls,
;;; written for the CHICKEN Scheme compiler:
;;;   http://www.call-with-current-continuation.org/

;;; Example:
;;;  csi> (use udp)
;;;  csi> (define s (udp-open-socket))
;;;  csi> (udp-bind! s #f 0)
;;;  csi> (udp-connect! s "localhost" 13)  ; daytime service
;;;  csi> (udp-send s "\n")
;;;  csi> (receive (n data from-host from-port) (udp-recvfrom s 64)
;;;         (print* n " bytes from " from-host ":" from-port ": " data))
;;;  26 bytes from 127.0.0.1:13: Wed Dec 24 11:53:14 2003
;;;  csi> (udp-close-socket s)
;;;  csi>


; ----------------------------TERMS OF USE--------------------------------
; Copyright (c) 2003-2004, Category 5
; All rights reserved.
;
; Redistribution and use in source and binary forms, with or without
; modification, are permitted provided that the following conditions are
; met:
;
;   Redistributions of source code must retain the above copyright notice,
;   this list of conditions and the following disclaimer. Redistributions in
;   binary form must reproduce the above copyright notice, this list of
;   conditions and the following disclaimer in the documentation and/or
;   other materials provided with the distribution. Neither the name of the
;   author nor the names of its contributors may be used to endorse or
;   promote products derived from this software without specific prior
;   written permission.
;
; THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
; IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
; THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
; PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR
; CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
; EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
; PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
; PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
; LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
; NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
; SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
; ----------------------------TERMS OF USE--------------------------------


(declare
 (fixnum-arithmetic)
 (no-bound-checks)
 )

(register-feature! 'udp)

(module udp
  (io:event-dispatch io:descriptor io:read-handler io:write-handler
   io:exception-handler io:set-read-handler! io:set-write-handler!
   io:set-exception-handler!
   udp-socket? udp-bound? udp-connected? udp-open-socket
   udp-open-socket* udp-bind! udp-connect! udp-send udp-sendto
   udp-recv udp-recvfrom udp-close-socket udp-bound-port
   udp-set-multicast-interface udp-join-multicast-group)

(import scheme chicken extras foreign srfi-1 srfi-18)

;;; ------- copied from tcp.scm, more or less -------
#>

# include <errno.h>
#ifdef _WIN32
# if _MSC_VER > 1300
# include <winsock2.h>
# include <ws2tcpip.h>
# else
# include <winsock.h>
# endif
# define EWOULDBLOCK     0
# define socklen_t       int
static WSADATA wsa;
# define hstrerror strerror

const char *inet_ntop(int af, const void *src, char *dst, socklen_t cnt)
{
        if (af == AF_INET)
        {
                struct sockaddr_in in;
                memset(&in, 0, sizeof(in));
                in.sin_family = AF_INET;
                memcpy(&in.sin_addr, src, sizeof(struct in_addr));
                getnameinfo((struct sockaddr *)&in, sizeof(struct
sockaddr_in), dst, cnt, NULL, 0, NI_NUMERICHOST);
                return dst;
        }
        else if (af == AF_INET6)
        {
                struct sockaddr_in6 in;
                memset(&in, 0, sizeof(in));
                in.sin6_family = AF_INET6;
                memcpy(&in.sin6_addr, src, sizeof(struct in_addr6));
                getnameinfo((struct sockaddr *)&in, sizeof(struct
sockaddr_in6), dst, cnt, NULL, 0, NI_NUMERICHOST);
                return dst;
        }
        return NULL;
}

#else
# include <fcntl.h>
# include <sys/types.h>
# include <sys/socket.h>
# include <sys/time.h>
# include <netinet/in.h>
# include <arpa/inet.h>
# include <unistd.h>
# include <netdb.h>
# include <signal.h>
# define closesocket     close
# define INVALID_SOCKET  -1
#endif

#ifndef INET_ADDRSTRLEN
#define INET_ADDRSTRLEN 16
#endif

<#

(define-foreign-variable errno int "errno")
(define-foreign-variable h_errno int "h_errno")

(define-foreign-variable _af_inet int "AF_INET")
(define-foreign-variable _sock_dgram int "SOCK_DGRAM")
(define-foreign-variable _sockaddr_in_size int "sizeof(struct sockaddr_in)")
(define-foreign-variable _ipproto_udp int "IPPROTO_UDP")
(define-foreign-variable _invalid_socket int "INVALID_SOCKET")
(define-foreign-variable _ewouldblock int "EWOULDBLOCK")

(define net-socket (foreign-lambda int "socket" int int int))
(define net-bind (foreign-lambda int "bind" int scheme-pointer int))
(define net-close (foreign-lambda int "closesocket" int))
(define net-send (foreign-lambda int "send" int scheme-pointer int int))
(define net-sendto (foreign-lambda int "sendto" int scheme-pointer int int scheme-pointer int))
(define net-recv (foreign-lambda int "recv" int scheme-pointer int int))
(define net-recvfrom (foreign-lambda int "recvfrom" int scheme-pointer int int pointer c-pointer))
(define net-connect (foreign-lambda int "connect" int scheme-pointer int))

(define net-make-nonblocking
  (foreign-lambda* bool ([int fd])
#<<EOF
#ifdef _WIN32
  unsigned long mode = 1;
  if(ioctlsocket(fd, FIONBIO, &mode)) {
    errno = WSAGetLastError();
    fprintf(stderr, "ioctlsocket failed: errno=%d\n", errno);
    return(0);
  }else{
    return(1);
  }
#else
  int val = fcntl(fd, F_GETFL, 0);
  if(val == -1) return(0);
  return(fcntl(fd, F_SETFL, val | O_NONBLOCK) != -1);
#endif
EOF
))

(define net-getsockport
  (foreign-lambda* int ([int s])
    "struct sockaddr_in sa;"
    "int len = sizeof(struct sockaddr_in);"
    "if(getsockname(s, (struct sockaddr *)&sa, (socklen_t *)&len) != 0) return(-1);"
    "else return(ntohs(sa.sin_port));") )

(define net-gethostaddr
  (foreign-lambda* bool ((scheme-pointer saddr) (c-string host) (unsigned-short port))
    "struct hostent *he = gethostbyname(host);"
    "struct sockaddr_in *addr = (struct sockaddr_in *)saddr;"
    "if(he == NULL) return(0);"
    "memset(addr, 0, sizeof(struct sockaddr_in));"
    "addr->sin_family = AF_INET;"
    "addr->sin_port = htons((short)port);"
    "addr->sin_addr = *((struct in_addr *)he->h_addr);"
    "return(1);"))

(define net-startup
  (foreign-lambda* bool () #<<EOF
#ifdef _WIN32
     return(WSAStartup(MAKEWORD(1, 1), &wsa) == 0);
#else
     signal(SIGPIPE, SIG_IGN);
     return(1);
#endif
EOF
) )

(unless (net-startup)
  (##sys#signal-hook #:network-error "can not initialize Winsock") )
;;; ------- end of code from tcp.scm -------

(define net-hstrerror (foreign-lambda c-string "hstrerror" int))

;;; io-select : fd-vec fd-vec fd-vec timeout-secs timeout-usecs -> int
;;; take three vectors of fds we want to read from, write from, and
;;; handle exceptional events from, plus a timeout in seconds+microseconds,
;;; and call select(2).  Mutate the vector slots to be -1 if the relevant
;;; event didn't occur, otherwise leave them set to the fd number so we can
;;; reverse-map the fds back to their socket container structures later.
(define io-select
  (foreign-lambda* int ((scheme-object rv)
                        (scheme-object wv)
                        (scheme-object ev)
                        (int secs) (int usecs))
#<<EOF
  fd_set rfds, wfds, efds;
  int nrfds, nwfds, nefds;
  int *ra, *wa, *ea;
  int maxfd = -1;
  struct timeval tv;
  int i, ret;

  FD_ZERO(&rfds);
  FD_ZERO(&wfds);
  FD_ZERO(&efds);
  tv.tv_sec = secs;
  tv.tv_usec = usecs;

  nrfds = C_header_size(rv);
  nwfds = C_header_size(wv);
  nefds = C_header_size(ev);
  
  ra = (int *)malloc(nrfds*sizeof(int));
  wa = (int *)malloc(nwfds*sizeof(int));
  ea = (int *)malloc(nefds*sizeof(int));

  for (i=0; i < nrfds; i++) {
       ra[i] = C_unfix(C_block_item(rv, i));
       FD_SET(ra[i], &rfds);
       if (ra[i] > maxfd)
          maxfd = ra[i];
  }
  for (i=0; i < nwfds; i++) {
       wa[i] = C_unfix(C_block_item(wv, i));
       FD_SET(wa[i], &wfds);
       if (wa[i] > maxfd)
          maxfd = wa[i];
  }
  for (i=0; i < nefds; i++) {
       ea[i] = C_unfix(C_block_item(ev, i));
       FD_SET(ea[i], &efds);
       if (ea[i] > maxfd)
         maxfd = ea[i];
  }

  ret = select(maxfd+1, &rfds, &wfds, &efds,
                        tv.tv_sec == -1 ? NULL : &tv);

  if (ret > 0) {
    for (i=0; i < nrfds; i++) {
      if (!FD_ISSET(ra[i], &rfds))
        C_mutate(&C_block_item(rv, i), C_fix(-1));
    }
    for (i=0; i < nwfds; i++) {
      if (!FD_ISSET(wa[i], &wfds))
        C_mutate(&C_block_item(wv, i), C_fix(-1));
    }
    for (i=0; i < nefds; i++) {
      if (!FD_ISSET(ea[i], &efds))
        C_mutate(&C_block_item(ev, i), C_fix(-1));
    }
  }
  free(ra);
  free(wa);
  free(ea);
  return(ret);
EOF
))

;;; net-make-in-addr-any-addr : sockaddr-in-pointer port -> bool
;;; make a sockaddr_in structure with the address set to INADDR_ANY
;;; and the specified port.
(define net-make-in-addr-any-addr
  (foreign-lambda* bool ((scheme-pointer saddr) (int port))
#<<EOF
  struct sockaddr_in *addr = (struct sockaddr_in *)saddr;
  memset(addr, 0, sizeof(struct sockaddr_in));
  addr->sin_family = AF_INET;
  addr->sin_port = htons(port);
  addr->sin_addr.s_addr = INADDR_ANY;
  return(1);
EOF
))

;;; net-inaddr->string : sockaddr-in-pointer -> c-string
;;; Use inet_ntop(3) to turn a sockaddr_in address into a string.
(define net-inaddr->string
  (foreign-lambda* c-string ((scheme-pointer saddr))
#<<EOF
  static char s[INET_ADDRSTRLEN];
  struct sockaddr_in *addr = (struct sockaddr_in *)saddr;
  if (inet_ntop(AF_INET, &addr->sin_addr.s_addr, s, sizeof(s)) == NULL)
    return(NULL);
  return(s);
EOF
))
;;; net-inaddr-port : sockaddr-in-pointer -> int
;;; return the port number of a sockaddr_in structure.
(define net-inaddr-port
  (foreign-lambda* int ((scheme-pointer saddr))
#<<EOF
  struct sockaddr_in *addr = (struct sockaddr_in *)saddr;
  return(ntohs(addr->sin_port));
EOF
))


;;; error-signaling calls
(define net-error
  (lambda args
    (##sys#update-errno)
    (apply ##sys#signal-hook #:network-error args)))
(define net-herror
  (lambda (host)
    (net-error "hostname lookup failed" host (net-hstrerror h_errno))))
(define net-get-host-or-error
  (lambda (sa host port)
    (if (not (net-gethostaddr sa host port))
        (net-herror host))))

(define syscall-failed?
  (lambda (arg)
    (eq? arg -1)))

(define-foreign-variable error-message c-string "strerror(errno)")

(define restart-nonblocking
  (lambda (name fd i/o thunk)
    (let ((return-code (thunk)))
      (cond ((not (eq? return-code -1)) return-code)
            ((eq? errno _ewouldblock)
             (##sys#thread-block-for-i/o! ##sys#current-thread fd i/o)
             (yield)
             (restart-nonblocking name fd i/o thunk))
            (else (net-error error-message name))))))

(define (yield)
  (##sys#call-with-current-continuation
   (lambda (return)
     (let ((ct ##sys#current-thread))
       (##sys#setslot ct 1 (lambda () (return (##core#undefined))))
       (##sys#schedule) ) ) ) )

;;; io:event-dispatch : slist timeout-secs timeout-usecs -> bool
;;; high-level interface to io-select - take a list of descriptors
;;; packaged in records that have, among other things, slots for read,
;;; write, and exception handler callback procedures.  Extract the
;;; fd numbers we want to handle events from, call io-select, and
;;; use the results to run the handlers for the events that occurred.
(define io:event-dispatch
  (lambda (slist . args)
  (let-optionals args
     ((timeout-secs -1) (timeout-usecs -1) (timeout-handler #f))
       (let ((readers (map (lambda (s) (cons (##sys#slot s 1) s))
                           (filter (lambda (s) (##sys#slot s 4)) slist)))
             (writers (map (lambda (s) (cons (##sys#slot s 1) s))
                           (filter (lambda (s) (##sys#slot s 5)) slist)))
             (cepters (map (lambda (s) (cons (##sys#slot s 1) s))
                           (filter (lambda (s) (##sys#slot s 6)) slist))))
         (let ((rv (list->vector (map car readers)))
               (wv (list->vector (map car writers)))
               (ev (list->vector (map car cepters))))
           (let ((ret (io-select rv wv ev timeout-secs timeout-usecs)))
             (cond ((syscall-failed? ret) (net-error "select"))
                   ((fx= ret 0) (and (procedure? timeout-handler)
                                   (timeout-handler slist)))
                   (else
                    (let ((readable (map (lambda (fd) (cdr (assq fd readers)))
                                         (remove (lambda (fd) (fx= fd -1)) (vector->list rv))))
                          (writable (map (lambda (fd) (cdr (assq fd writers)))
                                         (remove (lambda (fd) (fx= fd -1)) (vector->list wv))))
                          (ceptable (map (lambda (fd) (cdr (assq fd cepters)))
                                         (remove (lambda (fd) (fx= fd -1)) (vector->list ev)))))
                      (for-each (lambda (s) ((##sys#slot s 4) s)) readable)
                      (for-each (lambda (s) ((##sys#slot s 5) s)) writable)
                      (for-each (lambda (s) ((##sys#slot s 6) s)) ceptable)
                      #t)))))))))

;;; udp-socket structure slots:
;;; 1  2      3          4            5             6
;;; fd bound? connected? read-handler write-handler except-handler
(define (udp-socket? x)
  (and (##core#inline "C_blockp" x)
       (##sys#structure? x 'udp-socket)))
(define (udp-bound? s)
  (and (udp-socket? s) (##sys#slot s 2)))
(define (udp-connected? s)
  (and (udp-socket? s) (##sys#slot s 3)))

;;; udp-open-socket : -> udp-socket
(define udp-open-socket
  (lambda ()
    (let ((s (net-socket _af_inet _sock_dgram 0)))
      (if (syscall-failed? s)
          (net-error "socket")
          (##sys#make-structure 'udp-socket s #f #f #f #f #f)))))

(define io:descriptor (lambda (s) (##sys#slot s 1)))
(define io:read-handler (lambda (s) (##sys#slot s 4)))
(define io:write-handler (lambda (s) (##sys#slot s 5)))
(define io:exception-handler (lambda (s) (##sys#slot s 6)))
(define io:set-read-handler! (lambda (s p) (##sys#setslot s 4 p)))
(define io:set-write-handler! (lambda (s p) (##sys#setslot s 5 p)))
(define io:set-exception-handler! (lambda (s p) (##sys#setslot s 6 p)))

;;; udp-open-socket* : -> udp-socket
;;; open a UDP socket and make it nonblocking
(define udp-open-socket*
  (lambda ()
    (let ((s (udp-open-socket)))
      (and (udp-socket? s) (net-make-nonblocking (io:descriptor s)) s))))
;;; udp-bind! : udp-socket host-string port-number -> unspecified
;;; bind a socket to a local address (possibly INADDR_ANY) and port
(define udp-bind!
  (lambda (sock host port)
    (let ((fd (io:descriptor sock))
          (addr (make-string _sockaddr_in_size)))
      (if host
          (net-get-host-or-error addr host port)
          (net-make-in-addr-any-addr addr port))
      (if (syscall-failed? (net-bind fd addr _sockaddr_in_size))
          (net-error "bind" host port)
          (##sys#setslot sock 2 #t)))))

(define udp-bound-port
  (lambda (sock)
    (let* ([fd (io:descriptor sock)]
           [port (net-getsockport fd)])
      (if (eq? -1 port)
        (net-error "getsockport"))
      port)))

;;; udp-connect! : udp-socket host-string port -> unspecified
;;; "connect" a socket.  In the case of UDP this does nothing more than
;;; store a peer address in the kernel socket structure for use with
;;; later calls to send(2).
(define udp-connect!
  (lambda (sock host port)
    (let ((fd (io:descriptor sock))
          (addr (make-string _sockaddr_in_size)))
      (net-get-host-or-error addr host port)
      (if (syscall-failed? (net-connect fd addr _sockaddr_in_size))
          (net-error "connect" host port)
          (##sys#setslot sock 3 #t)))))

;;; udp-send : udp-socket string -> unspecified
;;; send bytes in string to the peer for this socket as specified earlier
;;; with udp-connect!.  If the socket was not "connected", send(2) will
;;; raise an error.
(define udp-send
  (lambda (sock str)
    (let ((fd (io:descriptor sock)))
      (restart-nonblocking "send" fd #f
       (lambda ()
         (net-send fd str (string-length str) 0))))))

;;; udp-sendto : udp-socket host-string port-num string -> unspecified
;;; send bytes in string to host:port via udp-socket.
(define udp-sendto
  (lambda (sock host port str)
    (let ((fd (io:descriptor sock))
          (addr (make-string _sockaddr_in_size)))
      (net-get-host-or-error addr host port)
      (restart-nonblocking "sendto" fd #f
        (lambda () (net-sendto fd str (string-length str)
                                 0 addr _sockaddr_in_size))))))
                          
;;; udp-recv : udp-socket string -> [len packet]
;;; receive a packet and store the data in string, returning the
;;; length of the packet and the substring of len bytes.
(define udp-recv
  (lambda (sock len)
    (let ((fd (io:descriptor sock))
          (buf (make-string len)))
      (let ((result
             (restart-nonblocking "recv" fd #t
               (lambda () (net-recv fd buf len 0)))))
          (values result (substring buf 0 result))))))

      
;;; udp-recvfrom : udp-socket string -> [len packet host-string port-num]
;;; like recv but returns four values, including the length of the
;;; received packet and the host and port from which it was received.
(define udp-recvfrom
  (lambda (sock len)
    (let ((fd (io:descriptor sock))
          (buf (make-string len))
          (from (make-string _sockaddr_in_size)))
      (let-location ((fromlen int _sockaddr_in_size))
	(let ((result
	       (restart-nonblocking "recvfrom" fd #t
				    (lambda () (net-recvfrom fd buf len
							       0 from #$fromlen)))))
	  (values result (substring buf 0 result)
		  (net-inaddr->string from) (net-inaddr-port from)))))) )

;;; udp-close-socket : udp-socket -> bool
;;; close a socket.
(define udp-close-socket
  (lambda (sock)
    (let ((fd (io:descriptor sock)))
      (if (syscall-failed? (net-close fd)) #f #t))))

;;; multicast
(define net-set-multicast-interface
  (foreign-lambda* bool ((int s) (c-string host))
#<<EOF
#ifdef _WIN32
  DWORD ipaddr;
#else
  in_addr_t ipaddr;
#endif
  ipaddr = inet_addr(host);
  return (setsockopt(s, IPPROTO_IP, IP_MULTICAST_IF, &ipaddr, sizeof(ipaddr)) == 0);
EOF
))

(define udp-set-multicast-interface
  (lambda (sock interface-host)
    (let ([fd (io:descriptor sock)])
      (or (net-set-multicast-interface fd interface-host)
          (net-error "udp-set-multicast-interface" error-message)))))

(define net-join-multicast-group
  (foreign-lambda* bool ((int s) (c-string ihost) (c-string mhost) (bool join))
#<<EOF
  struct ip_mreq mreq;
  int optname;
  if (ihost) {
    mreq.imr_interface.s_addr = inet_addr(ihost);
  } else {
    mreq.imr_interface.s_addr = INADDR_ANY;
  }
  mreq.imr_multiaddr.s_addr = inet_addr(mhost);
  if (join) {
    optname = IP_ADD_MEMBERSHIP;
  } else {
    optname = IP_DROP_MEMBERSHIP;
  }
  return (setsockopt(s, IPPROTO_IP, optname, &mreq, sizeof(mreq)) == 0);
EOF
))

(define udp-join-multicast-group
  (lambda (sock interface-host multicast-host . opt)
    (let ([fd (io:descriptor sock)]
          [join (and (pair? opt) (car opt))])
      (or (net-join-multicast-group fd interface-host multicast-host join)
          (net-error "udp-join-multicast-group" error-message)))))

)

;;; END