<commit_before>//! Shadowsocks UDP server

use std::{io, net::SocketAddr, sync::Arc, time::Duration};

use bytes::Bytes;
use futures::future::{self, AbortHandle};
use io::ErrorKind;
use lfu_cache::TimedLfuCache;
use log::{debug, error, info, trace, warn};
use shadowsocks::{
    lookup_then,
    net::UdpSocket as OutboundUdpSocket,
    relay::{
        socks5::Address,
        udprelay::{ProxySocket, MAXIMUM_UDP_PAYLOAD_SIZE},
    },
    ServerConfig,
};
use spin::Mutex as SpinMutex;
use tokio::{
    sync::{mpsc, Mutex},
    time,
};

use crate::net::MonProxySocket;

use super::context::ServiceContext;

type AssociationMap = TimedLfuCache<SocketAddr, UdpAssociation>;
type SharedAssociationMap = Arc<Mutex<AssociationMap>>;

pub struct UdpServer {
    context: Arc<ServiceContext>,
    assoc_map: SharedAssociationMap,
    cleanup_abortable: AbortHandle,
}

impl Drop for UdpServer {
    fn drop(&mut self) {
        self.cleanup_abortable.abort();
    }
}

impl UdpServer {
    pub fn new(context: Arc<ServiceContext>, time_to_live: Option<Duration>, capacity: Option<usize>) -> UdpServer {
        let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
        let assoc_map = Arc::new(Mutex::new(match capacity {
            Some(capacity) => TimedLfuCache::with_capacity_and_expiration(capacity, time_to_live),
            None => TimedLfuCache::with_expiration(time_to_live),
        }));

        let cleanup_abortable = {
            let assoc_map = assoc_map.clone();
            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
                loop {
                    time::sleep(time_to_live).await;

                    // cleanup expired associations
                    assoc_map.lock().await.evict_expired();
                }
            });
            tokio::spawn(cleanup_task);
            cleanup_abortable
        };

        UdpServer {
            context,
            assoc_map,
            cleanup_abortable,
        }
    }

    pub async fn run(mut self, svr_cfg: &ServerConfig) -> io::Result<()> {
        let socket = ProxySocket::bind(self.context.context(), svr_cfg).await?;

        info!(
            "shadowsocks udp server listening on {}",
            socket.local_addr().expect("listener.local_addr"),
        );

        let socket = MonProxySocket::from_socket(socket, self.context.flow_stat());
        let listener = Arc::new(socket);

        let mut buffer = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        loop {
            let (n, peer_addr, target_addr) = match listener.recv_from(&mut buffer).await {
                Ok(s) => s,
                Err(err) => {
                    error!("udp server recv_from failed with error: {}", err);
                    continue;
                }
            };

            if self.context.check_outbound_blocked(&target_addr).await {
                error!("udp client {} outbound {} blocked by ACL rules", peer_addr, target_addr);
                continue;
            }

            let data = &buffer[..n];
            if let Err(err) = self.send_packet(&listener, peer_addr, target_addr, data).await {
                error!(
                    "udp packet relay {} with {} bytes failed, error: {}",
                    peer_addr,
                    data.len(),
                    err
                );
            }
        }
    }

    async fn send_packet(
        &mut self,
        listener: &Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        target_addr: Address,
        data: &[u8],
    ) -> io::Result<()> {
        let mut assoc_map = self.assoc_map.lock().await;

        if let Some(assoc) = assoc_map.get(&peer_addr) {
            return assoc.try_send((target_addr, Bytes::copy_from_slice(data)));
        }

        let assoc = UdpAssociation::new(
            self.context.clone(),
            listener.clone(),
            peer_addr,
            self.assoc_map.clone(),
        );

        trace!("created udp association for {}", peer_addr);

        assoc.try_send((target_addr, Bytes::copy_from_slice(data)))?;
        assoc_map.insert(peer_addr, assoc);

        Ok(())
    }
}

struct UdpAssociation {
    assoc: Arc<UdpAssociationContext>,
    sender: mpsc::Sender<(Address, Bytes)>,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.assoc.outbound_ipv4_socket.lock().abort();
        self.assoc.outbound_ipv6_socket.lock().abort();
    }
}

impl UdpAssociation {
    fn new(
        context: Arc<ServiceContext>,
        inbound: Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        assoc_map: SharedAssociationMap,
    ) -> UdpAssociation {
        let (assoc, sender) = UdpAssociationContext::new(context, inbound, peer_addr, assoc_map);
        UdpAssociation { assoc, sender }
    }

    fn try_send(&self, data: (Address, Bytes)) -> io::Result<()> {
        if let Err(..) = self.sender.try_send(data) {
            let err = io::Error::new(ErrorKind::Other, "udp relay channel full");
            return Err(err);
        }
        Ok(())
    }
}

enum UdpAssociationState {
    Empty,
    Connected {
        socket: Arc<OutboundUdpSocket>,
        abortable: AbortHandle,
    },
    Aborted,
}

impl Drop for UdpAssociationState {
    fn drop(&mut self) {
        if let UdpAssociationState::Connected { ref abortable, .. } = *self {
            abortable.abort();
        }
    }
}

impl UdpAssociationState {
    fn empty() -> UdpAssociationState {
        UdpAssociationState::Empty
    }

    fn set_connected(&mut self, socket: Arc<OutboundUdpSocket>, abortable: AbortHandle) {
        *self = UdpAssociationState::Connected { socket, abortable };
    }

    fn abort(&mut self) {
        *self = UdpAssociationState::Aborted;
    }
}

struct UdpAssociationContext {
    context: Arc<ServiceContext>,
    inbound: Arc<MonProxySocket>,
    peer_addr: SocketAddr,
    outbound_ipv4_socket: SpinMutex<UdpAssociationState>,
    outbound_ipv6_socket: SpinMutex<UdpAssociationState>,
    assoc_map: SharedAssociationMap,
    target_cache: Mutex<TimedLfuCache<SocketAddr, Address>>,
}

impl Drop for UdpAssociationContext {
    fn drop(&mut self) {
        trace!("udp association for {} is closed", self.peer_addr);
    }
}

impl UdpAssociationContext {
    fn new(
        context: Arc<ServiceContext>,
        inbound: Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        assoc_map: SharedAssociationMap,
    ) -> (Arc<UdpAssociationContext>, mpsc::Sender<(Address, Bytes)>) {
        // Pending packets 1024 should be good enough for a server.
        // If there are plenty of packets stuck in the channel, dropping exccess packets is a good way to protect the server from
        // being OOM.
        let (sender, receiver) = mpsc::channel(1024);

        let assoc = Arc::new(UdpAssociationContext {
            context,
            inbound,
            peer_addr,
            outbound_ipv4_socket: SpinMutex::new(UdpAssociationState::empty()),
            outbound_ipv6_socket: SpinMutex::new(UdpAssociationState::empty()),
            assoc_map,
            // Cache for remembering the original Address of target,
            // when recv_from a SocketAddr, we have to know whch Address that client was originally requested.
            //
            // XXX: 128 target addresses should be enough for __one__ client.
            //      1 hours should be enough for caching the address mapping. Most of the DNS records' TTL won't last that long.
            target_cache: Mutex::new(TimedLfuCache::with_capacity_and_expiration(
                128,
                Duration::from_secs(1 * 60 * 60),
            )),
        });

        let l2r_task = {
            let assoc = assoc.clone();
            assoc.copy_l2r(receiver)
        };
        tokio::spawn(l2r_task);

        (assoc, sender)
    }

    async fn copy_l2r(self: Arc<Self>, mut receiver: mpsc::Receiver<(Address, Bytes)>) {
        while let Some((target_addr, data)) = receiver.recv().await {
            trace!(
                "udp relay {} -> {} with {} bytes",
                self.peer_addr,
                target_addr,
                data.len()
            );

            let assoc = self.clone();
            if let Err(err) = assoc.copy_l2r_dispatch(&target_addr, &data).await {
                error!(
                    "udp relay {} -> {} with {} bytes, error: {}",
                    self.peer_addr,
                    target_addr,
                    data.len(),
                    err
                );
            }
        }
    }

