<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::io::{Error, ErrorKind};
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};
use std::{mem, thread};
use std::net::{TcpStream, TcpListener};
use std::os::unix::io::{RawFd, AsRawFd, IntoRawFd};
use std::collections::LinkedList;

use libc;
use errno::errno;
use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use libc::{c_int, c_void};
use config::Config;
use openssl::ssl::{SslStream, SslContext};

use stats;
use types::*;
use resources::ResourcePool;
use ss::nonblocking::plain::Plain;
use ss::nonblocking::secure::Secure;
use ss::{Socket, Stream, SRecv, SSend, TcpOptions, SocketOptions};

// We need to be able to access our resource pool from several methods
static mut pool: *mut ResourcePool = 0 as *mut ResourcePool;

// Global SslContext
static mut ssl_context: *mut SslContext = 0 as *mut SslContext;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLET  - Fd is in EdgeTriggered mode (notification on state changes)
// EPOLLIN  - Data is available in kerndl buffer
const EVENTS: u32 = event_type::EPOLLET | event_type::EPOLLIN;

/// Starts the epoll wait and incoming connection listener threads.
pub fn begin(config: Config, handler: Box<EventHandler>) {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Stream>::new()));

    // Resource pool
    let mut rp = ResourcePool::new(config.workers);
    unsafe {
        pool = &mut rp;
    }

    // Wrap our event handler into something that can be safely shared
    // between threads.
    let e_handler = Handler(Box::into_raw(handler));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    let streams2 = sockets.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            event_loop(epfd2, streams2, e_handler);
        })
        .unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let streams3 = sockets.clone();
    let prox = thread::Builder::new()
                   .name("TCP Incoming Listener".to_string())
                   .spawn(move || {
                       listen(config, epfd3, streams3);
                   })
                   .unwrap();

    // Stay alive forever, or at least we hope
    let _ = prox.join();
}

/// Incoming connection listening thread
fn listen(config: Config, epfd: RawFd, streams: StreamList) {
    // Setup server and listening port
    let listener_result = try_setup_tcp_listener(&config);
    if listener_result.is_err() {
        error!("Setting up server: {}", listener_result.unwrap_err());
        return;
    }

    // If we're using SSL, setup our context reference
    if config.ssl.is_some() {
        setup_ssl_context(&config);
    }

    // Begin listening for new connections
    let listener = listener_result.unwrap();
    for accept_result in listener.incoming() {
        match accept_result {
            Ok(tcp_stream) => handle_new_connection(tcp_stream, &config, epfd, streams.clone()),
            Err(e) => error!("Accepting connection: {}", e)
        }
    }

    drop(listener);
}

fn try_setup_tcp_listener(config: &Config) -> Result<TcpListener, Error> {
    let create_result = TcpListener::bind((&config.addr[..], config.port));
    if create_result.is_err() {
        return create_result;
    }

    let listener = create_result.unwrap();
    let server_fd = listener.as_raw_fd();

    // Set the SO_REUSEADDR so that restarts after crashes do not take 5min to unbind
    // the initial port
    unsafe {
        let optval: c_int = 1;
        let opt_result = libc::setsockopt(server_fd,
                                          libc::SOL_SOCKET,
                                          libc::SO_REUSEADDR,
                                          &optval as *const _ as *const c_void,
                                          mem::size_of::<c_int>() as u32);
        if opt_result < 0 {
            return Err(Error::from_raw_os_error(errno().0 as i32));
        }
    }

    Ok(listener)
}

fn setup_ssl_context(config: &Config) {
    unsafe {
        ssl_context = Box::into_raw(Box::new(config.ssl.clone().unwrap()));
    }
}

fn handle_new_connection(tcp_stream: TcpStream, config: &Config, epfd: RawFd, streams: StreamList) {
    // Update our total opened file descriptors
    stats::fd_opened();

    // Create and configure a new socket
    let mut socket = Socket::new(tcp_stream.into_raw_fd());
    let result = setup_new_socket(&mut socket);
    if result.is_err() {
        close_fd(socket.as_raw_fd());
        return;
    }

    // Setup our stream
    let stream = match config.ssl {
        Some(_) => {
            let ssl_result = unsafe { SslStream::accept(&(*ssl_context), socket) };
            match ssl_result {
                Ok(ssl_stream) => {
                    let secure_stream = Secure::new(ssl_stream);
                    Stream::new(Box::new(secure_stream))
                }
                Err(ssl_error) => {
                    error!("Creating SslStream: {}", ssl_error);
                    return;
                }
            }
        }
        None => {
            let plain_text = Plain::new(socket);
            Stream::new(Box::new(plain_text))
        }
    };

    // Add stream to our server
    add_to_epoll(epfd, stream.as_raw_fd(), streams.clone());
    add_stream_to_master_list(stream, streams.clone());
}

fn setup_new_socket(socket: &mut Socket) -> Result<(), ()> {
    let result = socket.set_nonblocking();
    if result.is_err() {
        error!("Setting fd to nonblocking: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_nodelay(true);
    if result.is_err() {
        error!("Setting tcp_nodelay: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_keepalive(true);
    if result.is_err() {
        error!("Setting tcp_keepalive: {}", result.unwrap_err());
        return Err(());
    }

    Ok(())
}

/// Event loop for handling all epoll events
fn event_loop(epfd: RawFd, streams: StreamList, handler: Handler) {
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe {
        events.set_len(100);
    }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    handle_epoll_event(epfd, &events[x], streams.clone(), handler.clone());
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        };
    }
}

/// Finds the stream the epoll event is associated with and parses the event type
/// to hand off to specific handlers
fn handle_epoll_event(epfd: RawFd, event: &EpollEvent, streams: StreamList, handler: Handler) {
    const READ_EVENT: u32 = event_type::EPOLLIN;

    // Locate the stream the event was for
    let mut stream;
    {
        // Mutex lock
        // Find the stream the event was for
        let mut guard = match streams.lock() {
            Ok(guard) => guard,
            Err(poisoned) => {
                warn!("StreamList Mutex was poisoned, using anyway");
                poisoned.into_inner()
            }
        };
        let list = guard.deref_mut();

        let mut found = false;
        let mut index = 1usize;
        for s in list.iter() {
            if s.as_raw_fd() == event.data as RawFd {
                found = true;
                break;
            }
            index += 1;
        }

        if !found {
            error!("Could not find fd in stream list? Removing from epoll and closing fd");

            let fd = event.data as RawFd;
            remove_fd_from_epoll(epfd, fd);
            close_fd(fd);
            return;
        }

        if index == 1 {
            stream = list.pop_front().unwrap();
        } else {
            let mut split = list.split_off(index - 1);
            stream = split.pop_front().unwrap();
            list.append(&mut split);
        }
    } // Mutex unlock

    if (event.events & READ_EVENT) > 0 {
        trace!("event was read event");
        let _ = handle_read_event(epfd, &mut stream, handler).map(|_| {
            add_stream_to_master_list(stream, streams.clone());
        });
    } else {
        trace!("event was drop event");

        let fd = stream.as_raw_fd();
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        let stream_fd = stream.as_raw_fd();
        unsafe {
            (*pool).run(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(stream_fd);
            });
        }
    }
}

/// Reads all available data on the stream.
///
/// If a complete message(s) is available, each message will be routed through the
/// resource pool.
///
/// If an error occurs during the read, the stream is dropped from the server.
fn handle_read_event(epfd: RawFd, stream: &mut Stream, handler: Handler) -> Result<(), ()> {
    trace!("handle read event");

    match stream.recv() {
        Ok(_) => {
            let mut rx_queue = stream.drain_rx_queue();
            for payload in rx_queue.iter_mut() {
                // Check if this is a request for stats
                if payload.len() == 6 && payload[0] == 0x04 && payload[1] == 0x04 {
                    let u8ptr: *const u8 = &payload[2] as *const _;
                    let f32ptr: *const f32 = u8ptr as *const _;
                    let sec = unsafe { *f32ptr };
                    let stream_cpy = stream.clone();
                    unsafe {
                        (*pool).run(move || {
                            let mut s = stream_cpy.clone();
                            let result = stats::as_serialized_buffer(sec);
                            if result.is_ok() {
                                let _ = s.send(&result.unwrap()[..]);
                            }
                        });
                    }
                    return Ok(());
                }

                // TODO - Refactor once better function passing traits are available in stable.
                let handler_cpy = handler.clone();
                let stream_cpy = stream.clone();
                let payload_cpy = payload.clone();
                unsafe {
                    (*pool).run(move || {
                        let Handler(ptr) = handler_cpy;
                        (*ptr).on_data_received(stream_cpy.clone(), payload_cpy.clone());
                    });
                }
            }
            Ok(())
        }
        Err(e) => {
            if e.kind() == ErrorKind::WouldBlock {
                return Ok(());
            }

            error!("stream.fd: {} during read: {}", stream.as_raw_fd(), e);

            remove_fd_from_epoll(epfd, stream.as_raw_fd());
            close_fd(stream.as_raw_fd());

            let stream_fd = stream.as_raw_fd();
            unsafe {
                (*pool).run(move || {
                    let Handler(ptr) = handler;
                    (*ptr).on_stream_closed(stream_fd.clone());
                });
            }

            Err(())
        }
    }
}

/// Inserts the stream back into the master list of streams
fn add_stream_to_master_list(stream: Stream, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("StreamList Mutex failed, using anyway...");
            poisoned.into_inner()
        }
    };
    let stream_list = guard.deref_mut();
    stream_list.push_back(stream);
    stats::conn_recv();
}

/// Adds a new fd to the epoll instance
fn add_to_epoll(epfd: RawFd, fd: RawFd, streams: StreamList) {
    let result = epoll::ctl(epfd,
                            ctl_op::ADD,
                            fd,
                            &mut EpollEvent {
                                data: fd as u64,
                                events: EVENTS,
                            });

    if result.is_ok() {
        trace!("New fd added to epoll");
    } else {
        let e = result.unwrap_err();
        error!("CtrlError during add: {}", e);
        remove_fd_from_list(fd, streams.clone());
        close_fd(fd);
    }
}

/// Removes a fd from the epoll instance
fn remove_fd_from_epoll(epfd: RawFd, fd: RawFd) {
    debug!("removing fd from epoll");

    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL. We'll be as backwards compatible as possible.
    let _ = epoll::ctl(epfd,
                       ctl_op::DEL,
                       fd,
                       &mut EpollEvent {
                           data: 0 as u64,
                           events: 0 as u32,
                       })
                .map_err(|e| warn!("Epoll CtrlError during del: {}", e));
}

/// Removes stream with id from master list
fn remove_fd_from_list(fd: RawFd, streams: StreamList) {
    trace!("removing fd: {} from master list", fd);

    let mut guard = match streams.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("StreamList Mutex was poisoned, using anyway");
            poisoned.into_inner()
        }
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for s in list.iter() {
        if s.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    if !found {
        trace!("fd: {} not found in list", fd);
        return;
    }

    if index == 1 {
        list.pop_front();
    } else {
        let mut split = list.split_off(index - 1);
        split.pop_front();
        list.append(&mut split);
    }

    stats::conn_lost();
}

/// Closes a fd with the kernel
fn close_fd(fd: RawFd) {
    unsafe {
        let result = libc::close(fd);
        if result < 0 {
            error!("Error closing fd: {}",
                   Error::from_raw_os_error(result as i32));
            return;
        }
    }
    stats::fd_closed();
}
<commit_msg>Fix race condition on fast network joins<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::io::{Error, ErrorKind};
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};
use std::{mem, thread};
use std::net::{TcpStream, TcpListener};
use std::os::unix::io::{RawFd, AsRawFd, IntoRawFd};
use std::collections::LinkedList;

use libc;
use errno::errno;
use epoll;
use epoll::util::*;
use epoll::EpollEvent;
use libc::{c_int, c_void};
use config::Config;
use openssl::ssl::{SslStream, SslContext};

use stats;
use types::*;
use resources::ResourcePool;
use ss::nonblocking::plain::Plain;
use ss::nonblocking::secure::Secure;
use ss::{Socket, Stream, SRecv, SSend, TcpOptions, SocketOptions};

// We need to be able to access our resource pool from several methods
static mut pool: *mut ResourcePool = 0 as *mut ResourcePool;

// Global SslContext
static mut ssl_context: *mut SslContext = 0 as *mut SslContext;

// When added to epoll, these will be the conditions of kernel notification:
//
// EPOLLET  - Fd is in EdgeTriggered mode (notification on state changes)
// EPOLLIN  - Data is available in kerndl buffer
const EVENTS: u32 = event_type::EPOLLET | event_type::EPOLLIN;

/// Starts the epoll wait and incoming connection listener threads.
pub fn begin(config: Config, handler: Box<EventHandler>) {
    // Master socket list
    let sockets = Arc::new(Mutex::new(LinkedList::<Stream>::new()));

    // Resource pool
    let mut rp = ResourcePool::new(config.workers);
    unsafe {
        pool = &mut rp;
    }

    // Wrap our event handler into something that can be safely shared
    // between threads.
    let e_handler = Handler(Box::into_raw(handler));

    // Epoll instance
    let result = epoll::create1(0);
    if result.is_err() {
        let err = result.unwrap_err();
        error!("Unable to create epoll instance: {}", err);
        panic!()
    }
    let epfd = result.unwrap();

    // Epoll wait thread
    let epfd2 = epfd.clone();
    let streams2 = sockets.clone();
    thread::Builder::new()
        .name("Epoll Wait".to_string())
        .spawn(move || {
            event_loop(epfd2, streams2, e_handler);
        })
        .unwrap();

    // New connection thread
    let epfd3 = epfd.clone();
    let streams3 = sockets.clone();
    let prox = thread::Builder::new()
        .name("TCP Incoming Listener".to_string())
        .spawn(move || {
           listen(config, epfd3, streams3);
        })
        .unwrap();

    // Stay alive forever, or at least we hope
    let _ = prox.join();
}

/// Incoming connection listening thread
fn listen(config: Config, epfd: RawFd, streams: StreamList) {
    // Setup server and listening port
    let listener_result = try_setup_tcp_listener(&config);
    if listener_result.is_err() {
        error!("Setting up server: {}", listener_result.unwrap_err());
        return;
    }

    // If we're using SSL, setup our context reference
    if config.ssl.is_some() {
        setup_ssl_context(&config);
    }

    // Begin listening for new connections
    let listener = listener_result.unwrap();
    for accept_result in listener.incoming() {
        match accept_result {
            Ok(tcp_stream) => handle_new_connection(tcp_stream, &config, epfd, streams.clone()),
            Err(e) => error!("Accepting connection: {}", e)
        }
    }

    drop(listener);
}

fn try_setup_tcp_listener(config: &Config) -> Result<TcpListener, Error> {
    let create_result = TcpListener::bind((&config.addr[..], config.port));
    if create_result.is_err() {
        return create_result;
    }

    let listener = create_result.unwrap();
    let server_fd = listener.as_raw_fd();

    // Set the SO_REUSEADDR so that restarts after crashes do not take 5min to unbind
    // the initial port
    unsafe {
        let optval: c_int = 1;
        let opt_result = libc::setsockopt(server_fd,
                                          libc::SOL_SOCKET,
                                          libc::SO_REUSEADDR,
                                          &optval as *const _ as *const c_void,
                                          mem::size_of::<c_int>() as u32);
        if opt_result < 0 {
            return Err(Error::from_raw_os_error(errno().0 as i32));
        }
    }

    Ok(listener)
}

fn setup_ssl_context(config: &Config) {
    unsafe {
        ssl_context = Box::into_raw(Box::new(config.ssl.clone().unwrap()));
    }
}

fn handle_new_connection(tcp_stream: TcpStream, config: &Config, epfd: RawFd, streams: StreamList) {
    // Update our total opened file descriptors
    stats::fd_opened();

    // Create and configure a new socket
    let mut socket = Socket::new(tcp_stream.into_raw_fd());
    let result = setup_new_socket(&mut socket);
    if result.is_err() {
        close_fd(socket.as_raw_fd());
        return;
    }

    // Setup our stream
    let stream = match config.ssl {
        Some(_) => {
            let ssl_result = unsafe { SslStream::accept(&(*ssl_context), socket) };
            match ssl_result {
                Ok(ssl_stream) => {
                    let secure_stream = Secure::new(ssl_stream);
                    Stream::new(Box::new(secure_stream))
                }
                Err(ssl_error) => {
                    error!("Creating SslStream: {}", ssl_error);
                    return;
                }
            }
        }
        None => {
            let plain_text = Plain::new(socket);
            Stream::new(Box::new(plain_text))
        }
    };

    // Add stream to our server
    let fd = stream.as_raw_fd();
    add_stream_to_master_list(stream, streams.clone());
    add_to_epoll(epfd, fd, streams.clone());
}