    async fn copy_l2r_dispatch(self: Arc<Self>, target_addr: &Address, data: &[u8]) -> io::Result<()> {
        match *target_addr {
            Address::SocketAddress(sa) => match sa {
                SocketAddr::V4(..) => self.copy_ipv4_l2r_dispatch(sa, data).await,
                SocketAddr::V6(..) => self.copy_ipv6_l2r_dispatch(sa, data).await,
            },
            Address::DomainNameAddress(ref dname, port) => {
                let sa = lookup_then!(self.context.context_ref(), dname, port, |sa| {
                    match sa {
                        SocketAddr::V4(..) => self.clone().copy_ipv4_l2r_dispatch(sa, data).await,
                        SocketAddr::V6(..) => self.clone().copy_ipv6_l2r_dispatch(sa, data).await,
                    }
                })?
                .0;

                // Record resolved address as reverse index
                self.target_cache.lock().await.insert(sa, target_addr.clone());

                Ok(())
            }
        }
    }

    async fn copy_ipv4_l2r_dispatch(self: Arc<Self>, target_addr: SocketAddr, data: &[u8]) -> io::Result<()> {
        let outbound = {
            let mut handle = self.outbound_ipv4_socket.lock();

            match *handle {
                UdpAssociationState::Empty => {
                    // Create a new connection to proxy server

                    let socket =
                        OutboundUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                    let socket: Arc<OutboundUdpSocket> = Arc::new(socket);

                    let (r2l_fut, r2l_abortable) = {
                        let assoc = self.clone();
                        future::abortable(assoc.copy_r2l(socket.clone()))
                    };

                    // CLIENT <- REMOTE
                    tokio::spawn(r2l_fut);
                    debug!(
                        "created udp association for {} with {:?}",
                        self.peer_addr,
                        self.context.connect_opts_ref()
                    );

                    handle.set_connected(socket.clone(), r2l_abortable);
                    socket
                }
                UdpAssociationState::Connected { ref socket, .. } => socket.clone(),
                UdpAssociationState::Aborted => {
                    debug!(
                        "udp association for {} (bypassed) have been aborted, dropped packet {} bytes to {}",
                        self.peer_addr,
                        data.len(),
                        target_addr
                    );
                    return Ok(());
                }
            }
        };

        let n = outbound.send_to(data, target_addr).await?;
        if n != data.len() {
            warn!(
                "{} -> {} sent {} bytes != expected {} bytes",
                self.peer_addr,
                target_addr,
                n,
                data.len()
            );
        }

        Ok(())
    }

    async fn copy_ipv6_l2r_dispatch(self: Arc<Self>, target_addr: SocketAddr, data: &[u8]) -> io::Result<()> {
        let outbound = {
            let mut handle = self.outbound_ipv6_socket.lock();

            match *handle {
                UdpAssociationState::Empty => {
                    // Create a new connection to proxy server

                    let socket =
                        OutboundUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                    let socket: Arc<OutboundUdpSocket> = Arc::new(socket);

                    let (r2l_fut, r2l_abortable) = {
                        let assoc = self.clone();
                        future::abortable(assoc.copy_r2l(socket.clone()))
                    };

                    // CLIENT <- REMOTE
                    tokio::spawn(r2l_fut);
                    debug!(
                        "created udp association for {} with {:?}",
                        self.peer_addr,
                        self.context.connect_opts_ref()
                    );

                    handle.set_connected(socket.clone(), r2l_abortable);
                    socket
                }
                UdpAssociationState::Connected { ref socket, .. } => socket.clone(),
                UdpAssociationState::Aborted => {
                    debug!(
                        "udp association for {} (bypassed) have been aborted, dropped packet {} bytes to {}",
                        self.peer_addr,
                        data.len(),
                        target_addr
                    );
                    return Ok(());
                }
            }
        };

        let n = outbound.send_to(data, target_addr).await?;
        if n != data.len() {
            warn!(
                "{} -> {} sent {} bytes != expected {} bytes",
                self.peer_addr,
                target_addr,
                n,
                data.len()
            );
        }

        Ok(())
    }

    async fn copy_r2l(self: Arc<Self>, outbound: Arc<OutboundUdpSocket>) -> io::Result<()> {
        let mut buffer = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        loop {
            let (n, addr) = match outbound.recv_from(&mut buffer).await {
                Ok(r) => {
                    // Keep association alive in map
                    let _ = self.assoc_map.lock().await.get(&self.peer_addr);
                    r
                }
                Err(err) => {
                    error!(
                        "udp failed to receive from outbound socket, peer_addr: {}, error: {}",
                        self.peer_addr, err
                    );
                    time::sleep(Duration::from_secs(1)).await;
                    continue;
                }
            };

            trace!("udp relay {} <- {} received {} bytes", self.peer_addr, addr, n);

            let data = &buffer[..n];

            let target_addr = match self.target_cache.lock().await.get(&addr) {
                Some(a) => a.clone(),
                None => Address::from(addr),
            };

            trace!(
                "udp relay {} <- {} ({}) with {} bytes",
                self.peer_addr,
                target_addr,
                addr,
                data.len()
            );

            // Send back to client
            if let Err(err) = self.inbound.send_to(self.peer_addr, &target_addr, data).await {
                warn!(
                    "udp failed to send back to client {}, from target {} ({}), error: {}",
                    self.peer_addr, target_addr, addr, err
                );
            }

            trace!(
                "udp relay {} <- {} ({}) with {} bytes",
                self.peer_addr,
                target_addr,
                addr,
                data.len()
            );
        }
    }
}
<commit_msg>record every resolved target SocketAddr<commit_after>//! Shadowsocks UDP server

use std::{io, net::SocketAddr, sync::Arc, time::Duration};

use bytes::Bytes;
use futures::future::{self, AbortHandle};
use io::ErrorKind;
use lfu_cache::TimedLfuCache;
use log::{debug, error, info, trace, warn};
use shadowsocks::{
    lookup_then,
    net::UdpSocket as OutboundUdpSocket,
    relay::{
        socks5::Address,
        udprelay::{ProxySocket, MAXIMUM_UDP_PAYLOAD_SIZE},
    },
    ServerConfig,
};
use spin::Mutex as SpinMutex;
use tokio::{
    sync::{mpsc, Mutex},
    time,
};

use crate::net::MonProxySocket;

use super::context::ServiceContext;

type AssociationMap = TimedLfuCache<SocketAddr, UdpAssociation>;
type SharedAssociationMap = Arc<Mutex<AssociationMap>>;

pub struct UdpServer {
    context: Arc<ServiceContext>,
    assoc_map: SharedAssociationMap,
    cleanup_abortable: AbortHandle,
}

impl Drop for UdpServer {
    fn drop(&mut self) {
        self.cleanup_abortable.abort();
    }
}

impl UdpServer {
    pub fn new(context: Arc<ServiceContext>, time_to_live: Option<Duration>, capacity: Option<usize>) -> UdpServer {
        let time_to_live = time_to_live.unwrap_or(crate::DEFAULT_UDP_EXPIRY_DURATION);
        let assoc_map = Arc::new(Mutex::new(match capacity {
            Some(capacity) => TimedLfuCache::with_capacity_and_expiration(capacity, time_to_live),
            None => TimedLfuCache::with_expiration(time_to_live),
        }));

        let cleanup_abortable = {
            let assoc_map = assoc_map.clone();
            let (cleanup_task, cleanup_abortable) = future::abortable(async move {
                loop {
                    time::sleep(time_to_live).await;

                    // cleanup expired associations
                    assoc_map.lock().await.evict_expired();
                }
            });
            tokio::spawn(cleanup_task);
            cleanup_abortable
        };

        UdpServer {
            context,
            assoc_map,
            cleanup_abortable,
        }
    }