fn setup_new_socket(socket: &mut Socket) -> Result<(), ()> {
    let result = socket.set_nonblocking();
    if result.is_err() {
        error!("Setting fd to nonblocking: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_nodelay(true);
    if result.is_err() {
        error!("Setting tcp_nodelay: {}", result.unwrap_err());
        return Err(());
    }

    let result = socket.set_tcp_keepalive(true);
    if result.is_err() {
        error!("Setting tcp_keepalive: {}", result.unwrap_err());
        return Err(());
    }

    Ok(())
}

/// Event loop for handling all epoll events
fn event_loop(epfd: RawFd, streams: StreamList, handler: Handler) {
    let mut events = Vec::<EpollEvent>::with_capacity(100);
    unsafe {
        events.set_len(100);
    }

    loop {
        match epoll::wait(epfd, &mut events[..], -1) {
            Ok(num_events) => {
                for x in 0..num_events as usize {
                    handle_epoll_event(epfd, &events[x], streams.clone(), handler.clone());
                }
            }
            Err(e) => {
                error!("Error on epoll::wait(): {}", e);
                panic!()
            }
        };
    }
}

/// Finds the stream the epoll event is associated with and parses the event type
/// to hand off to specific handlers
fn handle_epoll_event(epfd: RawFd, event: &EpollEvent, streams: StreamList, handler: Handler) {
    const READ_EVENT: u32 = event_type::EPOLLIN;

    // Locate the stream the event was for
    let mut stream;
    {
        // Mutex lock
        // Find the stream the event was for
        let mut guard = match streams.lock() {
            Ok(guard) => guard,
            Err(poisoned) => {
                warn!("StreamList Mutex was poisoned, using anyway");
                poisoned.into_inner()
            }
        };
        let list = guard.deref_mut();

        let mut found = false;
        let mut index = 1usize;
        for s in list.iter() {
            if s.as_raw_fd() == event.data as RawFd {
                found = true;
                break;
            }
            index += 1;
        }

        if !found {
            error!("Could not find fd in stream list? Removing from epoll and closing fd");

            let fd = event.data as RawFd;
            remove_fd_from_epoll(epfd, fd);
            close_fd(fd);
            return;
        }

        if index == 1 {
            stream = list.pop_front().unwrap();
        } else {
            let mut split = list.split_off(index - 1);
            stream = split.pop_front().unwrap();
            list.append(&mut split);
        }
    } // Mutex unlock

    if (event.events & READ_EVENT) > 0 {
        trace!("event was read event");
        let _ = handle_read_event(epfd, &mut stream, handler).map(|_| {
            add_stream_to_master_list(stream, streams.clone());
        });
    } else {
        trace!("event was drop event");

        let fd = stream.as_raw_fd();
        remove_fd_from_epoll(epfd, fd);
        close_fd(fd);

        let stream_fd = stream.as_raw_fd();
        unsafe {
            (*pool).run(move || {
                let Handler(ptr) = handler;
                (*ptr).on_stream_closed(stream_fd);
            });
        }
    }
}

/// Reads all available data on the stream.
///
/// If a complete message(s) is available, each message will be routed through the
/// resource pool.
///
/// If an error occurs during the read, the stream is dropped from the server.
fn handle_read_event(epfd: RawFd, stream: &mut Stream, handler: Handler) -> Result<(), ()> {
    trace!("handle read event");

    match stream.recv() {
        Ok(_) => {
            let mut rx_queue = stream.drain_rx_queue();
            for payload in rx_queue.iter_mut() {
                // Check if this is a request for stats
                if payload.len() == 6 && payload[0] == 0x04 && payload[1] == 0x04 {
                    let u8ptr: *const u8 = &payload[2] as *const _;
                    let f32ptr: *const f32 = u8ptr as *const _;
                    let sec = unsafe { *f32ptr };
                    let stream_cpy = stream.clone();
                    unsafe {
                        (*pool).run(move || {
                            let mut s = stream_cpy.clone();
                            let result = stats::as_serialized_buffer(sec);
                            if result.is_ok() {
                                let _ = s.send(&result.unwrap()[..]);
                            }
                        });
                    }
                    return Ok(());
                }

                // TODO - Refactor once better function passing traits are available in stable.
                let handler_cpy = handler.clone();
                let stream_cpy = stream.clone();
                let payload_cpy = payload.clone();
                unsafe {
                    (*pool).run(move || {
                        let Handler(ptr) = handler_cpy;
                        (*ptr).on_data_received(stream_cpy.clone(), payload_cpy.clone());
                    });
                }
            }
            Ok(())
        }
        Err(e) => {
            if e.kind() == ErrorKind::WouldBlock {
                return Ok(());
            }

            error!("stream.fd: {} during read: {}", stream.as_raw_fd(), e);

            remove_fd_from_epoll(epfd, stream.as_raw_fd());
            close_fd(stream.as_raw_fd());

            let stream_fd = stream.as_raw_fd();
            unsafe {
                (*pool).run(move || {
                    let Handler(ptr) = handler;
                    (*ptr).on_stream_closed(stream_fd.clone());
                });
            }

            Err(())
        }
    }
}

/// Inserts the stream back into the master list of streams
fn add_stream_to_master_list(stream: Stream, streams: StreamList) {
    let mut guard = match streams.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("StreamList Mutex failed, using anyway...");
            poisoned.into_inner()
        }
    };
    let stream_list = guard.deref_mut();
    stream_list.push_back(stream);
    stats::conn_recv();
}

/// Adds a new fd to the epoll instance
fn add_to_epoll(epfd: RawFd, fd: RawFd, streams: StreamList) {
    let result = epoll::ctl(epfd,
                            ctl_op::ADD,
                            fd,
                            &mut EpollEvent {
                                data: fd as u64,
                                events: EVENTS,
                            });

    if result.is_ok() {
        trace!("New fd added to epoll");
    } else {
        let e = result.unwrap_err();
        error!("CtrlError during add: {}", e);
        remove_fd_from_list(fd, streams.clone());
        close_fd(fd);
    }
}

/// Removes a fd from the epoll instance
fn remove_fd_from_epoll(epfd: RawFd, fd: RawFd) {
    debug!("removing fd from epoll");

    // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
    // a non-null pointer in event, even though this argument is ignored.
    // Since Linux 2.6.9, event can be specified as NULL when using
    // EPOLL_CTL_DEL. We'll be as backwards compatible as possible.
    let _ = epoll::ctl(epfd,
                       ctl_op::DEL,
                       fd,
                       &mut EpollEvent {
                           data: 0 as u64,
                           events: 0 as u32,
                       })
                .map_err(|e| warn!("Epoll CtrlError during del: {}", e));
}

/// Removes stream with id from master list
fn remove_fd_from_list(fd: RawFd, streams: StreamList) {
    trace!("removing fd: {} from master list", fd);

    let mut guard = match streams.lock() {
        Ok(guard) => guard,
        Err(poisoned) => {
            warn!("StreamList Mutex was poisoned, using anyway");
            poisoned.into_inner()
        }
    };
    let list = guard.deref_mut();

    let mut found = false;
    let mut index = 1usize;
    for s in list.iter() {
        if s.as_raw_fd() == fd {
            found = true;
            break;
        }
        index += 1;
    }

    if !found {
        trace!("fd: {} not found in list", fd);
        return;
    }

    if index == 1 {
        list.pop_front();
    } else {
        let mut split = list.split_off(index - 1);
        split.pop_front();
        list.append(&mut split);
    }

    stats::conn_lost();
}

/// Closes a fd with the kernel
fn close_fd(fd: RawFd) {
    unsafe {
        let result = libc::close(fd);
        if result < 0 {
            error!("Error closing fd: {}",
                   Error::from_raw_os_error(result as i32));
            return;
        }
    }
    stats::fd_closed();
}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate env_logger;
extern crate futures;
extern crate futures_cpupool;
extern crate grpc;
extern crate protobuf;
#[macro_use]
extern crate log;

extern crate grpc_interop;
use grpc_interop::*;

use std::thread;

use bytes::Bytes;

use futures::stream;
use futures::stream::Stream;

use grpc::*;
use futures::Async;

static DICTIONARY: &'static str = "ABCDEFGHIJKLMNOPQRSTUVabcdefghijklmnoqprstuvwxyz0123456789";
// Note: due to const restrictions, this is calculated by hand.
static DICTIONARY_SIZE: usize = 58;

/**
 * Returns a Vec<u8> with a given size containing printable u8s.
 */
fn make_string(size: usize) -> Vec<u8> {
    let dict = DICTIONARY.to_owned().into_bytes();
    let mut result = Vec::<u8>::with_capacity(size);

    for n in 0..size {
        result.push(dict[n % DICTIONARY_SIZE]);
    }

    return result;
}

fn echo_custom_metadata(req_metadata: &Metadata) -> Metadata {
    static TEST_ECHO_KEY: &'static str = "x-grpc-test-echo-initial";
    let mut metadata = Metadata::new();
    if let Some(value) = req_metadata.get(TEST_ECHO_KEY) {
        metadata.add(MetadataKey::from(TEST_ECHO_KEY), Bytes::from(value));
    }
    metadata
}

fn echo_custom_trailing(req_metadata: &Metadata) -> Metadata {
    static TEST_ECHO_TRAILING_KEY: &'static str = "x-grpc-test-echo-trailing-bin";
    let mut metadata = Metadata::new();
    if let Some(value) = req_metadata.get(TEST_ECHO_TRAILING_KEY) {
        metadata.add(
            MetadataKey::from(TEST_ECHO_TRAILING_KEY),
            Bytes::from(value),
        );
    }
    metadata
}

struct TestServerImpl {}

impl TestService for TestServerImpl {
    // One empty request followed by one empty response.
    fn empty_call(&self, _ctx: ServerHandlerContext, _req: ServerRequestSingle<Empty>, resp: ServerResponseUnarySink<Empty>) -> grpc::Result<()> {
        resp.finish(Empty::new())
    }

    // One request followed by one response.
    fn unary_call(&self, _ctx: ServerHandlerContext, req: ServerRequestSingle<SimpleRequest>, mut resp: ServerResponseUnarySink<SimpleResponse>) -> grpc::Result<()> {
        if req.message.get_response_status().get_code() != 0 {
            info!("requested to send grpc error {}", req.message.get_response_status().get_code());
            return resp.send_grpc_error(
                GrpcStatus::from_code_or_unknown(req.message.get_response_status().get_code() as u32),
                req.message.get_response_status().message.clone(),
            );
        }

        let mut payload = Payload::new();
        payload.set_body(make_string(req.message.get_response_size() as usize));
        let mut response = SimpleResponse::new();
        response.set_payload(payload);
        resp.send_metadata(echo_custom_metadata(&req.metadata))?;
        resp.finish_with_trailers(response, echo_custom_trailing(&req.metadata))
    }