    pub async fn run(mut self, svr_cfg: &ServerConfig) -> io::Result<()> {
        let socket = ProxySocket::bind(self.context.context(), svr_cfg).await?;

        info!(
            "shadowsocks udp server listening on {}",
            socket.local_addr().expect("listener.local_addr"),
        );

        let socket = MonProxySocket::from_socket(socket, self.context.flow_stat());
        let listener = Arc::new(socket);

        let mut buffer = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        loop {
            let (n, peer_addr, target_addr) = match listener.recv_from(&mut buffer).await {
                Ok(s) => s,
                Err(err) => {
                    error!("udp server recv_from failed with error: {}", err);
                    continue;
                }
            };

            if self.context.check_outbound_blocked(&target_addr).await {
                error!("udp client {} outbound {} blocked by ACL rules", peer_addr, target_addr);
                continue;
            }

            let data = &buffer[..n];
            if let Err(err) = self.send_packet(&listener, peer_addr, target_addr, data).await {
                error!(
                    "udp packet relay {} with {} bytes failed, error: {}",
                    peer_addr,
                    data.len(),
                    err
                );
            }
        }
    }

    async fn send_packet(
        &mut self,
        listener: &Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        target_addr: Address,
        data: &[u8],
    ) -> io::Result<()> {
        let mut assoc_map = self.assoc_map.lock().await;

        if let Some(assoc) = assoc_map.get(&peer_addr) {
            return assoc.try_send((target_addr, Bytes::copy_from_slice(data)));
        }

        let assoc = UdpAssociation::new(
            self.context.clone(),
            listener.clone(),
            peer_addr,
            self.assoc_map.clone(),
        );

        trace!("created udp association for {}", peer_addr);

        assoc.try_send((target_addr, Bytes::copy_from_slice(data)))?;
        assoc_map.insert(peer_addr, assoc);

        Ok(())
    }
}

struct UdpAssociation {
    assoc: Arc<UdpAssociationContext>,
    sender: mpsc::Sender<(Address, Bytes)>,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.assoc.outbound_ipv4_socket.lock().abort();
        self.assoc.outbound_ipv6_socket.lock().abort();
    }
}

impl UdpAssociation {
    fn new(
        context: Arc<ServiceContext>,
        inbound: Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        assoc_map: SharedAssociationMap,
    ) -> UdpAssociation {
        let (assoc, sender) = UdpAssociationContext::new(context, inbound, peer_addr, assoc_map);
        UdpAssociation { assoc, sender }
    }

    fn try_send(&self, data: (Address, Bytes)) -> io::Result<()> {
        if let Err(..) = self.sender.try_send(data) {
            let err = io::Error::new(ErrorKind::Other, "udp relay channel full");
            return Err(err);
        }
        Ok(())
    }
}

enum UdpAssociationState {
    Empty,
    Connected {
        socket: Arc<OutboundUdpSocket>,
        abortable: AbortHandle,
    },
    Aborted,
}

impl Drop for UdpAssociationState {
    fn drop(&mut self) {
        if let UdpAssociationState::Connected { ref abortable, .. } = *self {
            abortable.abort();
        }
    }
}

impl UdpAssociationState {
    fn empty() -> UdpAssociationState {
        UdpAssociationState::Empty
    }

    fn set_connected(&mut self, socket: Arc<OutboundUdpSocket>, abortable: AbortHandle) {
        *self = UdpAssociationState::Connected { socket, abortable };
    }

    fn abort(&mut self) {
        *self = UdpAssociationState::Aborted;
    }
}

struct UdpAssociationContext {
    context: Arc<ServiceContext>,
    inbound: Arc<MonProxySocket>,
    peer_addr: SocketAddr,
    outbound_ipv4_socket: SpinMutex<UdpAssociationState>,
    outbound_ipv6_socket: SpinMutex<UdpAssociationState>,
    assoc_map: SharedAssociationMap,
    target_cache: Mutex<TimedLfuCache<SocketAddr, Address>>,
}

impl Drop for UdpAssociationContext {
    fn drop(&mut self) {
        trace!("udp association for {} is closed", self.peer_addr);
    }
}

impl UdpAssociationContext {
    fn new(
        context: Arc<ServiceContext>,
        inbound: Arc<MonProxySocket>,
        peer_addr: SocketAddr,
        assoc_map: SharedAssociationMap,
    ) -> (Arc<UdpAssociationContext>, mpsc::Sender<(Address, Bytes)>) {
        // Pending packets 1024 should be good enough for a server.
        // If there are plenty of packets stuck in the channel, dropping exccess packets is a good way to protect the server from
        // being OOM.
        let (sender, receiver) = mpsc::channel(1024);

        let assoc = Arc::new(UdpAssociationContext {
            context,
            inbound,
            peer_addr,
            outbound_ipv4_socket: SpinMutex::new(UdpAssociationState::empty()),
            outbound_ipv6_socket: SpinMutex::new(UdpAssociationState::empty()),
            assoc_map,
            // Cache for remembering the original Address of target,
            // when recv_from a SocketAddr, we have to know whch Address that client was originally requested.
            //
            // XXX: 128 target addresses should be enough for __one__ client.
            //      1 hours should be enough for caching the address mapping. Most of the DNS records' TTL won't last that long.
            target_cache: Mutex::new(TimedLfuCache::with_capacity_and_expiration(
                128,
                Duration::from_secs(1 * 60 * 60),
            )),
        });

        let l2r_task = {
            let assoc = assoc.clone();
            assoc.copy_l2r(receiver)
        };
        tokio::spawn(l2r_task);

        (assoc, sender)
    }

    async fn copy_l2r(self: Arc<Self>, mut receiver: mpsc::Receiver<(Address, Bytes)>) {
        while let Some((target_addr, data)) = receiver.recv().await {
            trace!(
                "udp relay {} -> {} with {} bytes",
                self.peer_addr,
                target_addr,
                data.len()
            );

            let assoc = self.clone();
            if let Err(err) = assoc.copy_l2r_dispatch(&target_addr, &data).await {
                error!(
                    "udp relay {} -> {} with {} bytes, error: {}",
                    self.peer_addr,
                    target_addr,
                    data.len(),
                    err
                );
            }
        }
    }

    async fn copy_l2r_dispatch(self: Arc<Self>, target_addr: &Address, data: &[u8]) -> io::Result<()> {
        match *target_addr {
            Address::SocketAddress(sa) => match sa {
                SocketAddr::V4(..) => self.copy_ipv4_l2r_dispatch(sa, data).await,
                SocketAddr::V6(..) => self.copy_ipv6_l2r_dispatch(sa, data).await,
            },
            Address::DomainNameAddress(ref dname, port) => {
                lookup_then!(self.context.context_ref(), dname, port, |sa| {
                    // Record resolved address as reverse index
                    self.target_cache.lock().await.insert(sa, target_addr.clone());

                    match sa {
                        SocketAddr::V4(..) => self.clone().copy_ipv4_l2r_dispatch(sa, data).await,
                        SocketAddr::V6(..) => self.clone().copy_ipv6_l2r_dispatch(sa, data).await,
                    }
                })?;

                Ok(())
            }
        }
    }

    async fn copy_ipv4_l2r_dispatch(self: Arc<Self>, target_addr: SocketAddr, data: &[u8]) -> io::Result<()> {
        let outbound = {
            let mut handle = self.outbound_ipv4_socket.lock();

            match *handle {
                UdpAssociationState::Empty => {
                    // Create a new connection to proxy server

                    let socket =
                        OutboundUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                    let socket: Arc<OutboundUdpSocket> = Arc::new(socket);

                    let (r2l_fut, r2l_abortable) = {
                        let assoc = self.clone();
                        future::abortable(assoc.copy_r2l(socket.clone()))
                    };

                    // CLIENT <- REMOTE
                    tokio::spawn(r2l_fut);
                    debug!(
                        "created udp association for {} with {:?}",
                        self.peer_addr,
                        self.context.connect_opts_ref()
                    );

                    handle.set_connected(socket.clone(), r2l_abortable);
                    socket
                }
                UdpAssociationState::Connected { ref socket, .. } => socket.clone(),
                UdpAssociationState::Aborted => {
                    debug!(
                        "udp association for {} (bypassed) have been aborted, dropped packet {} bytes to {}",
                        self.peer_addr,
                        data.len(),
                        target_addr
                    );
                    return Ok(());
                }
            }
        };

        let n = outbound.send_to(data, target_addr).await?;
        if n != data.len() {
            warn!(
                "{} -> {} sent {} bytes != expected {} bytes",
                self.peer_addr,
                target_addr,
                n,
                data.len()
            );
        }

        Ok(())
    }