    // One request followed by one response. Response has cache control
    // headers set such that a caching HTTP proxy (such as GFE) can
    // satisfy subsequent requests.
    fn cacheable_unary_call(&self, _ctx: ServerHandlerContext, _req: ServerRequestSingle<SimpleRequest>, resp: ServerResponseUnarySink<SimpleResponse>) -> grpc::Result<()> {
        // TODO: implement fully
        resp.finish(SimpleResponse::new())
    }

    // One request followed by a sequence of responses (streamed download).
    // The server returns the payload with client desired type and sizes.
    fn streaming_output_call(&self, o: ServerHandlerContext, mut req: ServerRequestSingle<StreamingOutputCallRequest>, resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        let sizes = req
            .message
            .take_response_parameters()
            .into_iter()
            .map(|res| res.get_size() as usize);
        let output = stream::iter_ok(sizes).map(|size| {
            let mut response = StreamingOutputCallResponse::new();
            let mut payload = Payload::new();
            payload.set_body(make_string(size));
            response.set_payload(payload);
            response
        });
        o.pump(output, resp);
        Ok(())
    }

    // A sequence of requests followed by one response (streamed upload).
    // The server returns the aggregated size of client payload as the result.
    fn streaming_input_call(&self, _ctx: ServerHandlerContext, req: ServerRequest<StreamingInputCallRequest>, resp: ServerResponseUnarySink<StreamingInputCallResponse>) -> grpc::Result<()> {
        let mut resp = Some(resp);
        let mut aggregate_size = 0;
        req.register_stream_handler_basic(move |message| {
            Ok(match message {
                Some(m) => {
                    aggregate_size += m.get_payload().body.len() as i32;
                }
                None => {
                    let mut response = StreamingInputCallResponse::new();
                    response.set_aggregated_payload_size(aggregate_size);
                    resp.take().unwrap().finish(response)?;
                }
            })
        });
        Ok(())
    }

    // A sequence of requests with each request served by the server immediately.
    // As one request could lead to multiple responses, this interface
    // demonstrates the idea of full duplexing.
    fn full_duplex_call(&self, o: ServerHandlerContext, req: ServerRequest<StreamingOutputCallRequest>, mut resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        let metadata = req.metadata();
        info!("sending custom metadata");
        resp.send_metadata(echo_custom_metadata(&metadata))?;
        let mut req = req.into_stream();
        o.spawn_poll_fn(move || {
            loop {
                if let Async::NotReady = resp.poll()? {
                    return Ok(Async::NotReady);
                }
                match req.poll()? {
                    Async::Ready(Some(m)) => {
                        if m.get_response_status().get_code() != 0 {
                            info!("requested to send grpc error {}", m.get_response_status().get_code());
                            resp.send_grpc_error(
                                GrpcStatus::from_code_or_unknown(m.get_response_status().code as u32),
                                m.get_response_status().message.clone(),
                            )?;
                            return Ok(Async::Ready(()));
                        }

                        for p in &m.response_parameters {
                            info!("requested to send data of size {}", p.size);
                            resp.send_data({
                                let mut response = StreamingOutputCallResponse::new();
                                let mut payload = Payload::new();
                                payload.set_body(make_string(p.size as usize));
                                response.set_payload(payload);
                                response
                            })?;
                        }
                    }
                    Async::Ready(None) => {
                        info!("sending custom trailers");
                        resp.send_trailers(echo_custom_trailing(&metadata))?;
                        return Ok(Async::Ready(()));
                    }
                    Async::NotReady => {
                        return Ok(Async::NotReady)
                    }
                }
            }
        });
        Ok(())
    }

    // TODO: implement this if we find an interop client that needs it.
    // A sequence of requests followed by a sequence of responses.
    // The server buffers all the client requests and then serves them in order. A
    // stream of responses are returned to the client when the server starts with
    // first request.
    fn half_duplex_call(&self, _ctx: ServerHandlerContext, _req: ServerRequest<StreamingOutputCallRequest>, mut resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        resp.send_trailers(Metadata::new())
    }
}

fn main() {
    env_logger::init();

    let mut server = ServerBuilder::new_plain();
    server.http.set_port(DEFAULT_PORT);
    server.add_service(TestServiceServer::new_service_def(TestServerImpl {}));
    let _server = server.build().expect("server");

    loop {
        thread::park();
    }
}
<commit_msg>Interop server cleanup<commit_after>extern crate bytes;
extern crate env_logger;
extern crate futures;
extern crate futures_cpupool;
extern crate grpc;
extern crate protobuf;
#[macro_use]
extern crate log;

extern crate grpc_interop;
use grpc_interop::*;

use std::thread;

use bytes::Bytes;

use futures::stream;
use futures::stream::Stream;

use grpc::*;
use futures::Async;

static DICTIONARY: &'static str = "ABCDEFGHIJKLMNOPQRSTUVabcdefghijklmnoqprstuvwxyz0123456789";

/**
 * Returns a Vec<u8> with a given size containing printable u8s.
 */
fn make_string(size: usize) -> Vec<u8> {
    let dict = DICTIONARY.to_owned().into_bytes();
    let mut result = Vec::<u8>::with_capacity(size);

    for n in 0..size {
        result.push(dict[n % DICTIONARY.len()]);
    }

    return result;
}

fn echo_custom_metadata(req_metadata: &Metadata) -> Metadata {
    static TEST_ECHO_KEY: &'static str = "x-grpc-test-echo-initial";
    let mut metadata = Metadata::new();
    if let Some(value) = req_metadata.get(TEST_ECHO_KEY) {
        metadata.add(MetadataKey::from(TEST_ECHO_KEY), Bytes::from(value));
    }
    metadata
}

fn echo_custom_trailing(req_metadata: &Metadata) -> Metadata {
    static TEST_ECHO_TRAILING_KEY: &'static str = "x-grpc-test-echo-trailing-bin";
    let mut metadata = Metadata::new();
    if let Some(value) = req_metadata.get(TEST_ECHO_TRAILING_KEY) {
        metadata.add(
            MetadataKey::from(TEST_ECHO_TRAILING_KEY),
            Bytes::from(value),
        );
    }
    metadata
}

struct TestServerImpl {}

impl TestService for TestServerImpl {
    // One empty request followed by one empty response.
    fn empty_call(&self, _ctx: ServerHandlerContext, _req: ServerRequestSingle<Empty>, resp: ServerResponseUnarySink<Empty>) -> grpc::Result<()> {
        resp.finish(Empty::new())
    }

    // One request followed by one response.
    fn unary_call(&self, _ctx: ServerHandlerContext, req: ServerRequestSingle<SimpleRequest>, mut resp: ServerResponseUnarySink<SimpleResponse>) -> grpc::Result<()> {
        if req.message.get_response_status().get_code() != 0 {
            info!("requested to send grpc error {}", req.message.get_response_status().get_code());
            return resp.send_grpc_error(
                GrpcStatus::from_code_or_unknown(req.message.get_response_status().get_code() as u32),
                req.message.get_response_status().message.clone(),
            );
        }

        let mut payload = Payload::new();
        payload.set_body(make_string(req.message.get_response_size() as usize));
        let mut response = SimpleResponse::new();
        response.set_payload(payload);
        resp.send_metadata(echo_custom_metadata(&req.metadata))?;
        resp.finish_with_trailers(response, echo_custom_trailing(&req.metadata))
    }

    // One request followed by one response. Response has cache control
    // headers set such that a caching HTTP proxy (such as GFE) can
    // satisfy subsequent requests.
    fn cacheable_unary_call(&self, _ctx: ServerHandlerContext, _req: ServerRequestSingle<SimpleRequest>, resp: ServerResponseUnarySink<SimpleResponse>) -> grpc::Result<()> {
        // TODO: implement fully
        resp.finish(SimpleResponse::new())
    }

    // One request followed by a sequence of responses (streamed download).
    // The server returns the payload with client desired type and sizes.
    fn streaming_output_call(&self, o: ServerHandlerContext, mut req: ServerRequestSingle<StreamingOutputCallRequest>, resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        let sizes = req
            .message
            .take_response_parameters()
            .into_iter()
            .map(|res| res.get_size() as usize);
        let output = stream::iter_ok(sizes).map(|size| {
            let mut response = StreamingOutputCallResponse::new();
            let mut payload = Payload::new();
            payload.set_body(make_string(size));
            response.set_payload(payload);
            response
        });
        o.pump(output, resp);
        Ok(())
    }

    // A sequence of requests followed by one response (streamed upload).
    // The server returns the aggregated size of client payload as the result.
    fn streaming_input_call(&self, _ctx: ServerHandlerContext, req: ServerRequest<StreamingInputCallRequest>, resp: ServerResponseUnarySink<StreamingInputCallResponse>) -> grpc::Result<()> {
        let mut resp = Some(resp);
        let mut aggregate_size = 0;
        req.register_stream_handler_basic(move |message| {
            Ok(match message {
                Some(m) => {
                    aggregate_size += m.get_payload().body.len() as i32;
                }
                None => {
                    let mut response = StreamingInputCallResponse::new();
                    response.set_aggregated_payload_size(aggregate_size);
                    resp.take().unwrap().finish(response)?;
                }
            })
        });
        Ok(())
    }

    // A sequence of requests with each request served by the server immediately.
    // As one request could lead to multiple responses, this interface
    // demonstrates the idea of full duplexing.
    fn full_duplex_call(&self, o: ServerHandlerContext, req: ServerRequest<StreamingOutputCallRequest>, mut resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        let metadata = req.metadata();
        info!("sending custom metadata");
        resp.send_metadata(echo_custom_metadata(&metadata))?;
        let mut req = req.into_stream();
        o.spawn_poll_fn(move || {
            loop {
                if let Async::NotReady = resp.poll()? {
                    return Ok(Async::NotReady);
                }
                match req.poll()? {
                    Async::Ready(Some(m)) => {
                        if m.get_response_status().get_code() != 0 {
                            info!("requested to send grpc error {}", m.get_response_status().get_code());
                            resp.send_grpc_error(
                                GrpcStatus::from_code_or_unknown(m.get_response_status().code as u32),
                                m.get_response_status().message.clone(),
                            )?;
                            return Ok(Async::Ready(()));
                        }

                        for p in &m.response_parameters {
                            info!("requested to send data of size {}", p.size);
                            resp.send_data({
                                let mut response = StreamingOutputCallResponse::new();
                                let mut payload = Payload::new();
                                payload.set_body(make_string(p.size as usize));
                                response.set_payload(payload);
                                response
                            })?;
                        }
                    }
                    Async::Ready(None) => {
                        info!("sending custom trailers");
                        resp.send_trailers(echo_custom_trailing(&metadata))?;
                        return Ok(Async::Ready(()));
                    }
                    Async::NotReady => {
                        return Ok(Async::NotReady)
                    }
                }
            }
        });
        Ok(())
    }

    // TODO: implement this if we find an interop client that needs it.
    // A sequence of requests followed by a sequence of responses.
    // The server buffers all the client requests and then serves them in order. A
    // stream of responses are returned to the client when the server starts with
    // first request.
    fn half_duplex_call(&self, _ctx: ServerHandlerContext, _req: ServerRequest<StreamingOutputCallRequest>, mut resp: ServerResponseSink<StreamingOutputCallResponse>) -> grpc::Result<()> {
        resp.send_trailers(Metadata::new())
    }
}

fn main() {
    env_logger::init();

    let mut server = ServerBuilder::new_plain();
    server.http.set_port(DEFAULT_PORT);
    server.add_service(TestServiceServer::new_service_def(TestServerImpl {}));
    let _server = server.build().expect("server");

    loop {
        thread::park();
    }
}
<|endoftext|>


<commit_before>use std::ops::Not;

use crate::types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
#[serde(untagged)]
pub enum ReplyMarkup {
    InlineKeyboardMarkup(InlineKeyboardMarkup),
    ReplyKeyboardMarkup(ReplyKeyboardMarkup),
    ReplyKeyboardRemove(ReplyKeyboardRemove),
    ForceReply(ForceReply),
}

impl From<InlineKeyboardMarkup> for ReplyMarkup {
    fn from(value: InlineKeyboardMarkup) -> ReplyMarkup {
        ReplyMarkup::InlineKeyboardMarkup(value)
    }
}

impl From<ReplyKeyboardMarkup> for ReplyMarkup {
    fn from(value: ReplyKeyboardMarkup) -> ReplyMarkup {
        ReplyMarkup::ReplyKeyboardMarkup(value)
    }
}

impl From<ReplyKeyboardRemove> for ReplyMarkup {
    fn from(value: ReplyKeyboardRemove) -> ReplyMarkup {
        ReplyMarkup::ReplyKeyboardRemove(value)
    }
}

impl From<ForceReply> for ReplyMarkup {
    fn from(value: ForceReply) -> ReplyMarkup {
        ReplyMarkup::ForceReply(value)
    }
}

/// This object represents a custom keyboard with reply options.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ReplyKeyboardMarkup {
    keyboard: Vec<Vec<KeyboardButton>>,
    #[serde(skip_serializing_if = "Not::not")]
    resize_keyboard: bool,
    #[serde(skip_serializing_if = "Not::not")]
    one_time_keyboard: bool,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ReplyKeyboardMarkup {
    pub fn new() -> Self {
        ReplyKeyboardMarkup {
            keyboard: Vec::new(),
            resize_keyboard: false,
            one_time_keyboard: false,
            selective: false,
        }
    }

    fn init(rows: Vec<Vec<KeyboardButton>>) -> Self {
        let mut keyboard = Self::new();
        keyboard.keyboard = rows;
        keyboard
    }

    /// Requests clients to resize the keyboard vertically for
    /// optimal fit (e.g., make the keyboard smaller if there
    /// are just two rows of buttons). Defaults to false, in which case
    /// the custom keyboard is always of the same height as the app's standard keyboard.
    pub fn resize_keyboard(&mut self) -> &mut Self {
        self.resize_keyboard = true;
        self
    }

    /// Requests clients to hide the keyboard as soon as it's been used.
    /// The keyboard will still be available, but clients will automatically
    /// display the usual letter-keyboard in the chat – the user can
    /// press a special button in the input field to see the custom
    /// keyboard again. Defaults to false.
    pub fn one_time_keyboard(&mut self) -> &mut Self {
        self.one_time_keyboard = true;
        self
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }

    pub fn add_row(&mut self, row: Vec<KeyboardButton>) -> &mut Vec<KeyboardButton> {
        self.keyboard.push(row);
        self.keyboard.last_mut().unwrap()
    }

    pub fn add_empty_row(&mut self) -> &mut Vec<KeyboardButton> {
        self.add_row(Default::default())
    }
}

impl From<Vec<Vec<KeyboardButton>>> for ReplyKeyboardMarkup {
    fn from(value: Vec<Vec<KeyboardButton>>) -> Self {
        Self::init(value)
    }
}

/// This object represents one button of the reply keyboard.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct KeyboardButton {
    text: String,
    #[serde(skip_serializing_if = "Not::not")]
    request_contact: bool,
    #[serde(skip_serializing_if = "Not::not")]
    request_location: bool,
}

impl KeyboardButton {
    pub fn new<S: AsRef<str>>(text: S) -> Self {
        Self {
            text: text.as_ref().to_string(),
            request_contact: false,
            request_location: false,
        }
    }

    /// The user's phone number will be sent as a contact when the
    /// button is pressed. Available in private chats only
    pub fn request_contact(&mut self) -> &mut Self {
        self.request_location = false;
        self.request_contact = true;
        self
    }

    /// The user's current location will be sent when the
    /// button is pressed. Available in private chats only
    pub fn request_location(&mut self) -> &mut Self {
        self.request_contact = false;
        self.request_location = true;
        self
    }
}

impl<'a> From<&'a str> for KeyboardButton {
    fn from(value: &'a str) -> KeyboardButton {
        KeyboardButton::new(value)
    }
}

impl From<String> for KeyboardButton {
    fn from(value: String) -> KeyboardButton {
        KeyboardButton::new(value)
    }
}

/// Upon receiving a message with this object, Telegram clients will remove
/// the current custom keyboard and display the default letter-keyboard.
/// By default, custom keyboards are displayed until a new keyboard is sent
/// by a bot. An exception is made for one-time keyboards that are hidden
/// immediately after the user presses a button (see ReplyKeyboardMarkup).
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ReplyKeyboardRemove {
    remove_keyboard: True,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ReplyKeyboardRemove {
    pub fn new() -> Self {
        Self {
            remove_keyboard: True,
            selective: false,
        }
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }
}

/// This object represents an inline keyboard that appears right next to the message it belongs to.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct InlineKeyboardMarkup {
    inline_keyboard: Vec<Vec<InlineKeyboardButton>>,
}

impl InlineKeyboardMarkup {
    pub fn new() -> Self {
        Self {
            inline_keyboard: Default::default(),
        }
    }

    fn init(inline_keyboard: Vec<Vec<InlineKeyboardButton>>) -> Self {
        Self { inline_keyboard }
    }

    pub fn add_row(&mut self, row: Vec<InlineKeyboardButton>) -> &mut Vec<InlineKeyboardButton> {
        self.inline_keyboard.push(row);
        self.inline_keyboard.last_mut().unwrap()
    }

    pub fn add_empty_row(&mut self) -> &mut Vec<InlineKeyboardButton> {
        self.add_row(Default::default())
    }
}

impl From<Vec<Vec<InlineKeyboardButton>>> for InlineKeyboardMarkup {
    fn from(value: Vec<Vec<InlineKeyboardButton>>) -> Self {
        Self::init(value)
    }
}

/// This object represents one button of an inline keyboard.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct InlineKeyboardButton {
    text: String,
    #[serde(flatten)]
    kind: InlineKeyboardButtonKind,
}