    async fn copy_ipv6_l2r_dispatch(self: Arc<Self>, target_addr: SocketAddr, data: &[u8]) -> io::Result<()> {
        let outbound = {
            let mut handle = self.outbound_ipv6_socket.lock();

            match *handle {
                UdpAssociationState::Empty => {
                    // Create a new connection to proxy server

                    let socket =
                        OutboundUdpSocket::connect_any_with_opts(&target_addr, self.context.connect_opts_ref()).await?;
                    let socket: Arc<OutboundUdpSocket> = Arc::new(socket);

                    let (r2l_fut, r2l_abortable) = {
                        let assoc = self.clone();
                        future::abortable(assoc.copy_r2l(socket.clone()))
                    };

                    // CLIENT <- REMOTE
                    tokio::spawn(r2l_fut);
                    debug!(
                        "created udp association for {} with {:?}",
                        self.peer_addr,
                        self.context.connect_opts_ref()
                    );

                    handle.set_connected(socket.clone(), r2l_abortable);
                    socket
                }
                UdpAssociationState::Connected { ref socket, .. } => socket.clone(),
                UdpAssociationState::Aborted => {
                    debug!(
                        "udp association for {} (bypassed) have been aborted, dropped packet {} bytes to {}",
                        self.peer_addr,
                        data.len(),
                        target_addr
                    );
                    return Ok(());
                }
            }
        };

        let n = outbound.send_to(data, target_addr).await?;
        if n != data.len() {
            warn!(
                "{} -> {} sent {} bytes != expected {} bytes",
                self.peer_addr,
                target_addr,
                n,
                data.len()
            );
        }

        Ok(())
    }

    async fn copy_r2l(self: Arc<Self>, outbound: Arc<OutboundUdpSocket>) -> io::Result<()> {
        let mut buffer = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        loop {
            let (n, addr) = match outbound.recv_from(&mut buffer).await {
                Ok(r) => {
                    // Keep association alive in map
                    let _ = self.assoc_map.lock().await.get(&self.peer_addr);
                    r
                }
                Err(err) => {
                    error!(
                        "udp failed to receive from outbound socket, peer_addr: {}, error: {}",
                        self.peer_addr, err
                    );
                    time::sleep(Duration::from_secs(1)).await;
                    continue;
                }
            };

            trace!("udp relay {} <- {} received {} bytes", self.peer_addr, addr, n);

            let data = &buffer[..n];

            let target_addr = match self.target_cache.lock().await.get(&addr) {
                Some(a) => a.clone(),
                None => Address::from(addr),
            };

            // Send back to client
            if let Err(err) = self.inbound.send_to(self.peer_addr, &target_addr, data).await {
                warn!(
                    "udp failed to send back to client {}, from target {} ({}), error: {}",
                    self.peer_addr, target_addr, addr, err
                );
            }

            trace!(
                "udp relay {} <- {} ({}) with {} bytes",
                self.peer_addr,
                target_addr,
                addr,
                data.len()
            );
        }
    }
}
<|endoftext|>


<commit_before>extern crate url;

use std::path::Path;
use self::url::Url;
use page::Page;

use util;

fn build_github_url(user: &str, repo: &str, branch: &Option<String>, page: &Page) -> String {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            format!("https://github.com/{}/{}/tree/{}", user, repo, b)
        } else {
            format!("https://github.com/{}/{}", user, repo)
        },
        &Page::Diff {ref lhs, ref rhs} => format!("https://github.com/{}/{}/compare/{}...{}", user, repo, lhs, rhs),
        &Page::Commit {ref hash} => format!("https://github.com/{}/{}/commit/{}", user, repo, hash),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => format!("https://github.com/{}/{}/blob/{}/{}", user, repo, hash, relative_path),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => format!("https://github.com/{}/{}/blob/{}/{}#L{}", user, repo, hash, relative_path, line),
    }
}

fn build_bitbucket_url(user: &str, repo: &str, branch: &Option<String>, page: &Page) -> util::Result<String> {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            Ok(format!("https://bitbucket.org/{}/{}/branch/{}", user, repo, b))
        } else {
            Ok(format!("https://bitbucket.org/{}/{}", user, repo))
        },
        &Page::Diff {..} => Err("BitBucket does not support diff between commits (see https://bitbucket.org/site/master/issues/4779/ability-to-diff-between-any-two-commits)".to_string()),
        &Page::Commit {ref hash} => Ok(format!("https://bitbucket.org/{}/{}/commits/{}", user, repo, hash)),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => Ok(format!("https://bitbucket.org/{}/{}/src/{}/{}", user, repo, hash, relative_path)),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => {
            let file = Path::new(relative_path)
                .file_name()
                .ok_or(format!("Cannot get file name from path: {}", relative_path))?
                .to_str()
                .ok_or(format!("Cannot convert path to UTF8 string: {}", relative_path))?;
            Ok(format!("https://bitbucket.org/{}/{}/src/{}/{}#{}-{}", user, repo, hash, relative_path, file, line))
        },
    }
}

fn build_github_enterprise_url(host: &str, user: &str, repo: &str, branch: &Option<String>, page: &Page) -> String {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            format!("https://{}/{}/{}/tree/{}", host, user, repo, b)
        } else {
            format!("https://{}/{}/{}", host, user, repo)
        },
        &Page::Diff {ref lhs, ref rhs} => format!("https://{}/{}/{}/compare/{}...{}", host, user, repo, lhs, rhs),
        &Page::Commit {ref hash} => format!("https://{}/{}/{}/commit/{}", host, user, repo, hash),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => format!("https://{}/{}/{}/blob/{}/{}", host, user, repo, hash, relative_path),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => format!("https://{}/{}/{}/blob/{}/{}#L{}", host, user, repo, hash, relative_path, line),
    }
}

fn build_gitlab_url(user: &str, repo: &str, branch: &Option<String>, page: &Page) -> String {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            format!("https://gitlab.com/{}/{}/tree/{}", user, repo, b)
        } else {
            format!("https://gitlab.com/{}/{}", user, repo)
        },
        &Page::Diff {ref lhs, ref rhs} => format!("https://gitlab.com/{}/{}/compare/{}...{}", user, repo, lhs, rhs),
        &Page::Commit {ref hash} => format!("https://gitlab.com/{}/{}/commit/{}", user, repo, hash),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => format!("https://gitlab.com/{}/{}/blob/{}/{}", user, repo, hash, relative_path),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => format!("https://gitlab.com/{}/{}/blob/{}/{}#L{}", user, repo, hash, relative_path, line),
    }
}

// Note: Parse '/user/repo.git' or '/user/repo' or 'user/repo' into 'user' and 'repo'
fn user_and_repo_from_path<'a>(path: &'a str) -> util::Result<(&'a str, &'a str)> {
    let mut split = path.split('/').skip_while(|s| s.is_empty());
    let user = split.next().ok_or(format!("Can't detect user name from path {}", path))?;
    let mut repo = split.next().ok_or(format!("Can't detect repository name from path {}", path))?;
    if repo.ends_with(".git") {
        // Slice '.git' from 'repo.git'
        repo = &repo[0 .. repo.len() - 4];
    }
    Ok((user, repo))
}

// Known URL formats
//  https://hosting_service.com/user/repo.git
//  git@hosting_service.com:user/repo.git (-> ssh://git@hosting_service.com:22/user/repo.git)
pub fn parse_and_build_page_url(repo: &String, page: &Page, branch: &Option<String>) -> util::Result<String> {
    let url = Url::parse(repo).map_err(|e| format!("{}", e))?;
    let path = url.path();
    let (user, repo_name) = user_and_repo_from_path(path)?;
    let host = url.host_str().ok_or(format!("Failed to parse host from {}", repo))?;
    match host {
        "github.com" => Ok(build_github_url(user, repo_name, branch, page)),
        "bitbucket.org" => build_bitbucket_url(user, repo_name, branch, page),
        "gitlab.com" => Ok(build_gitlab_url(user, repo_name, branch, page)),
        host => if host.starts_with("github.") {
            Ok(build_github_enterprise_url(host, user, repo_name, branch, page))
        } else {
            Err(format!("Unknown hosting service for URL {}", repo))
        },
    }
}
<commit_msg>refactor: reduce duplicates on building page URL<commit_after>extern crate url;

use std::path::Path;
use self::url::Url;
use page::Page;

use util;

fn build_github_like_url(host: &str, user: &str, repo: &str, branch: &Option<String>, page: &Page) -> String {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            format!("https://{}/{}/{}/tree/{}", host, user, repo, b)
        } else {
            format!("https://{}/{}/{}", host, user, repo)
        },
        &Page::Diff {ref lhs, ref rhs} => format!("https://{}/{}/{}/compare/{}...{}", host, user, repo, lhs, rhs),
        &Page::Commit {ref hash} => format!("https://{}/{}/{}/commit/{}", host, user, repo, hash),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => format!("https://{}/{}/{}/blob/{}/{}", host, user, repo, hash, relative_path),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => format!("https://{}/{}/{}/blob/{}/{}#L{}", host, user, repo, hash, relative_path, line),
    }
}

fn build_bitbucket_url(user: &str, repo: &str, branch: &Option<String>, page: &Page) -> util::Result<String> {
    match page {
        &Page::Open => if let &Some(ref b) = branch {
            Ok(format!("https://bitbucket.org/{}/{}/branch/{}", user, repo, b))
        } else {
            Ok(format!("https://bitbucket.org/{}/{}", user, repo))
        },
        &Page::Diff {..} => Err("BitBucket does not support diff between commits (see https://bitbucket.org/site/master/issues/4779/ability-to-diff-between-any-two-commits)".to_string()),
        &Page::Commit {ref hash} => Ok(format!("https://bitbucket.org/{}/{}/commits/{}", user, repo, hash)),
        &Page::FileOrDir {ref relative_path, ref hash, line: None} => Ok(format!("https://bitbucket.org/{}/{}/src/{}/{}", user, repo, hash, relative_path)),
        &Page::FileOrDir {ref relative_path, ref hash, line: Some(line)} => {
            let file = Path::new(relative_path)
                .file_name()
                .ok_or(format!("Cannot get file name from path: {}", relative_path))?
                .to_str()
                .ok_or(format!("Cannot convert path to UTF8 string: {}", relative_path))?;
            Ok(format!("https://bitbucket.org/{}/{}/src/{}/{}#{}-{}", user, repo, hash, relative_path, file, line))
        },
    }
}

// Note: Parse '/user/repo.git' or '/user/repo' or 'user/repo' into 'user' and 'repo'
fn user_and_repo_from_path<'a>(path: &'a str) -> util::Result<(&'a str, &'a str)> {
    let mut split = path.split('/').skip_while(|s| s.is_empty());
    let user = split.next().ok_or(format!("Can't detect user name from path {}", path))?;
    let mut repo = split.next().ok_or(format!("Can't detect repository name from path {}", path))?;
    if repo.ends_with(".git") {
        // Slice '.git' from 'repo.git'
        repo = &repo[0 .. repo.len() - 4];
    }
    Ok((user, repo))
}

// Known URL formats
//  1. https://hosting_service.com/user/repo.git
//  2. git@hosting_service.com:user/repo.git (-> ssh://git@hosting_service.com:22/user/repo.git)
pub fn parse_and_build_page_url(repo: &String, page: &Page, branch: &Option<String>) -> util::Result<String> {
    let url = Url::parse(repo).map_err(|e| format!("{}", e))?;
    let path = url.path();
    let (user, repo_name) = user_and_repo_from_path(path)?;
    let host = url.host_str().ok_or(format!("Failed to parse host from {}", repo))?;
    match host {
        "github.com" | "gitlab.com" => Ok(build_github_like_url(host, user, repo_name, branch, page)),
        "bitbucket.org" => build_bitbucket_url(user, repo_name, branch, page),
        _ => if host.starts_with("github.") {
            Ok(build_github_like_url(host, user, repo_name, branch, page))
        } else {
            Err(format!("Unknown hosting service for URL {}", repo))
        },
    }
}
<|endoftext|>


<commit_before>// In a chat window add dozens of nicks, each printing some random lines.

extern crate libc;
extern crate mio;
extern crate term_input;
extern crate termbox_simple;
extern crate time;
extern crate tiny;

use mio::Events;
use mio::Poll;
use mio::PollOpt;
use mio::Ready;
use mio::Token;
use mio::unix::EventedFd;
use std::fs::File;
use std::io::Read;
use std::rc::Rc;

use term_input::{Event, Input};
use tiny::config::Colors;
use tiny::tui::{MsgTarget, TUIRet, Timestamp, TUI};

fn main() {
    let chan_target = MsgTarget::Chan {
        serv_name: "debug",
        chan_name: "chan",
    };

    let mut tui = TUI::new(Colors::default());
    tui.new_server_tab("debug");
    tui.new_chan_tab("debug", "chan");
    tui.show_topic("This is channel topic", Timestamp::now(), &chan_target);
    tui.draw();

    {
        let mut text = String::new();
        let mut file = File::open("test/lipsum.txt").unwrap();
        file.read_to_string(&mut text).unwrap();

        for (line_idx, line) in text.lines().enumerate() {
            let now = Timestamp::now();
            let nick = format!("nick_{}", line_idx);
            tui.add_nick(&nick, Some(now), &chan_target);
            tui.add_privmsg(&nick, line, now, &chan_target, false);
        }
    }

    tui.set_nick("debug", Rc::new("some_long_nick_name____".to_owned()));
    tui.draw();

    let poll = Poll::new().unwrap();
    poll.register(
        &EventedFd(&libc::STDIN_FILENO),
        Token(libc::STDIN_FILENO as usize),
        Ready::readable(),
        PollOpt::level(),
    ).unwrap();

    let mut input = Input::new();
    let mut ev_buffer: Vec<Event> = Vec::new();
    let mut events = Events::with_capacity(10);
    'mainloop: loop {
        match poll.poll(&mut events, None) {
            Err(_) => {
                // usually SIGWINCH, which is caught by term_input
                tui.resize();
                tui.draw();
            }
            Ok(_) => {
                input.read_input_events(&mut ev_buffer);
                for ev in ev_buffer.drain(0..) {
                    match tui.handle_input_event(ev) {
                        TUIRet::Input { msg, from } =>
                            if msg == "/clear".chars().collect::<Vec<char>>() {
                                tui.clear(&from.to_target())
                            } else {
                                tui.add_msg(
                                    &msg.into_iter().collect::<String>(),
                                    Timestamp::now(),
                                    &MsgTarget::Server { serv_name: "debug" },
                                );
                            },
                        TUIRet::Abort => {
                            break 'mainloop;
                        }
                        TUIRet::EventIgnored(Event::FocusGained) => {
                            tui.add_msg(
                                "focus gained",
                                Timestamp::now(),
                                &MsgTarget::Server { serv_name: "debug" },
                            );
                        }
                        TUIRet::EventIgnored(Event::FocusLost) => {
                            tui.add_msg(
                                "focus lost",
                                Timestamp::now(),
                                &MsgTarget::Server { serv_name: "debug" },
                            );
                        }
                        _ =>
                            {}
                    }
                }
                tui.draw();
            }
        }
    }
}
<commit_msg>Fix example program after 643d17b<commit_after>// In a chat window add dozens of nicks, each printing some random lines.

extern crate libc;
extern crate mio;
extern crate term_input;
extern crate termbox_simple;
extern crate time;
extern crate tiny;