impl InlineKeyboardButton {
    /// Data to be sent in a callback query to the bot when button is pressed, 1-64 bytes
    pub fn callback<T: AsRef<str>, C: AsRef<str>>(text: T, callback: C) -> Self {
        Self {
            text: text.as_ref().to_string(),
            kind: InlineKeyboardButtonKind::CallbackData(callback.as_ref().to_string()),
        }
    }

    pub fn url<T: AsRef<str>, U: AsRef<str>>(text: T, url: U) -> Self {
        Self {
            text: text.as_ref().to_string(),
            kind: InlineKeyboardButtonKind::Url(url.as_ref().to_string()),
        }
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub enum InlineKeyboardButtonKind {
    #[serde(rename = "url")]
    Url(String),          // TODO(knsd): Url?
    #[serde(rename = "callback_data")]
    CallbackData(String), // TODO(knsd) Validate size?
//  SwitchInlineQuery(String),
//  SwitchInlineQueryCurrentChat(String),
//  CallbackGame(CallbackGame),
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
struct InlineKeyboardButtonRaw<'a> {
    text: &'a str,
    #[serde(skip_serializing_if = "Option::is_none")]
    url: Option<&'a str>, // TODO(knsd): Url?
    #[serde(skip_serializing_if = "Option::is_none")]
    callback_data: Option<&'a str>, //TODO(knsd) Validate size?
    #[serde(skip_serializing_if = "Option::is_none")]
    switch_inline_query: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    switch_inline_query_current_chat: Option<&'a str>,
    //    callback_game: Option<CallbackGame>,
}

/// Upon receiving a message with this object, Telegram clients will
/// display a reply interface to the user (act as if the user has
/// selected the bot‘s message and tapped ’Reply'). This can be
/// extremely useful if you want to create user-friendly step-by-step
/// interfaces without having to sacrifice privacy mod
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ForceReply {
    force_reply: True,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ForceReply {
    pub fn new() -> Self {
        Self {
            force_reply: True,
            selective: false,
        }
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }
}
<commit_msg>impl From<Vec<Vec<InlineKeyboardButton>>>> for ReplyMarkup<commit_after>use std::ops::Not;

use crate::types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
#[serde(untagged)]
pub enum ReplyMarkup {
    InlineKeyboardMarkup(InlineKeyboardMarkup),
    ReplyKeyboardMarkup(ReplyKeyboardMarkup),
    ReplyKeyboardRemove(ReplyKeyboardRemove),
    ForceReply(ForceReply),
}

impl From<InlineKeyboardMarkup> for ReplyMarkup {
    fn from(value: InlineKeyboardMarkup) -> ReplyMarkup {
        ReplyMarkup::InlineKeyboardMarkup(value)
    }
}

impl From<Vec<Vec<InlineKeyboardButton>>> for ReplyMarkup {
    fn from(value: Vec<Vec<InlineKeyboardButton>>) -> ReplyMarkup {
        ReplyMarkup::InlineKeyboardMarkup(value.into())
    }
}

impl From<ReplyKeyboardMarkup> for ReplyMarkup {
    fn from(value: ReplyKeyboardMarkup) -> ReplyMarkup {
        ReplyMarkup::ReplyKeyboardMarkup(value)
    }
}

impl From<ReplyKeyboardRemove> for ReplyMarkup {
    fn from(value: ReplyKeyboardRemove) -> ReplyMarkup {
        ReplyMarkup::ReplyKeyboardRemove(value)
    }
}

impl From<ForceReply> for ReplyMarkup {
    fn from(value: ForceReply) -> ReplyMarkup {
        ReplyMarkup::ForceReply(value)
    }
}

/// This object represents a custom keyboard with reply options.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ReplyKeyboardMarkup {
    keyboard: Vec<Vec<KeyboardButton>>,
    #[serde(skip_serializing_if = "Not::not")]
    resize_keyboard: bool,
    #[serde(skip_serializing_if = "Not::not")]
    one_time_keyboard: bool,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ReplyKeyboardMarkup {
    pub fn new() -> Self {
        ReplyKeyboardMarkup {
            keyboard: Vec::new(),
            resize_keyboard: false,
            one_time_keyboard: false,
            selective: false,
        }
    }

    fn init(rows: Vec<Vec<KeyboardButton>>) -> Self {
        let mut keyboard = Self::new();
        keyboard.keyboard = rows;
        keyboard
    }

    /// Requests clients to resize the keyboard vertically for
    /// optimal fit (e.g., make the keyboard smaller if there
    /// are just two rows of buttons). Defaults to false, in which case
    /// the custom keyboard is always of the same height as the app's standard keyboard.
    pub fn resize_keyboard(&mut self) -> &mut Self {
        self.resize_keyboard = true;
        self
    }

    /// Requests clients to hide the keyboard as soon as it's been used.
    /// The keyboard will still be available, but clients will automatically
    /// display the usual letter-keyboard in the chat – the user can
    /// press a special button in the input field to see the custom
    /// keyboard again. Defaults to false.
    pub fn one_time_keyboard(&mut self) -> &mut Self {
        self.one_time_keyboard = true;
        self
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }

    pub fn add_row(&mut self, row: Vec<KeyboardButton>) -> &mut Vec<KeyboardButton> {
        self.keyboard.push(row);
        self.keyboard.last_mut().unwrap()
    }

    pub fn add_empty_row(&mut self) -> &mut Vec<KeyboardButton> {
        self.add_row(Default::default())
    }
}

impl From<Vec<Vec<KeyboardButton>>> for ReplyKeyboardMarkup {
    fn from(value: Vec<Vec<KeyboardButton>>) -> Self {
        Self::init(value)
    }
}

/// This object represents one button of the reply keyboard.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct KeyboardButton {
    text: String,
    #[serde(skip_serializing_if = "Not::not")]
    request_contact: bool,
    #[serde(skip_serializing_if = "Not::not")]
    request_location: bool,
}

impl KeyboardButton {
    pub fn new<S: AsRef<str>>(text: S) -> Self {
        Self {
            text: text.as_ref().to_string(),
            request_contact: false,
            request_location: false,
        }
    }

    /// The user's phone number will be sent as a contact when the
    /// button is pressed. Available in private chats only
    pub fn request_contact(&mut self) -> &mut Self {
        self.request_location = false;
        self.request_contact = true;
        self
    }

    /// The user's current location will be sent when the
    /// button is pressed. Available in private chats only
    pub fn request_location(&mut self) -> &mut Self {
        self.request_contact = false;
        self.request_location = true;
        self
    }
}

impl<'a> From<&'a str> for KeyboardButton {
    fn from(value: &'a str) -> KeyboardButton {
        KeyboardButton::new(value)
    }
}

impl From<String> for KeyboardButton {
    fn from(value: String) -> KeyboardButton {
        KeyboardButton::new(value)
    }
}

/// Upon receiving a message with this object, Telegram clients will remove
/// the current custom keyboard and display the default letter-keyboard.
/// By default, custom keyboards are displayed until a new keyboard is sent
/// by a bot. An exception is made for one-time keyboards that are hidden
/// immediately after the user presses a button (see ReplyKeyboardMarkup).
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ReplyKeyboardRemove {
    remove_keyboard: True,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ReplyKeyboardRemove {
    pub fn new() -> Self {
        Self {
            remove_keyboard: True,
            selective: false,
        }
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }
}

/// This object represents an inline keyboard that appears right next to the message it belongs to.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct InlineKeyboardMarkup {
    inline_keyboard: Vec<Vec<InlineKeyboardButton>>,
}

impl InlineKeyboardMarkup {
    pub fn new() -> Self {
        Self {
            inline_keyboard: Default::default(),
        }
    }

    fn init(inline_keyboard: Vec<Vec<InlineKeyboardButton>>) -> Self {
        Self { inline_keyboard }
    }

    pub fn add_row(&mut self, row: Vec<InlineKeyboardButton>) -> &mut Vec<InlineKeyboardButton> {
        self.inline_keyboard.push(row);
        self.inline_keyboard.last_mut().unwrap()
    }

    pub fn add_empty_row(&mut self) -> &mut Vec<InlineKeyboardButton> {
        self.add_row(Default::default())
    }
}

impl From<Vec<Vec<InlineKeyboardButton>>> for InlineKeyboardMarkup {
    fn from(value: Vec<Vec<InlineKeyboardButton>>) -> Self {
        Self::init(value)
    }
}

/// This object represents one button of an inline keyboard.
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct InlineKeyboardButton {
    text: String,
    #[serde(flatten)]
    kind: InlineKeyboardButtonKind,
}

impl InlineKeyboardButton {
    /// Data to be sent in a callback query to the bot when button is pressed, 1-64 bytes
    pub fn callback<T: AsRef<str>, C: AsRef<str>>(text: T, callback: C) -> Self {
        Self {
            text: text.as_ref().to_string(),
            kind: InlineKeyboardButtonKind::CallbackData(callback.as_ref().to_string()),
        }
    }

    pub fn url<T: AsRef<str>, U: AsRef<str>>(text: T, url: U) -> Self {
        Self {
            text: text.as_ref().to_string(),
            kind: InlineKeyboardButtonKind::Url(url.as_ref().to_string()),
        }
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub enum InlineKeyboardButtonKind {
    #[serde(rename = "url")]
    Url(String),          // TODO(knsd): Url?
    #[serde(rename = "callback_data")]
    CallbackData(String), // TODO(knsd) Validate size?
//  SwitchInlineQuery(String),
//  SwitchInlineQueryCurrentChat(String),
//  CallbackGame(CallbackGame),
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
struct InlineKeyboardButtonRaw<'a> {
    text: &'a str,
    #[serde(skip_serializing_if = "Option::is_none")]
    url: Option<&'a str>, // TODO(knsd): Url?
    #[serde(skip_serializing_if = "Option::is_none")]
    callback_data: Option<&'a str>, //TODO(knsd) Validate size?
    #[serde(skip_serializing_if = "Option::is_none")]
    switch_inline_query: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    switch_inline_query_current_chat: Option<&'a str>,
    //    callback_game: Option<CallbackGame>,
}

/// Upon receiving a message with this object, Telegram clients will
/// display a reply interface to the user (act as if the user has
/// selected the bot‘s message and tapped ’Reply'). This can be
/// extremely useful if you want to create user-friendly step-by-step
/// interfaces without having to sacrifice privacy mod
#[derive(Debug, Clone, PartialEq, PartialOrd, Serialize)]
pub struct ForceReply {
    force_reply: True,
    #[serde(skip_serializing_if = "Not::not")]
    selective: bool,
}

impl ForceReply {
    pub fn new() -> Self {
        Self {
            force_reply: True,
            selective: false,
        }
    }

    /// Use this method if you want to force reply from specific users only.
    /// Targets: 1) users that are @mentioned in the text of
    /// the Message object; 2) if the bot's message is a reply (has reply_to_message_id),
    /// sender of the original message.
    pub fn selective(&mut self) -> &mut Self {
        self.selective = true;
        self
    }
}
<|endoftext|>


<commit_before>//! A simple application that downloads the data dictionary and creates code or
//! data to reproduce it in the core library.
//!
//! ### How to use
//!
//! Simply run the application. It will automatically retrieve the dictionary
//! from the DCMTK data dictionary and store the result in "entries.rs".
//! Future versions will enable different kinds of outputs.
//!
//! Please use the `--help` flag for the full usage information.

use clap::{App, Arg};
use serde::Serialize;
use regex::Regex;

use heck::ShoutySnakeCase;
use std::{fs::{create_dir_all, File}, io::BufWriter};
use std::io::{BufRead, BufReader, Write};
use std::path::Path;

/// url to DCMTK dic file
const DEFAULT_LOCATION: &str = "https://raw.githubusercontent.com/DCMTK/dcmtk/master/dcmdata/data/dicom.dic";

fn main() {
    let matches = App::new("DICOM Dictionary Builder")
        .version("0.1.0")
        .arg(
            Arg::with_name("FROM")
                .default_value(DEFAULT_LOCATION)
                .help("Where to fetch the dictionary from"),
        )
        .arg(
            Arg::with_name("no-retired")
                .long("no-retired")
                .help("Whether to ignore retired tags")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("OUTPUT")
                .short("o")
                .help("The path to the output file")
                .default_value("tags.rs")
                .takes_value(true),
        )
        .get_matches();

    let ignore_retired = matches.is_present("no-retired");

    let src = matches.value_of("FROM").unwrap();

    let dst = Path::new(matches.value_of("OUTPUT").unwrap());

    if src.starts_with("http:") || src.starts_with("https:") {
        // read from URL
        println!("Downloading DICOM dictionary ...");
        let mut resp = reqwest::get(src).unwrap();
        let mut data = vec![];
        resp.copy_to(&mut data).unwrap();

        let preamble = data.split(|&b| b == b'\n')
            .filter_map(|l| std::str::from_utf8(l).ok())
            .find(|l| l.contains("Copyright"))
            .unwrap_or("");
        let preamble = format!(
            "Adapted from the DCMTK project.\nURL: {}\n{}",
            src,
            preamble,
        );

        let entries = parse_entries(&*data).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired, &preamble).expect("Failed to write file");
    } else {
        // read from File
        let file = File::open(src).unwrap();
        let entries = parse_entries(BufReader::new(file)).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired, "").expect("Failed to write file");
    }
}

type DynResult<T> = Result<T, Box<dyn std::error::Error>>;

fn parse_entries<R: BufRead>(source: R) -> DynResult<Vec<Entry>> {
    let mut result = vec![];

    for line in source.lines() {
        let line = line?;
        if line.starts_with('#') {
            continue;
        }

        // (0010,0010)	PN	PatientName	1	DICOM
        let parts: Vec<&str> = line.split('\t').collect();

        if parts[4] == "ILLEGAL" || parts[4] == "PRIVATE" || parts[4] == "GENERIC" {
            continue;
        }

        let mut vr = parts[1].to_string();

        // Some fields are dependent upon context.
        // We may want to support this at some point, but for now,
        // let's just choose some good defaults.
        if vr == "up" {
            vr = "UL".to_string();
        }
        if vr == "xs" {
            vr = "US".to_string();
        }
        if vr == "ox" {
            vr = "OB".to_string();
        }
        if vr == "px" {
            vr = "OB".to_string();
        }
        if vr == "lt" {
            vr = "OW".to_string();
        }
        if vr == "na" {
            // These are the "Item", "ItemDelimitationItem", "SequenceDelimitationItem", etc, values.
            // Question, should we generate const values for these, and use them internally?
            continue;
        }

        let mut alias = parts[2].to_string();
        if alias.starts_with("RETIRED_") {
            alias = alias.trim_start_matches("RETIRED_").to_string();
        }

        let tag = parts[0].to_string();

        let regex_tag = Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{4})\)$")?;
        let regex_tag_group100 = Regex::new(r"^\(([0-9A-F]{2})00-[0-9A-F]{2}FF,([0-9A-F]{4})\)$")?;
        let regex_tag_element100 =
            Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{2})00-[0-9A-F]{2}FF\)$")?;

        let cap = regex_tag.captures(tag.as_str());
        let tag_type;
        let tag_declaration = if let Some(cap) = cap {
            // single tag
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Single;
            format!("Tag(0x{}, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_group100.captures(tag.as_str()) {
            // tag range over groups: (ggxx, eeee)
            let group = cap.get(1).expect("capture group 1: group portion").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Group100;
            format!("Tag(0x{}00, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_element100.captures(tag.as_str()) {
            // tag range over elements: (gggg, eexx)
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap
                .get(2)
                .expect("capture group 2: element portion")
                .as_str();
            tag_type = TagType::Element100;
            format!("Tag(0x{}, 0x{}00)", group, elem)
        } else {
            panic!("invalid tag: {}", alias);
        };

        result.push(Entry {
            tag,
            vr,
            alias,
            vm: parts[3].to_string(),
            obs: parts[4].to_string(),
            is_retired: parts[4].contains("retired"),
            tag_declaration,
            tag_type,
        });
    }

    Ok(result)
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
enum TagType {
    Single,
    Group100,
    Element100,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
struct Entry {
    /// (0010,0010)
    tag: String,
    /// PN
    vr: String,
    /// PatientName
    alias: String,
    /// 1
    vm: String,
    /// DICOM
    obs: String,
    /// Retired field?
    is_retired: bool,
    /// Tag declaration "Tag(0x6000, 0x1102)"
    tag_declaration: String,
    /// The type the tag represents
    tag_type: TagType,
}

/// Write the tag dictionary as Rust code.
fn to_code_file<P>(
    dest_path: P,
    entries: Vec<Entry>,
    include_retired: bool,
    preamble: &str,
) -> DynResult<()>
where
    P: AsRef<Path>,
{
    if let Some(p_dir) = dest_path.as_ref().parent() {
        create_dir_all(&p_dir)?;
    }
    let mut f = BufWriter::new(File::create(&dest_path)?);

    f.write_all(
        b"//! Automatically generated. Edit at your own risk.\n"
    )?;

    for line in preamble.split('\n') {
        write!(f, "//! {}\n", line)?;
    }

    f.write_all(b"\n\
    use dicom_core::dictionary::{DictionaryEntryRef, TagRange, TagRange::*};\n\
    use dicom_core::Tag;\n\
    use dicom_core::VR::*;\n\n",
    )?;

    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        writeln!(f, "/// {} {} {} {} {}", e.alias, e.tag, e.vr, e.vm, e.obs)?;

        let tag_type = match e.tag_type {
            TagType::Single => "Tag",
            _ => "TagRange",
        };

        let tag_set = match e.tag_type {
            TagType::Single => e.tag_declaration.clone(),
            TagType::Group100 => format!("Group100({})", e.tag_declaration),
            TagType::Element100 => format!("Element100({})", e.tag_declaration),
        };

        writeln!(
            f,
            "#[rustfmt::skip]\npub const {}: {} = {};",
            e.alias.to_shouty_snake_case(),
            tag_type,
            tag_set,
        )?;
    }

    f.write_all(
        b"\n\n\
    type E = DictionaryEntryRef<'static>;\n\n\
    #[rustfmt::skip]\n\
    pub(crate) const ENTRIES: &[E] = &[\n",
    )?;
    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        let (vr1, vr2) = e.vr.split_at(2);

        let second_vr = if vr2 != "" {
            format!(" /*{} */", vr2)
        } else {
            vr2.to_string()
        };

        let tag_set = match e.tag_type {
            TagType::Single => format!("Single({})", e.alias.to_shouty_snake_case()),
            _ => e.alias.to_shouty_snake_case(),
        };

        writeln!(
            f,
            "    E {{ tag: {}, alias: \"{}\", vr: {}{} }}, // {}",
            tag_set, e.alias, vr1, second_vr, e.obs
        )?;
    }
    f.write_all(b"];\n")?;

    Ok(())
}
<commit_msg>[dictionary-builder] add link to license in preamble<commit_after>//! A simple application that downloads the data dictionary and creates code or
//! data to reproduce it in the core library.
//!
//! ### How to use
//!
//! Simply run the application. It will automatically retrieve the dictionary
//! from the DCMTK data dictionary and store the result in "entries.rs".
//! Future versions will enable different kinds of outputs.
//!
//! Please use the `--help` flag for the full usage information.

use clap::{App, Arg};
use serde::Serialize;
use regex::Regex;

use heck::ShoutySnakeCase;
use std::{fs::{create_dir_all, File}, io::BufWriter};
use std::io::{BufRead, BufReader, Write};
use std::path::Path;

/// url to DCMTK dic file
const DEFAULT_LOCATION: &str = "https://raw.githubusercontent.com/DCMTK/dcmtk/master/dcmdata/data/dicom.dic";

fn main() {
    let matches = App::new("DICOM Dictionary Builder")
        .version("0.1.0")
        .arg(
            Arg::with_name("FROM")
                .default_value(DEFAULT_LOCATION)
                .help("Where to fetch the dictionary from"),
        )
        .arg(
            Arg::with_name("no-retired")
                .long("no-retired")
                .help("Whether to ignore retired tags")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("OUTPUT")
                .short("o")
                .help("The path to the output file")
                .default_value("tags.rs")
                .takes_value(true),
        )
        .get_matches();

    let ignore_retired = matches.is_present("no-retired");

    let src = matches.value_of("FROM").unwrap();

    let dst = Path::new(matches.value_of("OUTPUT").unwrap());

    if src.starts_with("http:") || src.starts_with("https:") {
        // read from URL
        println!("Downloading DICOM dictionary ...");
        let mut resp = reqwest::get(src).unwrap();
        let mut data = vec![];
        resp.copy_to(&mut data).unwrap();

        let preamble = data.split(|&b| b == b'\n')
            .filter_map(|l| std::str::from_utf8(l).ok())
            .find(|l| l.contains("Copyright"))
            .unwrap_or("");
        let preamble = format!(
            "Adapted from the DCMTK project.\nURL: {}\nLicense: {}\n{}",
            src,
            "https://github.com/DCMTK/dcmtk/blob/master/COPYRIGHT",
            preamble,
        );

        let entries = parse_entries(&*data).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired, &preamble).expect("Failed to write file");
    } else {
        // read from File
        let file = File::open(src).unwrap();
        let entries = parse_entries(BufReader::new(file)).unwrap();
        println!("Writing to file ...");
        to_code_file(dst, entries, !ignore_retired, "").expect("Failed to write file");
    }
}