use mio::Events;
use mio::Poll;
use mio::PollOpt;
use mio::Ready;
use mio::Token;
use mio::unix::EventedFd;
use std::fs::File;
use std::io::Read;

use term_input::{Event, Input};
use tiny::config::Colors;
use tiny::tui::{MsgTarget, TUIRet, Timestamp, TUI};

fn main() {
    let chan_target = MsgTarget::Chan {
        serv_name: "debug",
        chan_name: "chan",
    };

    let mut tui = TUI::new(Colors::default());
    tui.new_server_tab("debug");
    tui.new_chan_tab("debug", "chan");
    tui.show_topic("This is channel topic", Timestamp::now(), &chan_target);
    tui.draw();

    {
        let mut text = String::new();
        let mut file = File::open("test/lipsum.txt").unwrap();
        file.read_to_string(&mut text).unwrap();

        for (line_idx, line) in text.lines().enumerate() {
            let now = Timestamp::now();
            let nick = format!("nick_{}", line_idx);
            tui.add_nick(&nick, Some(now), &chan_target);
            tui.add_privmsg(&nick, line, now, &chan_target, false);
        }
    }

    tui.set_nick("debug", "some_long_nick_name____");
    tui.draw();

    let poll = Poll::new().unwrap();
    poll.register(
        &EventedFd(&libc::STDIN_FILENO),
        Token(libc::STDIN_FILENO as usize),
        Ready::readable(),
        PollOpt::level(),
    ).unwrap();

    let mut input = Input::new();
    let mut ev_buffer: Vec<Event> = Vec::new();
    let mut events = Events::with_capacity(10);
    'mainloop: loop {
        match poll.poll(&mut events, None) {
            Err(_) => {
                // usually SIGWINCH, which is caught by term_input
                tui.resize();
                tui.draw();
            }
            Ok(_) => {
                input.read_input_events(&mut ev_buffer);
                for ev in ev_buffer.drain(0..) {
                    match tui.handle_input_event(ev) {
                        TUIRet::Input { msg, from } =>
                            if msg == "/clear".chars().collect::<Vec<char>>() {
                                tui.clear(&from.to_target())
                            } else {
                                tui.add_msg(
                                    &msg.into_iter().collect::<String>(),
                                    Timestamp::now(),
                                    &MsgTarget::Server { serv_name: "debug" },
                                );
                            },
                        TUIRet::Abort => {
                            break 'mainloop;
                        }
                        TUIRet::EventIgnored(Event::FocusGained) => {
                            tui.add_msg(
                                "focus gained",
                                Timestamp::now(),
                                &MsgTarget::Server { serv_name: "debug" },
                            );
                        }
                        TUIRet::EventIgnored(Event::FocusLost) => {
                            tui.add_msg(
                                "focus lost",
                                Timestamp::now(),
                                &MsgTarget::Server { serv_name: "debug" },
                            );
                        }
                        _ =>
                            {}
                    }
                }
                tui.draw();
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;

use serde::Deserialize;
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, StringWithSeparator, TryFromInto};
use shakmaty::{
    fen::Fen,
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode, Color,
};

use crate::{
    api::{Error, LaxFen, LilaVariant},
    model::{Mode, Month, RatingGroup, Speed, UserName, Year},
    opening::{Opening, Openings},
};

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct MastersQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde_as(as = "TryFromInto<u16>")]
    #[serde(default)]
    pub since: Year,
    #[serde_as(as = "TryFromInto<u16>")]
    #[serde(default = "Year::max_value")]
    pub until: Year,
    #[serde(flatten)]
    pub limits: Limits,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct LichessQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde(flatten)]
    pub limits: Limits,
    #[serde(flatten)]
    pub filter: LichessQueryFilter,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct LichessQueryFilter {
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Speed>>")]
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, RatingGroup>>")]
    #[serde(default)]
    pub ratings: Option<Vec<RatingGroup>>,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default)]
    pub since: Month,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "Month::max_value")]
    pub until: Month,
}

impl LichessQueryFilter {
    pub fn contains_speed(&self, speed: Speed) -> bool {
        self.speeds
            .as_ref()
            .map_or(true, |speeds| speeds.contains(&speed))
    }

    pub fn contains_rating_group(&self, rating_group: RatingGroup) -> bool {
        self.ratings.as_ref().map_or(true, |ratings| {
            ratings.contains(&min(rating_group, RatingGroup::Group2500))
        })
    }

    pub fn top_group(&self) -> Option<RatingGroup> {
        let mut top_group = None;
        for group in RatingGroup::ALL.into_iter().rev() {
            if !self.contains_rating_group(group) || group < RatingGroup::Group2000 {
                break;
            } else {
                top_group = Some(group);
            }
        }
        top_group
    }
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PlayerQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde_as(as = "DisplayFromStr")]
    pub player: UserName,
    #[serde_as(as = "DisplayFromStr")]
    pub color: Color,
    #[serde(flatten)]
    pub filter: PlayerQueryFilter,
    #[serde(flatten)]
    pub limits: Limits,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PlayerQueryFilter {
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Mode>>")]
    #[serde(default)]
    pub modes: Option<Vec<Mode>>,
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Speed>>")]
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default)]
    pub since: Month,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "Month::max_value")]
    pub until: Month,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct Play {
    #[serde(default)]
    pub variant: LilaVariant,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub fen: Option<LaxFen>,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    pub play: Vec<Uci>,
}

impl Play {
    pub fn position<'a>(
        self,
        openings: &'a Openings,
    ) -> Result<(Variant, Zobrist<VariantPosition, u128>, Option<&'a Opening>), Error> {
        let variant = Variant::from(self.variant);
        let mut pos = Zobrist::new(match self.fen {
            Some(fen) => {
                VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?
            }
            None => VariantPosition::new(variant),
        });
        let opening = openings.classify_and_play(&mut pos, self.play)?;
        Ok((variant, pos, opening))
    }
}

#[serde_as]
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Limits {
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "usize::max_value")]
    pub top_games: usize,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "usize::max_value")]
    pub recent_games: usize,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub moves: Option<usize>,
}
<commit_msg>elide a lifetime<commit_after>use std::cmp::min;

use serde::Deserialize;
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, StringWithSeparator, TryFromInto};
use shakmaty::{
    fen::Fen,
    uci::Uci,
    variant::{Variant, VariantPosition},
    zobrist::Zobrist,
    CastlingMode, Color,
};

use crate::{
    api::{Error, LaxFen, LilaVariant},
    model::{Mode, Month, RatingGroup, Speed, UserName, Year},
    opening::{Opening, Openings},
};

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct MastersQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde_as(as = "TryFromInto<u16>")]
    #[serde(default)]
    pub since: Year,
    #[serde_as(as = "TryFromInto<u16>")]
    #[serde(default = "Year::max_value")]
    pub until: Year,
    #[serde(flatten)]
    pub limits: Limits,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct LichessQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde(flatten)]
    pub limits: Limits,
    #[serde(flatten)]
    pub filter: LichessQueryFilter,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct LichessQueryFilter {
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Speed>>")]
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, RatingGroup>>")]
    #[serde(default)]
    pub ratings: Option<Vec<RatingGroup>>,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default)]
    pub since: Month,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "Month::max_value")]
    pub until: Month,
}

impl LichessQueryFilter {
    pub fn contains_speed(&self, speed: Speed) -> bool {
        self.speeds
            .as_ref()
            .map_or(true, |speeds| speeds.contains(&speed))
    }

    pub fn contains_rating_group(&self, rating_group: RatingGroup) -> bool {
        self.ratings.as_ref().map_or(true, |ratings| {
            ratings.contains(&min(rating_group, RatingGroup::Group2500))
        })
    }