type DynResult<T> = Result<T, Box<dyn std::error::Error>>;

fn parse_entries<R: BufRead>(source: R) -> DynResult<Vec<Entry>> {
    let mut result = vec![];

    for line in source.lines() {
        let line = line?;
        if line.starts_with('#') {
            continue;
        }

        // (0010,0010)	PN	PatientName	1	DICOM
        let parts: Vec<&str> = line.split('\t').collect();

        if parts[4] == "ILLEGAL" || parts[4] == "PRIVATE" || parts[4] == "GENERIC" {
            continue;
        }

        let mut vr = parts[1].to_string();

        // Some fields are dependent upon context.
        // We may want to support this at some point, but for now,
        // let's just choose some good defaults.
        if vr == "up" {
            vr = "UL".to_string();
        }
        if vr == "xs" {
            vr = "US".to_string();
        }
        if vr == "ox" {
            vr = "OB".to_string();
        }
        if vr == "px" {
            vr = "OB".to_string();
        }
        if vr == "lt" {
            vr = "OW".to_string();
        }
        if vr == "na" {
            // These are the "Item", "ItemDelimitationItem", "SequenceDelimitationItem", etc, values.
            // Question, should we generate const values for these, and use them internally?
            continue;
        }

        let mut alias = parts[2].to_string();
        if alias.starts_with("RETIRED_") {
            alias = alias.trim_start_matches("RETIRED_").to_string();
        }

        let tag = parts[0].to_string();

        let regex_tag = Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{4})\)$")?;
        let regex_tag_group100 = Regex::new(r"^\(([0-9A-F]{2})00-[0-9A-F]{2}FF,([0-9A-F]{4})\)$")?;
        let regex_tag_element100 =
            Regex::new(r"^\(([0-9A-F]{4}),([0-9A-F]{2})00-[0-9A-F]{2}FF\)$")?;

        let cap = regex_tag.captures(tag.as_str());
        let tag_type;
        let tag_declaration = if let Some(cap) = cap {
            // single tag
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Single;
            format!("Tag(0x{}, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_group100.captures(tag.as_str()) {
            // tag range over groups: (ggxx, eeee)
            let group = cap.get(1).expect("capture group 1: group portion").as_str();
            let elem = cap.get(2).expect("capture group 2: element").as_str();
            tag_type = TagType::Group100;
            format!("Tag(0x{}00, 0x{})", group, elem)
        } else if let Some(cap) = regex_tag_element100.captures(tag.as_str()) {
            // tag range over elements: (gggg, eexx)
            let group = cap.get(1).expect("capture group 1: group").as_str();
            let elem = cap
                .get(2)
                .expect("capture group 2: element portion")
                .as_str();
            tag_type = TagType::Element100;
            format!("Tag(0x{}, 0x{}00)", group, elem)
        } else {
            panic!("invalid tag: {}", alias);
        };

        result.push(Entry {
            tag,
            vr,
            alias,
            vm: parts[3].to_string(),
            obs: parts[4].to_string(),
            is_retired: parts[4].contains("retired"),
            tag_declaration,
            tag_type,
        });
    }

    Ok(result)
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
enum TagType {
    Single,
    Group100,
    Element100,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Serialize)]
struct Entry {
    /// (0010,0010)
    tag: String,
    /// PN
    vr: String,
    /// PatientName
    alias: String,
    /// 1
    vm: String,
    /// DICOM
    obs: String,
    /// Retired field?
    is_retired: bool,
    /// Tag declaration "Tag(0x6000, 0x1102)"
    tag_declaration: String,
    /// The type the tag represents
    tag_type: TagType,
}

/// Write the tag dictionary as Rust code.
fn to_code_file<P>(
    dest_path: P,
    entries: Vec<Entry>,
    include_retired: bool,
    preamble: &str,
) -> DynResult<()>
where
    P: AsRef<Path>,
{
    if let Some(p_dir) = dest_path.as_ref().parent() {
        create_dir_all(&p_dir)?;
    }
    let mut f = BufWriter::new(File::create(&dest_path)?);

    f.write_all(
        b"//! Automatically generated. Edit at your own risk.\n"
    )?;

    for line in preamble.split('\n') {
        write!(f, "//! {}\n", line)?;
    }

    f.write_all(b"\n\
    use dicom_core::dictionary::{DictionaryEntryRef, TagRange, TagRange::*};\n\
    use dicom_core::Tag;\n\
    use dicom_core::VR::*;\n\n",
    )?;

    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        writeln!(f, "/// {} {} {} {} {}", e.alias, e.tag, e.vr, e.vm, e.obs)?;

        let tag_type = match e.tag_type {
            TagType::Single => "Tag",
            _ => "TagRange",
        };

        let tag_set = match e.tag_type {
            TagType::Single => e.tag_declaration.clone(),
            TagType::Group100 => format!("Group100({})", e.tag_declaration),
            TagType::Element100 => format!("Element100({})", e.tag_declaration),
        };

        writeln!(
            f,
            "#[rustfmt::skip]\npub const {}: {} = {};",
            e.alias.to_shouty_snake_case(),
            tag_type,
            tag_set,
        )?;
    }

    f.write_all(
        b"\n\n\
    type E = DictionaryEntryRef<'static>;\n\n\
    #[rustfmt::skip]\n\
    pub(crate) const ENTRIES: &[E] = &[\n",
    )?;
    for e in &entries {
        if !include_retired && e.is_retired {
            continue;
        }

        let (vr1, vr2) = e.vr.split_at(2);

        let second_vr = if vr2 != "" {
            format!(" /*{} */", vr2)
        } else {
            vr2.to_string()
        };

        let tag_set = match e.tag_type {
            TagType::Single => format!("Single({})", e.alias.to_shouty_snake_case()),
            _ => e.alias.to_shouty_snake_case(),
        };

        writeln!(
            f,
            "    E {{ tag: {}, alias: \"{}\", vr: {}{} }}, // {}",
            tag_set, e.alias, vr1, second_vr, e.obs
        )?;
    }
    f.write_all(b"];\n")?;

    Ok(())
}
<|endoftext|>