    pub fn top_group(&self) -> Option<RatingGroup> {
        let mut top_group = None;
        for group in RatingGroup::ALL.into_iter().rev() {
            if !self.contains_rating_group(group) || group < RatingGroup::Group2000 {
                break;
            } else {
                top_group = Some(group);
            }
        }
        top_group
    }
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PlayerQuery {
    #[serde(flatten)]
    pub play: Play,
    #[serde_as(as = "DisplayFromStr")]
    pub player: UserName,
    #[serde_as(as = "DisplayFromStr")]
    pub color: Color,
    #[serde(flatten)]
    pub filter: PlayerQueryFilter,
    #[serde(flatten)]
    pub limits: Limits,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PlayerQueryFilter {
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Mode>>")]
    #[serde(default)]
    pub modes: Option<Vec<Mode>>,
    #[serde_as(as = "Option<StringWithSeparator<CommaSeparator, Speed>>")]
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default)]
    pub since: Month,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "Month::max_value")]
    pub until: Month,
}

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct Play {
    #[serde(default)]
    pub variant: LilaVariant,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub fen: Option<LaxFen>,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    pub play: Vec<Uci>,
}

impl Play {
    pub fn position(
        self,
        openings: &Openings,
    ) -> Result<(Variant, Zobrist<VariantPosition, u128>, Option<&Opening>), Error> {
        let variant = Variant::from(self.variant);
        let mut pos = Zobrist::new(match self.fen {
            Some(fen) => {
                VariantPosition::from_setup(variant, &Fen::from(fen), CastlingMode::Chess960)?
            }
            None => VariantPosition::new(variant),
        });
        let opening = openings.classify_and_play(&mut pos, self.play)?;
        Ok((variant, pos, opening))
    }
}

#[serde_as]
#[derive(Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Limits {
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "usize::max_value")]
    pub top_games: usize,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default = "usize::max_value")]
    pub recent_games: usize,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub moves: Option<usize>,
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate test;

#[macro_use]
extern crate nom;

use nom::digit;
use nom::alphanumeric;
use std::str::FromStr;
use std::str;

named!(symbol_name, call!(alphanumeric));

named!(number<u64>,
    map_res!(
        map_res!(
            digit,
            str::from_utf8
        ),
        FromStr::from_str
    )
);

// ========================================================================== //

#[derive(Debug, PartialEq)]
pub enum UnaryOperator {
    Not,
    Minus,
    Negate,
}

#[derive(Debug, PartialEq)]
pub enum BinaryOperator {
    LogicOr,
    LogicAnd,
    Or,
    And,
    Equals,
    NotEquals,
    Greater,
    Less,
    LessOrEquals,
    GreaterOrEquals,
    ShiftRight,
    ShiftLeft,
    Plus,
    Minus,
    Multiply,
    Divide,
    Remainder,
}

#[derive(Debug, PartialEq)]
pub enum Expression<'a> {
    Ident(&'a str),
    Number(u64),
    Call {
        function: &'a str,
        args: Vec<Expression<'a>>,
    },
    UnaryOp {
        operator: UnaryOperator,
        right: Box<Expression<'a>>,
    },
    BinaryOp {
        left: Box<Expression<'a>>,
        operator: BinaryOperator,
        right: Box<Expression<'a>>,
    },
    TernaryOp {
        condition: Box<Expression<'a>>,
        left: Box<Expression<'a>>,
        right: Box<Expression<'a>>,
    },
}

named!(value_ident<Expression>, map!(
    map_res!(
        symbol_name,
        str::from_utf8
    ),
    |x| Expression::Ident(x)
));

named!(value_number<Expression>, map!(
    number,
    |x| Expression::Number(x)
));

named!(value_nested<Expression>, delimited!(
    tag!("("),
    ws!(expression),
    tag!(")")
));

named!(value_call<Expression>, do_parse!(
    func: map_res!(symbol_name, str::from_utf8)
    >>
    ws!(tag!("("))
    >>
    args: separated_list!(
        ws!(tag!(",")),
        expression
    )
    >>
    ws!(tag!(")"))
    >>
    (Expression::Call{
        function: func,
        args: args
    })
));

named!(pub value<Expression>, alt_complete!(
    value_call | value_nested | value_number | value_ident
));

named!(expr_unary_op<Expression>, do_parse!(
    op: ws!(alt_complete!(
        tag!("-") | tag!("!") | tag!("~")
    ))
    >>
    right: expr_level_1
    >>
    (Expression::UnaryOp{
        operator: match op[0] as char {
            '-' => UnaryOperator::Minus,
            '!' => UnaryOperator::Not,
            '~' => UnaryOperator::Negate,
            _ => panic!("Invalid operator"),
        },
        right: Box::new(right),
    })
));

named!(expr_level_1<Expression>, alt_complete!(
    expr_unary_op | value
));

named!(expr_level_2<Expression>, do_parse!(
    first: expr_level_1
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("*") | tag!("/") | tag!("%")
            )),
            expr_level_1
        ),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0[0] as char {
                    '*' => BinaryOperator::Multiply,
                    '/' => BinaryOperator::Divide,
                    '%' => BinaryOperator::Remainder,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_3<Expression>, do_parse!(
    first: expr_level_2
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("+") | tag!("-")
            )),
            expr_level_2
        ),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0[0] as char {
                    '+' => BinaryOperator::Plus,
                    '-' => BinaryOperator::Minus,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_4<Expression>, do_parse!(
    first: expr_level_3
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("<<") | tag!(">>")
            )),
            expr_level_3
        ),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0 {
                    b"<<" => BinaryOperator::ShiftLeft,
                    b">>" => BinaryOperator::ShiftRight,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_5<Expression>, do_parse!(
    first: expr_level_4
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("==") | tag!("!=") | tag!("<=") | tag!(">=") | tag!("<") | tag!(">")
            )),
            expr_level_4
        ),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0 {
                    b"==" => BinaryOperator::Equals,
                    b"!=" => BinaryOperator::NotEquals,
                    b"<=" => BinaryOperator::LessOrEquals,
                    b">=" => BinaryOperator::GreaterOrEquals,
                    b"<" => BinaryOperator::Less,
                    b">" => BinaryOperator::Greater,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_6<Expression>, do_parse!(
    first: expr_level_5
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("&")), expr_level_5),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::And,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_7<Expression>, do_parse!(
    first: expr_level_6
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("|")), expr_level_6),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::Or,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_8<Expression>, do_parse!(
    first: expr_level_7
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("&&")), expr_level_7),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::LogicAnd,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_9<Expression>, do_parse!(
    first: expr_level_8
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("||")), expr_level_8),
        first,
        |prev, new: (&'a [u8], Expression<'a>)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::LogicOr,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_ternary_op<Expression>, do_parse!(
    cond: expr_level_9
    >>
    ws!(tag!("?"))
    >>
    left: expression
    >>
    ws!(tag!(":"))
    >>
    right: expression
    >>
    (Expression::TernaryOp{
        condition: Box::new(cond),
        left: Box::new(left),
        right: Box::new(right),
    })
));

named!(pub expression<Expression>, alt_complete!(
    expr_ternary_op | expr_level_9
));

macro_rules! assert_done {
    ($res:expr) => (
        match $res {
            nom::IResult::Done(b"", v @ _) => {
                println!("ok: {:?}", v);
            }
            r @ _ => panic!("fail: {:?}", r),
        }
    )
}

#[test]
fn test_ws() {
    let x = b"a ( b ( d , ( 0 ) ) , c )";
    assert_done!(expression(x));
    let y = b"a(b(d,(0)),c)";
    assert_eq!(expression(x), expression(y));
}

#[test]
fn test_ternary() {
    let x = b"a ( b ) ? c ( d ) : e";
    assert_done!(expression(x));
}

#[test]
fn test_logic_or() {
    let x = b"a || b";
    assert_done!(expression(x));
}

#[bench]
fn bench_some(b: &mut test::Bencher) {
    b.iter(|| {
               let x = test::black_box(b"a ( b ( d , ( 0 < 3 ) ) , c | d )");
               expression(x)
           })
}

fn main() {
    let x = b"a - ! - b";
    println!("{:?}", expression(x));
}
<commit_msg>Use owned strings<commit_after>#![feature(test)]
extern crate test;

#[macro_use]
extern crate nom;

use nom::digit;
use nom::alphanumeric;
use std::str::FromStr;
use std::str;

named!(symbol_name, call!(alphanumeric));

named!(number<u64>,
    map_res!(
        map_res!(
            digit,
            str::from_utf8
        ),
        FromStr::from_str
    )
);

// ========================================================================== //

#[derive(Debug, PartialEq)]
pub enum UnaryOperator {
    LogicNot,
    Minus,
    BitwiseNot,
}

#[derive(Debug, PartialEq)]
pub enum BinaryOperator {
    LogicOr,
    LogicAnd,
    BitwiseOr,
    BitwiseAnd,
    Equals,
    NotEquals,
    Lesser,
    Greater,
    LesserOrEquals,
    GreaterOrEquals,
    ShiftRight,
    ShiftLeft,
    Plus,
    Minus,
    Multiply,
    Divide,
    Remainder,
}

#[derive(Debug, PartialEq)]
pub enum Expression {
    Ident(String),
    Number(u64),
    Call {
        function: String,
        arguments: Vec<Expression>,
    },
    UnaryOp {
        operator: UnaryOperator,
        right: Box<Expression>,
    },
    BinaryOp {
        left: Box<Expression>,
        operator: BinaryOperator,
        right: Box<Expression>,
    },
    TernaryOp {
        condition: Box<Expression>,
        left: Box<Expression>,
        right: Box<Expression>,
    },
}

named!(value_ident<Expression>, map!(
    map_res!(
        symbol_name,
        str::from_utf8
    ),
    |x: &str| Expression::Ident(x.into())
));

named!(value_number<Expression>, map!(
    number,
    |x| Expression::Number(x)
));

named!(value_nested<Expression>, delimited!(
    tag!("("),
    ws!(expression),
    tag!(")")
));

named!(value_call<Expression>, do_parse!(
    func: map_res!(symbol_name, str::from_utf8)
    >>
    ws!(tag!("("))
    >>
    args: separated_list!(
        ws!(tag!(",")),
        expression
    )
    >>
    ws!(tag!(")"))
    >>
    (Expression::Call{
        function: func.into(),
        arguments: args
    })
));

named!(pub value<Expression>, alt_complete!(
    value_nested | value_call | value_number | value_ident
));

named!(expr_unary_op<Expression>, do_parse!(
    op: ws!(alt_complete!(
        tag!("-") | tag!("!") | tag!("~")
    ))
    >>
    right: expr_level_1
    >>
    (Expression::UnaryOp{
        operator: match op[0] as char {
            '-' => UnaryOperator::Minus,
            '!' => UnaryOperator::LogicNot,
            '~' => UnaryOperator::BitwiseNot,
            _ => panic!("Invalid operator"),
        },
        right: Box::new(right),
    })
));

named!(expr_level_1<Expression>, alt_complete!(
    expr_unary_op | value
));

named!(expr_level_2<Expression>, do_parse!(
    first: expr_level_1
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("*") | tag!("/") | tag!("%")
            )),
            expr_level_1
        ),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0[0] as char {
                    '*' => BinaryOperator::Multiply,
                    '/' => BinaryOperator::Divide,
                    '%' => BinaryOperator::Remainder,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_3<Expression>, do_parse!(
    first: expr_level_2
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("+") | tag!("-")
            )),
            expr_level_2
        ),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0[0] as char {
                    '+' => BinaryOperator::Plus,
                    '-' => BinaryOperator::Minus,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_4<Expression>, do_parse!(
    first: expr_level_3
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("<<") | tag!(">>")
            )),
            expr_level_3
        ),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0 {
                    b"<<" => BinaryOperator::ShiftLeft,
                    b">>" => BinaryOperator::ShiftRight,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_5<Expression>, do_parse!(
    first: expr_level_4
    >>
    fold: fold_many0!(pair!(
            ws!(alt_complete!(
                tag!("==") | tag!("!=") | tag!("<=") | tag!(">=") | tag!("<") | tag!(">")
            )),
            expr_level_4
        ),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: match new.0 {
                    b"==" => BinaryOperator::Equals,
                    b"!=" => BinaryOperator::NotEquals,
                    b"<=" => BinaryOperator::LesserOrEquals,
                    b">=" => BinaryOperator::GreaterOrEquals,
                    b"<" => BinaryOperator::Lesser,
                    b">" => BinaryOperator::Greater,
                    _ => panic!("Invalid operator"),
                },
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_6<Expression>, do_parse!(
    first: expr_level_5
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("&")), expr_level_5),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::BitwiseAnd,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_7<Expression>, do_parse!(
    first: expr_level_6
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("|")), expr_level_6),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::BitwiseOr,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_8<Expression>, do_parse!(
    first: expr_level_7
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("&&")), expr_level_7),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::LogicAnd,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_level_9<Expression>, do_parse!(
    first: expr_level_8
    >>
    fold: fold_many0!(
        pair!(ws!(tag!("||")), expr_level_8),
        first,
        |prev, new: (&'a [u8], Expression)| {
            Expression::BinaryOp {
                left: Box::new(prev),
                operator: BinaryOperator::LogicOr,
                right: Box::new(new.1)
            }
        }
    )
    >>
    (fold)
));

named!(expr_ternary_op<Expression>, do_parse!(
    cond: expr_level_9
    >>
    ws!(tag!("?"))
    >>
    left: expression
    >>
    ws!(tag!(":"))
    >>
    right: expression
    >>
    (Expression::TernaryOp{
        condition: Box::new(cond),
        left: Box::new(left),
        right: Box::new(right),
    })
));

named!(pub expression<Expression>, alt_complete!(
    expr_ternary_op | expr_level_9
));

macro_rules! assert_done {
    ($res:expr) => (
        match $res {
            nom::IResult::Done(b"", v @ _) => {
                println!("ok: {:?}", v);
            }
            r @ _ => panic!("fail: {:?}", r),
        }
    )
}

#[test]
fn test_ws() {
    let x = b"a ( b ( d , ( 0 ) ) , c )";
    assert_done!(expression(x));
    let y = b"a(b(d,(0)),c)";
    assert_eq!(expression(x), expression(y));
}

#[test]
fn test_ternary() {
    let x = b"a ( b ) ? c ( d ) : e";
    assert_done!(expression(x));
}

#[test]
fn test_logic_or() {
    let x = b"a || b";
    assert_done!(expression(x));
}

#[bench]
fn bench_some(b: &mut test::Bencher) {
    b.iter(|| {
               let x = test::black_box(b"((a((1) - - - (2) - - - (3)) + 1))");
               expression(x)
           })
}

fn main() {
    let x = b"((a((1) - - - (2) - - - (3)) + 1))";
    println!("{:#?}", expression(x));
}
<|endoftext|>


<commit_before><commit_msg>Add the 7 optional FSL intents<commit_after><|endoftext|>


<commit_before>mod names;

///Search for a mnemonic, synonym or other name in the ascii name
///table and return any one that matches.
pub fn lookup_by_name(name: &str) -> Option<char> {
    for &(ch_name, ch) in names::NAMES {
        if name.to_lowercase() == ch_name {
            return Some(ch);
        }
    }
    None
}

pub fn additional_names(ch: char) -> Option<names::Information> {
    let code = ch as u8;
    if code < 128 {
        Some(names::PRINTABLE_CHARS[code as usize].clone())
    } else {
        None
    }
}
<commit_msg>Correctly search for mnemonic names<commit_after>mod names;

///Search for a mnemonic, synonym or other name in the ascii name
///table and return any one that matches.
pub fn lookup_by_name(name: &str) -> Option<char> {
    for &(ch_name, ch) in names::NAMES {
        if name.to_lowercase() == ch_name.to_lowercase() {
            return Some(ch);
        }
    }
    None
}

pub fn additional_names(ch: char) -> Option<names::Information> {
    let code = ch as u8;
    if code < 128 {
        Some(names::PRINTABLE_CHARS[code as usize].clone())
    } else {
        None
    }
}
<|endoftext|>